From ea0f8ce62f64f290614d73a7311ea2e28289dbc4 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Thu, 23 Jul 2026 18:55:54 +1200 Subject: [PATCH 01/21] feat: state machine init --- CHANGELOG.md | 22 + src/compact.rs | 203 ++++---- src/compact/types.rs | 169 +++++- src/config.rs | 51 +- src/engine.rs | 2 + src/engine/bare.rs | 2 +- src/engine/machine.rs | 1130 +++++++++++++++++++++++++++++++++++++++++ src/error.rs | 2 +- src/presets.rs | 2 +- 9 files changed, 1467 insertions(+), 116 deletions(-) create mode 100644 src/engine/machine.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index ec9774e..61579a0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. ### Added +- `engine::machine::LoopMachine` and supporting types — a sans-IO, serializable + state machine that owns every agent-loop decision (turn counting, max-turn + enforcement, tool-call validity, stop-reason routing, compaction trigger, + history, cancellation). `Serialize + Deserialize`, with no `async`, no + `tokio`, and no `ApiClient` in its surface. Includes `RunConfig`, + `MachineStep` (`CallLLM`/`CallTools`/`Compact`/`Done`), `ModelTurn`, + `PendingToolCall`, `MachineOutcome`, and `MachineState`. The machine is not + yet wired into `BareLoop` in this release; it is public so a driver can be + written against it. +- `LoopError` now derives `Serialize`, `Deserialize`, `PartialEq`, and `Eq`. +- `compact::types::CompactReason` now derives `Serialize` and `Deserialize`. - `DeltaPart::Thinking { text }` variant + `on_thinking_delta` observer event (`ThinkingDeltaContext`): reasoning-model tokens (Claude extended-thinking, DeepSeek-R1, OpenAI o-series, Gemini 2.5+) are now routed as their own @@ -185,6 +196,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. ### Changed +- **Breaking (compaction thresholds → basis points):** the compaction trigger + threshold and compaction-target fraction are now integers in basis points + (0–10_000; 10_000 = 100%) instead of `f64` fractions. This removes lossy + `u64 → f64` casts from the threshold arithmetic. Affected APIs: + `ContextManager::with_threshold(u16)` and `with_compact_target_pct(u16)` + (were `f64`); `ContextManager::threshold() -> u16` and + `compact_target_pct() -> u16` (were `f64`); + `LoopConfig::with_compact_threshold(u16)` and the `compact_threshold` field + (were `f64`). The default is `8_000` (was `0.80`); the clamp range is + `[1_000, 10_000]` (was `[0.1, 1.0]`). Migration: multiply existing `f64` + values by `10_000` and round — `0.80 → 8_000`, `0.50 → 5_000`, `0.70 → 7_000`. - **Breaking (renames):** consuming builder methods that return `Self` are now uniformly prefixed `with_`, matching the crate-wide convention. The old no-prefix names are removed. Affected types and methods (old → new): diff --git a/src/compact.rs b/src/compact.rs index 9e9f2aa..7d40a05 100644 --- a/src/compact.rs +++ b/src/compact.rs @@ -42,10 +42,10 @@ //! //! let manager = ContextManager::new(Arc::new(compactor)) //! .with_context_window(200_000) -//! .with_threshold(0.80); +//! .with_threshold(8_000); //! //! assert_eq!(manager.context_window(), 200_000); -//! assert!((manager.threshold() - 0.80).abs() < f64::EPSILON); +//! assert_eq!(manager.threshold(), 8_000); //! ``` //! //! # Integration with `BareLoop` @@ -177,9 +177,9 @@ pub trait ContextCompactor: Send + Sync { /// 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_threshold(8_000) // triggers at 160k tokens /// .with_compact_target(CompactBase::Context) // target = % of 200k -/// .with_compact_target_pct(0.50); // compact to 50% of 200k = 100k +/// .with_compact_target_pct(5_000); // compact to 50% of 200k = 100k /// ``` #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub enum CompactBase { @@ -193,11 +193,11 @@ pub enum CompactBase { /// Target is a percentage of the trigger threshold. /// - /// `target = compact_threshold_tokens × compact_target_pct` + /// `target = compact_threshold_tokens × compact_target_pct / 10_000` /// - /// 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). + /// This is the default. With the default `threshold = 8_000` bps (80%) and + /// `compact_target_pct = 7_000` bps (70%), compaction targets 56% of the + /// context window (`8000 × 7000 / 10_000² = 0.56`). #[default] Threshold, } @@ -233,7 +233,7 @@ pub enum CompactBase { /// /// let manager = ContextManager::new(Arc::new(compactor)) /// .with_context_window(1_000) -/// .with_threshold(0.80); +/// .with_threshold(8_000); /// /// // Short conversation — no compaction needed. /// let messages = vec![ @@ -245,17 +245,54 @@ pub enum CompactBase { /// ``` pub struct ContextManager { /// The compaction strategy. + /// + /// The [`ContextCompactor`](super::compact::ContextCompactor) implementation + /// that performs the actual summarization or truncation when compaction runs. + /// The manager decides *when* and *how much* to compact; this field decides + /// *how* the messages are rewritten. compactor: Arc, + /// Model context window size in tokens. + /// + /// The hard upper bound on tokens the model accepts in one request. Every + /// threshold and target is expressed as a fraction of this value, so it is + /// the denominator for [`compact_threshold_tokens`](Self::compact_threshold_tokens) + /// and the emergency-zone and compaction-target calculations. context_window: u64, - /// Threshold (0.0–1.0) at which compaction triggers. - threshold: f64, + + /// Threshold in basis points (`0–10_000`; `10_000` = 100%) at which compaction triggers. + /// + /// When estimated token usage reaches `threshold_bps * context_window / 10_000`, + /// [`should_compact`](Self::should_compact) returns `true` and a compaction + /// pass runs before the next turn. Set via [`with_threshold`](Self::with_threshold), + /// which clamps to `[1_000, 10_000]`; defaults to `8_000` (80%). + threshold_bps: u16, + /// Whether auto-compaction is enabled. + /// + /// When `false`, [`should_compact`](Self::should_compact) always returns + /// `false` and [`ensure_context_fits`](Self::ensure_context_fits) never + /// triggers compaction — the host must manage context size manually (useful + /// in tests or fixed-length sessions). Toggle via + /// [`with_auto_compact`](Self::with_auto_compact). auto_compact: bool, + /// The base used to compute the compaction target. + /// + /// Whether [`compact_target_tokens`](Self::compact_target_tokens) measures + /// the post-compaction size against the trigger threshold + /// ([`CompactBase::Threshold`], the default) or the full context window + /// ([`CompactBase::Context`]). Set via + /// [`with_compact_target`](Self::with_compact_target). compact_base: CompactBase, - /// The fraction (0.0–1.0) of the target base to compact down to. - compact_target: f64, + + /// Fraction of the target base to compact down to, in basis points (`0–10_000`). + /// + /// The post-compaction size aimed for: `compact_target_bps * base / 10_000`, + /// where `base` is determined by [`compact_base`](Self::compact_base). + /// Defaults to `7_000` (70%); set and clamped to `[1_000, 10_000]` via + /// [`with_compact_target_pct`](Self::with_compact_target_pct). + compact_target_bps: u16, } impl ContextManager { @@ -266,19 +303,19 @@ impl ContextManager { /// | Setting | Default | /// |----------------------|------------------------------| /// | `context_window` | 200_000 | - /// | `threshold` | 0.80 | + /// | `threshold` | 8_000 bps (80%) | /// | `auto_compact` | `true` | /// | `compact_target` | [`CompactBase::Threshold`] | - /// | `compact_target_pct` | 0.70 | + /// | `compact_target_pct` | 7_000 bps (70%) | #[must_use] pub fn new(compactor: Arc) -> Self { Self { compactor, context_window: 200_000, - threshold: 0.80, + threshold_bps: 8_000, auto_compact: true, compact_base: CompactBase::Threshold, - compact_target: 0.70, + compact_target_bps: 7_000, } } @@ -292,12 +329,12 @@ impl ContextManager { self } - /// Set the compaction threshold (0.0–1.0). + /// Set the compaction threshold in basis points (`0–10_000`; `10_000` = 100%). /// - /// Clamped to `[0.1, 1.0]` to prevent degenerate configurations. + /// Clamped to `[1_000, 10_000]` (10%–100%) to prevent degenerate configurations. #[must_use] - pub fn with_threshold(mut self, threshold: f64) -> Self { - self.threshold = threshold.clamp(0.1, 1.0); + pub fn with_threshold(mut self, threshold_bps: u16) -> Self { + self.threshold_bps = threshold_bps.clamp(1_000, 10_000); self } @@ -322,13 +359,14 @@ impl ContextManager { self } - /// Set the fraction (0.0–1.0) of the target base to compact down to. + /// Set the fraction of the target base to compact down to, in basis points + /// (`0–10_000`; `10_000` = 100%). /// - /// Clamped to `[0.1, 1.0]` to prevent degenerate configurations. - /// Defaults to `0.70` (70%). + /// Clamped to `[1_000, 10_000]` (10%–100%) to prevent degenerate configurations. + /// Defaults to `7_000` (70%). #[must_use] - pub fn with_compact_target_pct(mut self, pct: f64) -> Self { - self.compact_target = pct.clamp(0.1, 1.0); + pub fn with_compact_target_pct(mut self, pct_bps: u16) -> Self { + self.compact_target_bps = pct_bps.clamp(1_000, 10_000); self } @@ -341,13 +379,13 @@ impl ContextManager { self.context_window } - /// The compaction threshold (0.0–1.0). + /// The compaction threshold in basis points (`0–10_000`; `10_000` = 100%). /// /// Compaction triggers when estimated tokens reach - /// `context_window * threshold`. See [`with_threshold`](Self::with_threshold). + /// `context_window * threshold / 10_000`. See [`with_threshold`](Self::with_threshold). #[must_use] - pub fn threshold(&self) -> f64 { - self.threshold + pub fn threshold(&self) -> u16 { + self.threshold_bps } /// Whether auto-compaction is enabled. @@ -367,27 +405,25 @@ impl ContextManager { self.compact_base } - /// The fraction (0.0–1.0) of the target base to compact down to. + /// The fraction of the target base to compact down to, in basis points + /// (`0–10_000`; `10_000` = 100%). /// /// See [`with_compact_target_pct`](Self::with_compact_target_pct). #[must_use] - pub fn compact_target_pct(&self) -> f64 { - self.compact_target + pub fn compact_target_pct(&self) -> u16 { + self.compact_target_bps } /// 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. + /// Equal to `context_window * threshold_bps / 10_000`, computed in + /// saturating integer arithmetic (widened to `u128` to avoid overflow). #[must_use] - #[allow( - clippy::cast_possible_truncation, - clippy::cast_sign_loss, - clippy::cast_precision_loss - )] pub fn compact_threshold_tokens(&self) -> u64 { - (self.threshold * self.context_window as f64) as u64 + const BASIS: u128 = 10_000; + let product = + u128::from(self.threshold_bps).saturating_mul(u128::from(self.context_window)) / BASIS; + u64::try_from(product).unwrap_or(u64::MAX) } /// The token count to compact down to. @@ -395,20 +431,17 @@ impl ContextManager { /// 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` + /// - [`CompactBase::Threshold`]: `compact_threshold_tokens × pct_bps / 10_000` + /// - [`CompactBase::Context`]: `context_window × pct_bps / 10_000` #[must_use] - #[allow( - clippy::cast_possible_truncation, - clippy::cast_sign_loss, - clippy::cast_precision_loss - )] pub fn compact_target_tokens(&self) -> u64 { + const BASIS: u128 = 10_000; 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 + let product = u128::from(self.compact_target_bps).saturating_mul(u128::from(base)) / BASIS; + u64::try_from(product).unwrap_or(u64::MAX) } /// Estimate the token count for a slice of messages. @@ -686,7 +719,7 @@ impl fmt::Debug for ContextManager { .field("threshold", &self.threshold()) .field("auto_compact", &self.auto_compact) .field("compact_target", &self.compact_base) - .field("compact_target_pct", &self.compact_target) + .field("compact_target_pct", &self.compact_target_pct()) .finish_non_exhaustive() } } @@ -733,7 +766,7 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)); assert_eq!(manager.context_window(), 200_000); - assert!((manager.threshold() - 0.80).abs() < f64::EPSILON); + assert_eq!(manager.threshold(), 8_000); assert!(manager.auto_compact()); } @@ -742,22 +775,22 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(100_000) - .with_threshold(0.50) + .with_threshold(5_000) .with_auto_compact(false); assert_eq!(manager.context_window(), 100_000); - assert!((manager.threshold() - 0.50).abs() < f64::EPSILON); + assert_eq!(manager.threshold(), 5_000); assert!(!manager.auto_compact()); } #[test] fn test_threshold_clamped() { let compactor = TruncatingCompactor::new(); - let manager = ContextManager::new(Arc::new(compactor)).with_threshold(0.01); - assert!((manager.threshold() - 0.1).abs() < f64::EPSILON); + let manager = ContextManager::new(Arc::new(compactor)).with_threshold(100); + assert_eq!(manager.threshold(), 1_000); let compactor2 = TruncatingCompactor::new(); - let manager = ContextManager::new(Arc::new(compactor2)).with_threshold(2.0); - assert!((manager.threshold() - 1.0).abs() < f64::EPSILON); + let manager = ContextManager::new(Arc::new(compactor2)).with_threshold(20_000); + assert_eq!(manager.threshold(), 10_000); } #[test] @@ -765,8 +798,8 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(1_000) - .with_threshold(0.80); - // 800 is the threshold for 1000 * 0.80 + .with_threshold(8_000); + // 800 is the threshold for 1000 * 8000 / 10_000 assert!(!manager.should_compact(799)); } @@ -775,7 +808,7 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(1_000) - .with_threshold(0.80); + .with_threshold(8_000); assert!(manager.should_compact(800)); } @@ -784,7 +817,7 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(1_000) - .with_threshold(0.50); // threshold at 500, but emergency at 950 + .with_threshold(5_000); // threshold at 500, but emergency at 950 assert!(manager.should_compact(950)); } @@ -808,7 +841,7 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(1_000) - .with_threshold(0.80); + .with_threshold(8_000); assert_eq!( manager.compact_reason(800), CompactReason::ThresholdExceeded @@ -827,58 +860,58 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(200_000) - .with_threshold(0.80); + .with_threshold(8_000); 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 + // Default: CompactBase::Threshold with pct 7_000 bps (70%) + // threshold = 200_000 * 8_000 / 10_000 = 160_000 + // target = 160_000 * 7_000 / 10_000 = 112_000 let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(200_000) - .with_threshold(0.80); + .with_threshold(8_000); 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 + // CompactBase::Context with pct 5_000 bps (50%) + // target = 200_000 * 5_000 / 10_000 = 100_000 let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(200_000) - .with_threshold(0.80) + .with_threshold(8_000) .with_compact_target(CompactBase::Context) - .with_compact_target_pct(0.50); + .with_compact_target_pct(5_000); 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 + // CompactBase::Threshold with pct 5_000 bps (50%) + // threshold = 200_000 * 8_000 / 10_000 = 160_000 + // target = 160_000 * 5_000 / 10_000 = 80_000 let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(200_000) - .with_threshold(0.80) + .with_threshold(8_000) .with_compact_target(CompactBase::Threshold) - .with_compact_target_pct(0.50); + .with_compact_target_pct(5_000); 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 manager = ContextManager::new(Arc::new(compactor)).with_compact_target_pct(100); + assert_eq!(manager.compact_target_pct(), 1_000); 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); + let manager2 = ContextManager::new(Arc::new(compactor2)).with_compact_target_pct(20_000); + assert_eq!(manager2.compact_target_pct(), 10_000); } #[test] @@ -886,7 +919,7 @@ mod tests { 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); + assert_eq!(manager.compact_target_pct(), 7_000); } #[tokio::test] @@ -971,7 +1004,7 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(1_000_000) // very large window - .with_threshold(0.80); + .with_threshold(8_000); let msgs = make_conversation(3); let result = manager.ensure_context_fits(msgs.clone(), 1).await; match result { @@ -987,12 +1020,12 @@ mod tests { let compactor = TruncatingCompactor::new() .with_min_messages(4) .with_preserve_recent(2); - // Use a window that triggers compaction at 50% but still fits + // Use a window that triggers compaction at 10% but still fits // the preserved 2 messages after compaction. // 2 messages ≈ 18 tokens, so 200 tokens is plenty of headroom. let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(200) - .with_threshold(0.10); + .with_threshold(1_000); let msgs = make_conversation(20); // 40 messages let result = manager.ensure_context_fits(msgs, 1).await; match result { diff --git a/src/compact/types.rs b/src/compact/types.rs index daf47ff..d6e339c 100644 --- a/src/compact/types.rs +++ b/src/compact/types.rs @@ -11,6 +11,7 @@ //! - [`EnsureContextResult`] — result of [`ContextManager::ensure_context_fits`](super::ContextManager::ensure_context_fits). use crate::message::Message; +use serde::{Deserialize, Serialize}; use std::fmt; // =================================================== @@ -22,13 +23,30 @@ use std::fmt; /// 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)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum CompactReason { /// Token usage exceeded the configured threshold percentage. + /// + /// This is the routine, expected trigger: estimated context size crossed the + /// [`ContextManager`](super::ContextManager)'s threshold (80% by default), + /// so a compaction pass runs proactively before the next turn to keep the + /// context comfortably below the window. ThresholdExceeded, + /// Token usage is dangerously close to the context window limit. + /// + /// This is the fallback safety trigger, firing when usage reaches the + /// emergency zone (95% of the window) regardless of the configured + /// threshold. An emergency compaction should compact more aggressively + /// than a routine threshold pass because the conversation is on the verge + /// of overflowing the model's window. Emergency, + /// Compaction was explicitly requested (e.g. by the agent or a tool). + /// + /// Compaction was forced by an explicit caller rather than by a size-based + /// trigger — for example a host application compacting on demand, or a tool + /// that wants to free context before producing a large result. Manual, } @@ -55,12 +73,31 @@ impl fmt::Display for CompactReason { #[derive(Debug, Clone)] pub struct CompactionContext { /// Estimated token count before compaction. + /// + /// The compactor's input size, computed with the crate's standard + /// 4-chars-per-token heuristic. Compaction aims to bring the post-compaction + /// size below this so the conversation fits with headroom for the next turn. pub tokens_before: u64, + /// Why compaction was triggered. + /// + /// Compactors may use the trigger to pick a strategy — an + /// [`Emergency`](CompactReason::Emergency) trigger warrants more aggressive + /// summarization than a routine [`ThresholdExceeded`](CompactReason::ThresholdExceeded). pub reason: CompactReason, + /// The model's context window size. + /// + /// The hard upper bound on tokens the model accepts in one request. This is + /// the denominator every compaction threshold and target is expressed + /// against, so the compactor can decide how much to keep. pub context_window: u64, + /// The current turn number in the session. + /// + /// Zero-indexed within the run. Useful for compaction strategies that weight + /// recent turns more heavily, or for correlating a compaction pass back to + /// the turn that triggered it in logs. pub turn: usize, } @@ -76,14 +113,42 @@ pub struct CompactionContext { #[derive(Debug, Clone)] pub struct CompactionOutcome { /// The compacted message list. + /// + /// The messages that remain after compaction — typically a summary or + /// truncation of the original conversation. The caller feeds this list back + /// into the loop as the new history. May equal the input when compaction + /// decided no change was needed (see [`CompactionOutcome::no_change`]). pub messages: Vec, + /// Estimated token count after compaction. + /// + /// The post-compaction size of [`messages`](Self::messages), estimated with + /// the same heuristic used for the pre-compaction count, so before/after + /// values are directly comparable. pub tokens_after: u64, + /// Estimated tokens saved by compaction. + /// + /// The difference between the pre-compaction token count and + /// [`tokens_after`](Self::tokens_after). Zero when compaction made no change + /// or when it enlarged the conversation (e.g. injecting a summary that + /// outweighs the messages it replaced). pub tokens_saved: u64, + /// Whether compaction succeeded. + /// + /// `true` when the compactor produced a usable message list, even if that + /// list is unchanged. `false` only when the compactor itself failed — in + /// that case [`error`](Self::error) describes the failure and + /// [`messages`](Self::messages) typically holds the original input. pub success: bool, + /// Error message if compaction failed. + /// + /// `Some(description)` when [`success`](Self::success) is `false`, carrying + /// the compactor's human-readable failure reason. `None` on success. Typed + /// as a [`String`] because it surfaces to observers/logs, not to programmatic + /// control flow (the loop treats any failed compaction uniformly). pub error: Option, } @@ -141,40 +206,107 @@ impl CompactionOutcome { #[derive(Debug, Clone)] pub struct CompactTelemetry { /// Why compaction was triggered. + /// + /// The [`CompactReason`] that caused this pass, useful for distinguishing + /// routine threshold compactions from emergency or manual ones when reading + /// the telemetry. pub trigger: CompactReason, + /// Conversation stats before compaction. + /// + /// A snapshot of the conversation as it was when compaction began — message + /// counts broken down by role and token estimate. See [`PreCompactStats`]. pub pre_compact: PreCompactStats, + /// Conversation stats after compaction. + /// + /// A snapshot of the conversation after compaction completed, plus the + /// savings achieved. Compare against [`pre_compact`](Self::pre_compact) to + /// measure the effect of the pass. See [`PostCompactStats`]. pub post_compact: PostCompactStats, + /// Wall-clock duration of the compaction. + /// + /// Time spent inside the compactor's `compact` call for this pass, measured + /// from just before the call to just after. Excludes the token-estimate + /// bookkeeping done before and after the call itself. pub duration: std::time::Duration, } /// Conversation statistics captured before compaction. +/// +/// A breakdown of the conversation's shape at the moment compaction begins: +/// how many messages there are, their estimated token cost, and how they split +/// across user/assistant/tool roles. Captured by +/// [`ContextManager::build_telemetry`](super::ContextManager::build_telemetry) +/// and bundled into [`CompactTelemetry::pre_compact`]. #[derive(Debug, Clone)] pub struct PreCompactStats { /// Total number of messages in the conversation. + /// + /// Every message in the history about to be compacted, regardless of role + /// or content. This is the input size the compactor operates on. pub total_messages: usize, + /// Estimated token count. + /// + /// The pre-compaction token estimate of the whole conversation, using the + /// standard 4-chars-per-token heuristic. This is the number compared against + /// the threshold to decide whether compaction was needed. pub estimated_tokens: u64, + /// Number of user-role messages. + /// + /// Messages whose role is [`User`](crate::message::Role::User). Includes + /// both genuine user turns and tool-result messages, which are conventionally + /// sent with the user role. pub user_messages: usize, + /// Number of assistant-role messages. + /// + /// Messages whose role is [`Assistant`](crate::message::Role::Assistant) — + /// the model's own responses, including any that carried tool-call requests. pub assistant_messages: usize, + /// Number of messages containing tool calls or results. + /// + /// Messages with at least one tool-call or tool-result part, regardless of + /// role. These are often worth preserving across compaction because they + /// carry the intermediate state of the tool loop. pub tool_messages: usize, } /// Conversation statistics captured after compaction. +/// +/// A summary of the conversation's shape after compaction completes, together +/// with how much the pass reclaimed. Captured by +/// [`ContextManager::build_telemetry`](super::ContextManager::build_telemetry) +/// and bundled into [`CompactTelemetry::post_compact`]. #[derive(Debug, Clone)] pub struct PostCompactStats { /// Total number of messages after compaction. + /// + /// The size of the compacted message list. Smaller than the pre-compaction + /// [`total_messages`](PreCompactStats::total_messages) when compaction + /// removed or summarized messages; equal when it made no change. pub total_messages: usize, /// Estimated token count after compaction. + /// + /// The post-compaction token estimate, comparable to + /// [`estimated_tokens`](PreCompactStats::estimated_tokens) from the + /// pre-compaction snapshot. The difference is [`tokens_saved`](Self::tokens_saved). pub estimated_tokens: u64, /// Tokens removed by compaction. + /// + /// How many tokens the pass reclaimed: the pre-compaction estimate minus + /// the post-compaction estimate. Saturates at zero, so it never goes + /// negative even if a summary injection made the conversation larger. pub tokens_saved: u64, /// Percentage of tokens saved (0–100). + /// + /// [`tokens_saved`](Self::tokens_saved) as a share of the pre-compaction + /// estimate, expressed as a whole-number percentage. Clamped to `0..=100`; + /// `0` when nothing was saved or when the pre-compaction estimate was zero. pub percent_saved: u8, } @@ -190,14 +322,39 @@ pub struct PostCompactStats { #[derive(Debug, Clone)] pub struct ContextOverflow { /// Estimated token count of the conversation. + /// + /// How many tokens the conversation occupies when it overflows — the same + /// heuristic estimate used everywhere else in the subsystem. Compare + /// against [`context_window`](Self::context_window) (or use + /// [`overflow`](Self::overflow)) to see by how much it exceeded the limit. pub tokens_used: u64, + /// The model's context window size. + /// + /// The hard token limit the conversation failed to fit under, even after a + /// compaction pass. The denominator [`utilization`](Self::utilization) is + /// measured against. pub context_window: u64, + /// How many messages were in the conversation. + /// + /// The message count at the point of overflow, useful for diagnosing + /// whether the overflow came from many small messages or a few large ones. pub message_count: usize, + /// The reason compaction was attempted. + /// + /// The [`CompactReason`] that triggered the (failed) compaction attempt. + /// An [`Emergency`](CompactReason::Emergency) trigger here means even an + /// aggressive compaction could not bring the conversation back under the + /// window. pub trigger: CompactReason, + /// Error from the compactor, if compaction was attempted. + /// + /// `Some(description)` when a compactor ran but returned an error that + /// prevented recovery; `None` when the conversation was simply too large + /// to reduce (compaction succeeded but the result still overflowed). pub compactor_error: Option, } @@ -245,8 +402,18 @@ impl std::error::Error for ContextOverflow {} #[derive(Debug, Clone)] pub enum EnsureContextResult { /// Compaction occurred and produced a shorter message list. + /// + /// The wrapped [`CompactionOutcome`] carries the compacted messages, the + /// token savings, and whether the pass succeeded. Feed + /// [`outcome.messages`](CompactionOutcome::messages) back into the loop as + /// the new history. Compacted(CompactionOutcome), + /// No compaction was needed; messages returned as-is. + /// + /// The conversation fit comfortably within the threshold, so no compaction + /// pass ran. The wrapped message list is the original input, unchanged; use + /// it directly as the next-turn history. NoAction(Vec), } diff --git a/src/config.rs b/src/config.rs index 00ab6e7..c22034c 100644 --- a/src/config.rs +++ b/src/config.rs @@ -79,14 +79,14 @@ pub struct LoopConfig { /// ([`validate`](LoopConfig::validate)). pub context_window: u64, - /// Threshold to trigger auto-compaction, as a fraction of the context - /// window (0.0–1.0). + /// Threshold to trigger auto-compaction, in basis points (`0–10_000`; + /// `10_000` = 100% of the context window). /// /// When the estimated token usage of the conversation exceeds - /// `compact_threshold * context_window`, the compactor runs before the next - /// turn to avoid an overflow. Defaults to `0.80`. Must be finite and in - /// `[0.0, 1.0]` ([`validate`](LoopConfig::validate)). - pub compact_threshold: f64, + /// `compact_threshold * context_window / 10_000`, the compactor runs before + /// the next turn to avoid an overflow. Defaults to `8_000` (80%). Must be in + /// `[0, 10_000]` ([`validate`](LoopConfig::validate)). + pub compact_threshold: u16, /// Whether auto-compaction is enabled. /// @@ -121,7 +121,7 @@ impl Default for LoopConfig { /// | [`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` | + /// | [`compact_threshold`](LoopConfig::compact_threshold) | `8_000` (80%) | /// | [`auto_compact`](LoopConfig::auto_compact) | `true` | /// | [`parallel_tool_dispatch`](LoopConfig::parallel_tool_dispatch) | `Sequential`, `max_concurrency: 8` | /// @@ -142,7 +142,7 @@ impl Default for LoopConfig { max_turns: 200, max_tokens: 16_384, context_window: 200_000, - compact_threshold: 0.80, + compact_threshold: 8_000, auto_compact: true, parallel_tool_dispatch: ParallelDispatchConfig::default(), } @@ -218,16 +218,16 @@ impl LoopConfig { self } - /// Set the auto-compaction trigger threshold as a fraction of the context - /// window (0.0–1.0). + /// Set the auto-compaction trigger threshold in basis points (`0–10_000`; + /// `10_000` = 100% of the context window). /// - /// When estimated token usage exceeds `compact_threshold * context_window`, - /// compaction runs before the next turn. This builder stores the value - /// verbatim — it does **not** clamp; the range is enforced by - /// [`validate`](LoopConfig::validate), which is the single source of truth - /// for the `[0.0, 1.0]` bound. + /// When estimated token usage exceeds + /// `compact_threshold * context_window / 10_000`, compaction runs before the + /// next turn. This builder stores the value verbatim — it does **not** clamp; + /// the range is enforced by [`validate`](LoopConfig::validate), which is the + /// single source of truth for the `[0, 10_000]` bound. #[must_use] - pub fn with_compact_threshold(mut self, compact_threshold: f64) -> Self { + pub fn with_compact_threshold(mut self, compact_threshold: u16) -> Self { self.compact_threshold = compact_threshold; self } @@ -257,7 +257,7 @@ impl LoopConfig { /// Validate the configuration fields. /// /// Checks that: - /// - `compact_threshold` is in the range `[0.0, 1.0]` (not `NaN`). + /// - `compact_threshold` is in the range `[0, 10_000]`. /// - `max_turns` is greater than zero. /// - `context_window` is greater than zero. /// - `max_tokens` is greater than zero. @@ -276,7 +276,7 @@ impl LoopConfig { /// let config = LoopConfig::default(); /// assert!(config.validate().is_ok()); /// - /// let bad = LoopConfig::default().with_compact_threshold(1.5); + /// let bad = LoopConfig::default().with_compact_threshold(15_000); /// assert!(bad.validate().is_err()); /// ``` #[must_use = "validation errors should not be silently ignored"] @@ -301,12 +301,9 @@ impl LoopConfig { "model must not be empty".to_string(), )); } - if self.compact_threshold.is_nan() - || self.compact_threshold < 0.0 - || self.compact_threshold > 1.0 - { + if self.compact_threshold > 10_000 { return Err(crate::error::LoopError::Config(format!( - "compact_threshold must be in [0.0, 1.0], got {}", + "compact_threshold must be in [0, 10_000], got {}", self.compact_threshold ))); } @@ -504,8 +501,8 @@ mod tests { #[test] fn with_compact_threshold_stores_value_without_clamping() { - let config = LoopConfig::default().with_compact_threshold(1.5); - assert!((config.compact_threshold - 1.5).abs() < f64::EPSILON); + let config = LoopConfig::default().with_compact_threshold(15_000); + assert_eq!(config.compact_threshold, 15_000); } #[test] @@ -540,7 +537,7 @@ mod tests { let defaults = LoopConfig::default(); assert_eq!(config.max_tokens, defaults.max_tokens); assert_eq!(config.context_window, defaults.context_window); - assert!((config.compact_threshold - defaults.compact_threshold).abs() < f64::EPSILON); + assert_eq!(config.compact_threshold, defaults.compact_threshold); assert_eq!(config.auto_compact, defaults.auto_compact); } @@ -553,7 +550,7 @@ mod tests { #[test] fn validate_still_rejects_builder_built_invalid_config() { - let config = LoopConfig::default().with_compact_threshold(1.5); + let config = LoopConfig::default().with_compact_threshold(15_000); assert!(config.validate().is_err()); } } diff --git a/src/engine.rs b/src/engine.rs index 5b0c09c..cdbbca9 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -25,6 +25,8 @@ mod bare; pub mod contributor; pub mod loop_core; +pub mod machine; pub use bare::*; pub use contributor::*; +pub use machine::*; diff --git a/src/engine/bare.rs b/src/engine/bare.rs index af73bc8..a8f344a 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -535,7 +535,7 @@ impl BareLoop { /// .with_min_messages(6); /// let manager = ContextManager::new(Arc::new(compactor)) /// .with_context_window(200_000) - /// .with_threshold(0.80); + /// .with_threshold(8_000); /// /// let mut agent = BareLoop::new(client, registry, config); /// agent.set_context_manager(Arc::new(manager)); diff --git a/src/engine/machine.rs b/src/engine/machine.rs new file mode 100644 index 0000000..df99f28 --- /dev/null +++ b/src/engine/machine.rs @@ -0,0 +1,1130 @@ +//! Sans-IO agent-loop state machine. +//! +//! [`LoopMachine`] owns every *decision* the agent loop makes — turn counting, +//! max-turn enforcement, tool-call validity, stop-reason routing, compaction +//! triggering, history accumulation, and the cancellation flag — and performs +//! zero IO. It is advanced by feeding it outcomes via its methods; a thin IO +//! driver (the `BareLoop` in [`crate::engine`]) runs a `match +//! machine.next_step()` loop, performs the side effect each step requests, and +//! feeds the result back. +//! +//! Because the machine is pure and [`Serialize`] + [`Deserialize`], a run can be +//! serialized mid-flight and resumed in another process. +//! +//! [`Serialize`]: serde::Serialize +//! [`Deserialize`]: serde::Deserialize + +use serde::{Deserialize, Serialize}; + +use crate::compact::types::CompactReason; +use crate::config::ParallelDispatchConfig; +use crate::engine::loop_core::{StopReason, ToolCall}; +use crate::error::LoopError; +use crate::message::{Message, MessagePart, Role, ToolContent}; + +/// Per-run configuration. +/// +/// The slice of agent configuration that varies across `run()` calls on the +/// same agent (turn/token budgets, compaction policy, dispatch mode). It is +/// owned by the machine so that the machine's decisions are self-contained and +/// a serialized machine carries its configuration with it. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RunConfig { + /// Maximum number of turns before forcing completion. + /// + /// A safety cap on runaway loops: once the machine has taken this many + /// turns without the model finishing, the run ends with + /// [`MachineOutcome::MaxTurnsExceeded`]. Defaults to `200`. Set it per-run + /// to bound long-running tool loops. + pub max_turns: usize, + + /// Maximum tokens for each API response. + /// + /// The per-response cap the driver sends to the provider, limiting the + /// length of a single model reply. Defaults to `16_384`. + pub max_tokens: u32, + + /// Context window size in tokens. + /// + /// The hard upper bound on tokens the model accepts in one request. The + /// machine uses it as the denominator for the compaction trigger: the + /// [`MachineStep::Compact`] step fires once estimated context size reaches + /// [`compact_threshold`](Self::compact_threshold) of this window. Defaults + /// to `200_000`; should match the configured model's real window. + pub context_window: u64, + + /// Threshold to trigger auto-compaction, in basis points (`0–10_000`; + /// `10_000` = 100% of the context window). + /// + /// The machine emits [`MachineStep::Compact`] once the current context size + /// exceeds this fraction of [`context_window`](Self::context_window). + /// Defaults to `8_000` (80%). Only consulted when + /// [`auto_compact`](Self::auto_compact) is `true`. + pub compact_threshold: u16, + + /// Whether auto-compaction is enabled. + /// + /// When `false`, the machine never emits [`MachineStep::Compact`] on its + /// own — the host must manage context size manually (useful in tests or + /// fixed-length runs). When `true` (the default), the machine gates + /// compaction with [`compact_threshold`](Self::compact_threshold). + pub auto_compact: bool, + + /// How independent tool calls within a single turn are dispatched. + /// + /// Controls whether the driver runs a turn's independent tool calls + /// sequentially or in parallel, up to a concurrency limit. See + /// [`ParallelDispatchConfig`]. + pub parallel_tool_dispatch: ParallelDispatchConfig, +} + +impl Default for RunConfig { + fn default() -> Self { + Self { + max_turns: 200, + max_tokens: 16_384, + context_window: 200_000, + compact_threshold: 8_000, + auto_compact: true, + parallel_tool_dispatch: ParallelDispatchConfig::default(), + } + } +} + +/// One unit of work the driver must perform before feeding an outcome back. +/// +/// Returned by [`LoopMachine::next_step`]. The driver matches on this and, for +/// each non-terminal variant, performs the requested work and feeds the result +/// back into the machine via the matching method ([`LoopMachine::model_response`], +/// [`LoopMachine::tool_results`], [`LoopMachine::compaction_result`]). +#[derive(Debug, Clone, Serialize, Deserialize)] +#[non_exhaustive] +pub enum MachineStep { + /// Request an LLM call over the machine's current [`history`](LoopMachine::history). + /// + /// The driver builds the feed (the messages actually sent to the LLM) from + /// [`LoopMachine::history`], calls the provider, and feeds the completed + /// [`ModelTurn`] back via [`LoopMachine::model_response`]. `turn` is the + /// 1-indexed number of the turn being requested. + CallLLM { + /// The 1-indexed turn number being requested. + /// + /// Starts at `1` on the first call after construction and increments + /// with each completed model response. The driver can use it to tag + /// observer events, logs, and rate-limit bookkeeping so each turn is + /// correlatable back to its request. + turn: usize, + }, + + /// Request that the driver dispatch these tool calls. + /// + /// Each [`PendingToolCall`] may carry a [`preresolved_result`](PendingToolCall::preresolved_result): + /// when set, the driver returns that content as the tool result *without* + /// dispatching (used for invalid-tool-call recovery). After dispatch, the + /// driver feeds the tool-result [`Message`]s back via + /// [`LoopMachine::tool_results`]. + CallTools { + /// The tool calls awaiting dispatch, with any preresolved results. + /// + /// Exactly the calls the model requested in the preceding + /// [`MachineStep::CallLLM`], in order. Entries whose + /// [`preresolved_result`](PendingToolCall::preresolved_result) is set + /// should be fed straight back rather than dispatched to a tool. + calls: Vec, + }, + + /// Request a context compaction pass over the history. + /// + /// The machine decides *when* to compact (this step); the driver's + /// `ContextCompactor` decides *how*. After compacting, the driver feeds the + /// compacted history back via [`LoopMachine::compaction_result`]. + Compact { + /// Why compaction was triggered. + /// + /// The [`CompactReason`] that caused the machine to emit this step, which + /// a driver may use to pick a compaction strategy. + reason: CompactReason, + }, + + /// Terminal — the run is complete. + /// + /// The wrapped [`MachineOutcome`] describes *how* the run ended: normal + /// completion, the max-turn limit, cancellation, or a failure. Once the + /// machine emits this, every subsequent [`LoopMachine::next_step`] repeats + /// the same outcome — the machine is finished and accepts no further input. + Done(MachineOutcome), +} + +/// A tool call awaiting dispatch, with an optional preresolved result. +/// +/// The machine classifies each model-emitted tool call against the names the +/// driver reported as available ([`ModelTurn::available_tools`]). For a name the +/// driver did not advertise, the machine sets +/// [`preresolved_result`](Self::preresolved_result) to a synthetic error +/// [`Message`]; the driver feeds that back as the tool result without +/// dispatching, so the model learns the name is unknown and can correct itself. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[non_exhaustive] +pub struct PendingToolCall { + /// The tool call requested by the model. + /// + /// Exactly as the model emitted it — call id, tool name, and JSON input — + /// to be dispatched against the driver's [`ToolRegistry`](crate::tool::ToolRegistry). + pub call: ToolCall, + + /// A pre-resolved result for calls suppressed by invalid-tool-call recovery. + /// + /// `Some(message)` when the tool name was not in the call's + /// [`available_tools`](ModelTurn::available_tools): the machine has built a + /// synthetic error [`Message`] and the driver should feed it back as the + /// tool result *without* dispatching, so the model learns the name is + /// unknown and can correct itself. `None` for valid calls, which the driver + /// dispatches normally. + pub preresolved_result: Option, +} + +/// A completed model call, fed back to the machine by the driver. +/// +/// The driver consumes the provider's stream, accumulates the final assistant +/// [`Message`], and packages it with its usage and stop reason into this +/// struct. Streaming stays driver-side; the machine only ever sees a completed +/// turn. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[non_exhaustive] +pub struct ModelTurn { + /// The assistant message the model produced. + /// + /// The fully-accumulated response, including any text and tool-call + /// [`MessagePart`] parts. + pub message: Message, + + /// Input tokens reported for this call. + /// + /// The prompt-side token count the provider returned for this request. + pub input_tokens: u64, + + /// Output tokens reported for this call. + /// + /// The completion-side token count the provider returned. + pub output_tokens: u64, + + /// Why the model stopped. + /// + /// The provider's [`StopReason`] for this response. + pub stop_reason: StopReason, + + /// Tool names the driver advertised as available for this call. + /// + /// The set of registered tool names the driver sent to the provider for + /// this turn (e.g. from [`ToolRegistry::tool_names`](crate::tool::ToolRegistry::tool_names)). + /// A tool call whose name is not in this list is treated as unknown and + /// given a preresolved error result (see + /// [`PendingToolCall::preresolved_result`]). + pub available_tools: Vec, +} + +/// The terminal outcome of a run. +/// +/// Carried by [`MachineStep::Done`] once [`LoopMachine::next_step`] decides the +/// run is over. Describes *how* it ended and carries the accounting a host needs +/// to log, bill, or report it. Each variant corresponds to one of the machine's +/// terminal conditions; once produced, the outcome is repeated unchanged on +/// every subsequent `next_step` call. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[non_exhaustive] +pub enum MachineOutcome { + /// The model finished without requesting further tools. + /// + /// The normal end state: the last model response carried no tool calls, so + /// the machine treated it as the final answer. The captured totals let a + /// host record exactly what the completed run cost. + Completed { + /// The final text the model produced. + /// + /// Concatenation of the text parts of the last assistant message. May be + /// empty if the model's final message carried only non-text parts. + final_text: String, + + /// Total input tokens consumed across the run. + /// + /// Sum of every turn's [`ModelTurn::input_tokens`]. + total_input_tokens: u64, + + /// Total output tokens consumed across the run. + /// + /// Sum of every turn's [`ModelTurn::output_tokens`]. + total_output_tokens: u64, + + /// Number of turns executed. + /// + /// How many model calls the run made before finishing. + turns_taken: usize, + }, + + /// The run was halted because `max_turns` was reached. + /// + /// The model kept requesting tools (or otherwise not finishing) until the + /// turn budget was exhausted, so the machine stopped it rather than loop + /// indefinitely. Usually a sign the agent isn't converging — raise + /// [`max_turns`](RunConfig::max_turns) or investigate the loop. + MaxTurnsExceeded { + /// Number of turns executed before the limit was hit. + /// + /// Equal to the configured [`RunConfig::max_turns`]. + turns_taken: usize, + }, + + /// The run was cancelled. + /// + /// A clean, cooperative termination caused by [`LoopMachine::cancel`] (the + /// driver calls it when its cancellation signal fires). Distinct from + /// [`Failed`](Self::Failed): cancellation is expected, not an error, and the + /// conversation may hold partial results in the history. Carries no extra + /// fields because there is nothing further to report. + Cancelled, + + /// The run failed with an error. + /// + /// A driver-side failure the machine was told about (for example a tool + /// dispatch or LLM call that errored terminally). The wrapped error is typed + /// so a host can match on it programmatically rather than parsing a string. + Failed { + /// The error that terminated the run. + /// + /// A host can pattern-match on the [`LoopError`] to recover or report + /// the specific failure. + error: LoopError, + }, +} + +/// Where the machine is in its request/respond cycle. +/// +/// One value is produced at each point in the loop and is also exposed by +/// [`LoopMachine::state`] for inspection. It is the machine's own step-protocol +/// bookkeeping — distinct from the public `LoopState` observation surface, which +/// a driver derives for observers. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub enum MachineState { + /// Ready to request the next model call. + /// + /// The starting state, and the state the machine returns to after tool + /// results or a compaction result are fed back — i.e. whenever it is ready + /// to call the model again. + Start, + + /// A model call has been requested; awaiting the response. + /// + /// Entered when the machine emits [`MachineStep::CallLLM`] and left when the + /// driver feeds the [`ModelTurn`] back via [`LoopMachine::model_response`]. + AwaitingModel { + /// The 1-indexed turn number in flight. + /// + /// Matches the `turn` carried on the outstanding [`MachineStep::CallLLM`]. + turn: usize, + }, + + /// Tool dispatch has been requested; awaiting the results. + /// + /// Entered when the machine emits [`MachineStep::CallTools`] and left when + /// the driver feeds the tool-result messages back via + /// [`LoopMachine::tool_results`]. + AwaitingTools { + /// The 1-indexed turn number the tool calls belong to. + /// + /// Lets a host correlate a dispatch back to the model response that + /// requested it. + turn: usize, + }, + + /// A compaction pass has been requested; awaiting the compacted history. + /// + /// Entered when the machine emits [`MachineStep::Compact`] and left when the + /// driver feeds the compacted history back via + /// [`LoopMachine::compaction_result`]. + AwaitingCompaction { + /// Why compaction was triggered. + /// + /// The [`CompactReason`] carried on the outstanding [`MachineStep::Compact`]. + reason: CompactReason, + }, + + /// The run has terminated. + /// + /// The wrapped [`MachineOutcome`] describes how it ended. Once in this + /// state, the machine accepts no further input. + Terminal(MachineOutcome), +} + +/// The sans-IO agent-loop state machine. +/// +/// Owns the append-only history and every decision the loop makes. Construct +/// one with [`LoopMachine::new`], then repeatedly call [`Self::next_step`] and +/// feed the result back. The machine performs no IO; a driver is required to +/// actually call the LLM, dispatch tools, and compact context. +/// +/// # Example +/// +/// Driving a machine to completion (illustrative — a real driver performs IO in +/// each arm): +/// +/// ```no_run +/// # use loopctl::engine::machine::{LoopMachine, RunConfig, ModelTurn, MachineStep, MachineOutcome}; +/// # fn build_turn(_: &LoopMachine) -> ModelTurn { unimplemented!() } +/// let mut machine = LoopMachine::new(RunConfig::default(), "Hello"); +/// loop { +/// match machine.next_step() { +/// MachineStep::CallLLM { .. } => { +/// let turn = build_turn(&machine); +/// machine.model_response(turn); +/// } +/// MachineStep::CallTools { .. } => { +/// machine.tool_results(Vec::new()); +/// } +/// MachineStep::Compact { .. } => { +/// machine.compaction_result(machine.history().to_vec(), 0); +/// } +/// MachineStep::Done(MachineOutcome::Completed { final_text, .. }) => { +/// assert_eq!(final_text, "Hi there"); +/// break; +/// } +/// _ => break, +/// } +/// } +/// ``` +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct LoopMachine { + /// The per-run configuration the machine decides against. + /// + /// The turn budget, compaction policy, and dispatch mode that govern this + /// run. Owned by the machine so it travels with a serialized instance. + config: RunConfig, + + /// The append-only conversation history. + /// + /// The complete record of the run: the opening user message, every + /// assistant response, and every tool result. The driver derives the LLM + /// feed from it on each call; compaction replaces it wholesale. + history: Vec, + + /// Where the machine is in its request/respond cycle. + /// + /// One of [`MachineState`]: ready to call the model, awaiting a model + /// response, awaiting tool results, awaiting a compaction result, or + /// terminal. Determines which step [`Self::next_step`] emits next. + state: MachineState, + + /// How many turns have completed. + /// + /// The count of model responses accepted so far. Bounded by + /// [`RunConfig::max_turns`]; reaching that cap ends the run with + /// [`MachineOutcome::MaxTurnsExceeded`]. + turns_taken: usize, + + /// Cumulative input tokens across the run, for accounting. + /// + /// Total prompt-side token consumption from every accepted turn. Reported + /// in [`MachineOutcome::Completed`] when the run finishes. Distinct from + /// [`context_tokens`](Self::context_tokens), which tracks current size for + /// the compaction trigger rather than lifetime consumption. + total_input_tokens: u64, + + /// Cumulative output tokens across the run, for accounting. + /// + /// Total generation-side token consumption from every accepted turn. + /// Reported in [`MachineOutcome::Completed`] when the run finishes. + total_output_tokens: u64, + + /// Current context-size estimate, in tokens. + /// + /// The machine's view of how large the conversation currently is, used to + /// decide when to compact. Compared against + /// [`compact_threshold`](RunConfig::compact_threshold) of + /// [`context_window`](RunConfig::context_window); crossing it triggers a + /// [`MachineStep::Compact`]. + context_tokens: u64, + + /// Whether [`Self::cancel`] has been called. + /// + /// A cancellation request from the driver. Once set, the next + /// [`Self::next_step`] ends the run with [`MachineOutcome::Cancelled`]. + cancelled: bool, + + /// Tool calls queued for the next [`MachineStep::CallTools`]. + /// + /// The classified calls (valid vs. unknown tool name) the most recent model + /// response requested, awaiting dispatch. Emitted as a `CallTools` step on + /// the next [`Self::next_step`]. + pending_tools: Vec, +} + +impl LoopMachine { + /// Construct a machine for a new run with the given `user_input`. + /// + /// The user's message is appended to the (initially empty) history as the + /// first entry, and the first [`Self::next_step`] call requests the initial + /// [`MachineStep::CallLLM`] at turn 1. + #[must_use] + pub fn new(config: RunConfig, user_input: impl Into) -> Self { + Self { + config, + history: vec![Message::user(user_input)], + state: MachineState::Start, + turns_taken: 0, + total_input_tokens: 0, + total_output_tokens: 0, + context_tokens: 0, + cancelled: false, + pending_tools: Vec::new(), + } + } + + /// The configuration the machine makes its decisions against. + /// + /// The driver reads it to build provider requests (e.g. + /// [`RunConfig::max_tokens`], [`RunConfig::parallel_tool_dispatch`]). + #[must_use] + pub fn config(&self) -> &RunConfig { + &self.config + } + + /// Return the next step the driver must perform. + /// + /// This is pure and idempotent: calling it twice with no intervening feed + /// method ([`Self::model_response`], [`Self::tool_results`], + /// [`Self::compaction_result`], [`Self::cancel`]) returns an equal step. + /// Once the machine is terminal, every subsequent call returns + /// [`MachineStep::Done`] with the same [`MachineOutcome`]. + /// + /// # Cancellation + /// + /// If [`Self::cancel`] has been called, the next call returns + /// [`MachineStep::Done`] with [`MachineOutcome::Cancelled`]. + pub fn next_step(&mut self) -> MachineStep { + if let MachineState::Terminal(outcome) = &self.state { + let outcome = outcome.clone(); + return MachineStep::Done(outcome); + } + + if self.cancelled { + let outcome = MachineOutcome::Cancelled; + self.state = MachineState::Terminal(outcome.clone()); + return MachineStep::Done(outcome); + } + + let turn = self.turns_taken.saturating_add(1); + match self.state.clone() { + MachineState::Start | MachineState::AwaitingModel { .. } => self.request_model(turn), + MachineState::AwaitingTools { .. } => { + let calls = std::mem::take(&mut self.pending_tools); + MachineStep::CallTools { calls } + } + MachineState::AwaitingCompaction { reason } => MachineStep::Compact { reason }, + MachineState::Terminal(outcome) => MachineStep::Done(outcome), + } + } + + /// Produce the step for a model-call request, honoring the turn budget and + /// the compaction trigger. + /// + /// Returns [`MachineOutcome::MaxTurnsExceeded`] when the budget is spent, + /// [`MachineStep::Compact`] when the context size has crossed the threshold + /// and auto-compaction is on, or [`MachineStep::CallLLM`] otherwise. + fn request_model(&mut self, turn: usize) -> MachineStep { + if self.turns_taken >= self.config.max_turns { + let turns_taken = self.turns_taken; + let outcome = MachineOutcome::MaxTurnsExceeded { turns_taken }; + self.state = MachineState::Terminal(outcome.clone()); + return MachineStep::Done(outcome); + } + if self.config.auto_compact && self.should_compact() { + let reason = CompactReason::ThresholdExceeded; + self.state = MachineState::AwaitingCompaction { reason }; + return MachineStep::Compact { reason }; + } + self.state = MachineState::AwaitingModel { turn }; + MachineStep::CallLLM { turn } + } + + /// Decide whether the current context size warrants a compaction pass. + /// + /// Compares the machine's current context-size estimate against the + /// configured [`compact_threshold`](RunConfig::compact_threshold) of the + /// [`context_window`](RunConfig::context_window): returns `true` when the + /// estimate has grown past that fraction of the window, signalling that the + /// next step should be [`MachineStep::Compact`] rather than another model + /// call. Returns `false` when compaction is disabled by a zero threshold or + /// window, or when the context still fits comfortably. + fn should_compact(&self) -> bool { + const BASIS: u128 = 10_000; + if self.config.context_window == 0 || self.config.compact_threshold == 0 { + return false; + } + let limit = u128::from(self.config.compact_threshold) + .saturating_mul(u128::from(self.config.context_window)); + u128::from(self.context_tokens).saturating_mul(BASIS) > limit + } + + /// Feed a completed model call back into the machine. + /// + /// The driver calls this after servicing a [`MachineStep::CallLLM`], passing + /// the model's response. If the response carried no tool calls the run is + /// complete and the machine becomes terminal with [`MachineOutcome::Completed`]; + /// otherwise the calls are classified (valid vs. unknown tool name) and the + /// next [`Self::next_step`] will request their dispatch via + /// [`MachineStep::CallTools`]. + /// + /// Has no effect once the machine is terminal. + pub fn model_response(&mut self, turn: ModelTurn) { + if self.is_terminal() { + return; + } + + let message = turn.message; + let tool_calls = Self::extract_tool_calls(&message); + self.history.push(message); + self.total_input_tokens = self.total_input_tokens.saturating_add(turn.input_tokens); + self.total_output_tokens = self.total_output_tokens.saturating_add(turn.output_tokens); + self.context_tokens = turn.input_tokens; + self.turns_taken = self.turns_taken.saturating_add(1); + + if tool_calls.is_empty() { + let final_text = self + .history + .last() + .map(Self::extract_text) + .unwrap_or_default(); + let outcome = MachineOutcome::Completed { + final_text, + total_input_tokens: self.total_input_tokens, + total_output_tokens: self.total_output_tokens, + turns_taken: self.turns_taken, + }; + self.state = MachineState::Terminal(outcome); + return; + } + + let turn_number = self.turns_taken; + self.pending_tools = tool_calls + .into_iter() + .map(|call| Self::classify(call, &turn.available_tools)) + .collect(); + self.state = MachineState::AwaitingTools { turn: turn_number }; + } + + /// Feed tool-result messages back into the machine. + /// + /// Each provided [`Message`] is appended to the history. After the results + /// are recorded, the next [`Self::next_step`] requests another + /// [`MachineStep::CallLLM`] (subject to the max-turn budget and compaction + /// trigger). Has no effect once the machine is terminal. + pub fn tool_results(&mut self, messages: Vec) { + if self.is_terminal() { + return; + } + self.history.extend(messages); + self.pending_tools.clear(); + self.state = MachineState::Start; + } + + /// Feed compacted history back into the machine. + /// + /// The driver calls this after servicing a [`MachineStep::Compact`], passing + /// the compacted history and `tokens_after` — its estimate of that history's + /// token size, which the machine adopts as the current context size so it + /// does not immediately request another compaction. The next + /// [`Self::next_step`] then requests the deferred [`MachineStep::CallLLM`]. + /// Has no effect once the machine is terminal. + pub fn compaction_result(&mut self, compacted: Vec, tokens_after: u64) { + if self.is_terminal() { + return; + } + self.history = compacted; + self.context_tokens = tokens_after; + self.state = MachineState::Start; + } + + /// Mark the run as cancelled. + /// + /// The next [`Self::next_step`] returns [`MachineStep::Done`] with + /// [`MachineOutcome::Cancelled`]. This is the state-correctness half of + /// cancellation: a driver that aborts an in-flight step calls this so the + /// machine's state reflects the abort rather than resuming the partial step. + pub fn cancel(&mut self) { + self.cancelled = true; + } + + /// Whether [`Self::cancel`] has been called. + /// + /// When `true`, the next [`Self::next_step`] goes terminal with + /// [`MachineOutcome::Cancelled`]. Reports the cancellation request, not + /// whether the machine is already terminal — use [`Self::is_terminal`] for + /// that. + #[must_use] + pub fn is_cancelled(&self) -> bool { + self.cancelled + } + + /// The append-only conversation history. + /// + /// Read-only access to the record the machine owns. The driver builds the + /// LLM feed from this on each [`MachineStep::CallLLM`]. + #[must_use] + pub fn history(&self) -> &[Message] { + &self.history + } + + /// The current state of the step protocol. + /// + /// A snapshot of where the machine is in its request/respond cycle — useful + /// for diagnostics, or for a driver that needs to know which feed method is + /// currently expected. + #[must_use] + pub fn state(&self) -> MachineState { + self.state.clone() + } + + /// Number of turns completed so far. + /// + /// How many model responses the machine has accepted. Also carried in + /// [`MachineOutcome::Completed`] and [`MachineOutcome::MaxTurnsExceeded`] + /// once the run ends. + #[must_use] + pub fn turns_taken(&self) -> usize { + self.turns_taken + } + + /// Total input tokens accumulated so far. + /// + /// The run's lifetime input consumption, for accounting. Also carried in + /// [`MachineOutcome::Completed`] once the run ends. + #[must_use] + pub fn total_input_tokens(&self) -> u64 { + self.total_input_tokens + } + + /// Total output tokens accumulated so far. + /// + /// The run's total generation cost, for accounting. Also carried in + /// [`MachineOutcome::Completed`] once the run ends. + #[must_use] + pub fn total_output_tokens(&self) -> u64 { + self.total_output_tokens + } + + /// Whether the machine has reached a terminal outcome. + /// + /// `true` once the machine has emitted a [`MachineOutcome`] (completion, + /// max-turns, cancellation, or failure). Once terminal, the feed methods + /// (`model_response`, `tool_results`, `compaction_result`) are no-ops and + /// [`Self::next_step`] keeps returning [`MachineStep::Done`] with the same + /// outcome. + #[must_use] + pub fn is_terminal(&self) -> bool { + matches!(self.state, MachineState::Terminal(_)) + } + + /// Classify a tool call against the advertised available tool names. + /// + /// Known names yield a plain [`PendingToolCall`] (`preresolved_result: + /// None`); unknown names get a synthetic error result so the driver can feed + /// it back without dispatching. + fn classify(call: ToolCall, available: &[String]) -> PendingToolCall { + let known = available.iter().any(|name| name == &call.tool); + if known { + return PendingToolCall { + call, + preresolved_result: None, + }; + } + let result = Self::unknown_tool_result(&call); + PendingToolCall { + call, + preresolved_result: Some(result), + } + } + + /// Build the tool-result message returned for an unknown tool call. + /// + /// Constructs a user-role [`Message`] carrying a single error tool-result + /// part that echoes the call's id and names the tool as unavailable. This is + /// the value placed on [`PendingToolCall::preresolved_result`] so a driver + /// can feed it straight back to the model — without dispatching — letting + /// the model learn the name is unknown and correct itself on the next turn. + fn unknown_tool_result(call: &ToolCall) -> Message { + let message = format!("tool '{}' is not available", call.tool); + Message::new( + Role::User, + vec![MessagePart::tool_result( + call.id.clone(), + ToolContent::Text(message), + true, + )], + ) + } + + /// Concatenate every text part of a message into a single string. + /// + /// Walks the message's parts in order and joins all text content, skipping + /// non-text parts (tool calls, tool results, images) entirely. Used to + /// derive the final answer text from the model's last response. + fn extract_text(message: &Message) -> String { + message + .parts + .iter() + .filter_map(|part| part.as_text()) + .collect::>() + .join("") + } + + /// Collect the tool calls a message requests, in order. + /// + /// Scans the message's parts for tool-call parts and maps each to a + /// [`ToolCall`] (carrying its id, tool name, and JSON input). Returns an + /// empty vector when the message contains no tool calls — the signal that + /// the turn is complete rather than a request to dispatch tools. + fn extract_tool_calls(message: &Message) -> Vec { + message + .parts + .iter() + .filter_map(|part| match part { + MessagePart::ToolCall { id, name, input } => Some(ToolCall { + id: id.clone(), + tool: name.clone(), + input: input.clone(), + }), + _ => None, + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::Value; + + fn small_machine(max_turns: usize) -> LoopMachine { + LoopMachine::new( + RunConfig { + max_turns, + ..RunConfig::default() + }, + "hello", + ) + } + + fn text_turn(text: &str, input_tokens: u64, output_tokens: u64) -> ModelTurn { + ModelTurn { + message: Message::assistant(text), + input_tokens, + output_tokens, + stop_reason: StopReason::EndTurn, + available_tools: Vec::new(), + } + } + + fn tool_turn(tool: &str, available: &[&str], input_tokens: u64) -> ModelTurn { + let part = MessagePart::tool_call("call_1", tool, Value::Object(serde_json::Map::new())); + ModelTurn { + message: Message::new(Role::Assistant, vec![part]), + input_tokens, + output_tokens: 10, + stop_reason: StopReason::ToolCall, + available_tools: available.iter().map(|s| (*s).to_string()).collect(), + } + } + + fn same_step(a: &MachineStep, b: &MachineStep) -> bool { + serde_json::to_string(a).unwrap_or_default() == serde_json::to_string(b).unwrap_or_default() + } + + #[test] + fn calling_llm_from_new_emits_call_llm_step() { + let mut machine = LoopMachine::new(RunConfig::default(), "hello"); + let step = machine.next_step(); + let MachineStep::CallLLM { turn } = step else { + panic!("expected CallLLM, got {step:?}"); + }; + assert_eq!(turn, 1); + assert_eq!(machine.state(), MachineState::AwaitingModel { turn: 1 }); + } + + #[test] + fn machine_api_has_no_async_no_tokio_no_apiclient() { + // If any driving method were async or needed a client/runtime, this + // body would not compile or would require a tokio context. It touches + // only the pure machine surface. + let mut machine = LoopMachine::new(RunConfig::default(), "hello"); + assert!(matches!(machine.next_step(), MachineStep::CallLLM { .. })); + machine.model_response(text_turn("hi", 5, 3)); + // A text-only turn completes the run, so the machine is terminal. + assert!(machine.is_terminal()); + assert_eq!(machine.turns_taken(), 1); + assert_eq!(machine.total_input_tokens(), 5); + assert_eq!(machine.total_output_tokens(), 3); + assert_eq!(machine.history().len(), 2); + assert!(matches!( + machine.state(), + MachineState::Terminal(MachineOutcome::Completed { .. }) + )); + assert_eq!(machine.config().max_turns, RunConfig::default().max_turns); + machine.cancel(); + assert!(machine.is_cancelled()); + } + + #[test] + fn resume_after_model_response_round_trips() { + let mut machine = small_machine(5); + let _ = machine.next_step(); + machine.model_response(tool_turn("echo", &["echo"], 10)); + // Now AwaitingTools. + let snapshot = serde_json::to_string(&machine).expect("serialize"); + let mut restored: LoopMachine = serde_json::from_str(&snapshot).expect("deserialize"); + let a = machine.next_step(); + let b = restored.next_step(); + assert!(same_step(&a, &b), "steps diverged after round-trip"); + } + + #[test] + fn resume_after_tool_results_round_trips() { + let mut machine = small_machine(5); + let _ = machine.next_step(); + machine.model_response(tool_turn("echo", &["echo"], 10)); + let step = machine.next_step(); + let MachineStep::CallTools { calls } = &step else { + panic!("expected CallTools, got {step:?}"); + }; + // Synthesize the tool-result message the driver would build. + let results: Vec = calls + .iter() + .map(|c| { + Message::new( + Role::User, + vec![MessagePart::tool_result( + c.call.id.clone(), + ToolContent::Text("ok".to_string()), + false, + )], + ) + }) + .collect(); + machine.tool_results(results); + let snapshot = serde_json::to_string(&machine).expect("serialize"); + let mut restored: LoopMachine = serde_json::from_str(&snapshot).expect("deserialize"); + let a = machine.next_step(); + let b = restored.next_step(); + assert!(same_step(&a, &b), "steps diverged after round-trip"); + } + + #[test] + fn resume_after_compaction_result_round_trips() { + // window = 100, threshold = 0.5 ⇒ compact once tokens > 50. + let mut machine = LoopMachine::new( + RunConfig { + max_turns: 5, + context_window: 100, + compact_threshold: 5_000, + auto_compact: true, + ..RunConfig::default() + }, + "hello", + ); + // Drive a tool-call turn past the threshold so the next step compacts. + let _ = machine.next_step(); // CallLLM + machine.model_response(tool_turn("echo", &["echo"], 60)); // AwaitingTools + let _ = machine.next_step(); // CallTools + machine.tool_results(vec![Message::user("tool-out")]); // → Start + assert!(matches!(machine.next_step(), MachineStep::Compact { .. })); + machine.compaction_result(vec![Message::user("compacted")], 0); + let snapshot = serde_json::to_string(&machine).expect("serialize"); + let mut restored: LoopMachine = serde_json::from_str(&snapshot).expect("deserialize"); + let a = machine.next_step(); + let b = restored.next_step(); + assert!(same_step(&a, &b), "steps diverged after round-trip"); + } + + #[test] + fn max_turns_enforced_by_machine() { + let mut machine = small_machine(2); + // Turn 1. + assert!(matches!( + machine.next_step(), + MachineStep::CallLLM { turn: 1 } + )); + machine.model_response(tool_turn("echo", &["echo"], 1)); + let _ = machine.next_step(); + machine.tool_results(vec![Message::user("r")]); + // Turn 2. + assert!(matches!( + machine.next_step(), + MachineStep::CallLLM { turn: 2 } + )); + machine.model_response(tool_turn("echo", &["echo"], 1)); + let _ = machine.next_step(); + machine.tool_results(vec![Message::user("r")]); + // Turn 3 must be denied. + match machine.next_step() { + MachineStep::Done(MachineOutcome::MaxTurnsExceeded { turns_taken }) => { + assert_eq!(turns_taken, 2); + } + other => panic!("expected MaxTurnsExceeded, got {other:?}"), + } + assert!(machine.is_terminal()); + } + + #[test] + fn cancel_returns_done_cancelled_at_next_step() { + let mut machine = small_machine(5); + let _ = machine.next_step(); + machine.cancel(); + match machine.next_step() { + MachineStep::Done(MachineOutcome::Cancelled) => {} + other => panic!("expected Done(Cancelled), got {other:?}"), + } + // Idempotent: stays terminal-cancelled. + assert!(matches!( + machine.next_step(), + MachineStep::Done(MachineOutcome::Cancelled) + )); + } + + #[test] + fn unknown_tool_call_gets_preresolved_result() { + let mut machine = small_machine(5); + let _ = machine.next_step(); + machine.model_response(tool_turn("ghost", &["echo", "ls"], 3)); + let step = machine.next_step(); + let MachineStep::CallTools { calls } = step else { + panic!("expected CallTools, got {step:?}"); + }; + let call = calls.first().expect("one call"); + let result = call + .preresolved_result + .as_ref() + .expect("unknown tool has a preresolved result"); + assert_eq!(result.role, Role::User); + assert!(result.parts.iter().any(|p| match p { + MessagePart::ToolResult { is_error, .. } => *is_error == Some(true), + _ => false, + })); + } + + #[test] + fn known_tool_call_emits_plain_pending_call() { + let mut machine = small_machine(5); + let _ = machine.next_step(); + machine.model_response(tool_turn("echo", &["echo", "ls"], 3)); + let step = machine.next_step(); + let MachineStep::CallTools { calls } = step else { + panic!("expected CallTools, got {step:?}"); + }; + let call = calls.first().expect("one call"); + assert!( + call.preresolved_result.is_none(), + "known tool has no preresolved result" + ); + } + + #[test] + fn compaction_triggered_when_tokens_exceed_threshold() { + // window = 100, threshold = 0.5 ⇒ compact once tokens > 50. + let mut machine = LoopMachine::new( + RunConfig { + max_turns: 5, + context_window: 100, + compact_threshold: 5_000, + auto_compact: true, + ..RunConfig::default() + }, + "hello", + ); + // First step: no tokens yet (0 > 50 is false) ⇒ CallLLM, not Compact. + assert!(matches!(machine.next_step(), MachineStep::CallLLM { .. })); + // A tool-call turn keeps the run going and accumulates 60 input tokens + // (exceeds 50). A text-only turn would complete the run instead. + machine.model_response(tool_turn("echo", &["echo"], 60)); + assert!(matches!(machine.next_step(), MachineStep::CallTools { .. })); + machine.tool_results(vec![Message::user("tool-out")]); + // The next step must compact before calling the model again. + match machine.next_step() { + MachineStep::Compact { reason } => { + assert_eq!(reason, CompactReason::ThresholdExceeded); + } + other => panic!("expected Compact, got {other:?}"), + } + } + + #[test] + fn compaction_result_replaces_history() { + // window = 100, threshold = 0.5 ⇒ compact once tokens > 50. + let mut machine = LoopMachine::new( + RunConfig { + max_turns: 5, + context_window: 100, + compact_threshold: 5_000, + auto_compact: true, + ..RunConfig::default() + }, + "hello", + ); + // Drive a tool-call turn past the threshold so the next step compacts. + let _ = machine.next_step(); // CallLLM + machine.model_response(tool_turn("echo", &["echo"], 60)); // AwaitingTools + let _ = machine.next_step(); // CallTools + machine.tool_results(vec![Message::user("tool-out")]); // → Start + assert!(matches!(machine.next_step(), MachineStep::Compact { .. })); + let compacted = vec![Message::user("compacted-only")]; + machine.compaction_result(compacted.clone(), 0); + // Compare by serialized form: Message is not PartialEq. + let got = serde_json::to_string(machine.history()).expect("serialize history"); + let want = serde_json::to_string(&compacted).expect("serialize expected"); + assert_eq!(got, want, "history must be replaced by the compacted slice"); + // After compaction the next step is the deferred CallLLM. + assert!(matches!(machine.next_step(), MachineStep::CallLLM { .. })); + } + + #[test] + fn history_accumulates_user_assistant_tool_round() { + let mut machine = small_machine(5); + let _ = machine.next_step(); + machine.model_response(tool_turn("echo", &["echo"], 1)); + let step = machine.next_step(); + let MachineStep::CallTools { calls } = step else { + panic!("expected CallTools, got {step:?}"); + }; + let result = Message::new( + Role::User, + calls + .iter() + .map(|c| { + MessagePart::tool_result( + c.call.id.clone(), + ToolContent::Text("ok".to_string()), + false, + ) + }) + .collect(), + ); + machine.tool_results(vec![result]); + + let roles: Vec = machine.history().iter().map(|m| m.role).collect(); + assert_eq!( + roles, + vec![Role::User, Role::Assistant, Role::User], + "history must be [user, assistant, tool_result(user)]" + ); + } + + #[test] + fn compact_reason_is_serde() { + for reason in [ + CompactReason::ThresholdExceeded, + CompactReason::Emergency, + CompactReason::Manual, + ] { + let text = serde_json::to_string(&reason).expect("serialize"); + let back: CompactReason = serde_json::from_str(&text).expect("deserialize"); + assert_eq!(back, reason); + } + } +} diff --git a/src/error.rs b/src/error.rs index a6349d3..50922d7 100644 --- a/src/error.rs +++ b/src/error.rs @@ -85,7 +85,7 @@ /// engine escalated to the model circuit breaker. /// - [`Internal`](LoopError::Internal) — A catch-all for unexpected /// or infrastructure-level errors. -#[derive(Debug, Clone, thiserror::Error)] +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error, serde::Serialize, serde::Deserialize)] #[non_exhaustive] pub enum LoopError { /// A tool was not found in the registry. diff --git a/src/presets.rs b/src/presets.rs index 9e2eedb..486c8fb 100644 --- a/src/presets.rs +++ b/src/presets.rs @@ -288,7 +288,7 @@ mod tests { assert_eq!(cfg.context_window, 120_000); assert_eq!(cfg.max_turns, 100); assert_eq!(cfg.max_tokens, 16_384); - assert_eq!(cfg.compact_threshold, 0.80); + assert_eq!(cfg.compact_threshold, 8_000); assert!(cfg.auto_compact); } From cdc0046331fc0d055ddfab3cb1fc52a498a5e852 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Fri, 24 Jul 2026 18:05:02 +1200 Subject: [PATCH 02/21] feat: xai alias --- Cargo.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Cargo.toml b/Cargo.toml index f7412f4..fc4235e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -50,6 +50,7 @@ anthropic = ["providers"] ollama = ["providers", "openai"] deepseek = ["providers", "openai"] grok = ["providers", "openai"] +xai = ["grok"] gemini = ["providers"] zai = ["providers", "anthropic"] grammar = ["providers"] From dd696c3685d2b335fb92e7123774be015a471ef4 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Fri, 24 Jul 2026 18:05:17 +1200 Subject: [PATCH 03/21] chore: gitignore upd --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index db7bb4d..0d49516 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ *.swo *~ .DS_Store +.env From 05f0dc5843656d7afb5c0bca83dd296ca0a89c43 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 27 Jul 2026 13:46:49 +1200 Subject: [PATCH 04/21] refactor: wip --- CHANGELOG.md | 70 +- Cargo.toml | 1 + README.md | 6 +- examples/chat.rs | 66 +- examples/echo-tool-cli.rs | 15 +- examples/hello-cli.rs | 13 +- examples/repl-cli.rs | 15 +- src/api.rs | 10 +- src/api/error.rs | 65 +- src/capabilities.rs | 70 +- src/compact.rs | 222 ++- src/compact/truncating.rs | 97 +- src/compact/types.rs | 43 +- src/config.rs | 492 ++----- src/detection.rs | 2 +- src/detection/convergence.rs | 112 +- src/detection/loop_detector.rs | 87 +- src/detection/manager.rs | 183 ++- src/engine.rs | 35 +- src/engine/bare.rs | 2154 +++++++++++++++++++----------- src/engine/bare/compact.rs | 113 +- src/engine/bare/dispatch.rs | 74 +- src/engine/bare/emission.rs | 107 +- src/engine/bare/message.rs | 43 +- src/engine/bare/stream.rs | 10 +- src/engine/core.rs | 14 + src/engine/core/lifecycle.rs | 684 ++++++++++ src/engine/{ => core}/machine.rs | 847 ++++++------ src/engine/loop_core.rs | 992 -------------- src/error.rs | 143 +- src/fallback.rs | 213 +-- src/hooks.rs | 55 +- src/hooks/builtin/auto_commit.rs | 341 ++++- src/hooks/context.rs | 319 ++++- src/hooks/executor.rs | 99 +- src/lib.rs | 6 +- src/{runtime.rs => managers.rs} | 666 ++++----- src/memory/builtin.rs | 8 - src/memory/entry.rs | 59 + src/message.rs | 159 +-- src/middleware.rs | 4 - src/middleware/memoize.rs | 12 - src/middleware/timeout.rs | 14 + src/middleware/unknown_tool.rs | 3 +- src/middleware/verify.rs | 8 - src/observer.rs | 187 +-- src/observer/context.rs | 11 +- src/presets.rs | 112 +- src/provider.rs | 144 +- src/provider/anthropic.rs | 175 +-- src/provider/gemini.rs | 190 +-- src/provider/openai.rs | 335 ++--- src/provider/sse.rs | 100 ++ src/reflection.rs | 36 - src/reflection/backoff.rs | 53 +- src/reflection/llm.rs | 26 +- src/stream.rs | 62 +- src/stream/handler.rs | 232 +++- src/stream/heartbeat.rs | 556 -------- src/stream/rate_limit.rs | 76 +- src/structured.rs | 14 + src/testing.rs | 113 +- src/tool.rs | 129 +- src/tool/health.rs | 234 +++- src/tool/permission.rs | 29 +- src/tool/registry.rs | 73 +- src/tool/shield.rs | 40 - tests/constrained_decode.rs | 111 ++ tests/examples_e2e.rs | 46 + tests/helpers.rs | 197 +++ tests/provider_e2e.rs | 160 +++ tests/provider_survival.rs | 60 + tests/structured_output.rs | 105 ++ 73 files changed, 6711 insertions(+), 5646 deletions(-) create mode 100644 src/engine/core.rs create mode 100644 src/engine/core/lifecycle.rs rename src/engine/{ => core}/machine.rs (59%) delete mode 100644 src/engine/loop_core.rs rename src/{runtime.rs => managers.rs} (52%) create mode 100644 src/provider/sse.rs delete mode 100644 src/stream/heartbeat.rs create mode 100644 tests/constrained_decode.rs create mode 100644 tests/examples_e2e.rs create mode 100644 tests/helpers.rs create mode 100644 tests/provider_e2e.rs create mode 100644 tests/provider_survival.rs create mode 100644 tests/structured_output.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 61579a0..65d76b3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,9 +15,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. history, cancellation). `Serialize + Deserialize`, with no `async`, no `tokio`, and no `ApiClient` in its surface. Includes `RunConfig`, `MachineStep` (`CallLLM`/`CallTools`/`Compact`/`Done`), `ModelTurn`, - `PendingToolCall`, `MachineOutcome`, and `MachineState`. The machine is not - yet wired into `BareLoop` in this release; it is public so a driver can be - written against it. + `PendingToolCall`, `MachineOutcome`, and `MachineState`. `BareLoop` now drives + a `LoopMachine` internally (`run()` is a `match machine.next_step()` loop); the + machine is exposed via `BareLoop::machine()` / `into_machine()` / + `from_machine()` for inspection and serialize-and-resume. +- `LoopMachine::inject(message)` — add an arbitrary message to the machine's + history (host steering, or `ContextContributor` goal re-injection). +- `Session`/`Run`/`Turn`/`Run` lifetime types (`engine::core`) — + the Session ⊃ [Run ⊃ [Turn]] hierarchy: one `Session` spans the process, + one `Run` per `run()` prompt, one `Turn` per loop iteration. `Session` + derives per-session totals (`total_turns`/`total_duration`/ + `total_input_tokens`/`total_output_tokens`) from its run list. `Run` + is the result of a `run()`; `Run::turn_count()`/`duration()`/`total_tokens()`. +- `SessionConfig` (`config`) — the session-scoped config slice (`session_id`, + `system_prompt`, `context_window`) with `with_*` builders, replacing the + session fields that lived on the old `LoopConfig`. - `LoopError` now derives `Serialize`, `Deserialize`, `PartialEq`, and `Eq`. - `compact::types::CompactReason` now derives `Serialize` and `Deserialize`. - `DeltaPart::Thinking { text }` variant + `on_thinking_delta` observer event @@ -196,17 +208,43 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. ### Changed -- **Breaking (compaction thresholds → basis points):** the compaction trigger - threshold and compaction-target fraction are now integers in basis points - (0–10_000; 10_000 = 100%) instead of `f64` fractions. This removes lossy - `u64 → f64` casts from the threshold arithmetic. Affected APIs: - `ContextManager::with_threshold(u16)` and `with_compact_target_pct(u16)` - (were `f64`); `ContextManager::threshold() -> u16` and - `compact_target_pct() -> u16` (were `f64`); - `LoopConfig::with_compact_threshold(u16)` and the `compact_threshold` field - (were `f64`). The default is `8_000` (was `0.80`); the clamp range is - `[1_000, 10_000]` (was `[0.1, 1.0]`). Migration: multiply existing `f64` - values by `10_000` and round — `0.80 → 8_000`, `0.50 → 5_000`, `0.70 → 7_000`. +- **Breaking (machine-driven engine):** `BareLoop::run()` is now a + `match machine.next_step()` loop driving a `LoopMachine`. The machine owns the + conversation history and every loop decision (turn count, max-turn, tool-call + validity, compaction trigger, cancellation); the driver owns IO (LLM call, + tool dispatch, compaction execution) and fires observers from the match-arms. + Observer event ordering is unchanged (pinned by golden tests). The + conversation is owned by the machine — read it via `BareLoop::conversation()` + (delegates to the machine) or `BareLoop::machine().history()`. +- **Breaking (Session/Run lifetime model):** the agent-loop lifetime is now + explicit. `LoopConfig` is **removed**; construction splits into + `SessionConfig` (session-scoped: `session_id`, `system_prompt`, + `context_window`) and `engine::RunConfig` (per-run: turn/token budgets, + compaction policy, dispatch mode). `SessionResult` is **removed** and unified + with the new `Run`/`Run` (`engine::core`): the per-run accumulator + is `Run`-shaped (`turns: Vec`, `error: Option`). The `model` + field is gone from config entirely — it lives on the `ApiClient` + (`ApiClient::model` / `set_model`). Migration: build `BareLoop::new` with a + `SessionConfig`; read `session_id`/`system_prompt`/`context_window` from + `SessionConfig`, run budgets from `RunConfig`, the model from the client. +- **Breaking (`Loop::run` signature + `initialize` removed):** `Loop::run` now + takes the per-run config — `run(&mut self, user_input: &str, run_config: + &RunConfig)` — and returns `Result`. Session + initialization happens once at construction (not per `run()`); each `run()` + receives a fresh `RunConfig`. `Loop::initialize` and `Loop::config()` are + removed. Migration: pass `&RunConfig::default()` (or a specific run config) + as the second `run()` argument; move any `initialize` setup into + construction. +- **Breaking (compaction thresholds → percentages):** the compaction trigger + threshold and compaction-target fraction are now `u8` percentages (0–100; + 100 = 100%) instead of `f64` fractions. Affected APIs: + `ContextManager::with_threshold(u8)` and `with_compact_target_pct(u8)` + (were `f64`); `ContextManager::threshold() -> u8` and + `compact_target_pct() -> u8` (were `f64`); + `SessionConfig::with_compact_threshold(u8)` and the `compact_threshold` field + (were `f64`). The default is `80` (was `0.80`); the clamp range is + `[1, 100]` (was `[0.1, 1.0]`). Migration: multiply existing `f64` + values by 100 and round — `0.80 → 80`, `0.50 → 50`, `0.70 → 70`. - **Breaking (renames):** consuming builder methods that return `Self` are now uniformly prefixed `with_`, matching the crate-wide convention. The old no-prefix names are removed. Affected types and methods (old → new): @@ -332,6 +370,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. ### Removed +- `Loop::process_turn` trait method and `BareLoop::run_turn_body` — the + machine-driven `run()` replaces the old per-turn execution path. The + `LoopMachine` is the new turn unit; drive it via `BareLoop::run()` (or + `LoopMachine::next_step()` directly for a custom driver). - `StreamTurnResult` (the handler no longer accumulates; the engine assembles the result from the event stream). - `StreamHandler::with_request_options` builder (options now flow via diff --git a/Cargo.toml b/Cargo.toml index fc4235e..a4a9640 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,7 @@ uuid = { version = "1", features = ["v4", "serde"] } tracing = "0.1" reqwest = { version = "0.13", default-features = false, features = ["json", "stream", "rustls"], optional = true } +bytes = "1" async-stream = "0.3" httpdate = { version = "1", optional = true } jsonschema = { version = "0.48", optional = true } diff --git a/README.md b/README.md index 05a3b08..507e3fb 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ and tool implementations; the framework handles the rest. | [`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 | +| [`runtime`](https://docs.rs/loopctl/latest/loopctl/runtime/index.html) | `LoopManagers` — 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.* | @@ -79,7 +79,7 @@ impl Tool for EchoTool { ```rust,no_run use loopctl::engine::BareLoop; -use loopctl::engine::loop_core::Loop; +use loopctl::engine::core::Loop; use loopctl::tool::ToolRegistry; use loopctl::config::LoopConfig; use std::sync::Arc; @@ -123,7 +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::engine::core::Loop; use loopctl::tool::ToolRegistry; use std::sync::Arc; diff --git a/examples/chat.rs b/examples/chat.rs index bebc4c1..9eb3f84 100644 --- a/examples/chat.rs +++ b/examples/chat.rs @@ -28,17 +28,14 @@ use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use loopctl::api::ApiClient; -use loopctl::config::LoopConfig; +use loopctl::config::SessionConfig; use loopctl::engine::BareLoop; -use loopctl::engine::loop_core::Loop; +use loopctl::engine::RunConfig; +use loopctl::engine::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< @@ -48,10 +45,6 @@ type ToolFuture = std::pin::Pin< >, >; -// ================================================== -// Observer -// ================================================== - /// A simple observer that prints tool calls and responses to stderr. struct PrintingObserver; @@ -74,10 +67,6 @@ impl LoopObserver for PrintingObserver { } } -// ================================================== -// Usage -// ================================================== - fn print_usage_and_exit() -> ! { eprintln!("No provider configured.\n"); eprintln!("Set one of:"); @@ -96,10 +85,6 @@ fn print_usage_and_exit() -> ! { std::process::exit(1); } -// ================================================== -// Tool functions -// ================================================== - fn echo_fn(input: serde_json::Value, _ctx: &ToolContext) -> ToolFuture { Box::pin(async move { let text = input @@ -180,11 +165,6 @@ fn build_tools() -> ToolRegistry { tools } -// ================================================== -// Minimal arithmetic expression evaluator -// (recursive descent: + - * / and parentheses) -// ================================================== - #[derive(Debug, Clone)] enum Token { Num(f64), @@ -342,17 +322,17 @@ fn parse_factor(tokens: &[Token], pos: &mut usize) -> Result { } } -// ================================================== -// 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::default().with_max_turns(10); + let config = SessionConfig::default(); + let run_config = RunConfig { + max_turns: 10, + ..RunConfig::default() + }; let mut agent = BareLoop::new(client, build_tools(), config); agent.register_observer(Arc::new(PrintingObserver)); @@ -377,8 +357,6 @@ async fn run_repl(client: Arc) { }); let stdin = io::stdin(); - let mut total_input: u64 = 0; - let mut total_output: u64 = 0; loop { print!("> "); @@ -397,25 +375,19 @@ async fn run_repl(client: Arc) { break; } - match agent.run(input).await { + match agent.run(input, &run_config).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()) - { + if result.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 + result.turn_count(), + result.input_tokens(), + result.output_tokens(), + agent.session().total_input_tokens(), + agent.session().total_output_tokens(), ); } Err(e) => { @@ -429,14 +401,6 @@ async fn run_repl(client: Arc) { } } -// ================================================== -// 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) => {{ diff --git a/examples/echo-tool-cli.rs b/examples/echo-tool-cli.rs index 0ab3219..fdafc59 100644 --- a/examples/echo-tool-cli.rs +++ b/examples/echo-tool-cli.rs @@ -11,9 +11,10 @@ use std::sync::Arc; -use loopctl::config::LoopConfig; +use loopctl::config::SessionConfig; use loopctl::engine::BareLoop; -use loopctl::engine::loop_core::Loop; +use loopctl::engine::RunConfig; +use loopctl::engine::core::Loop; use loopctl::testing::{MockApiClient, MockResponse, MockToolCall}; use loopctl::tool::{FnTool, ToolOutput, ToolRegistry}; use serde_json::json; @@ -78,7 +79,7 @@ async fn main() { ); // 3. Construct and run the loop. - let mut agent = BareLoop::new(Arc::new(client), tools, LoopConfig::default()); + let mut agent = BareLoop::new(Arc::new(client), tools, SessionConfig::default()); // Ctrl-C interrupts the in-flight turn via loopctl's CancelSignal. let cancel_signal = agent.cancel_signal(); @@ -89,11 +90,11 @@ async fn main() { }); let result = agent - .run("Please echo something.") + .run("Please echo something.", &RunConfig::default()) .await .expect("session should succeed"); - println!("Turns: {}", result.total_turns); - println!("Tool calls: {}", result.tool_calls); - println!("Output: {}", result.final_output.unwrap_or_default()); + println!("Turns: {}", result.turn_count()); + println!("Tool calls: {}", result.tool_call_count()); + println!("Output: {}", result.output.unwrap_or_default()); } diff --git a/examples/hello-cli.rs b/examples/hello-cli.rs index 6b91e33..dd4160b 100644 --- a/examples/hello-cli.rs +++ b/examples/hello-cli.rs @@ -13,9 +13,10 @@ use std::sync::Arc; -use loopctl::config::LoopConfig; +use loopctl::config::SessionConfig; use loopctl::engine::BareLoop; -use loopctl::engine::loop_core::Loop; +use loopctl::engine::RunConfig; +use loopctl::engine::core::Loop; use loopctl::testing::MockApiClient; use loopctl::tool::ToolRegistry; @@ -26,7 +27,7 @@ async fn main() { // 2. Build the components. let tools = ToolRegistry::new(); - let config = LoopConfig::default(); + let config = SessionConfig::default(); // 3. Construct the loop. let mut agent = BareLoop::new(Arc::new(client), tools, config); @@ -41,10 +42,10 @@ async fn main() { // 4. Run and print the result. let result = agent - .run("Say hello!") + .run("Say hello!", &RunConfig::default()) .await .expect("session should succeed"); - println!("Turns: {}", result.total_turns); - println!("Output: {}", result.final_output.unwrap_or_default()); + println!("Turns: {}", result.turn_count()); + println!("Output: {}", result.output.unwrap_or_default()); } diff --git a/examples/repl-cli.rs b/examples/repl-cli.rs index d0fbc4c..be02c3a 100644 --- a/examples/repl-cli.rs +++ b/examples/repl-cli.rs @@ -13,9 +13,10 @@ use std::io::{self, BufRead, Write}; use std::sync::Arc; -use loopctl::config::LoopConfig; +use loopctl::config::SessionConfig; use loopctl::engine::BareLoop; -use loopctl::engine::loop_core::Loop; +use loopctl::engine::RunConfig; +use loopctl::engine::core::Loop; use loopctl::testing::MockApiClient; use loopctl::tool::ToolRegistry; @@ -45,7 +46,11 @@ async fn main() { 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()); + let mut agent = BareLoop::new( + Arc::new(client), + ToolRegistry::new(), + SessionConfig::default(), + ); // A fresh agent per turn means a fresh CancelSignal per turn: Ctrl-C // interrupts the current turn only, and the next prompt gets a clean @@ -58,10 +63,10 @@ async fn main() { } }); - match agent.run(input).await { + match agent.run(input, &RunConfig::default()).await { Ok(result) => { listener.abort(); - let output = result.final_output.unwrap_or_default(); + let output = result.output.unwrap_or_default(); writeln!(&mut stdout, "{output}\n").unwrap(); } Err(e) => { diff --git a/src/api.rs b/src/api.rs index 8b57578..f5ff8ab 100644 --- a/src/api.rs +++ b/src/api.rs @@ -51,7 +51,7 @@ pub struct StreamRequest { /// (top-level `system` on Anthropic/Gemini, a leading `system` role message /// on OpenAI). When `None`, the provider receives no system prompt for the /// turn. The value is forwarded verbatim from - /// [`LoopConfig::system_prompt`](crate::config::LoopConfig::system_prompt); + /// [`SessionConfig::system_prompt`](crate::config::SessionConfig::system_prompt); /// set it there to drive this field. pub system: Option, @@ -67,7 +67,13 @@ pub struct StreamRequest { } impl StreamRequest { - /// Create a request with messages and no system prompt or tools. + /// Create a request carrying `messages` and no system prompt or tools. + /// + /// The minimal constructor: only the required conversation history + /// is set. Chain + /// [`with_system`](Self::with_system) and + /// [`with_tools`](Self::with_tools) to populate the optional fields, + /// or use struct-literal syntax to set every field at once. #[must_use] pub fn new(messages: Vec) -> Self { Self { diff --git a/src/api/error.rs b/src/api/error.rs index c23d341..3261c50 100644 --- a/src/api/error.rs +++ b/src/api/error.rs @@ -78,9 +78,6 @@ use thiserror::Error; #[repr(u16)] #[non_exhaustive] pub enum ErrorCode { - // ================================================== - // API errors (1000-1005) - // ================================================== /// A generic LLM API request failed. /// /// Returned when the API responds with an unexpected status or the @@ -130,9 +127,6 @@ pub enum ErrorCode { /// truncation) before retrying. Maps to **1005**. ApiContextOverflow = 1005, - // ================================================== - // Auth errors (1100-1101) - // ================================================== /// The API key is invalid or missing. /// /// The key may be expired, revoked, or not yet set. Check @@ -146,9 +140,6 @@ pub enum ErrorCode { /// Maps to **1101**. AuthFailed = 1101, - // ================================================== - // HTTP errors (1200-1202) - // ================================================== /// A low-level HTTP connection error. /// /// Typically a DNS failure, TCP timeout, or TLS handshake error. @@ -177,9 +168,6 @@ pub enum ErrorCode { /// the message for downstream log aggregation. HttpResponseError = 1202, - // ================================================== - // Tool errors (1300-1304) - // ================================================== /// The requested tool does not exist. /// /// The agent asked for a tool name that is not registered with the @@ -224,9 +212,6 @@ pub enum ErrorCode { /// but the caller should consider increasing the timeout first. ToolTimeout = 1304, - // ================================================== - // Config errors (1400-1403) - // ================================================== /// The configuration file could not be found. /// /// The path pointed to by the config setting does not exist or is @@ -250,9 +235,6 @@ pub enum ErrorCode { /// No value was provided and no default is available. Maps to **1403**. ConfigMissing = 1403, - // ================================================== - // I/O errors (1500-1502) - // ================================================== /// A required file was not found on disk. /// /// Similar to [`ConfigFileNotFound`] but for general I/O paths. @@ -271,27 +253,18 @@ pub enum ErrorCode { /// Disk full, permission denied, or path invalid. Maps to **1502**. IoWriteError = 1502, - // ================================================== - // JSON errors (1600) - // ================================================== /// A JSON payload could not be parsed. /// /// The response body or request payload contained malformed JSON. /// Maps to **1600**. JsonParseError = 1600, - // ================================================== - // Internal (1700) - // ================================================== /// An internal / unexpected error. /// /// Catch-all for bugs or unclassifiable failures. If you see this /// code in production please file a bug report. Maps to **1700**. InternalError = 1700, - // ================================================== - // Signal (1999) - // ================================================== /// The operation was interrupted by a user signal (e.g. Ctrl-C). /// /// Not an error per se — the agent loop checks for this code to @@ -366,9 +339,21 @@ pub enum ApiError { /// delay without re-parsing. #[error("Rate limit exceeded (retry after {retry_after:?}): {message}")] RateLimit { - /// The server-advised delay. + /// The server-advised delay before retrying. + /// + /// Carries the parsed `Retry-After` header (or equivalent hint) + /// so downstream layers — the stream handler's back-off and the + /// fallback manager's model switch — can honour the provider's + /// guidance without re-parsing the message. `None` when the + /// provider sent no header or it failed to parse. retry_after: Option, /// The error body or human-readable message. + /// + /// Preserved verbatim from the provider's response (for example + /// `"429 Too Many Requests"`), primarily for logging and + /// diagnostics. Classification does not depend on it — the + /// variant alone selects + /// [`ErrorCode::ApiRateLimited`]. message: String, }, @@ -448,10 +433,6 @@ pub enum ApiError { } impl ApiError { - // ================================================== - // Classifiers - // ================================================== - /// Derive the machine-readable [`ErrorCode`] for this error. /// /// Inspects the error variant and, for message-based variants @@ -606,6 +587,22 @@ impl ApiError { } /// Check a message string for context-overflow keywords. + /// + /// Shared keyword matcher behind [`is_context_overflow`](Self::is_context_overflow) + /// and the message inspection inside [`code`](Self::code). Returns + /// `true` when the lowercased message contains any of the phrases + /// providers use to signal that the prompt exceeded the model's + /// maximum context length: `"context"`, `"too many tokens"`, + /// `"exceeds maximum"`, or `"max tokens"`. + /// + /// The match is deliberately broad (substring on the lowercased + /// text) so that provider-specific phrasings — Anthropic's + /// `"prompt is too long"`, OpenAI's `"maximum context length"`, + /// and generic `"too many tokens"` renderings — all collapse to a + /// single signal. The trade-off is a small false-positive risk on + /// unrelated messages that happen to contain the word `"context"`, + /// which is acceptable given the downstream consumers retry + /// through compaction rather than aborting. fn is_context_overflow_internal(msg: &str) -> bool { let msg_lower = msg.to_lowercase(); msg_lower.contains("context") @@ -728,10 +725,6 @@ impl ApiError { matches!(self, Self::Tool(..)) } - // ================================================== - // Constructors - // ================================================== - /// Create a generic [`ApiError::Api`] variant. /// /// Use this for LLM API failures that don't warrant a more specific diff --git a/src/capabilities.rs b/src/capabilities.rs index 6ba4565..a6e2102 100644 --- a/src/capabilities.rs +++ b/src/capabilities.rs @@ -1,27 +1,27 @@ //! 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) +//! agent loop can depend on. [`LoopManagers`](crate::managers::LoopManagers) //! 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 *(requires `hooks` feature)* | -//! | [`PipelineAware`] | Middleware pipeline dispatch | -//! | `HealthTrackable` | Per-tool health monitoring *(requires `tool_health` feature)* | +//! | Trait | Purpose | +//! |---------------------|---------------------------------------------------------------| +//! | [`Observable`] | Lifecycle event observation | +//! | [`Detectable`] | Loop and convergence detection | +//! | [`FallbackCapable`] | Model fallback / circuit breaker | +//! | [`Compactable`] | Context compaction | +//! | [`StreamCapable`] | Resilient LLM streaming | +//! | [`PipelineAware`] | Middleware pipeline dispatch | +//! | `Hookable` | Bidirectional lifecycle hooks *(requires `hooks` feature)* | +//! | `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): +//! without pulling in the full [`LoopManagers`](crate::managers::LoopManagers): //! //! ```rust,ignore //! fn check_patterns(runtime: &impl Detectable) { @@ -42,10 +42,6 @@ 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 @@ -54,7 +50,7 @@ use crate::tool::health::ToolHealthRegistry; /// /// # Implementors /// -/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// - [`LoopManagers`](crate::managers::LoopManagers) — the framework's default implementation. /// /// # When to use /// @@ -62,7 +58,7 @@ use crate::tool::health::ToolHealthRegistry; /// events but don't need any other infrastructure capabilities. /// /// ```rust,ignore -/// fn process_turn(runtime: &impl Observable) { +/// fn run_step(runtime: &impl Observable) { /// runtime.observers().on_turn_start(&ctx); /// // ... do work ... /// runtime.observers().on_turn_end(&ctx); @@ -70,6 +66,10 @@ use crate::tool::health::ToolHealthRegistry; /// ``` pub trait Observable { /// Returns the observer host for lifecycle event fan-out. + /// + /// The observer host holds every registered observer and dispatches + /// all lifecycle events (turn start/end, stream deltas, tool dispatch, + /// compaction) to them. fn observers(&self) -> &ObserverHost; } @@ -82,7 +82,7 @@ pub trait Observable { /// /// # Implementors /// -/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// - [`LoopManagers`](crate::managers::LoopManagers) — the framework's default implementation. /// /// # When to use /// @@ -99,6 +99,9 @@ pub trait Observable { /// ``` pub trait Detectable { /// Returns the detection manager for loop and convergence detection. + /// + /// Use this to record tool calls or model responses and check whether + /// the agent is repeating itself or converging on a stable answer. fn detection(&self) -> &DetectionManager; } @@ -110,7 +113,7 @@ pub trait Detectable { /// /// # Implementors /// -/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// - [`LoopManagers`](crate::managers::LoopManagers) — the framework's default implementation. /// /// # When to use /// @@ -129,6 +132,9 @@ pub trait Detectable { /// ``` pub trait FallbackCapable { /// Returns the fallback manager (circuit breaker) for API model fallback. + /// + /// Use this to record API failures and check whether the circuit + /// breaker has tripped, indicating the primary model is unavailable. fn fallback(&self) -> &FallbackManager; } @@ -141,7 +147,7 @@ pub trait FallbackCapable { /// /// # Implementors /// -/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// - [`LoopManagers`](crate::managers::LoopManagers) — the framework's default implementation. /// /// # When to use /// @@ -150,6 +156,9 @@ pub trait FallbackCapable { /// implementations that need to manage the context window directly. pub trait Compactable { /// Returns the context manager, if compaction is configured. + /// + /// Returns `None` when no compactor is set — the loop will not + /// auto-compact, and the host must manage context size manually. fn context_manager(&self) -> Option<&Arc>; } @@ -162,7 +171,7 @@ pub trait Compactable { /// /// # Implementors /// -/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// - [`LoopManagers`](crate::managers::LoopManagers) — the framework's default implementation. /// /// # When to use /// @@ -190,7 +199,7 @@ pub trait StreamCapable { /// /// # Implementors /// -/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// - [`LoopManagers`](crate::managers::LoopManagers) — the framework's default implementation. /// /// # When to use /// @@ -199,6 +208,9 @@ pub trait StreamCapable { #[cfg(feature = "hooks")] pub trait Hookable { /// Returns the hook executor, if hooks are configured. + /// + /// Returns `None` when no hook executor is set — no pre/post hooks + /// will fire. Hooks can approve or reject actions, unlike observers. fn hook_executor(&self) -> Option<&HookExecutor>; } @@ -209,7 +221,7 @@ pub trait Hookable { /// /// # Implementors /// -/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// - [`LoopManagers`](crate::managers::LoopManagers) — the framework's default implementation. /// /// # When to use /// @@ -227,6 +239,10 @@ pub trait Hookable { /// ``` pub trait PipelineAware { /// Returns the tool middleware pipeline, if configured. + /// + /// Returns `None` when no pipeline is set — tool calls go directly + /// to the registry. When set, every tool call passes through the + /// pipeline's middleware layers before reaching the tool. fn pipeline(&self) -> Option<&ToolPipeline>; } @@ -240,9 +256,13 @@ pub trait PipelineAware { /// /// # Implementors /// -/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// - [`LoopManagers`](crate::managers::LoopManagers) — the framework's default implementation. #[cfg(feature = "tool_health")] pub trait HealthTrackable { /// Returns the tool health registry, if health tracking is configured. + /// + /// Returns `None` when no health registry is set — per-tool circuit + /// breakers will not fire. When set, records success/failure counts + /// and opens breakers for unhealthy tools. fn health_registry(&self) -> Option<&ToolHealthRegistry>; } diff --git a/src/compact.rs b/src/compact.rs index 7d40a05..d48523f 100644 --- a/src/compact.rs +++ b/src/compact.rs @@ -42,10 +42,10 @@ //! //! let manager = ContextManager::new(Arc::new(compactor)) //! .with_context_window(200_000) -//! .with_threshold(8_000); +//! .with_threshold(80); //! //! assert_eq!(manager.context_window(), 200_000); -//! assert_eq!(manager.threshold(), 8_000); +//! assert_eq!(manager.threshold(), 80); //! ``` //! //! # Integration with `BareLoop` @@ -70,10 +70,6 @@ pub use types::{ EnsureContextResult, PostCompactStats, PreCompactStats, }; -// =================================================== -// ContextCompactor trait -// =================================================== - /// Strategy trait for compacting a conversation's message history. /// /// Implementations define *how* to reduce a message list — truncation, @@ -150,10 +146,6 @@ pub trait ContextCompactor: Send + Sync { ) -> Pin + Send + '_>>; } -// =================================================== -// CompactBase -// =================================================== - /// Determines the base used to calculate the compaction target. /// /// When compaction triggers, the manager asks the compactor to reduce @@ -177,9 +169,9 @@ pub trait ContextCompactor: Send + Sync { /// let compactor = TruncatingCompactor::new(); /// let manager = ContextManager::new(Arc::new(compactor)) /// .with_context_window(200_000) -/// .with_threshold(8_000) // triggers at 160k tokens +/// .with_threshold(80) // triggers at 160k tokens /// .with_compact_target(CompactBase::Context) // target = % of 200k -/// .with_compact_target_pct(5_000); // compact to 50% of 200k = 100k +/// .with_compact_target_pct(50); // compact to 50% of 200k = 100k /// ``` #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub enum CompactBase { @@ -195,17 +187,13 @@ pub enum CompactBase { /// /// `target = compact_threshold_tokens × compact_target_pct / 10_000` /// - /// This is the default. With the default `threshold = 8_000` bps (80%) and - /// `compact_target_pct = 7_000` bps (70%), compaction targets 56% of the - /// context window (`8000 × 7000 / 10_000² = 0.56`). + /// This is the default. With the default `threshold = 80` (80%) and + /// `compact_target_pct = 70` (70%), compaction targets 56% of the + /// context window. #[default] Threshold, } -// =================================================== -// ContextManager -// =================================================== - /// Manages context window usage and triggers compaction when needed. /// /// Main entry point for context management. It monitors @@ -233,7 +221,7 @@ pub enum CompactBase { /// /// let manager = ContextManager::new(Arc::new(compactor)) /// .with_context_window(1_000) -/// .with_threshold(8_000); +/// .with_threshold(80); /// /// // Short conversation — no compaction needed. /// let messages = vec![ @@ -243,6 +231,7 @@ pub enum CompactBase { /// let tokens = ContextManager::estimate_tokens(&messages); /// assert!(!manager.should_compact(tokens)); /// ``` +#[derive(Clone)] pub struct ContextManager { /// The compaction strategy. /// @@ -260,13 +249,13 @@ pub struct ContextManager { /// and the emergency-zone and compaction-target calculations. context_window: u64, - /// Threshold in basis points (`0–10_000`; `10_000` = 100%) at which compaction triggers. + /// Threshold as a percentage (0–100) at which compaction triggers. /// - /// When estimated token usage reaches `threshold_bps * context_window / 10_000`, + /// When estimated token usage reaches `threshold * context_window / 100`, /// [`should_compact`](Self::should_compact) returns `true` and a compaction /// pass runs before the next turn. Set via [`with_threshold`](Self::with_threshold), - /// which clamps to `[1_000, 10_000]`; defaults to `8_000` (80%). - threshold_bps: u16, + /// which clamps to `[1, 100]`; defaults to `80`. + threshold: u8, /// Whether auto-compaction is enabled. /// @@ -286,13 +275,13 @@ pub struct ContextManager { /// [`with_compact_target`](Self::with_compact_target). compact_base: CompactBase, - /// Fraction of the target base to compact down to, in basis points (`0–10_000`). + /// Fraction of the target base to compact down to, as a percentage (0–100). /// - /// The post-compaction size aimed for: `compact_target_bps * base / 10_000`, + /// The post-compaction size aimed for: `compact_target * base / 100`, /// where `base` is determined by [`compact_base`](Self::compact_base). - /// Defaults to `7_000` (70%); set and clamped to `[1_000, 10_000]` via + /// Defaults to `70` (70%); set and clamped to `[1, 100]` via /// [`with_compact_target_pct`](Self::with_compact_target_pct). - compact_target_bps: u16, + compact_target: u8, } impl ContextManager { @@ -303,19 +292,19 @@ impl ContextManager { /// | Setting | Default | /// |----------------------|------------------------------| /// | `context_window` | 200_000 | - /// | `threshold` | 8_000 bps (80%) | + /// | `threshold` | 80 (80%) | /// | `auto_compact` | `true` | - /// | `compact_target` | [`CompactBase::Threshold`] | - /// | `compact_target_pct` | 7_000 bps (70%) | + /// | `compact_target` | [`CompactBase::Threshold`] | + /// | `compact_target_pct` | 70 (70%) | #[must_use] pub fn new(compactor: Arc) -> Self { Self { compactor, context_window: 200_000, - threshold_bps: 8_000, + threshold: 80, auto_compact: true, compact_base: CompactBase::Threshold, - compact_target_bps: 7_000, + compact_target: 70, } } @@ -329,12 +318,12 @@ impl ContextManager { self } - /// Set the compaction threshold in basis points (`0–10_000`; `10_000` = 100%). + /// Set the compaction threshold as a percentage (0–100). /// - /// Clamped to `[1_000, 10_000]` (10%–100%) to prevent degenerate configurations. + /// Clamped to `[1, 100]` to prevent degenerate configurations. #[must_use] - pub fn with_threshold(mut self, threshold_bps: u16) -> Self { - self.threshold_bps = threshold_bps.clamp(1_000, 10_000); + pub fn with_threshold(mut self, threshold: u8) -> Self { + self.threshold = threshold.clamp(1, 100); self } @@ -359,14 +348,14 @@ impl ContextManager { self } - /// Set the fraction of the target base to compact down to, in basis points - /// (`0–10_000`; `10_000` = 100%). + /// Set the fraction of the target base to compact down to, as a percentage + /// (`0–100`; `100` = 100%). /// - /// Clamped to `[1_000, 10_000]` (10%–100%) to prevent degenerate configurations. - /// Defaults to `7_000` (70%). + /// Clamped to `[1, 100]` to prevent degenerate configurations. + /// Defaults to `70`. #[must_use] - pub fn with_compact_target_pct(mut self, pct_bps: u16) -> Self { - self.compact_target_bps = pct_bps.clamp(1_000, 10_000); + pub fn with_compact_target_pct(mut self, pct: u8) -> Self { + self.compact_target = pct.clamp(1, 100); self } @@ -379,13 +368,13 @@ impl ContextManager { self.context_window } - /// The compaction threshold in basis points (`0–10_000`; `10_000` = 100%). + /// The compaction threshold as a percentage (0–100). /// /// Compaction triggers when estimated tokens reach /// `context_window * threshold / 10_000`. See [`with_threshold`](Self::with_threshold). #[must_use] - pub fn threshold(&self) -> u16 { - self.threshold_bps + pub fn threshold(&self) -> u8 { + self.threshold } /// Whether auto-compaction is enabled. @@ -405,25 +394,23 @@ impl ContextManager { self.compact_base } - /// The fraction of the target base to compact down to, in basis points - /// (`0–10_000`; `10_000` = 100%). + /// The fraction of the target base to compact down to, as a percentage + /// (`0–100`; `100` = 100%). /// /// See [`with_compact_target_pct`](Self::with_compact_target_pct). #[must_use] - pub fn compact_target_pct(&self) -> u16 { - self.compact_target_bps + pub fn compact_target_pct(&self) -> u8 { + self.compact_target } /// The token budget at which compaction triggers. /// - /// Equal to `context_window * threshold_bps / 10_000`, computed in - /// saturating integer arithmetic (widened to `u128` to avoid overflow). + /// Equal to `context_window * threshold`. #[must_use] pub fn compact_threshold_tokens(&self) -> u64 { - const BASIS: u128 = 10_000; - let product = - u128::from(self.threshold_bps).saturating_mul(u128::from(self.context_window)) / BASIS; - u64::try_from(product).unwrap_or(u64::MAX) + self.context_window + .saturating_mul(u64::from(self.threshold)) + / 100 } /// The token count to compact down to. @@ -431,17 +418,15 @@ impl ContextManager { /// Computed from [`compact_target`](Self::compact_target) and /// [`compact_target_pct`](Self::compact_target_pct): /// - /// - [`CompactBase::Threshold`]: `compact_threshold_tokens × pct_bps / 10_000` - /// - [`CompactBase::Context`]: `context_window × pct_bps / 10_000` + /// - [`CompactBase::Threshold`]: `compact_threshold_tokens × pct` + /// - [`CompactBase::Context`]: `context_window × pct` #[must_use] pub fn compact_target_tokens(&self) -> u64 { - const BASIS: u128 = 10_000; let base: u64 = match self.compact_base { CompactBase::Threshold => self.compact_threshold_tokens(), CompactBase::Context => self.context_window, }; - let product = u128::from(self.compact_target_bps).saturating_mul(u128::from(base)) / BASIS; - u64::try_from(product).unwrap_or(u64::MAX) + base.saturating_mul(u64::from(self.compact_target)) / 100 } /// Estimate the token count for a slice of messages. @@ -604,6 +589,27 @@ impl ContextManager { &self, messages: Vec, turn: usize, + ) -> Result { + self.compact_with_reason(messages, turn, CompactReason::Manual) + .await + } + + /// Trigger compaction with a specific reason, bypassing the threshold check. + /// + /// Use this when the decision to compact was already made (e.g. by the + /// driving state machine) and the `reason` should reach the compactor's + /// [`CompactionContext`] so it can vary strategy (e.g. more aggressive + /// summarization for [`CompactReason::Emergency`]). + /// + /// # Errors + /// + /// Returns a [`ContextOverflow`] if compaction fails or the result + /// still exceeds the context window. + pub async fn compact_with_reason( + &self, + messages: Vec, + turn: usize, + reason: CompactReason, ) -> Result { let tokens_before = Self::estimate_tokens(&messages); @@ -615,7 +621,7 @@ impl ContextManager { let target_tokens = self.compact_target_tokens(); let context = CompactionContext { tokens_before, - reason: CompactReason::Manual, + reason, context_window: self.context_window, turn, }; @@ -656,30 +662,19 @@ impl ContextManager { /// Note: Currently does not depend on instance state, but is a method /// by design to allow future configuration-aware telemetry. #[must_use] - #[allow( - clippy::unused_self, - clippy::arithmetic_side_effects, - clippy::cast_possible_truncation - )] pub fn build_telemetry( - &self, trigger: CompactReason, pre_messages: &[Message], post_messages: &[Message], - _outcome: &CompactionOutcome, start: Instant, ) -> CompactTelemetry { let pre_tokens = Self::estimate_tokens(pre_messages); let post_tokens = Self::estimate_tokens(post_messages); let tokens_saved = pre_tokens.saturating_sub(post_tokens); - let percent_saved = if pre_tokens > 0 { - let ratio = u128::from(tokens_saved) - .saturating_mul(100) - .saturating_div(u128::from(pre_tokens)); - u8::try_from(ratio.min(100)).unwrap_or(100) - } else { - 0 - }; + let percent_saved: u8 = tokens_saved + .checked_mul(100) + .and_then(|v| v.checked_div(pre_tokens)) + .map_or(0, |ratio| ratio.min(100) as u8); CompactTelemetry { trigger, @@ -724,10 +719,6 @@ impl fmt::Debug for ContextManager { } } -// =================================================== -// Tests -// =================================================== - #[cfg(test)] mod tests { use super::*; @@ -766,7 +757,7 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)); assert_eq!(manager.context_window(), 200_000); - assert_eq!(manager.threshold(), 8_000); + assert_eq!(manager.threshold(), 80); assert!(manager.auto_compact()); } @@ -775,22 +766,22 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(100_000) - .with_threshold(5_000) + .with_threshold(50) .with_auto_compact(false); assert_eq!(manager.context_window(), 100_000); - assert_eq!(manager.threshold(), 5_000); + assert_eq!(manager.threshold(), 50); assert!(!manager.auto_compact()); } #[test] fn test_threshold_clamped() { let compactor = TruncatingCompactor::new(); - let manager = ContextManager::new(Arc::new(compactor)).with_threshold(100); - assert_eq!(manager.threshold(), 1_000); + let manager = ContextManager::new(Arc::new(compactor)).with_threshold(1); + assert_eq!(manager.threshold(), 1); let compactor2 = TruncatingCompactor::new(); - let manager = ContextManager::new(Arc::new(compactor2)).with_threshold(20_000); - assert_eq!(manager.threshold(), 10_000); + let manager = ContextManager::new(Arc::new(compactor2)).with_threshold(200); + assert_eq!(manager.threshold(), 100); } #[test] @@ -798,7 +789,7 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(1_000) - .with_threshold(8_000); + .with_threshold(80); // 800 is the threshold for 1000 * 8000 / 10_000 assert!(!manager.should_compact(799)); } @@ -808,7 +799,7 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(1_000) - .with_threshold(8_000); + .with_threshold(80); assert!(manager.should_compact(800)); } @@ -817,7 +808,7 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(1_000) - .with_threshold(5_000); // threshold at 500, but emergency at 950 + .with_threshold(50); // threshold at 500, but emergency at 950 assert!(manager.should_compact(950)); } @@ -841,7 +832,7 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(1_000) - .with_threshold(8_000); + .with_threshold(80); assert_eq!( manager.compact_reason(800), CompactReason::ThresholdExceeded @@ -860,58 +851,50 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(200_000) - .with_threshold(8_000); + .with_threshold(80); assert_eq!(manager.compact_threshold_tokens(), 160_000); } #[test] fn test_compact_target_tokens_default() { - // Default: CompactBase::Threshold with pct 7_000 bps (70%) - // threshold = 200_000 * 8_000 / 10_000 = 160_000 - // target = 160_000 * 7_000 / 10_000 = 112_000 let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(200_000) - .with_threshold(8_000); + .with_threshold(80); assert_eq!(manager.compact_target_tokens(), 112_000); } #[test] fn test_compact_target_tokens_context_base() { - // CompactBase::Context with pct 5_000 bps (50%) - // target = 200_000 * 5_000 / 10_000 = 100_000 let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(200_000) - .with_threshold(8_000) + .with_threshold(80) .with_compact_target(CompactBase::Context) - .with_compact_target_pct(5_000); + .with_compact_target_pct(50); assert_eq!(manager.compact_target_tokens(), 100_000); } #[test] fn test_compact_target_tokens_threshold_base() { - // CompactBase::Threshold with pct 5_000 bps (50%) - // threshold = 200_000 * 8_000 / 10_000 = 160_000 - // target = 160_000 * 5_000 / 10_000 = 80_000 let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(200_000) - .with_threshold(8_000) + .with_threshold(80) .with_compact_target(CompactBase::Threshold) - .with_compact_target_pct(5_000); + .with_compact_target_pct(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(100); - assert_eq!(manager.compact_target_pct(), 1_000); + let manager = ContextManager::new(Arc::new(compactor)).with_compact_target_pct(1); + assert_eq!(manager.compact_target_pct(), 1); let compactor2 = TruncatingCompactor::new(); - let manager2 = ContextManager::new(Arc::new(compactor2)).with_compact_target_pct(20_000); - assert_eq!(manager2.compact_target_pct(), 10_000); + let manager2 = ContextManager::new(Arc::new(compactor2)).with_compact_target_pct(200); + assert_eq!(manager2.compact_target_pct(), 100); } #[test] @@ -919,7 +902,7 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)); assert_eq!(manager.compact_target(), CompactBase::Threshold); - assert_eq!(manager.compact_target_pct(), 7_000); + assert_eq!(manager.compact_target_pct(), 70); } #[tokio::test] @@ -1004,7 +987,7 @@ mod tests { let compactor = TruncatingCompactor::new(); let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(1_000_000) // very large window - .with_threshold(8_000); + .with_threshold(80); let msgs = make_conversation(3); let result = manager.ensure_context_fits(msgs.clone(), 1).await; match result { @@ -1025,7 +1008,7 @@ mod tests { // 2 messages ≈ 18 tokens, so 200 tokens is plenty of headroom. let manager = ContextManager::new(Arc::new(compactor)) .with_context_window(200) - .with_threshold(1_000); + .with_threshold(10); let msgs = make_conversation(20); // 40 messages let result = manager.ensure_context_fits(msgs, 1).await; match result { @@ -1089,11 +1072,11 @@ mod tests { #[test] fn test_build_telemetry() { let compactor = TruncatingCompactor::new(); - let manager = ContextManager::new(Arc::new(compactor)).with_context_window(1_000); + let _manager = ContextManager::new(Arc::new(compactor)).with_context_window(1_000); let pre = make_conversation(10); let post = make_conversation(2); - let outcome = CompactionOutcome { + let _outcome = CompactionOutcome { messages: post.clone(), tokens_after: 100, tokens_saved: 800, @@ -1102,13 +1085,8 @@ mod tests { }; let start = Instant::now(); - let telemetry = manager.build_telemetry( - CompactReason::ThresholdExceeded, - &pre, - &post, - &outcome, - start, - ); + let telemetry = + ContextManager::build_telemetry(CompactReason::ThresholdExceeded, &pre, &post, start); assert_eq!(telemetry.trigger, CompactReason::ThresholdExceeded); assert_eq!(telemetry.pre_compact.total_messages, 20); diff --git a/src/compact/truncating.rs b/src/compact/truncating.rs index 5545395..9b34d79 100644 --- a/src/compact/truncating.rs +++ b/src/compact/truncating.rs @@ -13,10 +13,6 @@ use std::collections::HashSet; 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 @@ -51,21 +47,25 @@ use std::pin::Pin; /// ``` #[derive(Debug, Clone)] pub struct TruncatingCompactor { - /// Number of recent messages to always preserve. + /// Number of recent messages to always preserve during compaction. + /// + /// This many messages from the end of the conversation are kept intact; + /// everything before them is dropped. Defaults to 4. preserve_recent: usize, - /// Minimum messages before compaction is considered. + + /// 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. Defaults to 6. min_messages: usize, } impl TruncatingCompactor { /// Create a new truncating compactor with sensible defaults. /// - /// Defaults: - /// - /// | Setting | Default | - /// |-------------------|---------| - /// | `preserve_recent` | 4 | - /// | `min_messages` | 6 | + /// Preserves the 4 most recent messages and requires at least 6 messages + /// before compaction is attempted. #[must_use] pub fn new() -> Self { Self { @@ -95,13 +95,17 @@ impl TruncatingCompactor { self } - /// Number of recent messages that will be preserved. + /// Number of recent messages that will be preserved during compaction. + /// + /// This many messages from the end of the conversation are always kept. #[must_use] pub fn preserve_recent(&self) -> usize { self.preserve_recent } - /// Minimum messages before compaction is attempted. + /// Minimum number of messages required before compaction is attempted. + /// + /// Conversations shorter than this are left untouched. #[must_use] pub fn min_messages(&self) -> usize { self.min_messages @@ -165,10 +169,6 @@ impl ContextCompactor for TruncatingCompactor { } } -// =================================================== -// Tool-call/result pair protection -// ================================================== - impl TruncatingCompactor { /// Adjust the split index to avoid orphaning tool-call/result pairs. /// @@ -235,10 +235,6 @@ impl TruncatingCompactor { } } -// =================================================== -// TokenSplitter -// =================================================== - /// Splits a conversation into "old" and "recent" at a turn boundary. /// /// Used by compactors (and agent-side code) that need to know which @@ -276,36 +272,59 @@ impl TruncatingCompactor { /// ``` #[derive(Debug, Clone)] pub struct TokenSplitter { - /// Number of recent messages to always preserve. + /// Number of recent messages to always preserve during a split. + /// + /// This many messages from the end of the conversation are kept in the + /// preserved portion. Defaults to 4. preserve_recent: usize, - /// Minimum messages before considering a split. + + /// Minimum number of messages before a split is attempted. + /// + /// Conversations shorter than this go entirely into the preserved + /// portion. Defaults to 6. 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). + /// Messages to compact or summarize (the old portion). + /// + /// These are the messages before the split point. They are candidates + /// for summarization, truncation, or removal by the compactor. pub to_compact: Vec, - /// Messages to preserve as-is (the "recent" part). + + /// Messages to preserve as-is (the recent portion). + /// + /// These are the messages after the split point. They are kept + /// untouched in the conversation history. pub preserved: Vec, - /// Estimated tokens in `to_compact`. + + /// Estimated token count of the old portion ([`to_compact`](Self::to_compact)). + /// + /// Approximate token usage of the messages eligible for summarization or + /// removal, computed from the pre-split history so callers can budget how + /// much context compaction should reclaim. pub compact_tokens: u64, - /// Estimated tokens in `preserved`. + + /// Estimated token count of the recent portion ([`preserved`](Self::preserved)). + /// + /// Approximate token usage of the messages kept intact after the split, + /// giving callers the residual context footprint that will carry into the + /// next model call. pub preserved_tokens: u64, + /// The index in the original message list where the split occurred. + /// + /// Zero when no split was needed (the entire conversation was preserved). pub split_index: usize, } impl TokenSplitter { /// Create a new splitter with sensible defaults. /// - /// Defaults: - /// - /// | Setting | Default | - /// |-------------------|---------| - /// | `preserve_recent` | 4 | - /// | `min_messages` | 6 | + /// Preserves the 4 most recent messages and requires at least 6 messages + /// before a split is attempted. #[must_use] pub fn new() -> Self { Self { @@ -314,14 +333,20 @@ impl TokenSplitter { } } - /// Set how many recent messages to preserve. + /// Set how many recent messages to preserve during a split. + /// + /// This many messages from the end of the conversation are kept in the + /// preserved portion. The rest are candidates for compaction. #[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. + /// Set the minimum number of messages before a split is attempted. + /// + /// Conversations shorter than this are returned entirely as preserved. + /// Prevents splitting very short conversations into meaningless fragments. #[must_use] pub fn with_min_messages(mut self, count: usize) -> Self { self.min_messages = count.max(2); diff --git a/src/compact/types.rs b/src/compact/types.rs index d6e339c..41c0a58 100644 --- a/src/compact/types.rs +++ b/src/compact/types.rs @@ -14,10 +14,6 @@ use crate::message::Message; use serde::{Deserialize, Serialize}; use std::fmt; -// =================================================== -// CompactReason -// =================================================== - /// Why compaction was triggered. /// /// Different triggers may warrant different compaction strategies. @@ -60,10 +56,6 @@ impl fmt::Display for CompactReason { } } -// =================================================== -// CompactionContext -// =================================================== - /// Metadata passed to [`ContextCompactor::compact`](super::ContextCompactor::compact) /// describing the compaction trigger and current state. /// @@ -101,10 +93,6 @@ pub struct CompactionContext { pub turn: usize, } -// =================================================== -// CompactionOutcome -// =================================================== - /// Result of a single compaction pass. /// /// Returned by [`ContextCompactor::compact`](super::ContextCompactor::compact), @@ -194,10 +182,6 @@ impl CompactionOutcome { } } -// =================================================== -// CompactTelemetry -// =================================================== - /// Telemetry data for a single compaction operation. /// /// Produced by [`ContextManager::ensure_context_fits`](super::ContextManager::ensure_context_fits) @@ -290,18 +274,21 @@ pub struct PostCompactStats { /// [`total_messages`](PreCompactStats::total_messages) when compaction /// removed or summarized messages; equal when it made no change. pub total_messages: usize, + /// Estimated token count after compaction. /// /// The post-compaction token estimate, comparable to /// [`estimated_tokens`](PreCompactStats::estimated_tokens) from the /// pre-compaction snapshot. The difference is [`tokens_saved`](Self::tokens_saved). pub estimated_tokens: u64, + /// Tokens removed by compaction. /// /// How many tokens the pass reclaimed: the pre-compaction estimate minus /// the post-compaction estimate. Saturates at zero, so it never goes /// negative even if a summary injection made the conversation larger. pub tokens_saved: u64, + /// Percentage of tokens saved (0–100). /// /// [`tokens_saved`](Self::tokens_saved) as a share of the pre-compaction @@ -310,10 +297,6 @@ pub struct PostCompactStats { pub percent_saved: u8, } -// =================================================== -// ContextOverflow error -// =================================================== - /// Error returned when the conversation cannot fit within the context /// window, even after compaction. /// @@ -359,13 +342,20 @@ pub struct ContextOverflow { } impl ContextOverflow { - /// How many tokens the conversation exceeds the window by. + /// How many tokens the conversation exceeds the context window by. + /// + /// Returns zero when the conversation fits within the window. Uses + /// saturating subtraction so an underflow never panics. #[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+). + /// The fraction of the context window currently consumed. + /// + /// Returns a value between `0.0` and `1.0` when the conversation fits, + /// or above `1.0` when it overflows. Returns infinity when the context + /// window is zero (division by zero). #[must_use] pub fn utilization(&self) -> f64 { if self.context_window == 0 { @@ -391,10 +381,6 @@ impl fmt::Display for ContextOverflow { 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 @@ -419,6 +405,11 @@ pub enum EnsureContextResult { impl EnsureContextResult { /// Extract the message list from this result, regardless of variant. + /// + /// Returns the compacted messages from [`Compacted`](Self::Compacted) or + /// the unchanged messages from [`NoAction`](Self::NoAction). Use this when + /// you only care about the resulting history and not whether compaction + /// actually occurred. #[must_use] pub fn into_messages(self) -> Vec { match self { diff --git a/src/config.rs b/src/config.rs index c22034c..a3832f3 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,326 +1,131 @@ //! Agent session configuration. //! -//! Defines [`LoopConfig`] — the configuration struct that controls agent -//! session parameters such as turn limits, model selection, and context -//! window size. +//! Defines the configuration types that control agent parameters: +//! [`SessionConfig`] for session-scoped settings and +//! [`RunConfig`](crate::engine::RunConfig) for per-run budgets. -use uuid::Uuid; - -/// Configuration for an agent session. +/// Session-scoped agent configuration. /// -/// 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. +/// The slice of agent configuration that is stable across `run()` calls on the +/// same agent: the session identity, system prompt, and the model's context +/// window. These do not vary from one prompt to the next within a session, so +/// they live here rather than on the per-run [`RunConfig`](crate::engine::core::RunConfig). /// /// # Construction /// -/// Use [`LoopConfig::default`] for sensible defaults, then override individual -/// fields either with the fluent `with_*()` builders or direct struct-literal -/// assignment — both are supported. +/// Use [`SessionConfig::default`] for sensible defaults, then override +/// individual fields with the `with_*()` builders. /// /// ``` -/// use loopctl::config::LoopConfig; +/// use loopctl::config::SessionConfig; /// -/// let config = LoopConfig::default() -/// .with_max_turns(50) -/// .with_model("default"); +/// let config = SessionConfig::default().with_context_window(128_000); +/// assert_eq!(config.context_window, 128_000); /// ``` #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -#[non_exhaustive] -pub struct LoopConfig { - /// Unique session identifier. - /// - /// A random UUID v4 generated on construction. Used to tag tool contexts, - /// session save/load paths, and observer events so a host app can correlate - /// turns across logs and persisted state. Two configs with the same - /// `session_id` are considered the same logical session. - pub session_id: Uuid, - - /// Model identifier passed to the API client on each request. - /// - /// Provider-specific (e.g. `"gpt-4o"`, `"claude-3.5-sonnet"`). The client - /// may override this at runtime via - /// [`ApiClient::set_model`](crate::api::ApiClient::set_model) when the - /// [`FallbackManager`](crate::fallback::FallbackManager) trips. Must be - /// non-empty ([`validate`](LoopConfig::validate) rejects whitespace-only). - pub model: String, - +pub struct SessionConfig { /// Optional system prompt override. /// /// When `Some`, this text is sent to the provider as the turn's system - /// prompt. When `None` (the default), the provider receives no system - /// prompt for the turn. Set this to inject a custom persona or instruction - /// set. + /// prompt. When `None` (the default), the provider receives no system prompt. pub system_prompt: Option, - /// Maximum number of turns before forcing completion. - /// - /// A safety cap: the agent loop halts with - /// [`LoopError::MaxTurnsExceeded`](crate::error::LoopError::MaxTurnsExceeded) - /// if it reaches this count without the model emitting a stop. Defaults to - /// `200`. Must be at least `1` ([`validate`](LoopConfig::validate)). - pub max_turns: usize, - - /// Maximum tokens for each API response. - /// - /// Sent to the provider as the `max_tokens` parameter on every request, - /// capping the length of a single model response. Defaults to `16_384`. - /// Must be at least `1` ([`validate`](LoopConfig::validate)). - pub max_tokens: u32, - /// Context window size in tokens. /// - /// Must match the actual window of [`model`](LoopConfig::model). Used by - /// the compaction subsystem to decide when conversation history exceeds - /// the budget (see [`compact_threshold`](LoopConfig::compact_threshold)). - /// Defaults to `200_000`. Must be at least `1` - /// ([`validate`](LoopConfig::validate)). + /// Must match the actual window of the configured model. Used by the + /// compaction subsystem to decide when the conversation exceeds the budget. pub context_window: u64, - /// Threshold to trigger auto-compaction, in basis points (`0–10_000`; - /// `10_000` = 100% of the context window). + /// Threshold to trigger auto-compaction, as a fraction of the context + /// window (0–100). Defaults to `80`. /// - /// When the estimated token usage of the conversation exceeds - /// `compact_threshold * context_window / 10_000`, the compactor runs before - /// the next turn to avoid an overflow. Defaults to `8_000` (80%). Must be in - /// `[0, 10_000]` ([`validate`](LoopConfig::validate)). - pub compact_threshold: u16, + /// When estimated token usage reaches this percentage of + /// [`context_window`](Self::context_window), the compaction subsystem is + /// invoked to summarize or truncate older messages before the next model + /// call. Lower it to compact more aggressively; raise it to defer + /// compaction and preserve more raw history. + pub compact_threshold: u8, - /// Whether auto-compaction is enabled. + /// Whether auto-compaction is enabled. Defaults to `true`. /// - /// When `true` (the default), the compactor runs automatically when - /// [`compact_threshold`](LoopConfig::compact_threshold) is reached. When - /// `false`, the agent never auto-compacts — the host app must manage - /// context size manually (useful for tests or fixed-length sessions). + /// When `true` the loop automatically runs the compactor once the + /// [`compact_threshold`](Self::compact_threshold) is reached. When `false` + /// the loop never compacts on its own, even under token pressure, leaving + /// context management entirely to the caller. pub auto_compact: bool, - - /// How independent tool calls within a single turn are dispatched. - /// - /// Defaults to [`ParallelMode::Sequential`] (one at a time). - /// Set [`ParallelMode::Parallel`] to run independent, - /// concurrency-safe calls concurrently up to - /// [`max_concurrency`](ParallelDispatchConfig::max_concurrency). - /// - /// Parallel mode changes observer/detection event ordering (PRE events are - /// batched, then POST events batched, rather than strictly paired per - /// call) and adds finer-grained cancellation. See - /// [`ParallelDispatchConfig`] and the dispatch module docs. - pub parallel_tool_dispatch: ParallelDispatchConfig, } -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) | `8_000` (80%) | - /// | [`auto_compact`](LoopConfig::auto_compact) | `true` | - /// | [`parallel_tool_dispatch`](LoopConfig::parallel_tool_dispatch) | `Sequential`, `max_concurrency: 8` | - /// - /// # Example - /// - /// ``` - /// use loopctl::config::LoopConfig; - /// - /// let config = LoopConfig::default(); - /// assert_eq!(config.max_turns, 200); - /// assert_eq!(config.model, "default"); - /// ``` +impl Default for SessionConfig { 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: 8_000, + compact_threshold: 80, auto_compact: true, - parallel_tool_dispatch: ParallelDispatchConfig::default(), } } } -impl LoopConfig { - /// Set the session identifier, overriding the random UUID v4 that - /// [`Default`](LoopConfig::default) generates. +impl SessionConfig { + /// Set the optional system prompt. /// - /// Use this to resume a prior session under its existing ID, or to pin a - /// deterministic ID in tests. + /// Stores `Some(prompt)` on the config; the provider receives it as + /// the turn's system prompt. Accepts any `Into` so string + /// literals work directly. Pass `None` (or skip this builder) to + /// leave the prompt unset, in which case the provider gets no + /// system prompt. #[must_use] - pub fn with_session_id(mut self, session_id: Uuid) -> Self { - self.session_id = session_id; + pub fn with_system_prompt(mut self, system_prompt: impl Into) -> Self { + self.system_prompt = Some(system_prompt.into()); self } - /// Set the model identifier passed to the API client on each request. + /// Set the model's context window size in tokens. /// - /// Accepts anything that converts into a `String` (for example `"gpt-4o"`), - /// so a `&str` literal works without an explicit `.to_string()`. The value - /// must be non-empty; see [`validate`](LoopConfig::validate). - #[must_use] - pub fn with_model(mut self, model: impl Into) -> Self { - self.model = model.into(); - self - } - - /// Set the optional system-prompt override. - /// - /// When `Some`, this text is sent to the provider as the turn's system - /// prompt (verbatim). When `None` (the default), the provider receives no - /// system prompt for the turn. Pass `Some(prompt)` to inject a custom - /// persona or instructions, or `None` to revert an earlier override. - #[must_use] - pub fn with_system_prompt(mut self, system_prompt: Option) -> Self { - self.system_prompt = system_prompt; - self - } - - /// Set the maximum number of turns before the loop forces completion. - /// - /// A safety cap: the loop halts with - /// [`LoopError::MaxTurnsExceeded`](crate::error::LoopError::MaxTurnsExceeded) - /// once it is reached. Must be at least `1`; see - /// [`validate`](LoopConfig::validate). - #[must_use] - pub fn with_max_turns(mut self, max_turns: usize) -> Self { - self.max_turns = max_turns; - self - } - - /// Set the maximum tokens for each API response, sent to the provider as - /// `max_tokens` on every request. - /// - /// Must be at least `1`; see [`validate`](LoopConfig::validate). - #[must_use] - pub fn with_max_tokens(mut self, max_tokens: u32) -> Self { - self.max_tokens = max_tokens; - self - } - - /// Set the context window size in tokens. - /// - /// Should match the actual window of the configured - /// [`model`](LoopConfig::model); the compaction subsystem uses it together - /// with [`compact_threshold`](LoopConfig::compact_threshold) to decide when - /// to run. Must be at least `1`; see [`validate`](LoopConfig::validate). + /// Must match the actual window of the configured model — the + /// compaction subsystem compares estimated token usage against this + /// value to decide when to compact. Setting it too low triggers + /// premature compaction; too high risks a provider-side context + /// overflow. #[must_use] pub fn with_context_window(mut self, context_window: u64) -> Self { self.context_window = context_window; self } - /// Set the auto-compaction trigger threshold in basis points (`0–10_000`; - /// `10_000` = 100% of the context window). + /// Set the compaction threshold as a percentage of the context + /// window (0–100). /// - /// When estimated token usage exceeds - /// `compact_threshold * context_window / 10_000`, compaction runs before the - /// next turn. This builder stores the value verbatim — it does **not** clamp; - /// the range is enforced by [`validate`](LoopConfig::validate), which is the - /// single source of truth for the `[0, 10_000]` bound. + /// When estimated token usage reaches this percentage of + /// [`context_window`](Self::context_window), the compaction + /// subsystem summarizes or truncates older messages before the next + /// model call. Lower it to compact more aggressively; raise it to + /// defer compaction and preserve more raw history. #[must_use] - pub fn with_compact_threshold(mut self, compact_threshold: u16) -> Self { - self.compact_threshold = compact_threshold; + pub fn with_compact_threshold(mut self, compact_threshold: u8) -> Self { + self.compact_threshold = compact_threshold.min(100); self } - /// Enable or disable automatic compaction. + /// Enable or disable auto-compaction. /// - /// When `true` (the default), compaction runs automatically once - /// [`compact_threshold`](LoopConfig::compact_threshold) is reached. When - /// `false`, the host application must manage context size itself. + /// When `true` (the default) the loop runs the compactor + /// automatically once [`compact_threshold`](Self::compact_threshold) + /// is reached. When `false` the loop never compacts on its own, + /// even under token pressure, leaving context management entirely + /// to the caller. #[must_use] pub fn with_auto_compact(mut self, auto_compact: bool) -> Self { self.auto_compact = auto_compact; self } - - /// Set how independent tool calls within a single turn are dispatched. - /// - /// Defaults to sequential. See [`ParallelDispatchConfig`] and - /// [`ParallelMode`] for the observer and loop-detection ordering - /// implications of enabling parallel dispatch. - #[must_use] - pub fn with_parallel_tool_dispatch(mut self, config: ParallelDispatchConfig) -> Self { - self.parallel_tool_dispatch = config; - self - } - - /// Validate the configuration fields. - /// - /// Checks that: - /// - `compact_threshold` is in the range `[0, 10_000]`. - /// - `max_turns` is greater than zero. - /// - `context_window` is greater than zero. - /// - `max_tokens` is greater than zero. - /// - `model` is not empty. - /// - /// # Errors - /// - /// Returns a [`Config`](crate::error::LoopError::Config) variant describing - /// or `Ok(())` if all fields are valid. - /// - /// # Example - /// - /// ``` - /// use loopctl::config::LoopConfig; - /// - /// let config = LoopConfig::default(); - /// assert!(config.validate().is_ok()); - /// - /// let bad = LoopConfig::default().with_compact_threshold(15_000); - /// assert!(bad.validate().is_err()); - /// ``` - #[must_use = "validation errors should not be silently ignored"] - pub fn validate(&self) -> Result<(), crate::error::LoopError> { - if self.max_turns == 0 { - return Err(crate::error::LoopError::Config( - "max_turns must be greater than 0".to_string(), - )); - } - if self.context_window == 0 { - return Err(crate::error::LoopError::Config( - "context_window must be greater than 0".to_string(), - )); - } - if self.max_tokens == 0 { - return Err(crate::error::LoopError::Config( - "max_tokens must be greater than 0".to_string(), - )); - } - if self.model.trim().is_empty() { - return Err(crate::error::LoopError::Config( - "model must not be empty".to_string(), - )); - } - if self.compact_threshold > 10_000 { - return Err(crate::error::LoopError::Config(format!( - "compact_threshold must be in [0, 10_000], got {}", - self.compact_threshold - ))); - } - if self.parallel_tool_dispatch.max_concurrency == 0 { - return Err(crate::error::LoopError::Config( - "parallel_tool_dispatch.max_concurrency must be at least 1".to_string(), - )); - } - Ok(()) - } } /// How independent tool calls within a single turn are dispatched. /// /// Defaults to [`Sequential`](ParallelMode::Sequential); /// opt into [`Parallel`](ParallelMode::Parallel) -/// via [`LoopConfig::parallel_tool_dispatch`]. +/// via [`ParallelDispatchConfig::mode`]. /// /// Parallel mode runs independent, concurrency-safe calls concurrently (up to /// [`ParallelDispatchConfig::max_concurrency`]). Sequential mode dispatches one @@ -350,6 +155,18 @@ pub enum ParallelMode { /// in input order) then batched POST events — see the dispatch module /// docs for the ordering invariant. Choose this for read-heavy, /// multi-call turns where latency is the sum of independent operations. + /// + /// # Side-effect divergence from Sequential + /// + /// In Sequential mode, loop detection and observer events fire on + /// **every** retry attempt — a tool that fails twice then succeeds + /// produces 3 detection operations and 3 observer pairs. + /// + /// In Parallel mode, detection and observer side-effects fire **once** + /// on the final result only (they are not thread-safe). Health + /// tracking fires on **every** attempt in both modes (it uses atomic + /// counters). Intermediate retries during parallel dispatch are + /// invisible to loop detection and observers but visible to health. Parallel, } @@ -362,11 +179,10 @@ pub enum ParallelMode { /// # Example /// /// ``` -/// use loopctl::config::{LoopConfig, ParallelDispatchConfig, ParallelMode}; +/// use loopctl::config::{ParallelDispatchConfig, ParallelMode}; /// -/// let config = LoopConfig::default().with_parallel_tool_dispatch( -/// ParallelDispatchConfig { mode: ParallelMode::Parallel, max_concurrency: 4 }, -/// ); +/// let dispatch = ParallelDispatchConfig { mode: ParallelMode::Parallel, max_concurrency: 4 }; +/// assert_eq!(dispatch.max_concurrency, 4); /// ``` #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct ParallelDispatchConfig { @@ -385,8 +201,7 @@ pub struct ParallelDispatchConfig { /// Defaults to `8`. Clamped to `[1, eligible_count]` at dispatch time (a /// 3-call batch never tries to acquire 8 permits). Setting this to `1` /// makes "parallel" mode behave like sequential — useful for debugging - /// (same code path, no concurrency). Must be at least 1 - /// ([`LoopConfig::validate`] rejects `0`). + /// (same code path, no concurrency). Must be at least 1. pub max_concurrency: usize, } @@ -404,153 +219,54 @@ 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(); - let msg = err.to_string(); - assert!( - msg.contains("max_tokens"), - "error should mention max_tokens: {msg}" - ); - } - - #[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(); - let msg = err.to_string(); - assert!(msg.contains("model"), "error should mention model: {msg}"); - } - - #[test] - fn validate_accepts_nonempty_model() { - let config = LoopConfig { - model: "gpt-4".to_string(), - ..LoopConfig::default() - }; - 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}"); - } - - #[test] - fn with_model_sets_field_and_accepts_str() { - let config = LoopConfig::default().with_model("gpt-4o"); - assert_eq!(config.model, "gpt-4o"); - } - - #[test] - fn with_system_prompt_sets_some() { - let config = LoopConfig::default().with_system_prompt(Some("p".to_string())); - assert_eq!(config.system_prompt, Some("p".to_string())); - } - - #[test] - fn with_system_prompt_none_clears_override() { - let config = LoopConfig::default().with_system_prompt(None); + fn session_config_default_context_window() { + let config = SessionConfig::default(); + assert_eq!(config.context_window, 200_000); assert!(config.system_prompt.is_none()); } #[test] - fn with_max_turns_sets_field() { - let config = LoopConfig::default().with_max_turns(7); - assert_eq!(config.max_turns, 7); + fn session_config_with_system_prompt_sets_some() { + let config = SessionConfig::default().with_system_prompt("be helpful"); + assert_eq!(config.system_prompt.as_deref(), Some("be helpful")); } #[test] - fn with_max_tokens_sets_field() { - let config = LoopConfig::default().with_max_tokens(123); - assert_eq!(config.max_tokens, 123); - } - - #[test] - fn with_context_window_sets_field() { - let config = LoopConfig::default().with_context_window(8192); + fn session_config_with_context_window_sets_field() { + let config = SessionConfig::default().with_context_window(8192); assert_eq!(config.context_window, 8192); } #[test] - fn with_compact_threshold_stores_value_without_clamping() { - let config = LoopConfig::default().with_compact_threshold(15_000); - assert_eq!(config.compact_threshold, 15_000); + fn session_config_builder_chain_composes_without_clobbering() { + let config = SessionConfig::default() + .with_system_prompt("p") + .with_context_window(128_000); + assert_eq!(config.system_prompt.as_deref(), Some("p")); + assert_eq!(config.context_window, 128_000); } #[test] - fn with_auto_compact_flips_default() { - let config = LoopConfig::default().with_auto_compact(false); - assert!(!config.auto_compact); + fn session_config_builder_does_not_mutate_source() { + let original = SessionConfig::default(); + let _modified = original.clone().with_context_window(1); + assert_eq!(original.context_window, 200_000); } #[test] - fn with_session_id_round_trips() { - let id = Uuid::new_v4(); - let config = LoopConfig::default().with_session_id(id); - assert_eq!(config.session_id, id); + fn parallel_dispatch_default_is_sequential_concurrency_8() { + let dispatch = ParallelDispatchConfig::default(); + assert_eq!(dispatch.mode, ParallelMode::Sequential); + assert_eq!(dispatch.max_concurrency, 8); } #[test] - fn with_parallel_tool_dispatch_round_trips() { + fn parallel_dispatch_struct_literal_round_trips() { let dispatch = ParallelDispatchConfig { mode: ParallelMode::Parallel, max_concurrency: 4, }; - let config = LoopConfig::default().with_parallel_tool_dispatch(dispatch); - assert_eq!(config.parallel_tool_dispatch.mode, ParallelMode::Parallel); - assert_eq!(config.parallel_tool_dispatch.max_concurrency, 4); - } - - #[test] - fn builder_chain_composes_without_clobbering() { - let config = LoopConfig::default().with_model("x").with_max_turns(5); - assert_eq!(config.model, "x"); - assert_eq!(config.max_turns, 5); - let defaults = LoopConfig::default(); - assert_eq!(config.max_tokens, defaults.max_tokens); - assert_eq!(config.context_window, defaults.context_window); - assert_eq!(config.compact_threshold, defaults.compact_threshold); - assert_eq!(config.auto_compact, defaults.auto_compact); - } - - #[test] - fn builder_does_not_mutate_source() { - let original = LoopConfig::default(); - let _modified = original.clone().with_max_turns(1); - assert_eq!(original.max_turns, 200); - } - - #[test] - fn validate_still_rejects_builder_built_invalid_config() { - let config = LoopConfig::default().with_compact_threshold(15_000); - assert!(config.validate().is_err()); + assert_eq!(dispatch.mode, ParallelMode::Parallel); + assert_eq!(dispatch.max_concurrency, 4); } } diff --git a/src/detection.rs b/src/detection.rs index 747e45c..ea67c5f 100644 --- a/src/detection.rs +++ b/src/detection.rs @@ -4,7 +4,7 @@ //! - **[`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`]. +//! For capability traits and the runtime bundle, see [`crate::managers`]. pub mod convergence; pub mod loop_detector; diff --git a/src/detection/convergence.rs b/src/detection/convergence.rs index 71ac4cf..749995a 100644 --- a/src/detection/convergence.rs +++ b/src/detection/convergence.rs @@ -84,10 +84,6 @@ use std::collections::VecDeque; use serde::{Deserialize, Serialize}; -// =================================================== -// ConvergenceAction -// =================================================== - /// Action to take when convergence is detected. /// /// When the [`ConvergenceDetector`] determines that the agent's recent responses @@ -182,10 +178,6 @@ pub enum ConvergenceAction { Compact, } -// =================================================== -// ConvergenceConfigError -// =================================================== - /// Error returned when [`ConvergenceConfig`] validation fails. /// /// [`ConvergenceDetector::new`] validates the configuration before @@ -214,6 +206,11 @@ pub enum ConvergenceConfigError { #[error("window_size must be at least 2, got {actual}")] WindowTooSmall { /// The invalid window size that was provided. + /// + /// Captured verbatim from + /// [`ConvergenceConfig::window_size`] so the diagnostic can + /// report the offending value (e.g. `0` or `1`) alongside the + /// required minimum of `2`. actual: usize, }, @@ -224,14 +221,15 @@ pub enum ConvergenceConfigError { #[error("similarity_threshold must be in [0.0, 1.0], got {actual}")] ThresholdOutOfRange { /// The invalid threshold value that was provided. + /// + /// Captured verbatim from + /// [`ConvergenceConfig::similarity_threshold`] so the diagnostic + /// can report the out-of-range value (e.g. `1.5` or `-0.2`) + /// alongside the required `[0.0, 1.0]` interval. actual: f32, }, } -// =================================================== -// ConvergenceConfig -// =================================================== - /// Configuration for convergence detection. /// /// Controls the sensitivity and behavior of the [`ConvergenceDetector`]. @@ -273,12 +271,32 @@ pub enum ConvergenceConfigError { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ConvergenceConfig { /// Whether convergence detection is active. Defaults to `true`. + /// + /// Master switch: when `false` the detector returns a "no convergence" + /// status for every response without inspecting the window, so the + /// subsystem can be disabled without removing it from the pipeline. pub enabled: bool, + /// Consecutive similar responses required to declare convergence. Must be ≥ 2. Defaults to `3`. + /// + /// The streak length that, once reached by consecutive similar responses, + /// sets [`ConvergenceStatus::detected`] to `true`. A larger value requires + /// longer repetition streaks before firing, reducing false positives. pub window_size: usize, + /// Jaccard similarity threshold (0.0–1.0). Defaults to `0.95`. + /// + /// Minimum Jaccard similarity a response must share with the immediately + /// preceding one to extend the consecutive-similar streak. Lower it to + /// catch paraphrased repetition; raise it toward `1.0` to demand + /// near-verbatim matches. pub similarity_threshold: f32, + /// Action on convergence. Defaults to [`ConvergenceAction::Stop`]. + /// + /// The [`ConvergenceAction`] returned inside [`ConvergenceStatus`] once + /// convergence is detected, telling the caller whether to halt, warn, + /// switch phase, ask the user, or compact the history. #[serde(default)] pub on_converge: ConvergenceAction, } @@ -314,10 +332,6 @@ impl Default for ConvergenceConfig { } } -// =================================================== -// ConvergenceStatus -// =================================================== - /// Convergence status with details. /// /// Returned by [`ConvergenceDetector::add_response`] and @@ -354,14 +368,39 @@ impl Default for ConvergenceConfig { #[derive(Debug, Clone, Default)] pub struct ConvergenceStatus { /// `true` when `consecutive_count >= window_size`. Other fields still valid when `false`. + /// + /// The headline result of a convergence check: `true` only once the + /// consecutive-similar streak has grown to the configured + /// [`ConvergenceConfig::window_size`]. The remaining fields remain + /// populated even when this is `false`, so callers can monitor trends. pub detected: bool, + /// Resets to `1` when similarity falls below threshold. + /// + /// Current length of the consecutive-similar streak. Grows by one each + /// time a response matches its predecessor above the threshold, and snaps + /// back to `1` as soon as a dissimilar response breaks the streak. pub consecutive_count: usize, + /// Highest Jaccard similarity (0.0–1.0). `0.0` when no comparison was made. + /// + /// Peak similarity observed between the latest response and any prior + /// response in the window during the most recent + /// [`add_response`](ConvergenceDetector::add_response) call, for + /// diagnostics and threshold tuning. pub similarity_score: f32, + /// Responses that exceeded the similarity threshold. Cleared on dissimilar response. + /// + /// Deduplicated set of responses that contributed to the current streak. + /// Emptied whenever a dissimilar response resets the streak, so it always + /// reflects the active run rather than historical matches. pub similar_responses: Vec, + /// Forwarded from [`ConvergenceConfig::on_converge`]; see [`ConvergenceAction`]. + /// + /// The action the caller should take when [`detected`](Self::detected) is + /// `true`, copied verbatim from the detector's configuration. pub action: ConvergenceAction, } @@ -395,10 +434,6 @@ impl ConvergenceStatus { } } -// =================================================== -// ConvergenceDetector -// =================================================== - /// Detects when agent responses have converged (become semantically similar). /// /// Maintains a sliding window of recent response strings and computes @@ -432,13 +467,38 @@ impl ConvergenceStatus { /// ``` #[derive(Debug)] pub struct ConvergenceDetector { - /// Immutable config set at construction. + /// Immutable configuration set at construction. + /// + /// Holds the validated [`ConvergenceConfig`] (window size, + /// similarity threshold, action) for the lifetime of the detector. + /// Exposed to callers via [`config`](ConvergenceDetector::config). config: ConvergenceConfig, - /// Bounded by `window_size`, ordered oldest→newest. - pub(super) window: VecDeque, - /// Resets to `1` on dissimilar response; convergence at `window_size`. - pub(super) consecutive_count: usize, - /// Deduplicated similar responses; cleared when streak breaks. + + /// Sliding window of recent responses, ordered oldest to newest. + /// + /// Bounded to [`ConvergenceConfig::window_size`] entries. Each + /// call to [`add_response`](ConvergenceDetector::add_response) + /// pushes the new response and evicts the oldest when full. The + /// front of the deque is the oldest kept response. Exposed to + /// callers via [`window`](ConvergenceDetector::window). + window: VecDeque, + + /// Current length of the ongoing similar-response streak. + /// + /// Resets to `1` whenever a new response is dissimilar to its + /// predecessor; otherwise it increments. Convergence is declared + /// once this counter reaches + /// [`window_size`](ConvergenceConfig::window_size). Exposed to + /// callers via + /// [`consecutive_count`](ConvergenceDetector::consecutive_count). + consecutive_count: usize, + + /// Deduplicated set of responses in the current similar streak. + /// + /// Populated while the streak continues and cleared whenever the + /// streak breaks. Drives the `similar_responses` field of + /// [`ConvergenceStatus`] so callers can inspect what the detector + /// considered "the same" before acting. similar_responses: Vec, } diff --git a/src/detection/loop_detector.rs b/src/detection/loop_detector.rs index 80d0f28..8a6ad2b 100644 --- a/src/detection/loop_detector.rs +++ b/src/detection/loop_detector.rs @@ -638,10 +638,24 @@ impl LoopDetectorConfig { #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct Operation { /// Tool name (e.g. `"Read"`, `"Edit"`, `"Bash"`). + /// + /// The identifier the tool was registered under, used as the first + /// component of the `(tool, primary_param, result_hash)` loop signature. pub tool: String, + /// Operation target (file path, command, etc.). Extracted by [`ToolSignature::extract_primary_param`]. + /// + /// A string that uniquely identifies what the operation acted on — + /// typically a file path or shell command — so two calls to the same tool + /// on different targets are treated as distinct operations. pub primary_param: String, + /// Hash of the tool result. `None` means results are ignored. Generated by [`hash_result`]. + /// + /// Compact digest of the tool's output. Two operations with identical + /// `tool` and `primary_param` but different result hashes are treated as + /// distinct (the agent is making progress); `None` disables that + /// result-aware comparison. pub result_hash: Option, } @@ -943,14 +957,38 @@ pub fn hash_result(content: &str) -> Option { #[derive(Debug, Clone, Default)] pub struct LoopStatus { /// `true` when an operation repeats beyond the configured threshold. + /// + /// The headline flag: set when at least one operation's occurrence count + /// reaches the applicable (generic or per-tool) threshold within the + /// sliding window, signalling a probable repetitive loop. pub is_looping: bool, + /// Operations tied for highest repetition. Empty when `is_looping` is `false`. + /// + /// Every operation sharing the maximum repetition count, so callers can + /// see the full set of culprits when several operations loop in lockstep. pub repeated_operations: Vec, + /// Repetitions of the most-repeated operation. Zero when no loop detected. + /// + /// The count attained by the entries in + /// [`repeated_operations`](Self::repeated_operations); compared against + /// [`LoopDetectorConfig::stop_threshold`] to decide on a forced stop. pub repetition_count: usize, + /// Loop description with count and suggestion. `None` when not looping or already warned. + /// + /// Human-readable message naming the looping operation, its repetition + /// count, and any tool-specific suggestion. Suppressed (set to `None`) + /// once a warning has been emitted for the same operation, to avoid + /// flooding the agent's context with duplicates. pub warning: Option, + /// `true` when `repetition_count >= stop_threshold` (non-zero). + /// + /// Escalation flag indicating the loop has persisted long enough that the + /// caller should halt the agent rather than merely warn. Inert when the + /// configured stop threshold is `0` (forced stops disabled). pub should_stop: bool, } @@ -988,15 +1026,48 @@ pub struct LoopStatus { /// operation rather than panicking). This means the detector degrades /// gracefully under contention but never blocks the agent loop. pub struct LoopDetector { - /// Sliding window of recent [`Operation`] records, bounded by [`LoopDetectorConfig::window_size`]. + /// Sliding window of recent [`Operation`] records. + /// + /// Bounded to [`LoopDetectorConfig::window_size`] entries, ordered + /// oldest to newest. Each recorded tool invocation pushes onto the + /// back and evicts from the front once the window is full; loop + /// detection scans this window to count repetitions of the same + /// operation. Guarded by a [`Mutex`] so the window can be updated + /// from the agent loop while being read by an observer thread. operations: Mutex>, - /// Detector configuration (thresholds and limits). Immutable after construction. + + /// Detector configuration: thresholds, window size, and limits. + /// + /// Held immutably for the lifetime of the detector (it is set in + /// the constructor and never mutated). Storing it by value keeps + /// threshold checks a single field access with no indirection. config: LoopDetectorConfig, - /// Per-turn tool call count. Reset by [`reset_turn`](LoopDetector::reset_turn). + + /// Number of tool calls recorded in the current turn. + /// + /// Incremented by every + /// [`record`](LoopDetector::record) call and reset to zero by + /// [`reset_turn`](LoopDetector::reset_turn). Drives the per-turn + /// tool-call budget check so a runaway turn can be flagged before + /// the sliding window fills. Guarded by a [`Mutex`]. turn_count: Mutex, - /// Operations that already triggered a warning. Cleared on result change or [`reset`](LoopDetector::reset). + + /// Operations that have already triggered a warning. + /// + /// Each entry is an [`Operation`] that crossed the warning + /// threshold, recorded so the detector does not re-warn on the + /// same pattern. Entries are removed when an operation's result + /// changes (progress observed) and the whole set is cleared by + /// [`reset`](LoopDetector::reset). Guarded by a [`Mutex`]. warned_operations: Mutex>, - /// Tool signature for extracting tool-specific parameters. + + /// Tool signature extractor used to derive each tool's primary parameter. + /// + /// A trait object held behind [`Arc`] so it can be shared cheaply + /// across clones of the detector. The signature decides which + /// argument identifies an operation as "the same" (e.g. a file + /// path for `Read`, a command for `Bash`), making repetition + /// detection tool-aware rather than string-comparing raw JSON. signature: Arc, } @@ -1413,6 +1484,7 @@ impl LoopDetector { .warned_operations .lock() .is_ok_and(|w| w.contains(first_op)); + if already_warned && !should_stop { return None; } @@ -1508,8 +1580,6 @@ 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() @@ -1539,6 +1609,7 @@ impl LoopDetector { if let Ok(mut ops) = self.operations.lock() { ops.clear(); } + if let Ok(mut warned) = self.warned_operations.lock() { warned.clear(); } @@ -1559,9 +1630,11 @@ impl LoopDetector { if let Ok(mut ops) = self.operations.lock() { ops.clear(); } + if let Ok(mut count) = self.turn_count.lock() { *count = 0; } + if let Ok(mut warned) = self.warned_operations.lock() { warned.clear(); } diff --git a/src/detection/manager.rs b/src/detection/manager.rs index d12c53e..af8c315 100644 --- a/src/detection/manager.rs +++ b/src/detection/manager.rs @@ -96,10 +96,6 @@ pub use super::convergence::ConvergenceAction; pub use super::convergence::ConvergenceConfigError; pub use super::loop_detector::LoopStatus; -// ================================================== -// Detected Pattern -// ================================================== - /// Represents a pattern detected in agent behavior. /// /// Returned by [`DetectionManager::record_operation`], @@ -166,11 +162,24 @@ pub enum DetectedPattern { /// } /// ``` LoopDetected { - /// Equal to [`LoopStatus::repetition_count`]. + /// How many times the repeated operation has been observed. + /// + /// Equal to [`LoopStatus::repetition_count`] at the moment the + /// loop crossed + /// [`DetectionConfig::loop_threshold`]. Callers can compare it + /// against the threshold to decide between a soft warning and a + /// hard stop. repetitions: usize, - /// Formatted as `"ToolName(primary_param)"`. + + /// Human-readable summary of the repeated operation. + /// + /// Formatted as `"ToolName(primary_param)"` (for example + /// `"Read(/etc/hosts)"` or `"Bash(ls -la)"`) using the + /// tool-signature extractor. Suitable for log lines and + /// user-facing diagnostics. pattern_description: String, }, + /// The agent's responses have become semantically similar. /// /// Emitted by [`DetectionManager::record_response`] when the internal @@ -191,11 +200,24 @@ pub enum DetectedPattern { /// - `consecutive_count` — how many consecutive response pairs exceeded /// the threshold. ConvergenceDetected { - /// Jaccard similarity (0.0–1.0) of the most recent response pair. + /// Jaccard similarity of the most recent response pair. + /// + /// A value in `[0.0, 1.0]` where `1.0` means identical + /// token sets. The detector fires once this score meets or + /// exceeds + /// [`DetectionConfig::convergence_threshold`]; callers can log + /// it to show *how* similar the converged responses were. similarity: f32, - /// Consecutive similar responses. + + /// Number of consecutive response pairs above the threshold. + /// + /// Reaches + /// [`DetectionConfig::convergence_count`] when convergence is + /// declared. A higher count means a longer run of near-identical + /// replies, which is a stronger signal that the agent is stuck. consecutive_count: usize, }, + /// Neither the loop detector nor the convergence detector has fired. /// /// The agent's tool calls are varying and/or its responses are diverging, @@ -204,10 +226,6 @@ pub enum DetectedPattern { NoPattern, } -// ================================================== -// Configuration -// ================================================== - /// Configuration for the [`DetectionManager`]. /// /// Groups all tunables for loop detection and convergence detection in a @@ -250,20 +268,62 @@ pub enum DetectedPattern { #[derive(Debug, Clone)] pub struct DetectionConfig { /// Consecutive similar operations before declaring a loop. Default: **3**. + /// + /// Number of times the same `(tool, primary_param, result_hash)` + /// signature must recur within the window before + /// [`DetectedPattern::LoopDetected`] is reported. Lower it to catch loops + /// sooner at the cost of more false positives. pub loop_threshold: usize, + /// Repetitions triggering forced stop (0 = disabled). Default: **10**. + /// + /// When a single operation's repetition count reaches this value, the + /// detector sets [`LoopStatus::should_stop`] so the caller halts the + /// agent. Setting it to `0` disables forced stops entirely while still + /// issuing warnings. pub stop_threshold: usize, + /// Whether loop detection is enabled. Default: **true**. + /// + /// Master switch for the tool-loop subsystem. When `false`, + /// [`DetectionManager::record_operation`] short-circuits and returns + /// [`DetectedPattern::NoPattern`] without consulting the loop detector. pub enable_loop_detection: bool, + /// Max operations kept in loop detector history. Default: **100**. + /// + /// Size of the sliding window the loop detector retains. A larger window + /// catches slower loops that span many interleaved operations; a smaller + /// one uses less memory and forgets stale patterns faster. pub max_history: usize, + /// Jaccard similarity threshold for convergence (0.0–1.0). Default: **0.95**. + /// + /// Minimum Jaccard similarity a response must share with its predecessor + /// to extend the convergence streak. Lower it to catch paraphrased + /// repetition; raise it toward `1.0` to demand near-verbatim matches. pub convergence_threshold: f32, + /// Consecutive similar responses for convergence. Default: **3**. + /// + /// Streak length of consecutive similar responses required before + /// [`DetectedPattern::ConvergenceDetected`] is reported. Maps to the + /// convergence detector's window size. pub convergence_count: usize, + /// Whether convergence detection is enabled. Default: **true**. + /// + /// Master switch for the response-convergence subsystem. When `false`, + /// [`DetectionManager::record_response`] short-circuits and returns + /// [`DetectedPattern::NoPattern`] without consulting the convergence + /// detector. pub enable_convergence_detection: bool, + /// Action on convergence. Default: [`ConvergenceAction::default()`]. + /// + /// The [`ConvergenceAction`] forwarded to callers once convergence is + /// detected — stop, warn, switch phase, ask the user, or compact the + /// history. Drives how the agent responds to a detected stall. pub on_converge: ConvergenceAction, } @@ -287,12 +347,12 @@ impl DetectionConfig { /// /// # Field Mapping /// - /// | `DetectionConfig` field | [`ConvergenceConfig`] field | - /// |--------------------------------|----------------------------------------| - /// | `enable_convergence_detection` | `enabled` | - /// | `convergence_count` | `window_size` | - /// | `convergence_threshold` | `similarity_threshold` | - /// | `on_converge` | `on_converge` | + /// | `DetectionConfig` field | [`ConvergenceConfig`] field | + /// |--------------------------------|-----------------------------| + /// | `enable_convergence_detection` | `enabled` | + /// | `convergence_count` | `window_size` | + /// | `convergence_threshold` | `similarity_threshold` | + /// | `on_converge` | `on_converge` | #[must_use] pub fn to_convergence_config(&self) -> ConvergenceConfig { ConvergenceConfig { @@ -323,10 +383,6 @@ impl DetectionConfig { } } -// ================================================== -// Statistics -// ================================================== - /// Cumulative statistics from the [`DetectionManager`]. /// /// Returned by [`DetectionManager::stats`] for observability and @@ -353,19 +409,34 @@ impl DetectionConfig { #[derive(Debug, Clone, Default)] pub struct DetectionStats { /// Tool-call operations recorded via `record_operation`. Reset by `reset`. + /// + /// Total count of operations fed into the loop detector since the last + /// [`DetectionManager::reset`], giving the denominator for loop-rate + /// calculations. pub turns_analyzed: usize, + /// Subset of `turns_analyzed` that returned `LoopDetected`. Reset by `reset`. + /// + /// How many of the recorded operations resulted in a + /// [`DetectedPattern::LoopDetected`] verdict, so callers can gauge how + /// frequently the agent gets stuck in tool loops. pub loops_detected: usize, + /// Times `record_response` returned `ConvergenceDetected`. Reset by `reset`. + /// + /// How many assistant responses triggered a + /// [`DetectedPattern::ConvergenceDetected`] verdict, indicating how often + /// the agent's replies collapsed into repetition. pub convergences_detected: usize, + /// Mirrors `LoopStatus::repetition_count`. Triggers `LoopDetected` at `loop_threshold`. Reset by `reset`. + /// + /// Live snapshot of the most-repeated operation's count, so observers can + /// watch a potential loop build up before it actually trips the + /// [`DetectionConfig::loop_threshold`]. pub current_streak: usize, } -// ================================================== -// Detection Manager -// ================================================== - /// Unified manager combining loop detection and convergence detection. /// /// [`DetectionManager`] is the primary entry point for the detection @@ -377,9 +448,9 @@ pub struct DetectionStats { /// /// Choose a constructor based on your needs: /// -/// | Constructor | Use case | -/// |------------------------------------------|-------------------------------------------| -/// | [`DetectionManager::new`] | Quick start with defaults | +/// | 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 | @@ -449,12 +520,37 @@ pub struct DetectionStats { /// - [`DetectionStats`] — cumulative observability counters. pub struct DetectionManager { /// Configuration thresholds and feature flags. + /// + /// Holds the validated [`DetectionConfig`] (loop threshold, stop + /// threshold, convergence settings) for the manager's lifetime. It + /// is set in the constructor and never mutated, so threshold checks + /// are a single field access with no locking. config: DetectionConfig, + /// Loop detector shared across compaction and analysis phases. + /// + /// Held behind [`Arc`] so callers that need direct access to the + /// tool-call window (for example, an observer that reports on + /// repetition independently of the manager) can clone the handle + /// rather than routing every query through the manager. The + /// detector's own internal state is `Mutex`-guarded. loop_detector: Arc, + /// Convergence detector guarded by interior mutability. + /// + /// [`ConvergenceDetector`] mutates its sliding window on every + /// recorded response, so it is wrapped in a [`Mutex`] to keep the + /// manager's public methods `&self`. This lets the manager be + /// shared across threads without an outer `mut` reference. convergence_detector: Mutex, + /// Cumulative detection statistics. + /// + /// Counters for turns analysed, loops detected, and convergences + /// observed since construction or the last + /// [`reset`](DetectionManager::reset). Updated under the + /// [`Mutex`] and snapshotted via + /// [`stats`](DetectionManager::stats) for observability dashboards. stats: Mutex, } @@ -590,10 +686,6 @@ impl DetectionManager { }) } - // ================================================== - // Loop detection (delegated to LoopDetector) - // ================================================== - /// Record an [`Operation`] for loop detection and return the current /// [`DetectedPattern`]. /// @@ -878,10 +970,6 @@ impl DetectionManager { &self.convergence_detector } - // ================================================== - // Convergence detection - // ================================================== - /// Record an assistant response for convergence detection. /// /// Forwards `response` to the [`ConvergenceDetector`] via @@ -992,10 +1080,6 @@ impl DetectionManager { .check_convergence() } - // ================================================== - // Pattern check - // ================================================== - /// Inspect the current detection state without recording new data. /// /// Checks both the loop detector and the convergence detector in @@ -1067,10 +1151,6 @@ impl DetectionManager { DetectedPattern::NoPattern } - // ================================================== - // Accessors - // ================================================== - /// Take a snapshot of the cumulative detection statistics. /// /// Clones the [`DetectionStats`] struct via the `Mutex`, @@ -1139,10 +1219,6 @@ impl DetectionManager { &self.config } - // ================================================== - // Reset - // ================================================== - /// Reset all internal detection state, as if the manager were freshly /// constructed. /// @@ -1214,10 +1290,6 @@ impl Default for DetectionManager { } } -// ================================================== -// Tests -// ================================================== - /// Tests for [`DetectionManager`]. /// /// Verifies loop detection, convergence detection, result-aware detection, @@ -1239,7 +1311,6 @@ mod tests { #[test] fn test_loop_detection() { let dm = DetectionManager::new().unwrap(); - // Same operation repeated many times for _ in 0..5 { let result = dm.record_tool_call("read_file", 42); if matches!(result, DetectedPattern::LoopDetected { .. }) { @@ -1256,7 +1327,7 @@ mod tests { for _ in 0..5 { let result = dm.record_response(response); if matches!(result, DetectedPattern::ConvergenceDetected { .. }) { - return; // test passes + return; } } panic!("Expected convergence detection after 5 identical responses"); @@ -1299,7 +1370,6 @@ mod tests { }; let dm = DetectionManager::new_with_config(config).unwrap(); - // Record 5 identical operations for _ in 0..5 { dm.record_tool_call("read_file", 42); } @@ -1315,7 +1385,6 @@ mod tests { fn test_record_operation_with_operation_struct() { let dm = DetectionManager::new().unwrap(); - // Record operations using the rich Operation API for _ in 0..5 { dm.record_operation(Operation::new("Read", "/test/file.txt")); } diff --git a/src/engine.rs b/src/engine.rs index cdbbca9..8e5a307 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -1,32 +1,41 @@ //! Engine module — the core agentic loop that ties all components together. //! //! The [`BareLoop`] is the framework's default implementation of an agent -//! execution loop. It orchestrates: +//! execution loop. It drives a sans-IO [`LoopMachine`] +//! (`engine::core`), matching on each [`MachineStep`] the machine emits +//! and performing the requested IO: //! -//! 1. **Initialize** — Set up the conversation with system prompt -//! 2. **Stream** — Send messages to the LLM API and accumulate the response -//! 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::engine::loop_core::SessionResult) +//! 1. **`CallLLM`** — Stream a model response over the machine's history. +//! 2. **`CallTools`** — Dispatch the tool calls the model requested. +//! 3. **`Compact`** — Run a context compactor over the history when the +//! machine's threshold is crossed. +//! 4. **`Done`** — The run ended; produce a [`Run`]. +//! +//! The machine owns the conversation history and every loop decision (turn +//! counting, max-turn, tool-call validity, compaction trigger, cancellation); +//! the driver owns the side effects (the LLM call, tool dispatch, observers, +//! the cancellation `select!`). //! //! # Example //! //! ```rust,ignore //! use loopctl::engine::BareLoop; +//! use loopctl::engine::RunConfig; //! use loopctl::tool::ToolRegistry; //! use loopctl::api::ApiClient; -//! use loopctl::config::LoopConfig; +//! use loopctl::config::SessionConfig; //! -//! let agent = BareLoop::new(client, registry, config); -//! let result = agent.run("Write a hello world program").await?; +//! let agent = BareLoop::new(client, registry, SessionConfig::default()); +//! let result = agent.run("Write a hello world program", &RunConfig::default()).await?; //! ``` +//! +//! [`LoopMachine`]: crate::engine::core::LoopMachine +//! [`MachineStep`]: crate::engine::core::MachineStep mod bare; pub mod contributor; -pub mod loop_core; -pub mod machine; +pub mod core; pub use bare::*; pub use contributor::*; -pub use machine::*; +pub use core::*; diff --git a/src/engine/bare.rs b/src/engine/bare.rs index a8f344a..8ec90cd 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -13,8 +13,8 @@ //! the LLM provider. //! - A [`ToolRegistry`](crate::tool::ToolRegistry) for dispatching tool //! calls the model requests. -//! - An [`LoopConfig`] governing session parameters (max turns, system -//! prompt, session ID). +//! - A [`SessionConfig`](crate::config::SessionConfig) governing session +//! parameters (system prompt, session ID, context window). //! - Optional [`LoopObserver`](crate::observer::LoopObserver) registrations for lifecycle instrumentation. //! //! # Key Design Decisions @@ -34,20 +34,21 @@ //! ```rust,ignore //! use loopctl::engine::BareLoop; //! use loopctl::tool::ToolRegistry; -//! use loopctl::config::LoopConfig; +//! use loopctl::config::SessionConfig; +//! use loopctl::engine::RunConfig; //! use std::sync::Arc; //! //! // 1. Build components //! let client = Arc::new(my_api_client); //! let registry = ToolRegistry::new(); -//! let config = LoopConfig::default(); +//! let config = SessionConfig::default(); //! //! // 2. Create the loop //! let mut agent = BareLoop::new(client, registry, config); //! //! // 3. Run -//! let result = agent.run("Hello, agent!").await?; -//! println!("Agent responded in {} turns", result.total_turns); +//! let result = agent.run("Hello, agent!", &RunConfig::default()).await?; +//! println!("Agent responded in {} turns", result.turn_count()); //! ``` use crate::api::ApiClient; @@ -58,11 +59,16 @@ use std::time::{Duration, Instant}; use crate::cancel::CancelSignal; use crate::compact::ContextManager; -use crate::config::LoopConfig; +use crate::config::SessionConfig; +use crate::engine::core::{ + LoopMachine, MachineOutcome, MachinePolicy, MachineState, MachineStep, ModelResponse, + PendingToolCall, Run, RunConfig, RunResult, Session, StopReason, ToolCall, +}; use crate::error::LoopError; -use crate::engine::loop_core::{LoopState, SessionResult, StopReason, ToolCall, TurnResult}; +use crate::capabilities::{Detectable, FallbackCapable}; +use crate::detection::{ConvergenceAction, DetectedPattern}; use crate::engine::{ContextContributor, ContributorContext}; #[cfg(all(test, feature = "hooks"))] use crate::hooks::Hook; @@ -74,6 +80,7 @@ use crate::hooks::context::{ use crate::hooks::context::{SessionEndContext as HookSessionEndContext, SessionEndReason}; #[cfg(feature = "hooks")] use crate::hooks::{HookAction, HookExecutor}; +use crate::managers::LoopManagers; use crate::message::{Message, MessagePart, Role, ToolContent}; use crate::middleware::{ToolDispatchContext, ToolPipeline, ToolPipelineBuilder}; use crate::observer::{ @@ -84,7 +91,6 @@ use crate::reflection::{ ExponentialBackoffRecovery, NoopReflector, RecoveryAction, RecoveryStrategy, ReflectionContext, Reflector, }; -use crate::runtime::LoopRuntime; use crate::stream::handler::StreamHandler; use crate::stream::{StreamAccumulator, StreamEvent, StreamStopReason, Usage}; use crate::structured::RequestOptions; @@ -98,10 +104,6 @@ mod emission; mod message; mod stream; -// ================================================== -// BareLoop -// ================================================== - /// The framework's default agent loop implementation. /// /// `BareLoop` ties together an [`ApiClient`], [`ToolRegistry`], @@ -121,18 +123,19 @@ mod stream; /// /// - [`new()`](BareLoop::new) — client + tools + config. /// - [`new_with_managers()`](BareLoop::new_with_managers) — full control, -/// including a [`LoopRuntime`]. +/// including a [`LoopManagers`]. /// /// # Example /// /// ```rust,ignore /// use loopctl::engine::BareLoop; /// use loopctl::tool::ToolRegistry; -/// use loopctl::config::LoopConfig; +/// use loopctl::config::SessionConfig; +/// use loopctl::engine::RunConfig; /// use std::sync::Arc; /// /// let registry = ToolRegistry::new(); -/// let config = LoopConfig::default(); +/// let config = SessionConfig::default(); /// /// let mut agent = BareLoop::new( /// Arc::new(my_client), @@ -140,8 +143,8 @@ mod stream; /// config, /// ); /// -/// let result = agent.run("Hello, agent!").await?; -/// println!("Agent responded in {} turns", result.total_turns); +/// let result = agent.run("Hello, agent!", &RunConfig::default()).await?; +/// println!("Agent responded in {} turns", result.turn_count()); /// ``` pub struct BareLoop { /// The LLM API client used to send conversation turns. @@ -156,19 +159,31 @@ pub struct BareLoop { /// by name in this registry and invokes it. tools: Arc, - /// Session parameters (max turns, model, system prompt). + /// Session-scoped state: id, config, start time, and run history. /// - /// See [`LoopConfig`] for the full set of options. - config: LoopConfig, + /// Owns the single source of truth for session identity, the session + /// config ([`SessionConfig`]), the session start instant, and the list + /// of [`Run`]s accumulated across `run()` calls. The in-flight run is + /// the last entry in `session.runs`, accessed via + /// [`current_run`](Self::current_run) / + /// [`current_run_mut`](Self::current_run_mut). + session: Session, - /// Conversation history (system + user + assistant + tool results). + /// The sans-IO agent-loop state machine. /// - /// 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, + /// Owns the conversation history and every loop decision (turn counting, + /// max-turn enforcement, compaction triggering, tool-call classification). + /// The driver advances it with `next_step()` and feeds outcomes back via + /// [`model_response`](LoopMachine::model_response), + /// [`tool_results`](LoopMachine::tool_results), + /// [`compaction_result`](LoopMachine::compaction_result), and + /// [`inject`](LoopMachine::inject). It is (re)created at the top of every + /// [`run()`](crate::engine::core::Loop::run) call from the run config + /// and user prompt; before that it holds an empty placeholder so the struct + /// is always valid. + machine: LoopMachine, - /// Framework runtime bundle — holds all cross-cutting infrastructure. + /// Framework managers bundle — holds all cross-cutting infrastructure. /// /// This is the single source of truth for: /// @@ -181,7 +196,7 @@ pub struct BareLoop { /// - Optional [`HookExecutor`] — bidirectional lifecycle hooks. /// - Optional [`ToolHealthRegistry`] — per-tool health tracking. /// - /// Reset at the start of every session via [`LoopRuntime::reset_all`]. + /// Reset at the start of every session via [`LoopManagers::reset_all`]. /// /// [`FallbackManager`]: crate::fallback::FallbackManager /// [`DetectionManager`]: crate::detection::DetectionManager @@ -191,7 +206,7 @@ pub struct BareLoop { /// [`StreamHandler`]: crate::stream::handler::StreamHandler /// [`HookExecutor`]: crate::hooks::HookExecutor /// [`ToolHealthRegistry`]: crate::tool::health::ToolHealthRegistry - managers: LoopRuntime, + managers: LoopManagers, /// Failure analyser for tool errors. /// @@ -215,34 +230,10 @@ pub struct BareLoop { /// will wake up mid-stream when cancelled. cancelled: Arc, - /// Current lifecycle state, exposed via the [`Loop`](crate::engine::loop_core::Loop) trait. - /// - /// Drives the engine's state machine - /// (`Idle → Processing → WaitingForTool → Processing → … → Completed` / - /// `Failed` / `Cancelled`, plus `Compacting`/`Reflecting` side-states). - /// Set throughout the turn loop and read by [`finalize`](crate::engine::loop_core::Loop::finalize) - /// to decide success vs. failure and to populate the - /// [`SessionResult`](crate::engine::loop_core::SessionResult) fields; also - /// returned from [`state`](crate::engine::loop_core::Loop::state) for - /// outside observers. The [`Idle`](LoopState::Idle) variant additionally - /// gates configuration setters via [`debug_assert_idle`](Self::debug_assert_idle). - state: LoopState, - - /// Session-level accumulator for turn counts, token usage, and tool calls. - /// - /// 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, - - /// Session start time, set by [`initialize`](crate::engine::loop_core::Loop::initialize). - /// - /// `None` only before the first `initialize()` call; `Some` thereafter. - /// Read in [`finalize`](crate::engine::loop_core::Loop::finalize) to - /// compute [`SessionResult::total_duration`](crate::engine::loop_core::SessionResult) - /// via `elapsed()`. Captured once per session (not per turn) so the - /// reported duration is the wall-clock lifetime of the whole session, - /// not a single turn. - session_start: Option, + /// Fallback config for [`run_config`](Self::run_config) before the first + /// `run()` call. Unreachable in practice (constructors seed a placeholder + /// run), but needed so the accessor returns `&RunConfig` without panicking. + fallback_config: RunConfig, /// Optional callback invoked for each text delta during streaming. /// @@ -270,10 +261,6 @@ pub struct BareLoop { request_options: RequestOptions, } -// ================================================== -// Run-loop helpers -// ================================================== - impl BareLoop { /// Maximum retry attempts for tool recovery before giving up. /// @@ -287,13 +274,13 @@ impl BareLoop { /// Create a new `BareLoop` with the given components. /// /// Initializes an empty conversation history and a - /// fresh [`LoopRuntime`]. The cancellation signal starts as non-cancelled. + /// fresh [`LoopManagers`]. The cancellation signal starts as non-cancelled. /// /// # Parameters /// /// - `client` — The LLM API client, wrapped in `Arc`. /// - `tools` — The [`ToolRegistry`] containing available tools. - /// - `config` — Session parameters (max turns, system prompt, etc.). + /// - `session_config` — Session parameters (session ID, system prompt, etc.). /// /// # Example /// @@ -301,48 +288,32 @@ impl BareLoop { /// let mut agent = BareLoop::new( /// Arc::new(my_client), /// ToolRegistry::new(), - /// LoopConfig::default(), + /// SessionConfig::default(), /// ); /// ``` - pub fn new(client: Arc, tools: ToolRegistry, config: LoopConfig) -> Self { - Self { - client, - tools: Arc::new(tools), - config, - conversation: Vec::new(), - managers: LoopRuntime::new(), - reflector: Arc::new(NoopReflector), - recovery: Arc::new(ExponentialBackoffRecovery::new(3)), - cancelled: Arc::new(CancelSignal::new()), - state: LoopState::Idle, - budget: SessionResult::default(), - session_start: None, - text_streamer: None, - contributors: Vec::new(), - request_options: RequestOptions::default(), - } + pub fn new(client: Arc, tools: ToolRegistry, session_config: SessionConfig) -> Self { + Self::new_with_managers(client, tools, session_config, LoopManagers::new()) } /// Create a new `BareLoop` with all components including managers. /// /// Use this constructor when you need to supply a pre-configured - /// [`LoopRuntime`] — for example, to enable loop detection or + /// [`LoopManagers`] — for example, to enable loop detection or /// circuit-breaker policies. /// /// # Parameters /// /// - `client` — The LLM API client, wrapped in `Arc`. /// - `tools` — The [`ToolRegistry`] containing available tools. - /// - `config` — Session parameters. - /// - `managers` — A pre-built [`LoopRuntime`]. + /// - `session_config` — Session parameters. + /// - `managers` — A pre-built [`LoopManagers`]. /// /// # Example /// /// ```rust,ignore - /// let managers = LoopRuntime::builder() + /// let managers = LoopManagers::new() /// .with_detection(DetectionManager::default()) - /// .with_fallback(FallbackManager::default()) - /// .build(); + /// .with_fallback(FallbackManager::default()); /// /// let mut agent = BareLoop::new_with_managers( /// Arc::new(my_client), @@ -354,47 +325,162 @@ impl BareLoop { pub fn new_with_managers( client: Arc, tools: ToolRegistry, - config: LoopConfig, - managers: LoopRuntime, + session_config: SessionConfig, + managers: LoopManagers, ) -> Self { Self { client, tools: Arc::new(tools), - config, - conversation: Vec::new(), + session: { + let mut s = Session::new(session_config); + s.runs.push(Run::default()); + s + }, + machine: LoopMachine::from_history(vec![Message::user("")]), managers, reflector: Arc::new(NoopReflector), recovery: Arc::new(ExponentialBackoffRecovery::new(3)), cancelled: Arc::new(CancelSignal::new()), - state: LoopState::Idle, - budget: SessionResult::default(), - session_start: None, + fallback_config: RunConfig::default(), text_streamer: None, contributors: Vec::new(), request_options: RequestOptions::default(), } } - // ================================================== - // Accessors - // ================================================== - /// Get the conversation history. /// /// Returns a slice of [`Message`] representing the full conversation - /// so far: system prompt (if applied), user messages, assistant - /// responses, and tool-result messages. + /// so far: the opening user message, contributor injections, assistant + /// responses, and tool-result messages. The history is owned by the + /// driving state machine; it is empty until the first + /// [`run()`](crate::engine::core::Loop::run) call mints a machine for + /// the run. pub fn conversation(&self) -> &[Message] { - &self.conversation + if self.session.session_start.is_none() { + return &[]; + } + self.machine.history() } - /// Get the agent configuration. + /// Get the session configuration. /// - /// Returns a reference to the [`LoopConfig`] that governs session - /// parameters such as max turns, system prompt, and session ID. + /// Returns a reference to the [`SessionConfig`] that holds session-scoped + /// parameters: the session ID, system prompt, and context window. /// The config is immutable for the lifetime of the loop. - pub fn config(&self) -> &LoopConfig { - &self.config + pub fn session_config(&self) -> &SessionConfig { + &self.session.config + } + + /// Get the session, including its config, start time, and completed runs. + /// + /// Returns a reference to the [`Session`] accumulated across `run()` + /// calls. Use this for cross-run accounting — for example + /// [`total_input_tokens`](Session::total_input_tokens) or + /// [`total_turns`](Session::total_turns) — without tracking totals + /// manually in the host. + #[must_use] + pub fn session(&self) -> &Session { + &self.session + } + + /// Borrow the in-flight run (the last entry in `session.runs`). + /// + /// Constructors seed `session.runs` with a placeholder, and `run()` + /// pushes a fresh [`Run`] before any access — so the last entry is + /// always present. + fn current_run(&self) -> Option<&Run> { + let len = self.session.runs.len(); + self.session.runs.get(len.saturating_sub(1)) + } + + /// Mutably borrow the in-flight run. + /// + /// Same contract as [`current_run`](Self::current_run) but `&mut`. + fn current_run_mut(&mut self) -> Option<&mut Run> { + let len = self.session.runs.len(); + self.session.runs.get_mut(len.saturating_sub(1)) + } + + /// Get the run configuration for the current run. + /// + /// Returns a reference to the [`RunConfig`] stored on the in-flight + /// [`Run`], governing per-run budgets (turn/token limits, compaction + /// policy, dispatch mode). This is the configuration applied to the + /// most recent `run()` call. + pub fn run_config(&self) -> &RunConfig { + self.current_run() + .map_or(&self.fallback_config, |run| &run.config) + } + + /// Build the policy struct the machine needs for `next_step()`. + /// + /// Combines the run's `max_turns` with the session's compaction knobs + /// into a single [`MachinePolicy`] passed fresh each call. + fn machine_policy(&self) -> MachinePolicy { + MachinePolicy { + max_turns: self.run_config().max_turns, + context_window: self.session.config.context_window, + compact_threshold: self.session.config.compact_threshold, + auto_compact: self.session.config.auto_compact, + } + } + + /// Borrow the driving state machine. + /// + /// Returns a reference to the [`LoopMachine`] that owns the current run's + /// history and decisions. Useful for inspecting the run in flight (e.g. the + /// accumulated history, turns taken, or the machine's internal state). The + /// machine is (re)created at the top of every + /// [`run()`](crate::engine::core::Loop::run) call; before the first run + /// it holds an empty placeholder. + #[must_use] + pub fn machine(&self) -> &LoopMachine { + &self.machine + } + + /// Consume the loop and take ownership of its state machine. + /// + /// Returns the [`LoopMachine`], dropping the rest of the loop (client, + /// tools, managers). Use this to checkpoint a run's machine for later + /// resumption via [`BareLoop::from_machine`]. + #[must_use] + pub fn into_machine(self) -> LoopMachine { + self.machine + } + + /// Build a loop around an existing state machine. + /// + /// Constructs a [`BareLoop`] whose machine is `machine` — for example to + /// resume a serialized run: deserialize the machine, wrap it in a loop with + /// the original client/tools, and continue driving it with + /// [`run()`](crate::engine::core::Loop::run). The session/run config + /// is taken from the machine. + #[must_use] + pub fn from_machine( + machine: LoopMachine, + session_config: SessionConfig, + client: Arc, + tools: ToolRegistry, + ) -> Self { + Self { + client, + tools: Arc::new(tools), + session: { + let mut s = Session::new(session_config); + s.runs.push(Run::default()); + s + }, + machine, + managers: LoopManagers::new(), + reflector: Arc::new(NoopReflector), + recovery: Arc::new(ExponentialBackoffRecovery::new(3)), + cancelled: Arc::new(CancelSignal::new()), + fallback_config: RunConfig::default(), + text_streamer: None, + contributors: Vec::new(), + request_options: RequestOptions::default(), + } } /// Get the tool registry. @@ -450,13 +536,9 @@ impl BareLoop { Arc::clone(&self.cancelled) } - // ================================================== - // Dependency setters - // ================================================== - /// Assert that the loop has not started running yet. /// - /// Configuration setters must be called before [`run()`](crate::engine::loop_core::Loop::run). + /// Configuration setters must be called before [`run()`](crate::engine::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. @@ -464,23 +546,24 @@ impl BareLoop { /// This check is only active in debug builds (`debug_assertions`). #[inline] fn debug_assert_idle(&self) { + let idle = self.machine.state() == MachineState::Start; debug_assert!( - matches!(self.state, LoopState::Idle), + idle, "BareLoop configuration setters must be called before run() — \ - current state is {:?}, expected Idle", - self.state + machine is {:?}, expected Start", + self.machine.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). + /// implementation. Must be called before [`run()`](crate::engine::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)). + /// (i.e., once the machine has advanced past [`MachineState::Start`]). /// /// # Example /// @@ -497,7 +580,7 @@ impl BareLoop { /// /// Replaces the default [`ExponentialBackoffRecovery`] with a /// caller-supplied implementation. Must be called before - /// [`run()`](crate::engine::loop_core::Loop::run). + /// [`run()`](crate::engine::core::Loop::run). /// /// # Panics (debug only) /// @@ -518,7 +601,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()`](crate::engine::loop_core::Loop::run). + /// Must be called before [`run()`](crate::engine::core::Loop::run). /// /// # Panics (debug only) /// @@ -535,14 +618,17 @@ impl BareLoop { /// .with_min_messages(6); /// let manager = ContextManager::new(Arc::new(compactor)) /// .with_context_window(200_000) - /// .with_threshold(8_000); + /// .with_threshold(80); /// /// let mut agent = BareLoop::new(client, registry, config); /// 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); + let synced = Arc::try_unwrap(manager) + .unwrap_or_else(|arc| (*arc).clone()) + .with_context_window(self.session.config.context_window); + self.managers.set_context_manager(Arc::new(synced)); } /// Set the [`StreamHandler`] for resilient streaming with retries, @@ -550,7 +636,7 @@ impl BareLoop { /// /// When set, the loop delegates streaming to the handler instead of /// using the inline streaming logic. Must be called before - /// [`run()`](crate::engine::loop_core::Loop::run). + /// [`run()`](crate::engine::core::Loop::run). /// /// # Panics (debug only) /// @@ -584,7 +670,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()`](crate::engine::loop_core::Loop::run). + /// Must be called before [`run()`](crate::engine::core::Loop::run). /// /// # Panics (debug only) /// @@ -613,7 +699,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()`](crate::engine::loop_core::Loop::run). + /// Must be called before [`run()`](crate::engine::core::Loop::run). /// /// # Panics (debug only) /// @@ -642,7 +728,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()`](crate::engine::loop_core::Loop::run). + /// Must be called before [`run()`](crate::engine::core::Loop::run). /// /// # Panics (debug only) /// @@ -685,7 +771,7 @@ impl BareLoop { /// in registration order. See [`LoopObserver`](crate::observer::LoopObserver) /// for the trait definition and available hooks. /// - /// Must be called before [`run()`](crate::engine::loop_core::Loop::run). + /// Must be called before [`run()`](crate::engine::core::Loop::run). /// /// # Panics (debug only) /// @@ -708,7 +794,7 @@ impl BareLoop { /// 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). + /// from the API during [`run`](crate::engine::core::Loop::run). /// This enables real-time display of the model's output without /// waiting for the full turn to complete. /// @@ -749,12 +835,12 @@ impl BareLoop { /// built without any — the turn-top consultation is a single cheap branch. /// /// Must be called before - /// [`run`](crate::engine::loop_core::Loop::run). + /// [`run`](crate::engine::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)). + /// (i.e., once the machine has advanced past [`MachineState::Start`]). /// /// # Example /// @@ -775,12 +861,12 @@ impl BareLoop { /// default ([`RequestOptions::default`]) for unconstrained behavior. /// /// Must be called before - /// [`run`](crate::engine::loop_core::Loop::run). + /// [`run`](crate::engine::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)). + /// (i.e., once the machine has advanced past [`MachineState::Start`]). /// /// # Example /// @@ -908,27 +994,27 @@ impl BareLoop { /// /// ```rust,ignore /// # use loopctl::engine::BareLoop; - /// # use loopctl::config::LoopConfig; + /// # use loopctl::config::SessionConfig; /// # 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()); + /// # let mut loop_ = BareLoop::new(client, tools, SessionConfig::default()); /// loop_.switch_model("model-b").with_context_window(8192).apply().unwrap(); - /// assert_eq!(loop_.config().model, "model-b"); - /// assert_eq!(loop_.config().context_window, 8192); + /// assert_eq!(loop_.client.model(), "model-b"); + /// assert_eq!(loop_.session_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::config::SessionConfig; /// # 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()); + /// # let mut loop_ = BareLoop::new(client, tools, SessionConfig::default()); /// loop_.switch_model("b").apply().unwrap(); /// ``` pub fn switch_model(&mut self, model: &str) -> ModelSwitch<'_, C> { @@ -936,14 +1022,9 @@ impl BareLoop { loop_: self, target_model: model.to_string(), context_window: None, - max_tokens: None, } } - // ================================================== - // Run helpers - // ================================================== - /// Pull per-turn `(input_tokens, output_tokens)` from optional [`Usage`]. /// /// Returns `(0, 0)` when the provider did not report usage for the turn. @@ -954,63 +1035,64 @@ impl BareLoop { } } - /// Dispatch a batch of tool calls, append the results to the conversation, - /// record the call count on `budget`, and fire `on_turn_end`. + /// Dispatch a batch of tool calls and return their aggregated result message. /// - /// `on_turn_end` fires on both success and error paths with the - /// corresponding `success` flag. On error the conversation is *not* - /// extended — the caller's error handling owns the terminal state. + /// Runs the calls through the configured dispatch path, records the call + /// count on `budget`, fires `on_turn_end` (on both success and error paths, + /// with the matching `success` flag), and returns the assembled tool-result + /// [`Message`] for the caller to feed into the driving machine via + /// [`LoopMachine::tool_results`]. The message is *not* pushed to the + /// history here — history is owned by the machine, so the caller decides + /// when to record it (alongside any preresolved results). /// /// Takes `budget` by mutable reference (rather than `&mut self`) because - /// the caller has already `mem::take`n it to avoid a double mutable borrow - /// while dispatch runs. + /// the caller has already split the borrow to dispatch against `self` while + /// accumulating into `budget`. /// /// # Errors /// /// Propagates [`LoopError::Cancelled`] if cancellation fired during /// dispatch, or any error the recovery system escalates to a hard failure - /// (e.g. loop detection aborts, exhaustion of retry budget). + /// (e.g. loop detection aborts, exhaustion of retry budget). On error the + /// tool-result message is meaningless and the caller's error handling owns + /// the terminal state. async fn dispatch_and_record( &mut self, tool_calls: &[ToolCall], turn_index: usize, - turn_duration: Duration, + turn_start: Instant, turn_input_tokens: u64, turn_output_tokens: u64, - budget: &mut SessionResult, - ) -> Result<(), LoopError> { - match self.dispatch_tools(tool_calls, turn_index).await { + ) -> Result { + let result = self.dispatch_tools(tool_calls, turn_index).await; + let turn_duration = turn_start.elapsed(); + match result { 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: turn_index, - success: true, - error: None, - duration_ms: Self::millis_u64(turn_duration), - input_tokens: turn_input_tokens, - output_tokens: turn_output_tokens, - }); - Ok(()) + self.notify_turn_end( + turn_index, + true, + None, + turn_duration, + turn_input_tokens, + turn_output_tokens, + ); + Ok(tool_result_msg) } Err(e) => { let err_str = e.to_string(); - self.managers.observers().on_turn_end(&TurnEndContext { - turn: turn_index, - success: false, - error: Some(err_str), - duration_ms: Self::millis_u64(turn_duration), - input_tokens: turn_input_tokens, - output_tokens: turn_output_tokens, - }); + self.notify_turn_end( + turn_index, + false, + Some(err_str), + turn_duration, + turn_input_tokens, + turn_output_tokens, + ); Err(e) } } } - // ================================================== - // Turn helpers (used by process_turn) - // ================================================== /// Stream one assistant response from the API and apply post-stream bookkeeping. /// @@ -1023,7 +1105,8 @@ impl BareLoop { /// sets the terminal state, and returns the error. /// /// `LoopError::Cancelled` short-circuits the failure bookkeeping: - /// cancellation is a clean termination, so it sets [`LoopState::Cancelled`] + /// cancellation is a clean termination, so the run loop records a + /// [`MachineOutcome::Cancelled`](crate::engine::core::MachineOutcome::Cancelled) /// without tripping the fallback or firing `on_stream_failure`. /// /// # Errors @@ -1052,12 +1135,11 @@ impl BareLoop { /// `(Message, Option, StreamStopReason)` and just needs the /// side-effects. fn record_stream_success(&mut self, usage: Option<&Usage>) { - self.managers.fallback.record_model_success(); + self.managers.fallback().record_model_success(); let (in_tok, out_tok) = Self::usage_tokens(usage); - // TODO: fire_stream_success self.managers.observers().on_stream_success(&StreamContext { - turn: self.budget.total_turns, + turn: self.current_run().map_or(0, Run::turn_count), model: self.client.model(), input_tokens: in_tok, output_tokens: out_tok, @@ -1077,25 +1159,28 @@ impl BareLoop { /// /// Then fires /// [`on_stream_failure`](crate::observer::LoopObserver::on_stream_failure) - /// (regardless of breaker outcome), sets the terminal - /// [`LoopState::Failed`], and returns the original error so the caller - /// can propagate it from [`do_stream`](Self::do_stream). + /// (regardless of breaker outcome) and returns the original error so the + /// caller can propagate it from [`do_stream`](Self::do_stream); the run + /// loop records the terminal + /// [`MachineOutcome::Failed`](crate::engine::core::MachineOutcome) on + /// the machine from the returned error. /// /// `LoopError::Cancelled` is **not** routed through here — cancellation - /// is a clean termination that sets [`LoopState::Cancelled`] without - /// tripping the breaker or firing `on_stream_failure`. The caller - /// ([`run_turn_body`](Self::run_turn_body)) handles cancellation before - /// reaching [`do_stream`](Self::do_stream). + /// is a clean termination that records + /// [`MachineOutcome::Cancelled`](crate::engine::core::MachineOutcome) + /// without tripping the breaker or firing `on_stream_failure`. The + /// [`run`](crate::engine::core::Loop::run) loop's `CallLLM` arm + /// handles cancellation before reaching [`do_stream`](Self::do_stream). fn record_stream_failure(&mut self, e: LoopError) -> LoopError { let tripped = if matches!(e, LoopError::RateLimitEscalation { .. }) { - self.managers.fallback.record_model_failure() + self.managers.fallback().record_model_failure() } else { - self.managers.fallback.record_api_failure() + self.managers.fallback().record_api_failure() }; if tripped { let from = self.client.model(); - if let Some(to) = self.managers.fallback.fallback_model() { + if let Some(to) = self.managers.fallback().fallback_model() { tracing::warn!(from = %from, to = %to, "fallback manager tripped"); self.managers .observers() @@ -1103,96 +1188,40 @@ impl BareLoop { } } - // TODO: fire_stream_failure self.managers .observers() .on_stream_failure(&StreamFailureContext { - turn: self.budget.total_turns, + turn: self.current_run().map_or(0, Run::turn_count), model: self.client.model(), error: e.clone(), }); - self.state = LoopState::Failed { - error: e.to_string(), - }; e } - /// Fold this turn's token counts into the running session totals on - /// `budget`. No-op when the provider did not report usage. + /// Fire [`on_turn_end`](crate::observer::LoopObserver::on_turn_end). /// - /// Uses saturating add so a runaway session cannot overflow the counters. - 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) for the - /// turn that just completed. - /// - /// Reports `total_turns - 1` as the turn number: callers invoke this after - /// the per-turn increment, so subtracting one recovers the 0-indexed turn - /// that just ran. - fn finish_turn(&mut self, turn_in: u64, turn_out: u64, duration: Duration) { + /// Single construction point for [`TurnEndContext`] — every turn-end + /// notification in the driver goes through here. + fn notify_turn_end( + &self, + turn: usize, + success: bool, + error: Option, + duration: Duration, + input_tokens: u64, + output_tokens: u64, + ) { self.managers.observers().on_turn_end(&TurnEndContext { - turn: self.budget.total_turns.saturating_sub(1), - success: true, - error: None, + turn, + success, + error, duration_ms: Self::millis_u64(duration), - input_tokens: turn_in, - output_tokens: turn_out, + input_tokens, + output_tokens, }); } - /// Build the [`TurnResult`] returned by a turn that completed the session - /// (`is_complete: true`, [`StopReason::EndTurn`]). Used by the no-tool-calls - /// and loop-detection success paths, both of which end the session from - /// inside `process_turn`. - 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 the - /// un-compacted history rather than failing the session. - 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" - ); - } - } - - /// Push the user's message onto the conversation (first turn only). - /// - /// Continuation turns receive `input == ""` because the previous turn's - /// tool results are already in the history. - fn record_user_input(&mut self, input: &str) { - if !input.is_empty() { - self.conversation.push(Message::user(input)); - } - } - /// Fire [`on_turn_start`](crate::observer::LoopObserver::on_turn_start) /// for the turn about to stream. /// @@ -1243,80 +1272,86 @@ impl BareLoop { } } - /// Consult the detection manager and, if a pattern fired, produce the - /// early-exit outcome for `process_turn` to return. - /// - /// Returns `None` when no pattern was detected — the caller continues with - /// tool extraction and dispatch. Returns `Some(Ok(..))` when detection - /// ended the turn softly (caller returns the `TurnResult`), or - /// `Some(Err(..))` when detection forced a hard failure (caller - /// propagates). Either `Some` arm is a terminal transition; both set the - /// appropriate [`LoopState`] before returning. /// Consult the detection manager and, if a pattern forced a hard stop, - /// produce the abort outcome for `process_turn` to return. + /// produce the abort outcome for the driver loop to act on. /// - /// Returns `None` when no pattern fired (caller continues with tool + /// Returns `None` when no pattern fired (the driver continues with tool /// extraction and dispatch), or `Some(Err(..))` with the propagated error /// when detection aborted the session. The terminal state is set via /// [`set_error_state`](Self::set_error_state) before returning. fn apply_loop_detection( &mut self, current_turn: usize, - pattern: &crate::detection::DetectedPattern, + pattern: &DetectedPattern, ) -> Option { - let e = self - .managers - .handle_detected_pattern(pattern, current_turn)?; + self.managers.notify_detected_pattern(pattern, current_turn); + let e = self.decide_detected_pattern(pattern)?; self.set_error_state(&e); Some(e) } - /// End the session because the model finished its turn without requesting - /// any tool calls. - /// - /// Fires [`on_turn_end`](crate::observer::LoopObserver::on_turn_end) - /// (via [`finish_turn`](Self::finish_turn)), transitions to - /// [`LoopState::Completed`], and returns a session-completing - /// [`TurnResult`] (`is_complete: true`) for `process_turn` to return. + /// Decide whether a detected pattern warrants aborting the loop. /// - /// Distinct from the success arm of - /// [`apply_loop_detection`](Self::apply_loop_detection), which transitions - /// to `Completed` *without* firing `on_turn_end` — that path ends the - /// session from detection, not from natural turn completion, so the - /// turn-end bookkeeping is intentionally skipped there. - fn complete_session( - &mut self, - text: String, - turn_in: u64, - turn_out: u64, - turn_start: Instant, - ) -> TurnResult { - // TODO: fire_complete_turn - self.finish_turn(turn_in, turn_out, turn_start.elapsed()); - self.state = LoopState::Completed { - summary: text.clone(), - }; - Self::turn_complete(text, turn_in, turn_out, turn_start.elapsed()) + /// Reads the detection config (`stop_threshold`, `on_converge`) to + /// determine if the pattern is severe enough to halt. Returns + /// `Some(LoopError)` to abort, `None` to continue. + fn decide_detected_pattern(&self, pattern: &DetectedPattern) -> Option { + let config = self.managers.detection().config(); + match pattern { + DetectedPattern::NoPattern => None, + DetectedPattern::LoopDetected { + repetitions, + pattern_description, + } => { + if *repetitions >= config.stop_threshold { + tracing::error!( + repetitions, + pattern = %pattern_description, + "stopping agent: loop threshold exceeded" + ); + Some(LoopError::LoopDetected { + message: format!("{pattern_description} repeated {repetitions} times"), + }) + } else { + None + } + } + DetectedPattern::ConvergenceDetected { .. } => match config.on_converge { + ConvergenceAction::Stop => Some(LoopError::LoopDetected { + message: "agent stopped: convergence detected".into(), + }), + ConvergenceAction::AskUser => Some(LoopError::LoopDetected { + message: "agent stopped: convergence detected, user input needed".into(), + }), + ConvergenceAction::Warn + | ConvergenceAction::Compact + | ConvergenceAction::SwitchPhase => None, + }, + } } - /// Set the terminal state for a propagated error. + /// Record the terminal outcome for a propagated error on the machine. /// - /// All errors from the turn body are recorded as [`LoopState::Failed`]. - /// Cancellation is handled separately by the `select!` in `process_turn`, - /// which sets [`LoopState::Cancelled`] directly and never reaches this - /// method. + /// Driver-loop errors are recorded as + /// [`MachineOutcome::Failed`](crate::engine::core::MachineOutcome::Failed) + /// on the machine. Cancellation that surfaces as a propagated + /// [`LoopError::Cancelled`] (for example when a retry loop observes the + /// cancel signal mid-dispatch) is recorded as + /// [`MachineOutcome::Cancelled`](crate::engine::core::MachineOutcome) + /// so a clean termination never reads as a failure. The machine is driven + /// to its terminal state so that [`state`](crate::engine::core::Loop::state) + /// reflects the outcome immediately. fn set_error_state(&mut self, e: &LoopError) { - self.state = LoopState::Failed { - error: e.to_string(), - }; + if matches!(e, LoopError::Cancelled) { + self.machine.cancel(); + let _ = self.machine.next_step(self.machine_policy()); + } else { + self.machine.fail(e.clone()); + } } } -// ================================================== -// ModelSwitch builder -// ================================================== - -/// Builder for a runtime model switch on [`BareLoop`]. +/// Builder for a model switch on [`BareLoop`]. /// /// Created by [`BareLoop::switch_model`]. Allows updating /// context-window and max-tokens alongside the model name, then applies @@ -1330,37 +1365,27 @@ 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. + /// If omitted, the existing 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 with_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 with_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`. + /// 3. Updates the session context window. /// 4. Resets the [`FallbackManager`](crate::fallback::FallbackManager) /// circuit breaker to `Primary` and updates the original-model /// tracker to the new model. @@ -1374,7 +1399,6 @@ impl ModelSwitch<'_, C> { loop_, target_model, context_window, - max_tokens, } = self; let trimmed = target_model.trim(); @@ -1384,22 +1408,17 @@ impl ModelSwitch<'_, C> { )); } - let from = loop_.config.model.clone(); + let from = loop_.client.model(); loop_.client.set_model(trimmed); - loop_.config.model = trimmed.to_string(); if let Some(cw) = context_window { - loop_.config.context_window = cw; + loop_.session.config.context_window = cw; } - if let Some(mt) = max_tokens { - loop_.config.max_tokens = mt; - } - - loop_.managers.fallback.reset(); + loop_.managers.fallback().reset(); loop_ .managers - .fallback + .fallback() .set_original_model(trimmed.to_string()); loop_ .managers @@ -1414,190 +1433,347 @@ impl ModelSwitch<'_, C> { } impl BareLoop { - /// Execute the turn body without cancellation awareness. + /// Inject any contributor messages ahead of the next model call. + /// + /// Each registered [`ContextContributor`] is consulted against the + /// machine-owned history snapshot; any message it returns is fed into the + /// machine via [`LoopMachine::inject`] in registration order. No-op when no + /// contributors are registered. + fn inject_contributors(&mut self, current_turn: usize) { + if self.contributors.is_empty() { + return; + } + let ctx = ContributorContext::new(current_turn, self.machine.history()); + let injected: Vec = self + .contributors + .iter() + .filter_map(|contributor| contributor.contribute(&ctx)) + .collect(); + for message in injected { + self.machine.inject(message); + } + } + + /// Handle a model-call request from the machine. /// - /// Cancellation is handled by the `select!` in `process_turn`, which - /// drops this future if `cancel.notified()` fires. + /// Fires the per-turn observer events in order, injects contributor + /// messages, streams the response, applies loop detection, feeds the + /// completed [`ModelResponse`] back to the machine, and keeps the run + /// turn count in sync. Cancellation races the stream via a biased + /// `select!`. /// /// # Errors /// - /// Returns [`LoopError`] on streaming failure, loop detection abort, tool - /// dispatch error, or compaction failure. - async fn run_turn_body( - &mut self, - current_turn: usize, - turn_start: Instant, - ) -> Result { - let (msg, usage, _stream_stop) = self.do_stream().await?; - self.accumulate_usage(usage.as_ref()); + /// Propagates [`LoopError::Cancelled`] when the cancel signal fires + /// mid-stream, or any streaming / loop-detection error. + async fn handle_call_llm(&mut self, turn: usize) -> Result<(), LoopError> { + let turn_start = Instant::now(); + let current_turn = turn.saturating_sub(1); + let is_first_turn = current_turn == 0; + let turn_input = if is_first_turn { + self.current_run() + .map_or(String::new(), |r| r.input.clone()) + } else { + self.machine + .history() + .last() + .map(|m| { + m.parts + .iter() + .filter_map(|p| p.as_text()) + .collect::>() + .join("") + }) + .unwrap_or_default() + }; - let text = Self::extract_text(&msg); - let (turn_in, turn_out) = Self::usage_tokens(usage.as_ref()); - let pattern = self.managers.detection.record_response(&text); - self.fire_response(current_turn, &text, usage); + self.fire_turn_start(current_turn, &turn_input); + self.inject_contributors(current_turn); + + let cancel = Arc::clone(&self.cancelled); + tokio::select! { + biased; + () = cancel.notified() => { + self.machine.cancel(); + self.notify_turn_end( + current_turn, + false, + Some("cancelled".into()), + turn_start.elapsed(), + 0, + 0, + ); + Err(LoopError::Cancelled) + } + stream_outcome = self.do_stream() => { + let (msg, usage, _stream_stop) = match stream_outcome { + Ok(triple) => triple, + Err(LoopError::Cancelled) => { + self.notify_turn_end( + current_turn, + false, + Some("cancelled".into()), + turn_start.elapsed(), + 0, + 0, + ); + return Err(LoopError::Cancelled); + } + Err(e) => return Err(e), + }; - if let Some(e) = self.apply_loop_detection(current_turn, &pattern) { - return Err(e); - } + let text = msg.text_content(); + let (turn_in, turn_out) = Self::usage_tokens(usage.as_ref()); + let pattern = self.managers.detection().record_response(&text); + self.fire_response(current_turn, &text, usage); + + if let Some(e) = self.apply_loop_detection(current_turn, &pattern) { + return Err(e); + } + + let tool_calls: Vec = msg + .tool_call_parts() + .into_iter() + .map(|(id, tool, input)| ToolCall { + id: id.to_string(), + tool: tool.to_string(), + input: input.clone(), + }) + .collect(); + let stop_reason = if tool_calls.is_empty() { + StopReason::EndTurn + } else { + StopReason::ToolCall + }; + let model_response = ModelResponse { + message: msg, + input_tokens: turn_in, + output_tokens: turn_out, + stop_reason, + available_tools: self.tools.tool_names(), + }; + self.machine.model_response(model_response); + + let turn_index = current_turn; + let is_empty = tool_calls.is_empty(); + if let Some(run) = self.current_run_mut() { + run.turns.push(crate::engine::core::Turn { + turn: turn_index, + input: turn_input, + output: text, + tool_calls, + input_tokens: turn_in, + output_tokens: turn_out, + }); + } - let tool_calls = Self::extract_tool_calls(&msg); - self.conversation.push(msg); - self.budget.total_turns = self.budget.total_turns.saturating_add(1); + if is_empty { + self.notify_turn_end( + current_turn, + true, + None, + turn_start.elapsed(), + turn_in, + turn_out, + ); + } + Ok(()) + } + } + } - if tool_calls.is_empty() { - // TODO: complete turn - return Ok(self.complete_session(text, turn_in, turn_out, turn_start)); + /// Handle a tool-dispatch request from the machine. + /// + /// Fires `on_tool_call_received`, dispatches the calls that are not + /// preresolved, feeds every tool-result message (dispatched plus + /// preresolved) back to the machine, and keeps the run budget in sync. + /// Cancellation races the dispatch via a biased `select!`. + /// + /// # Errors + /// + /// Propagates [`LoopError::Cancelled`] when the cancel signal fires during + /// dispatch, or any dispatch / loop-detection error. + async fn handle_call_tools( + &mut self, + turn: usize, + calls: &[PendingToolCall], + ) -> Result<(), LoopError> { + let turn_start = Instant::now(); + let current_turn = turn.saturating_sub(1); + let (turn_in, turn_out) = self + .current_run() + .and_then(|r| r.turns.last()) + .map_or((0, 0), |t| (t.input_tokens, t.output_tokens)); + let mut tool_calls: Vec = Vec::with_capacity(calls.len()); + let mut dispatch_calls: Vec = Vec::new(); + let mut preresolved: Vec = Vec::new(); + for pending in calls { + tool_calls.push(pending.call.clone()); + match &pending.preresolved_result { + Some(msg) => preresolved.push(msg.clone()), + None => dispatch_calls.push(pending.call.clone()), + } } self.fire_tool_calls_received(current_turn, &tool_calls); - self.state = LoopState::WaitingForTool { - tool: tool_calls - .first() - .map(|tc| tc.tool.clone()) - .unwrap_or_default(), - started_at: std::time::SystemTime::now(), - }; - let mut budget = std::mem::take(&mut self.budget); - let turn_duration = turn_start.elapsed(); - let dispatch_result = self - .dispatch_and_record( - &tool_calls, - current_turn, - turn_duration, - turn_in, - turn_out, - &mut budget, - ) - .await; - - self.budget = budget; - if let Err(e) = dispatch_result { - self.set_error_state(&e); - return Err(e); - } + let cancel = Arc::clone(&self.cancelled); + let dispatch = async { + self.dispatch_and_record(&dispatch_calls, current_turn, turn_start, turn_in, turn_out) + .await + }; - self.try_compact_context().await; - self.state = LoopState::Processing { - turn: self.budget.total_turns, + let dispatched: Message = tokio::select! { + biased; + () = cancel.notified() => { + self.machine.cancel(); + self.notify_turn_end( + current_turn, + false, + Some("cancelled".into()), + turn_start.elapsed(), + 0, + 0, + ); + return Err(LoopError::Cancelled); + } + result = dispatch => match result { + Ok(msg) => msg, + Err(e) => { + self.set_error_state(&e); + return Err(e); + } + }, }; - 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, - }) + preresolved.push(dispatched); + self.machine.tool_results(preresolved); + Ok(()) } -} - -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 { - config.validate()?; - 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(()) - }) + /// Handle a compaction request from the machine. + /// + /// Runs the configured [`ContextManager`](crate::compact::ContextManager) + /// over the machine-owned history (firing `on_compaction` and hooks), then + /// feeds the compacted history back to the machine. + /// + /// # Errors + /// + /// Propagates [`LoopError::ContextExceeded`] when compaction could not + /// reduce the history enough. + async fn handle_compact( + &mut self, + reason: crate::compact::types::CompactReason, + ) -> Result<(), LoopError> { + let turn = self.machine.turns_taken(); + // The machine is already `AwaitingCompaction` for this reason; the + // driver just performs the IO and feeds the result back. + let (compacted, tokens_after) = self.run_compaction(turn, reason).await?; + self.machine.compaction_result(compacted, tokens_after); + Ok(()) } +} - fn process_turn<'a>( +impl crate::engine::core::Loop for BareLoop { + fn run<'a>( &'a mut self, input: &'a str, - ) -> Pin> + Send + 'a>> { + run_config: &'a RunConfig, + ) -> Pin + Send + 'a>> { Box::pin(async move { - let turn_start = Instant::now(); - let current_turn = self.budget.total_turns; - - self.record_user_input(input); - self.fire_turn_start(current_turn, input); - - if !self.contributors.is_empty() { - let ctx = ContributorContext::new(current_turn, &self.conversation); - let injected: Vec = self - .contributors - .iter() - .filter_map(|contributor| contributor.contribute(&ctx)) - .collect(); - self.conversation.extend(injected); + let session_is_new = self.session.session_start.is_none(); + if session_is_new { + self.session.session_start = Some(Instant::now()); + self.notify_session_start(); } - let cancel = Arc::clone(&self.cancelled); - tokio::select! { - biased; - () = cancel.notified() => { - self.state = LoopState::Cancelled; - // TODO: fire_turn_cancelled - self.managers.observers().on_turn_end(&TurnEndContext { - turn: current_turn, - success: false, - error: Some("cancelled".into()), - duration_ms: Self::millis_u64(turn_start.elapsed()), - input_tokens: 0, - output_tokens: 0, - }); - Err(LoopError::Cancelled) + self.session.history.push(Message::user(input)); + self.session.runs.push(Run::new(input, run_config)); + self.managers.reset_all(); + self.machine = LoopMachine::from_history(self.session.history.clone()); + + loop { + let policy = self.machine_policy(); + match self.machine.next_step(policy) { + MachineStep::CallLLM { turn } => { + if let Err(e) = self.handle_call_llm(turn).await { + self.set_error_state(&e); + self.finalize(false).await?; + return Err(e); + } + } + MachineStep::CallTools { calls } => { + let turn = self.machine.turns_taken(); + if let Err(e) = self.handle_call_tools(turn, &calls).await { + self.set_error_state(&e); + self.finalize(false).await?; + return Err(e); + } + } + MachineStep::Compact { reason } => { + if let Err(e) = self.handle_compact(reason).await { + self.set_error_state(&e); + self.finalize(false).await?; + return Err(e); + } + } + MachineStep::Done(outcome) => match outcome { + MachineOutcome::Completed { final_text } => { + if let Some(run) = self.current_run_mut() { + run.output = Some(final_text.clone()); + } + + self.session.history.push(Message::assistant(final_text)); + + break; + } + MachineOutcome::MaxTurnsExceeded => { + let err = LoopError::MaxTurnsExceeded { + max: self.run_config().max_turns, + }; + self.finalize(false).await?; + return Err(err); + } + MachineOutcome::Cancelled => { + self.finalize(false).await?; + return Err(LoopError::Cancelled); + } + MachineOutcome::Failed { error } => { + self.finalize(false).await?; + return Err(error); + } + }, } - result = self.run_turn_body(current_turn, turn_start) => result, } + + self.finalize(true).await }) } fn should_continue(&self) -> bool { - if self.is_cancelled() { - return false; - } - self.budget.total_turns < self.config.max_turns + !self.machine.is_terminal() } fn finalize<'a>( &'a mut self, - ) -> Pin> + Send + 'a>> { + success: bool, + ) -> 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 { .. } | LoopState::Cancelled); - - // 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(), - LoopState::Cancelled => "session cancelled".to_string(), - _ => "session failed".to_string(), - }); + if let Some(run) = self.current_run_mut() { + run.end = Some(Instant::now()); } - // Notify observers/hooks using the final SessionResult. - self.notify_session_end(&self.budget, duration); + let run = self.current_run().cloned().unwrap_or_default(); + let duration = run.duration(); - Ok(self.budget.clone()) + self.notify_session_end(&run, success, duration); + + Ok(run) }) } - fn state(&self) -> LoopState { - self.state.clone() + fn state(&self) -> MachineState { + self.machine.state() } fn cancel(&self) { @@ -1608,24 +1784,20 @@ impl crate::engine::loop_core::Loop for BareLoop { 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, - }); + let max_turns = self.run_config().max_turns; + let turns = self.current_run().map_or(0, Run::turn_count); + if turns >= max_turns { + return Some(LoopError::MaxTurnsExceeded { max: max_turns }); } None } - - fn config(&self) -> &LoopConfig { - &self.config - } } #[cfg(test)] mod tests { use super::*; use crate::api::error::ApiError; - use crate::engine::loop_core::Loop; + use crate::engine::core::Loop; use crate::stream::{ DeltaPart, IndexedDelta, MessageDelta, MessageDeltaPayload, MessageMetadata, MessageStart, PartStart, Usage, @@ -1696,7 +1868,6 @@ mod tests { tool_input: Value, final_text: &str, ) { - // First response: tool_call let tool_events = vec![ StreamEvent::MessageStart(MessageStart { message: MessageMetadata { @@ -1720,7 +1891,60 @@ mod tests { ]; crate::error::recover_guard(self.responses.lock()).push(tool_events); - // Second response: end_turn with text + let text_events = vec![ + StreamEvent::MessageStart(MessageStart { + message: MessageMetadata { + id: "msg_final".into(), + role: "assistant".into(), + model: crate::error::recover_guard(self.model_name.lock()).clone(), + }, + }), + StreamEvent::PartStart(PartStart { + index: 0, + part: Some(MessagePart::text(final_text)), + }), + StreamEvent::IndexedDelta(IndexedDelta { + index: 0, + delta: DeltaPart::Text { + text: final_text.to_string(), + }, + }), + StreamEvent::PartStop, + StreamEvent::MessageDelta(MessageDelta { + delta: MessageDeltaPayload { + stop_reason: Some("end_turn".to_string()), + }, + usage: Some(Usage::new(30, 15)), + }), + StreamEvent::MessageStop, + ]; + crate::error::recover_guard(self.responses.lock()).push(text_events); + } + + fn add_multi_tool_then_text(&self, tools: &[(String, String, Value)], final_text: &str) { + let mut tool_events = vec![StreamEvent::MessageStart(MessageStart { + message: MessageMetadata { + id: "msg_tool".into(), + role: "assistant".into(), + model: crate::error::recover_guard(self.model_name.lock()).clone(), + }, + })]; + for (idx, (id, name, input)) in tools.iter().enumerate() { + tool_events.push(StreamEvent::PartStart(PartStart { + index: idx, + part: Some(MessagePart::tool_call(id, name, input.clone())), + })); + tool_events.push(StreamEvent::PartStop); + } + tool_events.push(StreamEvent::MessageDelta(MessageDelta { + delta: MessageDeltaPayload { + stop_reason: Some("tool_call".to_string()), + }, + usage: Some(Usage::new(50, 10)), + })); + tool_events.push(StreamEvent::MessageStop); + crate::error::recover_guard(self.responses.lock()).push(tool_events); + let text_events = vec![ StreamEvent::MessageStart(MessageStart { message: MessageMetadata { @@ -1829,7 +2053,6 @@ mod tests { } } - // Helper trait for Vec-like pop_front on Vec trait PopFront { fn pop_front(&mut self) -> Option; } @@ -2014,10 +2237,14 @@ mod tests { } } - fn make_config() -> LoopConfig { - LoopConfig { + fn make_config() -> SessionConfig { + SessionConfig::default() + } + + fn make_run_config() -> RunConfig { + RunConfig { max_turns: 10, - ..Default::default() + ..RunConfig::default() } } @@ -2028,11 +2255,10 @@ mod tests { let config = make_config(); let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); - let result = agent.run("Hi").await.unwrap(); + let result = agent.run("Hi", &RunConfig::default()).await.unwrap(); - assert!(result.success); - assert_eq!(result.total_turns, 1); - assert_eq!(result.final_output.as_deref(), Some("Hello! I'm done.")); + assert_eq!(result.turn_count(), 1); + assert_eq!(result.output.as_deref(), Some("Hello! I'm done.")); } #[tokio::test] @@ -2050,11 +2276,339 @@ mod tests { let config = make_config(); let mut agent = BareLoop::new(Arc::new(client), registry, config); - let result = agent.run("Echo hello").await.unwrap(); + let result = agent + .run("Echo hello", &RunConfig::default()) + .await + .unwrap(); + + assert_eq!(result.turn_count(), 2); // tool_call turn + end_turn + assert_eq!(result.tool_call_count(), 1); + } + + struct SequenceObserver { + log: Arc>>, + } + + impl SequenceObserver { + fn new(log: Arc>>) -> Self { + Self { log } + } + + fn record(&self, name: &str) { + crate::error::recover_guard(self.log.lock()).push(name.to_string()); + } + } + + impl crate::observer::LoopObserver for SequenceObserver { + fn name(&self) -> &'static str { + "sequence" + } + fn on_turn_start(&self, _ctx: &crate::observer::TurnStartContext) { + self.record("on_turn_start"); + } + fn on_text_delta(&self, _ctx: &crate::observer::TextDeltaContext) { + self.record("on_text_delta"); + } + fn on_stream_success(&self, _ctx: &crate::observer::StreamContext) { + self.record("on_stream_success"); + } + fn on_response(&self, _ctx: &crate::observer::ResponseContext) { + self.record("on_response"); + } + fn on_turn_end(&self, _ctx: &crate::observer::TurnEndContext) { + self.record("on_turn_end"); + } + fn on_tool_call_received(&self, _ctx: &crate::observer::ToolCallReceivedContext) { + self.record("on_tool_call_received"); + } + fn on_tool_pre(&self, _ctx: &crate::observer::ToolPreContext) { + self.record("on_tool_pre"); + } + fn on_tool_post(&self, _ctx: &crate::observer::ToolPostContext) { + self.record("on_tool_post"); + } + fn on_compaction(&self, _ctx: &crate::observer::CompactedContext) { + self.record("on_compaction"); + } + } + + fn sequence_log() -> Arc>> { + Arc::new(Mutex::new(Vec::new())) + } + + fn agent_with_sequence_observer( + client: MockClient, + registry: ToolRegistry, + log: Arc>>, + ) -> BareLoop { + let mut agent = BareLoop::new(Arc::new(client), registry, make_config()); + agent.register_observer(Arc::new(SequenceObserver::new(log))); + agent + } + + fn snapshot(log: &Arc>>) -> Vec { + crate::error::recover_guard(log.lock()).clone() + } + + #[tokio::test] + async fn observer_sequence_text_only_turn() { + let client = MockClient::new("test-model"); + client.add_text_response("Hi there."); + let log = sequence_log(); + let mut agent = agent_with_sequence_observer(client, ToolRegistry::new(), log.clone()); + agent.run("Hi", &RunConfig::default()).await.unwrap(); + + let events = snapshot(&log); + let turn_events: Vec<&String> = events + .iter() + .filter(|e| { + matches!( + e.as_str(), + "on_turn_start" + | "on_text_delta" + | "on_stream_success" + | "on_response" + | "on_turn_end" + ) + }) + .collect(); + let expected = [ + "on_turn_start", + "on_text_delta", + "on_stream_success", + "on_response", + "on_turn_end", + ]; + assert_eq!( + turn_events.iter().map(|s| s.as_str()).collect::>(), + expected + ); + } + + #[tokio::test] + async fn observer_sequence_tool_call_turn() { + let client = MockClient::new("test-model"); + client.add_tool_then_text("tool_1", "echo", json!({"message": "hi"}), "Done."); + let mut registry = ToolRegistry::new(); + registry.register(EchoTool); + let log = sequence_log(); + let mut agent = agent_with_sequence_observer(client, registry, log.clone()); + agent.run("echo hi", &RunConfig::default()).await.unwrap(); + + let events = snapshot(&log); + // The tool-call turn must announce the tool calls before dispatching. + assert!( + events.iter().any(|e| e == "on_tool_call_received"), + "tool-call turn fires on_tool_call_received" + ); + let pre = events.iter().position(|e| e == "on_tool_pre"); + let post = events.iter().position(|e| e == "on_tool_post"); + assert!( + pre.zip(post).is_some_and(|(p1, p2)| p1 < p2), + "on_tool_pre fires before on_tool_post" + ); + } + + #[tokio::test] + async fn observer_sequence_multi_tool_turn() { + let client = MockClient::new("test-model"); + // Two tool calls in one turn, then a final text turn. + client.add_multi_tool_then_text( + &[ + ( + "tool_a".to_string(), + "echo".to_string(), + json!({"message": "a"}), + ), + ( + "tool_b".to_string(), + "echo".to_string(), + json!({"message": "b"}), + ), + ], + "All done.", + ); + let mut registry = ToolRegistry::new(); + registry.register(EchoTool); + let log = sequence_log(); + let mut agent = agent_with_sequence_observer(client, registry, log.clone()); + agent + .run("echo twice", &RunConfig::default()) + .await + .unwrap(); - assert!(result.success); - assert_eq!(result.total_turns, 2); // tool_call turn + end_turn - assert_eq!(result.tool_calls, 1); + let events = snapshot(&log); + // Sequential dispatch: pre, post, pre, post — never interleaved. + let tool_seq: Vec<&String> = events + .iter() + .filter(|e| matches!(e.as_str(), "on_tool_pre" | "on_tool_post")) + .collect(); + assert_eq!( + tool_seq.iter().map(|s| s.as_str()).collect::>(), + ["on_tool_pre", "on_tool_post", "on_tool_pre", "on_tool_post"], + "multi-tool sequential dispatch keeps pre/post paired and ordered" + ); + } + + #[tokio::test] + async fn observer_sequence_compaction_turn() { + let client = MockClient::new("test-model"); + // Drive enough tokens to trip a low threshold, then finish. + client.add_text_response(&"x".repeat(200)); + client.add_text_response("compacted-and-done"); + let log = sequence_log(); + let mut agent = agent_with_sequence_observer(client, ToolRegistry::new(), log.clone()); + agent.set_context_manager(Arc::new( + crate::compact::ContextManager::new(Arc::new( + crate::compact::TruncatingCompactor::new(), + )) + .with_context_window(100) + .with_threshold(10), + )); + let run_config = RunConfig::default(); + let run_result = agent.run("fill it up", &run_config).await; + // The compaction scenario drives the run to completion; event placement is asserted below. + assert!(run_result.is_ok(), "compaction run completes"); + + let events = snapshot(&log); + // If compaction ran, on_compaction sits at a turn boundary (after a + // turn_end, before the next turn_start). If the estimate didn't trip, + // the scenario is N/A — assert placement only when present. + if let Some(idx) = events.iter().position(|e| e == "on_compaction") { + let before = events.get(idx.wrapping_sub(1)); + let after = events.get(idx + 1); + assert!( + before == Some(&"on_turn_end".to_string()) + || after == Some(&"on_turn_start".to_string()), + "on_compaction at idx {idx} sits at a turn boundary, got before={before:?} after={after:?}" + ); + } + } + + #[tokio::test] + async fn observer_sequence_cancelled_turn() { + let client = MockClient::new("test-model"); + // Never-ending tool calls so the loop is mid-flight when cancelled. + for _ in 0..5 { + client.add_tool_only_response("c1", "echo", json!({"message": "x"})); + } + let mut registry = ToolRegistry::new(); + registry.register(EchoTool); + let log = sequence_log(); + let mut agent = agent_with_sequence_observer(client, registry, log.clone()); + + let handle = agent.cancel_signal(); + let join = tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + handle.cancel(); + }); + let result = agent.run("go", &RunConfig::default()).await; + join.await.unwrap(); + assert!(result.is_err(), "cancelled run returns an error"); + + let events = snapshot(&log); + let started = events.iter().filter(|e| **e == "on_turn_start").count(); + let ended = events.iter().filter(|e| **e == "on_turn_end").count(); + assert!( + started >= 1 && ended >= 1, + "cancelled turn still fires on_turn_end (started={started}, ended={ended})" + ); + } + + struct ToolNameCapture { + captured: Arc>>, + } + impl crate::observer::LoopObserver for ToolNameCapture { + fn name(&self) -> &'static str { + "tool-name-capture" + } + fn on_tool_pre(&self, ctx: &crate::observer::ToolPreContext) { + *crate::error::recover_guard(self.captured.lock()) = Some(ctx.tool.clone()); + } + } + + #[tokio::test] + async fn dispatch_surfaces_tool_name_on_tool_pre() { + let client = MockClient::new("test-model"); + // A tool-call turn then a final text turn. The driver is dispatching the + // tool during `on_tool_pre`; the ToolNameCapture observer records the + // tool name carried on the context. + client.add_tool_then_text("tool_1", "echo", json!({"message": "hi"}), "done"); + let mut registry = ToolRegistry::new(); + registry.register(EchoTool); + + let captured = Arc::new(Mutex::new(None::)); + let mut agent = BareLoop::new(Arc::new(client), registry, make_config()); + agent.register_observer(Arc::new(ToolNameCapture { + captured: Arc::clone(&captured), + })); + agent.run("echo hi", &RunConfig::default()).await.unwrap(); + + let snapshot = crate::error::recover_guard(captured.lock()).clone(); + assert_eq!( + snapshot.as_deref(), + Some("echo"), + "tool name preserved on ToolPreContext during dispatch" + ); + } + + #[tokio::test] + async fn bareloop_machine_accessor_returns_machine() { + let client = MockClient::new("test-model"); + client.add_text_response("hi"); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + agent.run("hello", &RunConfig::default()).await.unwrap(); + // After a run, the machine is populated and history holds the turn. + let machine = agent.machine(); + assert!(machine.turns_taken() >= 1); + assert!(!machine.history().is_empty()); + } + + #[tokio::test] + async fn serialize_drop_deserialize_resume_preserves_history() { + let client = MockClient::new("test-model"); + client.add_text_response("first"); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + agent.run("prompt", &RunConfig::default()).await.unwrap(); + + // Take the machine and round-trip it through serde. + let machine = agent.into_machine(); + let serialized = serde_json::to_string(&machine).expect("serialize machine"); + let restored: LoopMachine = serde_json::from_str(&serialized).expect("deserialize machine"); + // Compare by serialized form: Message is not PartialEq. + let got = serde_json::to_string(restored.history()).expect("serialize history"); + let want = serde_json::to_string(machine.history()).expect("serialize history"); + assert_eq!(got, want, "history survives serialize/deserialize"); + + // Rebuild a loop around the restored machine. + let client2 = MockClient::new("test-model"); + let _rebuilt = BareLoop::from_machine( + restored, + make_config(), + Arc::new(client2), + ToolRegistry::new(), + ); + } + + #[tokio::test] + async fn session_id_stable_and_run_id_rotates_across_runs() { + let client = MockClient::new("test-model"); + client.add_text_response("first"); + client.add_text_response("second"); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + + let first = agent.run("one", &RunConfig::default()).await.unwrap(); + let first_session = agent.session().id; + let first_run = first.id; + + let second = agent.run("two", &RunConfig::default()).await.unwrap(); + let second_session = agent.session().id; + let second_run = second.id; + + // Session identity is stable across runs. + assert_eq!(first_session, second_session, "session_id is stable"); + // Each run mints a fresh id. + assert_ne!(first_run, second_run, "id rotates per run"); } #[tokio::test] @@ -2069,14 +2623,15 @@ mod tests { ); } - let mut config = make_config(); - config.max_turns = 3; - let mut registry = ToolRegistry::new(); registry.register(EchoTool); - let mut agent = BareLoop::new(Arc::new(client), registry, config); - let result = agent.run("Keep going").await; + let mut agent = BareLoop::new(Arc::new(client), registry, make_config()); + let run_config = RunConfig { + max_turns: 3, + ..RunConfig::default() + }; + let result = agent.run("Keep going", &run_config).await; assert!(result.is_err()); match result.unwrap_err() { LoopError::MaxTurnsExceeded { max } => assert_eq!(max, 3), @@ -2096,7 +2651,7 @@ mod tests { agent.cancel(); assert!(agent.is_cancelled()); - let result = agent.run("Hi").await; + let result = agent.run("Hi", &RunConfig::default()).await; assert!(result.is_err()); match result.unwrap_err() { LoopError::Cancelled => {} @@ -2110,7 +2665,7 @@ mod tests { let client = MockClient::new("test-model"); let config = make_config(); let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); - let result = agent.run("Hi").await; + let result = agent.run("Hi", &RunConfig::default()).await; assert!(result.is_err()); match result.unwrap_err() { LoopError::Api(msg) => assert!(msg.contains("No more mock responses"), "got: {msg}"), @@ -2126,12 +2681,14 @@ mod tests { // Empty registry — tool won't be found let config = make_config(); let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); - let result = agent.run("Use nonexistent tool").await.unwrap(); + let result = agent + .run("Use nonexistent tool", &RunConfig::default()) + .await + .unwrap(); // The tool-not-found should be returned as an error result in the conversation, // not as a hard error. The loop should continue and eventually get the end_turn. - assert!(result.success); - assert_eq!(result.total_turns, 2); + assert_eq!(result.turn_count(), 2); } #[tokio::test] @@ -2144,10 +2701,12 @@ mod tests { let config = make_config(); let mut agent = BareLoop::new(Arc::new(client), registry, config); - let result = agent.run("Use failing tool").await.unwrap(); + let result = agent + .run("Use failing tool", &RunConfig::default()) + .await + .unwrap(); - assert!(result.success); - assert_eq!(result.total_turns, 2); + assert_eq!(result.turn_count(), 2); } #[tokio::test] @@ -2160,9 +2719,8 @@ mod tests { let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); agent.register_observer(plugin.clone()); - let result = agent.run("Hi").await.unwrap(); + let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); - assert!(result.success); 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); @@ -2182,8 +2740,7 @@ mod tests { 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); + let _result = agent.run("Echo test", &RunConfig::default()).await.unwrap(); assert_eq!(plugin.tool_pres.load(Ordering::SeqCst), 1); assert_eq!(plugin.tool_posts.load(Ordering::SeqCst), 1); @@ -2207,14 +2764,25 @@ mod tests { let config = make_config(); let mut agent = BareLoop::new(Arc::new(client), registry, config); - // Pre-push the user message manually to inspect conversation after - agent.conversation.push(Message::user("Echo hello")); + // Driving the run builds the conversation in the machine-owned history. + agent + .run("Echo hello", &RunConfig::default()) + .await + .unwrap(); - // Can't run directly since we manually pushed, but let's verify the helpers - let msg = Message::assistant("test"); - let tool_calls = BareLoop::::extract_tool_calls(&msg); - assert!(tool_calls.is_empty()); + // History: [user, assistant(tool_call), user(tool_result), assistant(text)]. + let history = agent.conversation(); + assert_eq!( + history.len(), + 4, + "expected user, assistant, tool-result, final-answer" + ); + assert_eq!(history[0].role, Role::User); + assert_eq!(history[1].role, Role::Assistant); + assert_eq!(history[2].role, Role::User); + assert_eq!(history[3].role, Role::Assistant); + // The extract helpers still classify tool-call parts correctly. let msg_with_tools = Message::new( Role::Assistant, vec![ @@ -2222,7 +2790,15 @@ mod tests { MessagePart::tool_call("id1", "echo", json!({"message": "hi"})), ], ); - let tool_calls = BareLoop::::extract_tool_calls(&msg_with_tools); + let tool_calls: Vec = msg_with_tools + .tool_call_parts() + .into_iter() + .map(|(id, tool, input)| ToolCall { + id: id.to_string(), + tool: tool.to_string(), + input: input.clone(), + }) + .collect(); assert_eq!(tool_calls.len(), 1); assert_eq!(tool_calls[0].tool, "echo"); } @@ -2307,11 +2883,13 @@ mod tests { let config = make_config(); let mut agent = BareLoop::new(Arc::new(client), registry, config); - let result = agent.run("Echo twice").await.unwrap(); + let result = agent + .run("Echo twice", &RunConfig::default()) + .await + .unwrap(); - assert!(result.success); - assert_eq!(result.total_turns, 2); - assert_eq!(result.tool_calls, 2); + assert_eq!(result.turn_count(), 2); + assert_eq!(result.tool_call_count(), 2); } #[tokio::test] @@ -2326,8 +2904,7 @@ mod tests { crate::error::recover_guard(buf.lock()).push(delta.to_string()); })); - let result = agent.run("Hi").await.unwrap(); - assert!(result.success); + let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); let received = crate::error::recover_guard(received.lock()); assert!(!received.is_empty(), "streamer should have fired"); @@ -2353,8 +2930,7 @@ mod tests { crate::error::recover_guard(buf.lock()).push(delta.to_string()); })); - let result = agent.run("Hi").await.unwrap(); - assert!(result.success); + let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); let received = crate::error::recover_guard(received.lock()); assert!( @@ -2373,8 +2949,7 @@ mod tests { 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); + let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); } #[tokio::test] @@ -2426,7 +3001,7 @@ mod tests { crate::error::recover_guard(buf.lock()).push_str(delta); })); - agent.run("Use tool").await.unwrap(); + agent.run("Use tool", &RunConfig::default()).await.unwrap(); // The InputJson delta should NOT have triggered the streamer. // Only the "Done" text response in the second turn should. @@ -2495,8 +3070,7 @@ mod tests { }); agent.register_observer(recorder as Arc); - let result = agent.run("Hi").await.unwrap(); - assert!(result.success); + let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); let captured = crate::error::recover_guard(captured.lock()); assert_eq!(captured.len(), 3, "one on_text_delta per SSE text chunk"); @@ -2537,9 +3111,11 @@ mod tests { }); agent.register_observer(recorder as Arc); - let result = agent.run("Use echo then finish").await.unwrap(); - assert!(result.success); - assert_eq!(result.total_turns, 2); + let result = agent + .run("Use echo then finish", &RunConfig::default()) + .await + .unwrap(); + assert_eq!(result.turn_count(), 2); let response_turns = crate::error::recover_guard(response_turns.lock()); let deltas = crate::error::recover_guard(deltas.lock()); @@ -2626,7 +3202,7 @@ mod tests { }); agent.register_observer(recorder as Arc); - agent.run("Use tool").await.unwrap(); + agent.run("Use tool", &RunConfig::default()).await.unwrap(); assert_eq!( count.load(Ordering::SeqCst), @@ -2659,8 +3235,7 @@ mod tests { }); agent.register_observer(recorder as Arc); - let result = agent.run("Hi").await.unwrap(); - assert!(result.success); + let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); let captured = crate::error::recover_guard(captured.lock()); assert!( @@ -2702,8 +3277,7 @@ mod tests { }); agent.register_observer(recorder as Arc); - let result = agent.run("Hi").await.unwrap(); - assert!(result.success); + let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); let streamer_buf = crate::error::recover_guard(streamer_buf.lock()); let observer_buf = crate::error::recover_guard(observer_buf.lock()); @@ -2732,8 +3306,7 @@ mod tests { let observer = Arc::new(CountingObserver::new()); agent.register_observer(observer.clone()); - let result = agent.run("Use echo").await.unwrap(); - assert!(result.success); + let _result = agent.run("Use echo", &RunConfig::default()).await.unwrap(); assert_eq!( observer.tool_calls_received.load(Ordering::SeqCst), @@ -2790,8 +3363,10 @@ mod tests { let observer = Arc::new(CountingObserver::new()); agent.register_observer(observer.clone()); - let result = agent.run("Echo twice").await.unwrap(); - assert!(result.success); + let _result = agent + .run("Echo twice", &RunConfig::default()) + .await + .unwrap(); assert_eq!( observer.tool_calls_received.load(Ordering::SeqCst), @@ -2810,8 +3385,7 @@ mod tests { let observer = Arc::new(CountingObserver::new()); agent.register_observer(observer.clone()); - let result = agent.run("Hi").await.unwrap(); - assert!(result.success); + let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); assert_eq!( observer.tool_calls_received.load(Ordering::SeqCst), @@ -2860,8 +3434,7 @@ mod tests { }); agent.register_observer(recorder as Arc); - let result = agent.run("Use echo").await.unwrap(); - assert!(result.success); + let _result = agent.run("Use echo", &RunConfig::default()).await.unwrap(); let received = crate::error::recover_guard(received.lock()); let response = crate::error::recover_guard(response.lock()); @@ -2930,8 +3503,7 @@ mod tests { let observer = Arc::new(CountingObserver::new()); agent.register_observer(observer.clone()); - let result = agent.run("Use flaky").await.unwrap(); - assert!(result.success); + let _result = agent.run("Use flaky", &RunConfig::default()).await.unwrap(); assert_eq!( observer.tool_calls_received.load(Ordering::SeqCst), @@ -2948,10 +3520,9 @@ mod tests { fn test_accessors() { 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); - assert_eq!(agent.config().session_id, session_id); + assert_ne!(agent.session().id, uuid::Uuid::nil()); assert!(agent.conversation().is_empty()); assert!(!agent.is_cancelled()); } @@ -2970,18 +3541,18 @@ mod tests { } #[tokio::test] - async fn test_session_result_fields() { + async fn test_run_result_fields() { let client = MockClient::new("test-model"); client.add_text_response("Hello!"); let config = make_config(); - let session_id = config.session_id; let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); - let result = agent.run("Hi").await.unwrap(); + let result = agent.run("Hi", &RunConfig::default()).await.unwrap(); - assert_eq!(result.session_id, session_id); - assert!(result.total_duration > Duration::ZERO); - assert!(result.input_tokens > 0 || result.output_tokens > 0); // from mock usage + // Session identity lives on the loop, not the per-run result. + assert_ne!(agent.session().id, uuid::Uuid::nil()); + assert!(result.duration() > Duration::ZERO); + assert!(result.input_tokens() > 0 || result.output_tokens() > 0); // from mock usage } #[tokio::test] @@ -2989,14 +3560,14 @@ mod tests { let client = MockClient::new("test-model"); client.add_text_response("One and done."); - let mut config = make_config(); - config.max_turns = 1; - - let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); - let result = agent.run("Hi").await.unwrap(); + let run_config = RunConfig { + max_turns: 1, + ..RunConfig::default() + }; + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + let result = agent.run("Hi", &run_config).await.unwrap(); - assert!(result.success); - assert_eq!(result.total_turns, 1); + assert_eq!(result.turn_count(), 1); } #[tokio::test] @@ -3004,15 +3575,18 @@ mod tests { let client = MockClient::new("test-model"); client.add_text_response("Should not be reached."); - let mut config = make_config(); - config.max_turns = 0; - - let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); - let result = agent.run("Hi").await; + let run_config = RunConfig { + max_turns: 0, + ..RunConfig::default() + }; + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + let result = agent.run("Hi", &run_config).await; assert!(result.is_err()); + // With max_turns == 0 the loop never executes a turn and reports the + // budget as exhausted. match result.unwrap_err() { - LoopError::Config(msg) => assert!(msg.contains("max_turns")), - other => panic!("Expected Config error, got: {other}"), + LoopError::MaxTurnsExceeded { max } => assert_eq!(max, 0), + other => panic!("Expected MaxTurnsExceeded, got: {other}"), } } @@ -3049,14 +3623,16 @@ mod tests { let config = make_config(); let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); - let result = agent.run("Use missing tool").await.unwrap(); - assert!(result.success); + let _result = agent + .run("Use missing tool", &RunConfig::default()) + .await + .unwrap(); } #[tokio::test] async fn test_loop_detection_hard_stop_propagates_loop_error() { use crate::detection::{DetectionConfig, DetectionManager}; - use crate::runtime::LoopRuntime; + use crate::managers::LoopManagers; let mut registry = ToolRegistry::new(); registry.register(EchoTool); @@ -3065,7 +3641,7 @@ mod tests { client.add_tool_only_response(&format!("call_{i}"), "echo", json!({ "message": "hi" })); } - let runtime = LoopRuntime::new().with_detection( + let managers = LoopManagers::new().with_detection( DetectionManager::new_with_config(DetectionConfig { loop_threshold: 2, stop_threshold: 2, @@ -3075,8 +3651,8 @@ mod tests { ); let mut agent = - BareLoop::new_with_managers(Arc::new(client), registry, make_config(), runtime); - let result = agent.run("test").await; + BareLoop::new_with_managers(Arc::new(client), registry, make_config(), managers); + let result = agent.run("test", &RunConfig::default()).await; assert!( matches!(result, Err(LoopError::LoopDetected { .. })), @@ -3087,7 +3663,7 @@ mod tests { #[tokio::test] async fn test_loop_detection_soft_block_before_stop_threshold() { use crate::detection::{DetectionConfig, DetectionManager}; - use crate::runtime::LoopRuntime; + use crate::managers::LoopManagers; let mut registry = ToolRegistry::new(); registry.register(EchoTool); @@ -3096,7 +3672,7 @@ mod tests { client.add_tool_only_response("c2", "echo", json!({ "message": "hi" })); client.add_text_response("Done"); - let runtime = LoopRuntime::new().with_detection( + let managers = LoopManagers::new().with_detection( DetectionManager::new_with_config(DetectionConfig { loop_threshold: 2, stop_threshold: 10, @@ -3106,8 +3682,8 @@ mod tests { ); let mut agent = - BareLoop::new_with_managers(Arc::new(client), registry, make_config(), runtime); - let result = agent.run("test").await; + BareLoop::new_with_managers(Arc::new(client), registry, make_config(), managers); + let result = agent.run("test", &RunConfig::default()).await; assert!(result.is_ok(), "expected Ok, got {result:?}"); } @@ -3119,7 +3695,7 @@ mod tests { let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); agent.cancel(); - let result = agent.run("test").await; + let result = agent.run("test", &RunConfig::default()).await; assert!( matches!(result, Err(LoopError::Cancelled)), @@ -3136,10 +3712,9 @@ mod tests { client.add_tool_then_text("tool_1", "fail", json!({}), "Moving on"); let mut agent = BareLoop::new(Arc::new(client), registry, make_config()); - let result = agent.run("Test").await.unwrap(); + let result = agent.run("Test", &RunConfig::default()).await.unwrap(); - assert!(result.success); - assert_eq!(result.tool_calls, 1); + assert_eq!(result.tool_call_count(), 1); } #[tokio::test] @@ -3148,10 +3723,9 @@ mod tests { client.add_tool_then_text("tool_1", "nonexistent", json!({}), "OK"); let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); - let result = agent.run("Test").await.unwrap(); + let result = agent.run("Test", &RunConfig::default()).await.unwrap(); - assert!(result.success); - assert_eq!(result.tool_calls, 1); + assert_eq!(result.tool_call_count(), 1); } #[tokio::test] @@ -3163,10 +3737,9 @@ mod tests { client.add_tool_then_text("tool_1", "fail", json!({}), "OK"); let mut agent = BareLoop::new(Arc::new(client), registry, make_config()); - let result = agent.run("Test").await.unwrap(); + let result = agent.run("Test", &RunConfig::default()).await.unwrap(); - assert!(result.success); - assert_eq!(result.tool_calls, 1); + assert_eq!(result.tool_call_count(), 1); } #[tokio::test] @@ -3182,14 +3755,14 @@ mod tests { // Cancel before running agent.cancel(); - let result = agent.run("Test").await; + let result = agent.run("Test", &RunConfig::default()).await; assert!(result.is_err()); } #[tokio::test] async fn test_cancel_during_dispatch_lands_in_cancelled_state() { - // Cancellation fired after process_turn has begun dispatching flows - // through LoopState::Cancelled (not Failed). Uses AlwaysRecoverable so + // Cancellation fired after dispatch has begun flows through + // MachineOutcome::Cancelled (not Failed). Uses AlwaysRecoverable so // FailingTool's error triggers a retry; the retry loop polls // is_cancelled() at the top of each iteration (dispatch.rs), so the // cancel signal set here is observed on the next retry attempt. @@ -3245,15 +3818,15 @@ mod tests { signal.cancel(); }); - let result = agent.run("Test").await; + let result = agent.run("Test", &RunConfig::default()).await; match result { Err(LoopError::Cancelled) => {} other => panic!("expected Err(LoopError::Cancelled), got {other:?}"), } assert_eq!( agent.state(), - LoopState::Cancelled, - "cancellation must land in LoopState::Cancelled, not Failed", + MachineState::Terminal(MachineOutcome::Cancelled), + "cancellation must land in MachineOutcome::Cancelled, not Failed", ); } @@ -3340,7 +3913,7 @@ mod tests { let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); let signal = agent.cancel_signal(); - let handle = tokio::spawn(async move { agent.run("Hi").await }); + let handle = tokio::spawn(async move { agent.run("Hi", &RunConfig::default()).await }); for _ in 0..5 { tokio::task::yield_now().await; @@ -3368,9 +3941,8 @@ mod tests { let builder = ToolPipeline::builder(); agent.set_pipeline(builder).unwrap(); - let result = agent.run("Echo hello").await; - assert!(result.is_ok()); - assert!(result.unwrap().success); + let result = agent.run("Echo hello", &RunConfig::default()).await; + result.unwrap(); } struct TurnNumberCapture { @@ -3412,17 +3984,14 @@ mod tests { let mut registry = ToolRegistry::new(); registry.register(EchoTool); - let mut config = make_config(); - config.max_turns = 10; let capture = Arc::new(Mutex::new(Vec::::new())); - let mut agent = BareLoop::new(Arc::new(client), registry, config); + let mut agent = BareLoop::new(Arc::new(client), registry, make_config()); let builder = ToolPipeline::builder().with_middleware(TurnNumberCapture::new(Arc::clone(&capture))); agent.set_pipeline(builder).unwrap(); - let result = agent.run("test").await; - assert!(result.is_ok()); + let _result = agent.run("test", &make_run_config()).await; let turns = crate::error::recover_guard(capture.lock()).clone(); // Tool was called on turn 0 (first turn) and turn 1 (second turn). @@ -3444,17 +4013,15 @@ mod tests { 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); + let mut loop_ = BareLoop::new(client_arc.clone(), tools, SessionConfig::default()); loop_.switch_model("model-b").apply().unwrap(); - // Config was updated. - assert_eq!(loop_.config().model, "model-b"); + // Client was updated via set_model. + assert_eq!(loop_.client.model(), "model-b"); - // Client was also updated via set_model. + // The shared client handle sees the same update. assert_eq!(client_arc.model(), "model-b"); } @@ -3478,7 +4045,7 @@ mod tests { let client = std::sync::Arc::new(MockClient::new("m1")); let tools = ToolRegistry::new(); - let mut loop_ = BareLoop::new(client, tools, LoopConfig::default()); + let mut loop_ = BareLoop::new(client, tools, SessionConfig::default()); let obs = std::sync::Arc::new(RecordingObserver::default()); let obs_clone = obs.clone(); loop_.register_observer(obs); @@ -3489,7 +4056,7 @@ mod tests { // Observer should have received both switches. let recorded = crate::error::recover_guard(obs_clone.switches.lock()); assert_eq!(recorded.len(), 2, "should have 2 model-switch events"); - assert_eq!(recorded[0], ("default".to_string(), "m2".to_string())); + assert_eq!(recorded[0], ("m1".to_string(), "m2".to_string())); assert_eq!(recorded[1], ("m2".to_string(), "m3".to_string())); } @@ -3536,16 +4103,14 @@ mod tests { model_name: std::sync::Arc::new(std::sync::Mutex::new("static".to_string())), }); let tools = ToolRegistry::new(); - let mut loop_ = BareLoop::new(client, tools, LoopConfig::default()); + let mut loop_ = BareLoop::new(client, tools, SessionConfig::default()); // set_model returns false (unsupported), but apply() is best-effort - // and still updates config-level state. + // and still updates the session/client 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). + // The client is the source of truth for the model; an unsupported + // set_model leaves the client unchanged. assert_eq!(loop_.client.model(), "static"); } @@ -3553,19 +4118,17 @@ mod tests { 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); + let mut loop_ = BareLoop::new(client, tools, SessionConfig::default()); // Before switch, fallback manager has no original model set. - assert_eq!(loop_.managers.fallback.original_model(), None); + 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(), + loop_.managers.fallback().original_model(), Some("new-primary".to_string()) ); } @@ -3574,7 +4137,7 @@ mod tests { 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 mut loop_ = BareLoop::new(client, tools, SessionConfig::default()); let result = loop_.switch_model("").apply(); assert!(result.is_err()); @@ -3584,32 +4147,32 @@ mod tests { assert!(result.is_err()); // Model should remain unchanged. - assert_eq!(loop_.config().model, "default"); + assert_eq!(loop_.client.model(), "model"); } #[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()); + let mut loop_ = BareLoop::new(client, tools, SessionConfig::default()); loop_.switch_model("b").apply().unwrap(); - assert_eq!(loop_.config().model, "b"); + assert_eq!(loop_.client.model(), "b"); loop_.switch_model("c").apply().unwrap(); - assert_eq!(loop_.config().model, "c"); + assert_eq!(loop_.client.model(), "c"); loop_.switch_model("d").apply().unwrap(); - assert_eq!(loop_.config().model, "d"); + assert_eq!(loop_.client.model(), "d"); } #[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 mut loop_ = BareLoop::new(client, tools, SessionConfig::default()); - let original_cw = loop_.config().context_window; + let original_cw = loop_.session_config().context_window; assert_ne!(original_cw, 8192); loop_ @@ -3618,34 +4181,29 @@ mod tests { .apply() .unwrap(); - assert_eq!(loop_.config().model, "small-model"); - assert_eq!(loop_.config().context_window, 8192); + assert_eq!(loop_.client.model(), "small-model"); + assert_eq!(loop_.session_config().context_window, 8192); } #[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()); + let mut loop_ = BareLoop::new(client, tools, SessionConfig::default()); - loop_ - .switch_model("m2") - .with_max_tokens(4096) - .apply() - .unwrap(); + loop_.switch_model("m2").apply().unwrap(); - assert_eq!(loop_.config().model, "m2"); - assert_eq!(loop_.config().max_tokens, 4096); + assert_eq!(loop_.client.model(), "m2"); } #[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()); + let mut loop_ = BareLoop::new(client, tools, SessionConfig::default()); loop_.switch_model(" gpt-4o ").apply().unwrap(); - assert_eq!(loop_.config().model, "gpt-4o"); + assert_eq!(loop_.client.model(), "gpt-4o"); } #[tokio::test] @@ -3654,23 +4212,24 @@ mod tests { 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); + let mut loop_ = BareLoop::new(client, tools, SessionConfig::default()); // 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); + 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().state(), FallbackState::Primary); assert_eq!( - loop_.managers.fallback.original_model(), + loop_.managers.fallback().original_model(), Some("new-primary".to_string()) ); } @@ -3708,15 +4267,18 @@ mod tests { 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, + SessionConfig::default(), ); + loop_.session.runs.push(Run::new( + "", + &RunConfig { + max_turns: 5, + ..RunConfig::default() + }, + )); loop_.set_hook_executor(executor); (loop_, hook) } @@ -3725,11 +4287,31 @@ mod tests { #[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; + // Normal completion: success true, not cancelled, under max_turns. + loop_.current_run_mut().unwrap().turns = vec![ + crate::engine::core::Turn { + turn: 0, + input: String::new(), + output: String::new(), + tool_calls: vec![], + input_tokens: 0, + output_tokens: 0, + }, + crate::engine::core::Turn { + turn: 1, + input: String::new(), + output: String::new(), + tool_calls: vec![], + input_tokens: 0, + output_tokens: 0, + }, + ]; - loop_.notify_session_end(&loop_.budget.clone(), Duration::from_millis(100)); + loop_.notify_session_end( + &loop_.current_run().unwrap().clone(), + true, + Duration::from_millis(100), + ); assert_eq!(hook.captured(), Some(SessionEndReason::Complete)); } @@ -3739,11 +4321,31 @@ mod tests { 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_.current_run_mut().unwrap().turns = vec![ + crate::engine::core::Turn { + turn: 0, + input: String::new(), + output: String::new(), + tool_calls: vec![], + input_tokens: 0, + output_tokens: 0, + }, + crate::engine::core::Turn { + turn: 1, + input: String::new(), + output: String::new(), + tool_calls: vec![], + input_tokens: 0, + output_tokens: 0, + }, + ]; loop_.cancelled.cancel(); - loop_.notify_session_end(&loop_.budget.clone(), Duration::from_millis(100)); + loop_.notify_session_end( + &loop_.current_run().unwrap().clone(), + true, + Duration::from_millis(100), + ); assert_eq!(hook.captured(), Some(SessionEndReason::Cancelled)); } @@ -3752,11 +4354,23 @@ mod tests { #[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 + // Hit max_turns: turn count == max_turns, not cancelled, success true. + loop_.current_run_mut().unwrap().turns = (0..5) + .map(|i| crate::engine::core::Turn { + turn: i, + input: String::new(), + output: String::new(), + tool_calls: vec![], + input_tokens: 0, + output_tokens: 0, + }) + .collect(); - loop_.notify_session_end(&loop_.budget.clone(), Duration::from_millis(100)); + loop_.notify_session_end( + &loop_.current_run().unwrap().clone(), + true, + Duration::from_millis(100), + ); assert_eq!(hook.captured(), Some(SessionEndReason::MaxTurns)); } @@ -3764,12 +4378,13 @@ mod tests { #[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()); + let (loop_, hook) = loop_with_reason_hook(); - loop_.notify_session_end(&loop_.budget.clone(), Duration::from_millis(100)); + loop_.notify_session_end( + &loop_.current_run().unwrap().clone(), + false, + Duration::from_millis(100), + ); assert_eq!(hook.captured(), Some(SessionEndReason::Error)); } @@ -3777,18 +4392,19 @@ mod tests { #[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()); + let (loop_, hook) = loop_with_reason_hook(); - loop_.notify_session_end(&loop_.budget.clone(), Duration::from_millis(100)); + loop_.notify_session_end( + &loop_.current_run().unwrap().clone(), + false, + Duration::from_millis(100), + ); - assert_eq!(hook.captured(), Some(SessionEndReason::ContextOverflow)); + assert_eq!(hook.captured(), Some(SessionEndReason::Error)); } #[tokio::test] - async fn process_turn_cancel_during_streaming_returns_fast() { + async fn run_cancel_during_streaming_returns_fast() { let (client, tx) = StreamingMockClient::new("test-model"); tx.send(Ok(StreamEvent::MessageStart(MessageStart { message: MessageMetadata { @@ -3805,7 +4421,7 @@ mod tests { agent.register_observer(observer.clone()); let signal = agent.cancel_signal(); - let handle = tokio::spawn(async move { agent.run("Hi").await }); + let handle = tokio::spawn(async move { agent.run("Hi", &RunConfig::default()).await }); for _ in 0..5 { tokio::task::yield_now().await; @@ -3832,7 +4448,7 @@ mod tests { } #[tokio::test] - async fn process_turn_cancel_during_dispatch_fires_turn_end() { + async fn run_cancel_during_dispatch_fires_turn_end() { struct SlowTool { notify: Arc, } @@ -3878,7 +4494,8 @@ mod tests { agent.register_observer(observer.clone()); let signal = agent.cancel_signal(); - let handle = tokio::spawn(async move { agent.run("Use slow tool").await }); + let handle = + tokio::spawn(async move { agent.run("Use slow tool", &RunConfig::default()).await }); for _ in 0..10 { tokio::task::yield_now().await; @@ -3903,7 +4520,7 @@ mod tests { } #[tokio::test] - async fn process_turn_cancel_during_recovery_backoff_returns_fast() { + async fn run_cancel_during_recovery_backoff_returns_fast() { struct AlwaysRecoverable; impl crate::reflection::Reflector for AlwaysRecoverable { fn analyze( @@ -3952,7 +4569,8 @@ mod tests { )); let signal = agent.cancel_signal(); - let handle = tokio::spawn(async move { agent.run("Use failing tool").await }); + let handle = + tokio::spawn(async move { agent.run("Use failing tool", &RunConfig::default()).await }); for _ in 0..10 { tokio::task::yield_now().await; @@ -3976,7 +4594,7 @@ mod tests { #[tokio::test] async fn test_rate_limit_escalation_feeds_circuit_breaker() { use crate::fallback::FallbackManager; - use crate::runtime::LoopRuntime; + use crate::managers::LoopManagers; use crate::stream::handler::{ RateLimitConfig, StreamHandler, StreamRetryConfig, StreamTimeoutConfig, }; @@ -4025,33 +4643,29 @@ mod tests { // Circuit breaker: trips on a single model failure (threshold = 1) and // has a fallback model configured. - let mut runtime = LoopRuntime::new(); - runtime.fallback = FallbackManager::new_with_fallback("primary-model".to_string(), 1); - runtime.fallback.set_fallback_model("fallback-model"); - runtime.set_stream_handler(handler); + let mut managers = LoopManagers::new().with_fallback(FallbackManager::new_with_fallback( + "primary-model".to_string(), + 1, + )); + managers.fallback().set_fallback_model("fallback-model"); + managers.set_stream_handler(handler); let config = make_config(); let client = Arc::new(AlwaysRateLimitClient); - let mut agent = BareLoop::new_with_managers(client, ToolRegistry::new(), config, runtime); + let mut agent = BareLoop::new_with_managers(client, ToolRegistry::new(), config, managers); - let result = agent.run("Hi").await; + let result = agent.run("Hi", &RunConfig::default()).await; assert!(result.is_err(), "rate-limited turn should fail"); // The escalation arm called record_model_failure(); with threshold 1 the // breaker tripped into Fallback state. assert!( - agent.managers.fallback.is_using_fallback(), + agent.managers.fallback().is_using_fallback(), "escalation should trip the circuit breaker to the fallback model" ); } - // ========================================================== - // ContextContributor turn-boundary injection - // ========================================================== // - // A recording client that snapshots the inbound conversation on every - // stream_messages call, so the tests can assert exactly what the model - // received (including contributor-injected messages). #[derive(Clone)] struct RecordingClient { @@ -4240,7 +4854,6 @@ mod tests { } } - // A contributor that always returns the same System reminder. struct StaticReminder(String); impl ContextContributor for StaticReminder { fn contribute(&self, _ctx: &ContributorContext<'_>) -> Option { @@ -4251,7 +4864,6 @@ mod tests { } } - // A contributor that never injects. struct NeverContributor; impl ContextContributor for NeverContributor { fn contribute(&self, _ctx: &ContributorContext<'_>) -> Option { @@ -4259,8 +4871,6 @@ mod tests { } } - // A contributor that counts how many times it was consulted, injecting - // nothing. Shared by the cadence tests. struct CountingContributor { calls: Arc, } @@ -4271,7 +4881,6 @@ mod tests { } } - // A contributor that records the turn number it was consulted at. struct CapturingContributor { seen_turns: Arc>>, } @@ -4282,12 +4891,8 @@ mod tests { } } - // A contributor that counts how many times it was consulted. - fn contributor_config() -> LoopConfig { - LoopConfig { - max_turns: 10, - ..Default::default() - } + fn contributor_config() -> SessionConfig { + SessionConfig::default() } #[tokio::test] @@ -4297,8 +4902,7 @@ mod tests { let config = contributor_config(); let mut agent = BareLoop::new(Arc::new(client.clone()), ToolRegistry::new(), config); agent.add_contributor(Box::new(StaticReminder("stay on task".into()))); - let result = agent.run("Hi").await.unwrap(); - assert!(result.success); + let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); // The injected message reached the model: it's in the captured inbound // conversation after the user message. @@ -4330,8 +4934,7 @@ mod tests { let config = contributor_config(); let mut agent = BareLoop::new(Arc::new(client.clone()), ToolRegistry::new(), config); // No add_contributor call. - let result = agent.run("Hi").await.unwrap(); - assert!(result.success); + let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); let seen = client.first_seen(); // No System messages reached the model. @@ -4351,8 +4954,7 @@ mod tests { let config = contributor_config(); let mut agent = BareLoop::new(Arc::new(client.clone()), ToolRegistry::new(), config); agent.add_contributor(Box::new(NeverContributor)); - let result = agent.run("Hi").await.unwrap(); - assert!(result.success); + let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); let seen = client.first_seen(); assert!( @@ -4369,8 +4971,7 @@ mod tests { let mut agent = BareLoop::new(Arc::new(client.clone()), ToolRegistry::new(), config); agent.add_contributor(Box::new(StaticReminder("first".into()))); agent.add_contributor(Box::new(StaticReminder("second".into()))); - let result = agent.run("Hi").await.unwrap(); - assert!(result.success); + let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); // Find positions of the two reminders in the persisted history. let persisted = agent.conversation(); @@ -4396,14 +4997,22 @@ mod tests { let config = contributor_config(); let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); agent.add_contributor(Box::new(StaticReminder("remind".into()))); - agent.run("Hi").await.unwrap().total_turns + agent + .run("Hi", &RunConfig::default()) + .await + .unwrap() + .turn_count() }; let without_contrib = { let client = RecordingClient::new("test-model"); client.add_text_response("done"); let config = contributor_config(); let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); - agent.run("Hi").await.unwrap().total_turns + agent + .run("Hi", &RunConfig::default()) + .await + .unwrap() + .turn_count() }; assert_eq!( with_contrib, without_contrib, @@ -4421,12 +5030,12 @@ mod tests { let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); let c = Arc::clone(&counter); agent.add_contributor(Box::new(CountingContributor { calls: c })); - agent.run("Hi").await.unwrap(); + agent.run("Hi", &RunConfig::default()).await.unwrap(); // One turn ran; the contributor was consulted once. assert_eq!(counter.load(Ordering::Relaxed), 1); // And the model was called exactly once (proving the single turn). - assert_eq!(agent.budget.total_turns, 1); + assert_eq!(agent.current_run().unwrap().turn_count(), 1); } #[tokio::test] @@ -4443,8 +5052,8 @@ mod tests { let mut agent = BareLoop::new(Arc::new(client), registry, config); let c = Arc::clone(&counter); agent.add_contributor(Box::new(CountingContributor { calls: c })); - let result = agent.run("Echo hi").await.unwrap(); - assert_eq!(result.total_turns, 2, "tool_call turn + end_turn"); + let result = agent.run("Echo hi", &RunConfig::default()).await.unwrap(); + assert_eq!(result.turn_count(), 2, "tool_call turn + end_turn"); assert_eq!( counter.load(Ordering::Relaxed), 2, @@ -4457,20 +5066,21 @@ mod tests { #[should_panic(expected = "configuration setters must be called before run()")] fn test_add_contributor_panics_after_session_start() { let client = MockClient::new("test-model"); + client.add_text_response("ok"); let config = contributor_config(); let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); - // initialize() moves the loop out of Idle — subsequent add_contributor - // must panic in debug builds (matches set_reflector's contract). - let cfg = LoopConfig { - max_turns: 10, - ..Default::default() - }; - // Box the future so we can drop it without awaiting; initialize is the - // state transition under test. + // The first run() establishes the session (capturing the start time + // and firing on_session_start), moving the loop out of Idle. A + // subsequent add_contributor must panic in debug builds (matches + // set_reflector's contract). + // Box the future so we can drop it without awaiting; the session-init + // side effect is the state transition under test. { - let fut = agent.initialize(&cfg); + let run_config = RunConfig::default(); + let fut = agent.run("seed", &run_config); let mut fut = std::pin::pin!(fut); - futures::executor::block_on(fut.as_mut()).unwrap(); + let outcome = futures::executor::block_on(fut.as_mut()); + drop(outcome); } agent.add_contributor(Box::new(StaticReminder("late".into()))); } @@ -4489,23 +5099,17 @@ mod tests { let mut agent = BareLoop::new(Arc::new(client), registry, config); let s = Arc::clone(&seen_turns); agent.add_contributor(Box::new(CapturingContributor { seen_turns: s })); - agent.run("go").await.unwrap(); + agent.run("go", &RunConfig::default()).await.unwrap(); let turns = crate::error::recover_guard(seen_turns.lock()).clone(); assert_eq!(turns, vec![0, 1], "turn numbers are 0-indexed and per-turn"); } - // Suppress unused-warning for RecordingClient::call_count when no test - // references it; kept for future diagnostics. #[allow(dead_code)] fn _suppress_recording_client_dead_code(c: &RecordingClient) { let _ = c.call_count(); } - // ========================================================== - // RequestOptions engine plumbing - // ========================================================== - #[tokio::test] async fn test_request_options_default_is_unconstrained() { // A fresh BareLoop has default RequestOptions — the engine reproduces @@ -4515,7 +5119,7 @@ mod tests { let config = contributor_config(); let mut agent = BareLoop::new(Arc::new(client.clone()), ToolRegistry::new(), config); // No set_request_options call — default path. - agent.run("Hi").await.unwrap(); + agent.run("Hi", &RunConfig::default()).await.unwrap(); let opts = client.first_options(); assert!( @@ -4539,7 +5143,7 @@ mod tests { crate::structured::RequestOptions::new() .with_tool_constraint(crate::structured::ToolConstraint::Strict), ); - agent.run("Hi").await.unwrap(); + agent.run("Hi", &RunConfig::default()).await.unwrap(); let opts = client.first_options(); assert!( @@ -4556,24 +5160,21 @@ mod tests { #[should_panic(expected = "configuration setters must be called before run()")] fn test_set_request_options_panics_after_session_start() { let client = MockClient::new("test-model"); + client.add_text_response("ok"); let config = contributor_config(); let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); - let cfg = LoopConfig { - max_turns: 10, - ..Default::default() - }; + // The first run() establishes the session and moves the loop out of + // Idle; a subsequent set_request_options must panic in debug builds. { - let fut = agent.initialize(&cfg); + let run_config = RunConfig::default(); + let fut = agent.run("seed", &run_config); let mut fut = std::pin::pin!(fut); - futures::executor::block_on(fut.as_mut()).unwrap(); + let outcome = futures::executor::block_on(fut.as_mut()); + drop(outcome); } agent.set_request_options(crate::structured::RequestOptions::default()); } - // ========================================================== - // ConstrainedProfile::apply - // ========================================================== - #[tokio::test] async fn test_constrained_apply_wires_pipeline_and_contributor() { // Apply() sets the small-model pipeline and registers a GoalReminder. To prove @@ -4593,9 +5194,12 @@ mod tests { // Add a cadence-1 reminder so it fires this session. agent.add_contributor(Box::new(crate::presets::GoalReminder::new(1))); - let result = agent.run("ship the demo goal").await.unwrap(); + let result = agent + .run("ship the demo goal", &RunConfig::default()) + .await + .unwrap(); // Tool-call turn + end_turn = 2 turns. - assert!(result.total_turns >= 1); + assert!(result.turn_count() >= 1); // The contributor fired: a Role::System message carrying the first // user message text reached the provider on some turn's outbound @@ -4682,8 +5286,7 @@ mod tests { }); agent.register_observer(recorder as Arc); - let result = agent.run("Hi").await.unwrap(); - assert!(result.success); + let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); let captured = crate::error::recover_guard(captured.lock()); assert_eq!( @@ -4693,7 +5296,7 @@ mod tests { ); let joined: String = captured.iter().map(|(_, d)| d.as_str()).collect(); assert_eq!(joined, "First reasoning chunk"); - assert_eq!(captured[0].0, 0, "turn number matches budget.total_turns"); + assert_eq!(captured[0].0, 0, "turn number matches the run's turn count"); } #[tokio::test] @@ -4754,7 +5357,7 @@ mod tests { }); agent.register_observer(recorder as Arc); - agent.run("Hi").await.unwrap(); + agent.run("Hi", &RunConfig::default()).await.unwrap(); assert_eq!( *crate::error::recover_guard(text_calls.lock()), @@ -4781,9 +5384,8 @@ mod tests { .with_reflector(Arc::new(NoopReflector)) .with_request_options(RequestOptions::default()); - let result = agent.run("Hi").await.unwrap(); + let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); - assert!(result.success, "the fluent-built loop runs a turn"); assert_eq!( observer.turn_starts.load(Ordering::SeqCst), 1, diff --git a/src/engine/bare/compact.rs b/src/engine/bare/compact.rs index 08f13d5..6dc88eb 100644 --- a/src/engine/bare/compact.rs +++ b/src/engine/bare/compact.rs @@ -1,8 +1,12 @@ //! Context compaction for agent conversations. //! -//! When a [`ContextManager`](crate::compact::ContextManager) is configured, -//! checks token usage after each tool dispatch and triggers compaction if the -//! conversation exceeds the context window threshold. +//! The driving state machine decides *when* to compact by emitting +//! [`MachineStep::Compact`](crate::engine::core::MachineStep::Compact); the +//! [`ContextManager`](crate::compact::ContextManager) consulted here decides +//! *how*. [`run_compaction`](BareLoop::run_compaction) runs the configured +//! compactor over the machine-owned history, fires the compaction observer and +//! hook events, and returns the compacted messages plus a post-compaction token +//! estimate for the driver to feed back into the machine. use super::{ApiClient, BareLoop, Instant, LoopError}; #[cfg(feature = "hooks")] @@ -12,42 +16,56 @@ use crate::compact::EnsureContextResult; use crate::capabilities::Compactable; #[cfg(feature = "hooks")] use crate::capabilities::Hookable; +use crate::message::Message; use crate::observer::CompactedContext; impl BareLoop { - /// Check if context compaction is needed and perform it if so. + /// Compact the conversation context owned by the driving machine. /// - /// 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). + /// Run when the machine requests it via + /// [`MachineStep::Compact`](crate::engine::core::MachineStep::Compact). + /// Reads the history from [`LoopMachine::history`](crate::engine::core::LoopMachine::history), + /// asks the configured [`ContextManager`](crate::compact::ContextManager) + /// to reduce it, fires + /// [`on_compaction`](crate::observer::LoopObserver::on_compaction) and the + /// post-compact hook when compaction occurred, and returns the compacted + /// messages alongside the post-compaction token estimate. The caller feeds + /// both back via + /// [`LoopMachine::compaction_result`](crate::engine::core::LoopMachine::compaction_result). /// - /// When no `ContextManager` is set, this is a no-op. + /// When no `ContextManager` is set the history is returned unchanged with a + /// zero token estimate, so the machine can resume without compaction. When + /// a pre-compact hook aborts compaction the history is likewise returned + /// unchanged. /// /// # Errors /// - /// 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<(), LoopError> { + /// Returns [`LoopError::ContextExceeded`] when compaction was required but + /// the compactor could not reduce the history enough to fit the context + /// window. + pub(super) async fn run_compaction( + &mut self, + turn: usize, + reason: crate::compact::types::CompactReason, + ) -> Result<(Vec, u64), LoopError> { + let history = self.machine.history().to_vec(); let Some(ctx_manager) = self.managers.context_manager() else { - return Ok(()); + return Ok((history, 0)); }; - if self.pre_compact_hook_aborts() { - return Ok(()); + if self.pre_compact_hook_aborts(&history) { + return Ok((history, 0)); } - let messages_before = self.conversation.len(); + let messages_before = history.len(); let compact_start = Instant::now(); - let conversation = self.conversation.clone(); - let result = ctx_manager.ensure_context_fits(conversation, turn).await; + let result = ctx_manager.compact_with_reason(history, turn, reason).await; match result { Ok(EnsureContextResult::Compacted(outcome)) => { let tokens_after = outcome.tokens_after; let tokens_saved = outcome.tokens_saved; - self.conversation = outcome.messages; + let messages_after = outcome.messages.len(); let tokens_before = tokens_after.saturating_add(tokens_saved); self.managers.observers().on_compaction(&CompactedContext { tokens_before, @@ -56,16 +74,14 @@ impl BareLoop { }); self.notify_post_compact_hook( messages_before, + messages_after, tokens_after, tokens_saved, compact_start.elapsed(), ); - Ok(()) - } - Ok(EnsureContextResult::NoAction(messages)) => { - self.conversation = messages; - Ok(()) + Ok((outcome.messages, tokens_after)) } + Ok(EnsureContextResult::NoAction(messages)) => Ok((messages, 0)), Err(overflow) => Err(LoopError::ContextExceeded { used: overflow.tokens_used, limit: overflow.context_window, @@ -73,34 +89,58 @@ impl BareLoop { } } - /// Run the pre-compact hook. Returns `true` if the hook aborts compaction. + /// Run the pre-compact hook, returning `true` if any hook aborts. + /// + /// Builds a [`PreCompactContext`] from the current history and + /// session config, then asks the + /// [`HookExecutor`](crate::hooks::HookExecutor) to run every + /// registered `on_pre_compact` hook. If any hook returns + /// [`abort: true`](crate::hooks::context::CompactResult::abort), + /// this returns `true` and [`run_compaction`](Self::run_compaction) + /// skips compaction entirely for this trigger. Returns `false` + /// when there is no hook executor or no hook vetoes the compaction. #[cfg(feature = "hooks")] - fn pre_compact_hook_aborts(&self) -> bool { + fn pre_compact_hook_aborts(&self, history: &[Message]) -> bool { let Some(executor) = self.managers.hook_executor() else { return false; }; - let tokens_before = crate::compact::CompactionOutcome::estimate_tokens(&self.conversation); + let tokens_before = crate::compact::CompactionOutcome::estimate_tokens(history); let ctx = PreCompactContext { trigger: CompactTrigger::Auto, custom_instructions: None, - message_count: self.conversation.len(), + message_count: history.len(), tokens_before, - context_window: self.config.context_window, - session_id: self.config.session_id, + context_window: self.session.config.context_window, + session_id: self.session.id, }; executor.check_pre_compact(&ctx).abort } + /// No-op pre-compact hook check for builds without the `hooks` feature. + /// + /// Always returns `false` so [`run_compaction`](Self::run_compaction) + /// proceeds with compaction unconditionally — there are no hooks to + /// consult. #[cfg(not(feature = "hooks"))] - fn pre_compact_hook_aborts(&self) -> bool { + fn pre_compact_hook_aborts(&self, _history: &[Message]) -> bool { false } /// Notify the post-compact hook that compaction completed. + /// + /// Builds a [`PostCompactContext`] from the before/after message + /// counts, the post-compaction token estimate, the tokens saved, + /// and the wall-clock compaction duration, then fires every + /// registered `on_post_compact` hook via the + /// [`HookExecutor`](crate::hooks::HookExecutor). No-op when there + /// is no hook executor. Called only on the + /// [`EnsureContextResult::Compacted`] path — the `NoAction` path + /// skips it. #[cfg(feature = "hooks")] fn notify_post_compact_hook( &self, messages_before: usize, + messages_after: usize, tokens_after: u64, tokens_saved: u64, duration: std::time::Duration, @@ -108,22 +148,27 @@ impl BareLoop { let Some(executor) = self.managers.hook_executor() else { return; }; - let messages_after = self.conversation.len(); let ctx = PostCompactContext { trigger: CompactTrigger::Auto, messages_compacted: messages_before.saturating_sub(messages_after), tokens_saved, tokens_after, duration_ms: u64::try_from(duration.as_millis()).unwrap_or(0), - session_id: self.config.session_id, + session_id: self.session.id, }; executor.notify_post_compact(&ctx); } + /// No-op post-compact notification for builds without the `hooks` feature. + /// + /// Does nothing — there are no hooks to notify. Kept so + /// [`run_compaction`](Self::run_compaction) compiles identically + /// with and without the feature. #[cfg(not(feature = "hooks"))] fn notify_post_compact_hook( &self, _messages_before: usize, + _messages_after: usize, _tokens_after: u64, _tokens_saved: u64, _duration: std::time::Duration, diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index 3bb9a31..b23bf6f 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -13,6 +13,7 @@ use super::{ }; #[cfg(feature = "hooks")] use super::{PostToolUseContext, PreToolUseContext}; +use crate::capabilities::Detectable; #[cfg(feature = "tool_health")] use crate::capabilities::HealthTrackable; #[cfg(feature = "hooks")] @@ -45,10 +46,6 @@ enum RecoveryOutcome { Cancelled, } -// =================================================== -// Parallel dispatch: dependency graph + dispatch plan -// =================================================== - /// Analysis of a batch of tool calls, classifying each as parallelizable and /// grouping independent calls into waves. /// @@ -201,7 +198,7 @@ impl BareLoop { /// Execute a batch of tool calls and return results in input order. /// /// Routes to the sequential or parallel path based on - /// [`parallel_tool_dispatch`](crate::config::LoopConfig::parallel_tool_dispatch). + /// [`parallel_tool_dispatch`](crate::engine::RunConfig::parallel_tool_dispatch). /// Both paths honour mid-batch cancellation: if the cancel signal fires /// between (or during) calls, the method returns /// [`LoopError::Cancelled`] promptly. @@ -219,7 +216,7 @@ impl BareLoop { tool_calls: &[ToolCall], turn_idx: usize, ) -> Result, LoopError> { - match self.config.parallel_tool_dispatch.mode { + match self.run_config().parallel_tool_dispatch.mode { crate::config::ParallelMode::Parallel => { self.dispatch_tools_parallel(tool_calls, turn_idx).await } @@ -377,7 +374,7 @@ impl BareLoop { let graph = ToolDependencyGraph::from_calls(tool_calls, &self.tools); let plan = graph.plan(); let max_concurrency = self - .config + .run_config() .parallel_tool_dispatch .max_concurrency .clamp(1, eligible.len()); @@ -445,7 +442,8 @@ impl BareLoop { self.post_detection(tc, er); self.fire_tool_post(turn_idx, tc, er); self.notify_post_tool_use_hooks(tc, er, turn_idx); - self.record_tool_health(tc.tool.as_str(), er); + // Health was already recorded inside run_parallel_task on + // every attempt — don't double-count here. results.push(er.clone()); } else { let soft = ToolDispatchResult { @@ -555,7 +553,7 @@ impl BareLoop { { Ok((next_attempt, correction)) => { attempt = next_attempt; - Self::apply_correction_if_present(&mut tc, correction, &tool_result); + Self::apply_correction_if_present(&mut tc, correction); } Err(RecoveryOutcome::SoftError(returned_result)) => return Ok(returned_result), Err(RecoveryOutcome::Cancelled) => return Err(LoopError::Cancelled), @@ -594,6 +592,11 @@ impl BareLoop { Ok(r) => r, Err(e) => return Some(Err(e)), }; + // Record health on every attempt — the registry is thread-safe + // (atomic counters), so this is safe to call from concurrent + // tasks. Detection/observer side-effects stay in the POST phase + // (they are not thread-safe). + self.record_tool_health(tc.tool.as_str(), &tool_result); if !tool_result.is_error { return Some(Ok(tool_result)); } @@ -603,7 +606,7 @@ impl BareLoop { { Ok((next_attempt, correction)) => { attempt = next_attempt; - Self::apply_correction_if_present(&mut tc, correction, &tool_result); + Self::apply_correction_if_present(&mut tc, correction); } Err(RecoveryOutcome::SoftError(returned_result)) => { return Some(Ok(returned_result)); @@ -650,13 +653,9 @@ impl BareLoop { /// The correction modifies the tool call's input before the next retry /// attempt. If the correction cannot be applied, logs a warning and /// proceeds with the original input. - fn apply_correction_if_present( - tc: &mut ToolCall, - correction: Option, - tool_result: &ToolDispatchResult, - ) { + fn apply_correction_if_present(tc: &mut ToolCall, correction: Option) { let Some(correction) = correction else { return }; - if let CorrectionResult::Failed(msg) = tc.apply_correction(&correction, tool_result) { + if let CorrectionResult::Failed(msg) = tc.apply_correction(&correction) { tracing::warn!( tool = %tc.tool, error = %msg, @@ -682,14 +681,13 @@ impl BareLoop { let operation = Operation::from_input_with_signature( &tc.tool, &tc.input, - self.managers.detection.signature(), + self.managers.detection().signature(), ); - let pattern = self.managers.detection.record_operation(operation); + let pattern = self.managers.detection().record_operation(operation); - // Inline loop/convergence detection. Some(error) forces a hard stop — - // propagate it so the agent loop terminates. None means no detection - // fired and dispatch continues. - match self.managers.handle_detected_pattern(&pattern, turn_idx) { + // Notify observers, then decide whether to abort. + self.managers.notify_detected_pattern(&pattern, turn_idx); + match self.decide_detected_pattern(&pattern) { Some(e) => Err(e), None => Ok(None), } @@ -708,9 +706,9 @@ impl BareLoop { &tc.tool, &tc.input, result_hash, - self.managers.detection.signature(), + self.managers.detection().signature(), ); - self.managers.detection.record_operation(operation); + self.managers.detection().record_operation(operation); } /// Execute a single tool call. @@ -876,7 +874,7 @@ impl BareLoop { let ctx = PreToolUseContext { tool_name: tc.tool.clone(), input: tc.input.clone(), - session_id: self.config.session_id, + session_id: self.session.id, turn_number: turn_idx, }; match executor.check_pre_tool_use(&ctx) { @@ -944,7 +942,7 @@ impl BareLoop { .as_millis() .try_into() .unwrap_or(u64::MAX), - session_id: self.config.session_id, + session_id: self.session.id, turn_number: turn_idx, }; executor.notify_post_tool_use(&ctx); @@ -1098,8 +1096,9 @@ impl BareLoop { #[allow(clippy::unnecessary_literal_bound)] mod tests { use crate::api::error::ApiError; - use crate::config::LoopConfig; - use crate::engine::loop_core::ToolCall; + use crate::config::SessionConfig; + use crate::engine::core::ToolCall; + use crate::engine::{Run, RunConfig}; use crate::message::ToolContent; use crate::tool::{ Tool, ToolContext, ToolError, ToolOutput, ToolSchema, registry::ToolRegistry, @@ -1191,7 +1190,7 @@ mod tests { } fn make_loop(tools: ToolRegistry) -> BareLoop { - let config = LoopConfig::default(); + let config = SessionConfig::default(); let client = Arc::new(MockClient::new("test")); BareLoop::new(client, tools, config) } @@ -1425,10 +1424,17 @@ mod tests { } fn make_parallel_loop(tools: ToolRegistry) -> BareLoop { - let mut config = LoopConfig::default(); - config.parallel_tool_dispatch.mode = crate::config::ParallelMode::Parallel; let client = Arc::new(MockClient::new("test")); - BareLoop::new(client, tools, config) + let run_config = RunConfig { + parallel_tool_dispatch: crate::config::ParallelDispatchConfig { + mode: crate::config::ParallelMode::Parallel, + ..Default::default() + }, + ..RunConfig::default() + }; + let mut bare = BareLoop::new(client, tools, SessionConfig::default()); + bare.session.runs.push(Run::new("", &run_config)); + bare } #[tokio::test] @@ -1537,8 +1543,8 @@ mod tests { #[tokio::test] async fn parallel_sequential_fallback() { use crate::testing::MockTool; - let config = LoopConfig::default(); - // Default is Sequential — no change needed. + let config = SessionConfig::default(); + // Default dispatch mode is Sequential — no change needed. let mut registry = ToolRegistry::new(); registry.register( MockTool::new("a", "a") diff --git a/src/engine/bare/emission.rs b/src/engine/bare/emission.rs index 65ef617..49d7e5e 100644 --- a/src/engine/bare/emission.rs +++ b/src/engine/bare/emission.rs @@ -2,9 +2,9 @@ //! //! Fires observer callbacks and hook notifications when a session begins and //! ends. Other observer events (`on_turn_start`, `on_response`, etc.) are fired -//! directly at their call sites in `process_turn`. +//! directly at their call sites in the driver loop. -use super::{ApiClient, BareLoop, Duration, SessionResult}; +use super::{ApiClient, BareLoop, Duration, Run}; #[cfg(feature = "hooks")] use crate::capabilities::Hookable; #[cfg(feature = "hooks")] @@ -14,69 +14,83 @@ use crate::hooks::context::{ }; use crate::observer::{SessionEndContext, SessionStartContext}; -// ================================================== -// Session lifecycle notifications -// ================================================== - impl BareLoop { /// Notify all observers and hooks that the session has started. + /// + /// Fires [`on_session_start`](crate::observer::LoopObserver::on_session_start) + /// on every registered observer with the session id, then fires the + /// `on_session_start` hook (when the `hooks` feature is enabled and a + /// hook executor is configured). Called once at the beginning of a + /// run, before the first turn. pub(super) fn notify_session_start(&self) { self.managers .observers() .on_session_start(&SessionStartContext { - session_id: self.config.session_id, + session_id: self.session.id, }); self.notify_session_start_hook(); } /// Notify all observers and hooks that the session has ended. - pub(super) fn notify_session_end(&self, result: &SessionResult, duration: Duration) { + /// + /// Fires [`on_session_end`](crate::observer::LoopObserver::on_session_end) + /// on every registered observer, then fires the `on_session_end` hook + /// (when the `hooks` feature is enabled). The [`Run`] supplies the + /// per-run totals (turn count, tokens); `success` distinguishes a + /// normal exit from a failure so observers and hooks can branch on + /// the outcome, and `duration` is the wall-clock run length. + pub(super) fn notify_session_end(&self, result: &Run, success: bool, duration: Duration) { self.managers .observers() .on_session_end(&SessionEndContext { - success: result.success, - error: result.error.clone(), - total_turns: result.total_turns, + success, + error: if success { + None + } else { + Some("run failed".to_string()) + }, + total_turns: result.turn_count(), duration_ms: Self::millis_u64(duration), }); - self.notify_session_end_hook(result, duration); + self.notify_session_end_hook(result, success, duration); } /// Derive the structured [`SessionEndReason`] from the loop's /// terminal state. /// - /// Unlike a simple `success` boolean, this distinguishes - /// cancellation, max-turns exhaustion, and context overflow. + /// Maps the boolean `success` plus the cancellation flag and the + /// turn cap into the richer hook-level enum: cancellation wins over + /// every other outcome, then failure, then the turn cap, finally + /// normal completion. Used only to populate the hook end-context — + /// observers receive the simpler boolean `success` field. #[cfg(feature = "hooks")] fn session_end_reason(&self, success: bool) -> SessionEndReason { if self.is_cancelled() { SessionEndReason::Cancelled } else if !success { - if self - .budget - .error - .as_ref() - .is_some_and(|e| e.contains("context") || e.contains("overflow")) - { - SessionEndReason::ContextOverflow - } else { - SessionEndReason::Error - } - } else if self.budget.total_turns >= self.config.max_turns { + SessionEndReason::Error + } else if self.current_run().map_or(0, Run::turn_count) >= self.run_config().max_turns { SessionEndReason::MaxTurns } else { SessionEndReason::Complete } } + /// Fire the `on_session_start` hook when a hook executor is + /// configured. + /// + /// Builds a [`HookSessionStartContext`] from the session id, the + /// client's current model, and the process working directory, then + /// dispatches it to every registered session-start hook. No-op when + /// no hook executor is set. #[cfg(feature = "hooks")] fn notify_session_start_hook(&self) { let Some(executor) = self.managers.hook_executor() else { return; }; let ctx = HookSessionStartContext { - session_id: self.config.session_id, - model: self.config.model.clone(), + session_id: self.session.id, + model: self.client.model(), working_directory: std::env::current_dir() .map(|p| p.to_string_lossy().into_owned()) .unwrap_or_default(), @@ -84,29 +98,54 @@ impl BareLoop { executor.notify_session_start(&ctx); } + /// No-op session-start hook for builds without the `hooks` feature. + /// + /// Does nothing — there are no hooks to notify. Kept so + /// [`notify_session_start`](Self::notify_session_start) compiles + /// identically with and without the feature. #[cfg(not(feature = "hooks"))] fn notify_session_start_hook(&self) {} + /// Fire the `on_session_end` hook when a hook executor is + /// configured. + /// + /// Derives the [`SessionEndReason`] via + /// [`session_end_reason`](Self::session_end_reason), then builds a + /// [`HookSessionEndContext`] from the session id, reason, turn + /// count, total tokens, and run duration, and dispatches it to + /// every registered session-end hook. No-op when no hook executor + /// is set. #[cfg(feature = "hooks")] - fn notify_session_end_hook(&self, result: &SessionResult, duration: Duration) { + fn notify_session_end_hook(&self, result: &Run, success: bool, duration: Duration) { let Some(executor) = self.managers.hook_executor() else { return; }; - let reason = self.session_end_reason(result.success); + let reason = self.session_end_reason(success); let ctx = HookSessionEndContext { - session_id: result.session_id, + session_id: self.session.id, reason, - total_turns: result.total_turns, - total_tokens: result.input_tokens.saturating_add(result.output_tokens), + total_turns: result.turn_count(), + total_tokens: result.total_tokens(), duration_secs: duration.as_secs(), }; executor.notify_session_end(&ctx); } + /// No-op session-end hook for builds without the `hooks` feature. + /// + /// Does nothing — there are no hooks to notify. Kept so + /// [`notify_session_end`](Self::notify_session_end) compiles + /// identically with and without the feature. #[cfg(not(feature = "hooks"))] - fn notify_session_end_hook(&self, _result: &SessionResult, _duration: Duration) {} + fn notify_session_end_hook(&self, _result: &Run, _success: bool, _duration: Duration) {} - /// Convert a [`Duration`] to milliseconds as `u64`. + /// Convert a [`Duration`] to milliseconds as a `u64`. + /// + /// Saturates at `u64::MAX` if the duration exceeds the `u64` range + /// (only possible with a platform-specific `u128` millisecond count + /// far beyond any realistic run length). Used to populate the + /// observer/hook end-context `duration_ms` fields from a + /// [`Duration`]. pub(super) fn millis_u64(duration: Duration) -> u64 { u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) } diff --git a/src/engine/bare/message.rs b/src/engine/bare/message.rs index d9e5187..7e09690 100644 --- a/src/engine/bare/message.rs +++ b/src/engine/bare/message.rs @@ -5,49 +5,10 @@ //! assistant response, and computing token counts. use super::{ - ApiClient, BareLoop, Message, MessagePart, Role, ToolCall, ToolContext, ToolDispatchResult, - ToolSchema, + ApiClient, BareLoop, Message, MessagePart, Role, ToolContext, ToolDispatchResult, ToolSchema, }; impl BareLoop { - /// Extract all text content from a message. - /// - /// Iterates over the message's [`MessagePart`]s and concatenates - /// every text part into a single `String`. Non-text parts (e.g. - /// `ToolCall`, `ToolResult`) are silently skipped. - /// - /// Used to produce the [`SessionResult::final_output`] string when - /// the model ends its turn. - pub(super) fn extract_text(msg: &Message) -> String { - msg.parts - .iter() - .filter_map(|b| b.as_text()) - .collect::>() - .join("") - } - - /// Extract tool call information from a message. - /// - /// Scans the message's [`MessagePart`]s for `ToolCall` variants and - /// 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 { - msg.parts - .iter() - .filter_map(|part| match part { - MessagePart::ToolCall { id, name, input } => Some(ToolCall { - id: id.clone(), - tool: name.clone(), - input: input.clone(), - }), - _ => None, - }) - .collect() - } - /// Build the tool result message from executed tool results. /// /// Per API convention, tool results are sent as a **user** message @@ -96,7 +57,7 @@ impl BareLoop { /// enclosing session (e.g. for logging, tracing, or storage). pub(super) fn build_tool_context(&self) -> ToolContext { ToolContext { - session_id: self.config.session_id, + session_id: self.session.id, ..ToolContext::default() } } diff --git a/src/engine/bare/stream.rs b/src/engine/bare/stream.rs index 0a4c1b2..6c4d557 100644 --- a/src/engine/bare/stream.rs +++ b/src/engine/bare/stream.rs @@ -9,7 +9,7 @@ //! retry, timeout, fallback, and rate-limit handling. use super::{ - ApiClient, BareLoop, LoopError, Message, StreamAccumulator, StreamEvent, StreamStopReason, + ApiClient, BareLoop, LoopError, Message, Run, StreamAccumulator, StreamEvent, StreamStopReason, Usage, }; use crate::capabilities::StreamCapable; @@ -53,8 +53,8 @@ impl BareLoop { let handler = self.managers.stream_handler(); let mut stream = handler.stream_turn( &*self.client, - crate::api::StreamRequest::new(self.conversation.clone()) - .with_system(self.config.system_prompt.clone()) + crate::api::StreamRequest::new(self.machine.history().to_vec()) + .with_system(self.session.config.system_prompt.clone()) .with_tools(self.build_tool_schemas()), self.request_options.clone(), &self.cancelled, @@ -114,7 +114,7 @@ impl BareLoop { streamer(text); } self.managers.observers().on_text_delta(&TextDeltaContext { - turn: self.budget.total_turns, + turn: self.current_run().map_or(0, Run::turn_count), delta: text.clone(), }); } @@ -124,7 +124,7 @@ impl BareLoop { self.managers .observers() .on_thinking_delta(&ThinkingDeltaContext { - turn: self.budget.total_turns, + turn: self.current_run().map_or(0, Run::turn_count), delta: text.clone(), }); } diff --git a/src/engine/core.rs b/src/engine/core.rs new file mode 100644 index 0000000..b727ab1 --- /dev/null +++ b/src/engine/core.rs @@ -0,0 +1,14 @@ +//! Agent core: the [`Loop`] lifecycle trait, foundational +//! lifecycle data types, and the sans-IO [`LoopMachine`] +//! state machine. +//! +//! See [`lifecycle`] for the trait and value types, and [`machine`] for the +//! state machine. All public types are re-exported here so the paths +//! `crate::engine::core::Run`, `crate::engine::core::LoopMachine`, etc. remain +//! stable. + +pub mod lifecycle; +pub mod machine; + +pub use lifecycle::*; +pub use machine::*; diff --git a/src/engine/core/lifecycle.rs b/src/engine/core/lifecycle.rs new file mode 100644 index 0000000..90a0a49 --- /dev/null +++ b/src/engine/core/lifecycle.rs @@ -0,0 +1,684 @@ +//! Foundational lifecycle types and the core agent trait. +//! +//! Fundamental operations every agent must support, plus the data types that +//! flow through the agent lifecycle: per-run configuration +//! ([`RunConfig`]), lifecycle state ([`MachineState`]), +//! turn and run results ([`Run`]), and tool call representations +//! ([`ToolCall`]). +//! +//! # Lifecycle +//! +//! ```text +//! run(input, run_config) +//! → should_continue() → true [repeated] +//! → should_continue() → false +//! finalize() +//! ``` +//! +//! # Implementing +//! +//! At a minimum you must provide [`run`](Loop::run), +//! [`should_continue`](Loop::should_continue), +//! [`finalize`](Loop::finalize), +//! [`state`](Loop::state), and +//! [`cancel`](Loop::cancel). +//! +//! # Quick Start +//! +//! ``` +//! use loopctl::config::SessionConfig; +//! use loopctl::engine::core::{MachineState, Run}; +//! +//! let _config = SessionConfig::default(); +//! +//! let run = Run::new("Task done.", &Default::default()); +//! assert_eq!(run.turn_count(), 0); +//! ``` + +use std::future::Future; +use std::pin::Pin; +use std::time::{Duration, Instant}; + +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::config::ParallelDispatchConfig; +use crate::engine::core::machine::MachineState; +use crate::error::LoopError; +use crate::message::Message; +use crate::reflection::{Correction, CorrectionResult, CorrectionType}; + +/// Provide a fresh `Instant` for serde-deserialized fields. +/// +/// `Instant` has no `Default` (its zero point is unspecified), so +/// `#[serde(skip, default = "instant_now")]` is used on the +/// `Instant`-typed fields of [`Run`] to give deserialized runs a +/// sensible monotonic starting point. +fn instant_now() -> Instant { + Instant::now() +} + +/// Per-run configuration. +/// +/// The slice of agent configuration that varies across `run()` calls on the +/// same agent (turn/token budgets, compaction policy, dispatch mode). It is +/// owned by the machine so that the machine's decisions are self-contained and +/// a serialized machine carries its configuration with it. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RunConfig { + /// Maximum number of turns before forcing completion. + /// + /// A safety cap on runaway loops: once the machine has taken this many + /// turns without the model finishing, the run ends with + /// `MachineOutcome::MaxTurnsExceeded`. Defaults to `200`. Set it per-run + /// to bound long-running tool loops. + pub max_turns: usize, + + /// How independent tool calls within a single turn are dispatched. + /// + /// Controls whether the driver runs a turn's independent tool calls + /// sequentially or in parallel, up to a concurrency limit. See + /// [`ParallelDispatchConfig`]. + pub parallel_tool_dispatch: ParallelDispatchConfig, +} + +impl Default for RunConfig { + fn default() -> Self { + Self { + max_turns: 200, + parallel_tool_dispatch: ParallelDispatchConfig::default(), + } + } +} + +/// 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)] +#[non_exhaustive] +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 treats this as a completed run. + EndTurn, + + /// The model requested tool execution. + /// + /// The LLM response contains one or more tool calls. The framework should + /// dispatch them and feed the results back. + ToolCall, + + /// The maximum token limit was reached. + /// + /// The LLM hit the model's 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, +} + +/// 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 tool dispatch result with the output. +/// +/// # Serialization +/// +/// Implements `Serialize` and `Deserialize` for persistence and inter-process +/// communication. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolCall { + /// The call identifier assigned by the LLM API. + /// + /// Used to correlate this request with the tool-result message that the + /// framework produces after dispatching the tool. Unique within a single + /// model response. + pub id: String, + + /// The name of the tool to invoke. + /// + /// Must match a tool registered in the + /// [`ToolRegistry`](crate::tool::ToolRegistry). If the name is not + /// registered, the machine classifies the call as unknown and attaches a + /// preresolved error result instead of dispatching. + pub tool: String, + + /// The JSON input arguments for the tool. + /// + /// A free-form JSON value whose shape depends on the tool's input schema. + /// The framework passes this verbatim to the tool's `call` method; the + /// tool is responsible for validating and parsing it. + pub input: serde_json::Value, +} + +impl ToolCall { + /// Apply a [`Correction`] from the reflection system in place. + /// + /// Modifies `self` according to the correction strategy: + /// + /// - [`InputFix`](CorrectionType::InputFix) — replaces + /// `self.input` with the corrected input. + /// - [`ToolChange`](CorrectionType::ToolChange) — replaces + /// `self.tool` with an alternative tool name. + /// - Other types — no mutation; the retry proceeds with unchanged parameters. + /// + /// Returns a [`CorrectionResult`] indicating whether the correction + /// was applied, failed, or skipped. + pub fn apply_correction(&mut self, correction: &Correction) -> CorrectionResult { + match correction.correction_type { + CorrectionType::InputFix => { + if let Some(ref modified) = correction.modified_input { + if modified.is_object() { + tracing::debug!( + tool = %self.tool, + "applying InputFix correction from reflector" + ); + self.input = modified.clone(); + CorrectionResult::Applied + } else { + CorrectionResult::Failed( + "InputFix correction modified_input must be a JSON object".to_string(), + ) + } + } 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.tool, + new_tool = %alt, + "applying ToolChange correction from reflector" + ); + self.tool.clone_from(alt); + CorrectionResult::Applied + } else { + CorrectionResult::Failed( + "ToolChange correction missing alternative_tool".to_string(), + ) + } + } + CorrectionType::PrerequisiteFix | CorrectionType::ApproachChange => { + CorrectionResult::Skipped + } + CorrectionType::Escalate => CorrectionResult::Skipped, + } + } +} + +/// The result of one `run(prompt)` call. +/// +/// One prompt → tool loop → final answer. Fresh per call: `id` and the +/// accounting fields rotate with each run, while the owning +/// [`Session`] keeps its identity stable. +/// +/// `Run` is a type alias for [`Run`], so the two are interchangeable in +/// ``` +/// use loopctl::engine::core::Run; +/// +/// fn summarize(result: &Run) -> String { +/// format!("{} turns", result.turn_count()) +/// } +/// +/// let run: Run = Run::new("hello", &Default::default()); +/// assert_eq!(summarize(&run), "0 turns"); +/// ``` +/// One iteration of the agent loop — a single LLM call and any tools it +/// triggered. +/// +/// Each entry in [`Run::turns`] records what happened during one loop +/// iteration: the model's response text, any tool calls it requested, and the +/// token cost of that call. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Turn { + /// The 0-indexed position of this turn within its run. + /// + /// Counts from zero within the owning [`Run`]. Useful for + /// correlating a turn with observer events and logs. + pub turn: usize, + + /// What went into this turn's model call. + /// + /// For turn 0 this is the user's prompt. For subsequent turns it's the + /// tool-result text from the previous turn's dispatch — the context the + /// model sees when continuing the loop. + pub input: String, + + /// The text the model produced (concatenated text parts). + /// + /// May be empty if the turn was a pure tool-call response with no + /// accompanying text. On the final turn (no tool calls) this is the + /// run's answer. + pub output: String, + + /// The tool calls the model requested during this turn, if any. + /// + /// Empty when the turn ended with a plain text response. Each entry is + /// one tool invocation the loop dispatched (or would have dispatched). + pub tool_calls: Vec, + + /// Input tokens reported by the provider for this turn. + /// + /// The prompt-side token count — the size of the context the model read + /// to produce this turn's response. + pub input_tokens: u64, + + /// Output tokens reported by the provider for this turn. + /// + /// The completion-side token count — the size of the response the model + /// produced for this turn. + pub output_tokens: u64, +} + +/// One prompt → loop → final answer. +/// +/// Owned by a [`Session`]; a session accumulates one `Run` per `run()` call. +/// The turn list carries per-turn detail; aggregate totals are derived from it. +/// +/// # Example +/// +/// ``` +/// use loopctl::engine::core::Run; +/// +/// let mut run = Run::new("What is 2+2?", &Default::default()); +/// run.output = Some("4".to_string()); +/// +/// assert_eq!(run.output.as_deref(), Some("4")); +/// ``` +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Run { + /// A fresh identifier unique to this run. + /// + /// Minted when the run begins; differs across successive `run()` calls on + /// the same session, so a caller can tell runs apart. + pub id: Uuid, + + /// When this run started. + /// + /// Captured at the top of the run. Combined with [`end`](Self::end) by + /// [`duration`](Self::duration) to report the run's wall-clock span. + /// Skipped during serialization (monotonic `Instant` is not portable); + /// deserialized runs get a fresh `Instant::now()`. + #[serde(skip, default = "instant_now")] + pub start: Instant, + + /// When this run ended. + /// + /// `None` while the run is still in flight, `Some` once it terminates + /// (completed, cancelled, or failed). Skipped during serialization; + /// deserialized runs get `None` (treated as still in-flight). + #[serde(skip)] + pub end: Option, + + /// The turns executed during this run, in order. + /// + /// Each entry records one model call and any tools it triggered: the + /// input text the model saw, the output it produced, the tool calls it + /// requested, and the token cost of that call. Derived totals + /// ([`turn_count`](Self::turn_count), [`input_tokens`](Self::input_tokens), + /// [`output_tokens`](Self::output_tokens), + /// [`tool_call_count`](Self::tool_call_count)) aggregate over this list. + pub turns: Vec, + + /// The user prompt that started this run. + /// + /// The exact text passed to `run()`, recorded for replay and reporting. + /// This is the `input` of the first [`Turn`] in [`turns`](Self::turns). + pub input: String, + + /// The run's final output, if it completed. + /// + /// The assistant's final answer on a successful run; `None` while in flight + /// or when the run ended without producing a final message. For a + /// completed run this matches the `output` of the last + /// [`Turn`](Self::turns) that carried no tool calls. + pub output: Option, + + /// The per-run configuration that governed this run. + /// + /// The turn budget, dispatch mode, and compaction policy passed to the + /// `run()` call that started this run. Captured so a serialized run + /// carries its governing config. + pub config: RunConfig, +} + +impl Run { + /// Begin a fresh run for the given `input`, governed by `config`. + /// + /// Mints a new [`id`](Self::id), captures [`start`](Self::start) + /// as now, leaves [`end`](Self::end) as `None` (the run is in flight), + /// stores the per-run [`config`](Self::config), and zeroes all accounting + /// fields. The framework fills in turns, token totals, and the terminal + /// outcome as the run progresses. + /// + /// # Example + /// + /// ``` + /// use loopctl::engine::core::Run; + /// use loopctl::engine::RunConfig; + /// + /// let run = Run::new("Hello, agent!", &RunConfig::default()); + /// assert_eq!(run.input, "Hello, agent!"); + /// assert_eq!(run.turn_count(), 0); + /// assert!(run.end.is_none()); + /// ``` + #[must_use] + pub fn new(input: impl Into, config: &RunConfig) -> Self { + Self { + id: Uuid::new_v4(), + start: Instant::now(), + end: None, + turns: Vec::new(), + input: input.into(), + output: None, + config: config.clone(), + } + } + + /// Number of turns executed during this run. + /// + /// Equivalent to the length of [`turns`](Self::turns). Each turn + /// corresponds to one model call. + #[must_use] + pub fn turn_count(&self) -> usize { + self.turns.len() + } + + /// Total input tokens consumed across all turns. + /// + /// Sums the `input_tokens` field of every [`Turn`] in + /// [`turns`](Self::turns) — the total prompt-side cost of the run. + #[must_use] + pub fn input_tokens(&self) -> u64 { + self.turns.iter().map(|t| t.input_tokens).sum() + } + + /// Total output tokens consumed across all turns. + /// + /// Sums the `output_tokens` field of every [`Turn`] in + /// [`turns`](Self::turns) — the total completion-side cost of the run. + #[must_use] + pub fn output_tokens(&self) -> u64 { + self.turns.iter().map(|t| t.output_tokens).sum() + } + + /// Total number of tool calls dispatched across all turns. + /// + /// Sums the tool-call count of every [`Turn`] in + /// [`turns`](Self::turns). + #[must_use] + pub fn tool_call_count(&self) -> usize { + self.turns.iter().map(|t| t.tool_calls.len()).sum() + } + + /// Wall-clock duration of this run. + /// + /// Elapsed time from [`start`](Self::start) to [`end`](Self::end). + /// While the run is in flight (`end` is `None`), counts up from + /// `start` to now. + #[must_use] + pub fn duration(&self) -> Duration { + match self.end { + Some(end) => end.saturating_duration_since(self.start), + None => self.start.elapsed(), + } + } + + /// Total tokens (input + output) consumed by this run. + /// + /// Convenience: sums [`input_tokens`](Self::input_tokens) and + /// [`output_tokens`](Self::output_tokens) with saturating addition. + #[must_use] + pub fn total_tokens(&self) -> u64 { + self.input_tokens().saturating_add(self.output_tokens()) + } +} + +impl Default for Run { + fn default() -> Self { + Self::new(String::new(), &RunConfig::default()) + } +} + +/// One agent identity, stable across the process lifetime. +/// +/// A multi-turn REPL is one `Session` made of many [`Run`]s: the `session_id` +/// and start time are stable across `run()` calls, while each run rotates its +/// own identity and accounting. Per-session totals are derived from the run +/// list, never stored separately. +/// +/// # Example +/// +/// ``` +/// use loopctl::config::SessionConfig; +/// use loopctl::engine::core::{Run, Session}; +/// +/// let config = SessionConfig::default(); +/// let mut session = Session::new(config); +/// session.runs.push(Run::new("first prompt", &Default::default())); +/// session.runs.push(Run::new("second prompt", &Default::default())); +/// assert_eq!(session.runs.len(), 2); +/// ``` +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Session { + /// Unique session identifier. + /// + /// A random UUID v4 generated on construction. Two sessions with the same + /// `id` are considered the same logical session. Stable across `run()` + /// calls on the same agent. + pub id: Uuid, + + /// Session-scoped configuration. + /// + /// The system prompt, context window, and compaction settings — set once + /// at construction and unchanged across `run()` calls. See + /// [`SessionConfig`](crate::config::SessionConfig). + pub config: crate::config::SessionConfig, + + /// When the session started, or `None` before the first `run()` call. + /// + /// Captured once on the first run; unlike a run's start time it does not + /// rotate between `run()` calls. Skipped during serialization (monotonic + /// `Instant` is not portable); deserialized sessions get `None`. + #[serde(skip)] + pub session_start: Option, + + /// The runs executed in this session, in order. + /// + /// One [`Run`] per `run()` call. The session's derived totals + /// (turns, tokens, duration) sum over this list. The last entry is the + /// in-flight run (or the most recently completed one). + pub runs: Vec, + + /// TODO: add docs + pub history: Vec, +} + +impl Session { + /// Create a new session with a fresh random id and the given config. + /// + /// The session starts with `session_start = None` (set on the first + /// `run()` call) and an empty run list. + #[must_use] + pub fn new(config: crate::config::SessionConfig) -> Self { + Self { + id: Uuid::new_v4(), + config, + session_start: None, + runs: Vec::new(), + history: Vec::new(), + } + } + + /// Borrow the most recent run (the in-flight or just-completed one). + /// + /// Returns `None` when no `run()` call has been made yet. + #[must_use] + pub fn current_run(&self) -> Option<&Run> { + self.runs.last() + } + + /// Mutably borrow the most recent run (the in-flight or just-completed + /// one). + /// + /// Returns `None` when no `run()` call has been made yet. + #[must_use] + pub fn current_run_mut(&mut self) -> Option<&mut Run> { + self.runs.last_mut() + } + + /// Total turns across every run in this session. + /// + /// The sum of each [`Run`]'s turn count. Derived from `runs`, so it stays + /// correct as runs are added or mutated. + #[must_use] + pub fn total_turns(&self) -> usize { + self.runs.iter().map(Run::turn_count).sum() + } + + /// Total wall-clock duration across every run in this session. + /// + /// The sum of each [`Run::duration`]. + #[must_use] + pub fn total_duration(&self) -> Duration { + self.runs.iter().map(Run::duration).sum() + } + + /// Total input tokens across every run in this session. + /// + /// The sum of each [`Run::input_tokens`]. + #[must_use] + pub fn total_input_tokens(&self) -> u64 { + self.runs.iter().map(Run::input_tokens).sum() + } + + /// Total output tokens across every run in this session. + /// + /// The sum of each [`Run::output_tokens`]. + #[must_use] + pub fn total_output_tokens(&self) -> u64 { + self.runs.iter().map(Run::output_tokens).sum() + } +} + +/// 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::core::Loop; +/// use loopctl::engine::RunConfig; +/// use loopctl::engine::core::{MachineState, Run}; +/// use loopctl::error::LoopError; +/// +/// struct MyAgent { +/// state: MyState, +/// } +/// +/// impl Loop for MyAgent { +/// fn run<'a>( +/// &'a mut self, +/// input: &'a str, +/// run_config: &'a RunConfig, +/// ) -> Pin + Send + 'a>> { +/// Box::pin(async { Ok(loopctl::engine::core::Run::new(input, &Default::default())) }) +/// } +/// fn should_continue(&self) -> bool { +/// !self.state.is_complete +/// } +/// fn finalize<'a>(&'a mut self) -> Pin + Send + 'a>> { +/// Box::pin(async { Ok(loopctl::engine::core::Run::new("", &Default::default())) }) +/// } +/// fn state(&self) -> MachineState { +/// MachineState::Start +/// } +/// fn cancel(&self) {} +/// } +/// ``` +/// The result of a `run()` call — either the completed [`Run`] or a [`LoopError`]. +pub type RunResult = Result; + +/// The core agent lifecycle trait. +/// +/// Implement this trait to create a new type of agent. +pub trait Loop: Send + Sync { + /// Drive one run of the agent loop for the given user prompt. + /// + /// This is the main entry point for running an agent. Session-scoped state + /// (identity, start time, system prompt) is established once at + /// construction and stays stable across calls; each `run()` receives a + /// fresh [`RunConfig`] for the per-run budget and policy, mints a new + /// [`Run`], and drives the turn loop until + /// the model finishes or the budget is exhausted. + /// + /// `input` is the prompt for this run. It is passed to the first turn; + /// continuation turns receive an empty string (tool results are already in + /// the conversation history). + /// + /// # Errors + /// + /// - [`LoopError::Cancelled`] — if the session was cancelled. + /// - [`LoopError::MaxTurnsExceeded`] — if the turn limit was reached. + /// - Any error returned by the turn loop or [`finalize`](Self::finalize). + fn run<'a>( + &'a mut self, + input: &'a str, + run_config: &'a RunConfig, + ) -> 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 run and produce its result. + /// + /// Called once after the last turn. Use this to clean up resources + /// and produce a final [`Run`]. + fn finalize<'a>( + &'a mut self, + success: bool, + ) -> 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) -> MachineState; + + /// 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 + /// [`run`](Loop::run) and [`should_continue`](Loop::should_continue) + /// can observe it and return promptly across threads. + fn cancel(&self); + + /// Explain *why* [`should_continue`](Self::should_continue) returned `false`. + /// + /// Return `None` for normal completion (the model finished) or `Some(err)` + /// when the session was forced to stop (`Cancelled`, `MaxTurnsExceeded`, + /// etc.). The default implementation returns `None` (normal completion). + fn stop_reason(&self) -> Option { + None + } +} diff --git a/src/engine/machine.rs b/src/engine/core/machine.rs similarity index 59% rename from src/engine/machine.rs rename to src/engine/core/machine.rs index df99f28..7440f47 100644 --- a/src/engine/machine.rs +++ b/src/engine/core/machine.rs @@ -1,96 +1,19 @@ -//! Sans-IO agent-loop state machine. +//! The sans-IO agent-loop state machine. //! -//! [`LoopMachine`] owns every *decision* the agent loop makes — turn counting, +//! [`LoopMachine`] owns every *decision* the loop makes — turn counting, //! max-turn enforcement, tool-call validity, stop-reason routing, compaction //! triggering, history accumulation, and the cancellation flag — and performs -//! zero IO. It is advanced by feeding it outcomes via its methods; a thin IO -//! driver (the `BareLoop` in [`crate::engine`]) runs a `match -//! machine.next_step()` loop, performs the side effect each step requests, and -//! feeds the result back. -//! -//! Because the machine is pure and [`Serialize`] + [`Deserialize`], a run can be -//! serialized mid-flight and resumed in another process. -//! -//! [`Serialize`]: serde::Serialize -//! [`Deserialize`]: serde::Deserialize +//! zero IO. The machine is policy-free and serializes as pure state; the turn +//! budget and compaction thresholds are passed in fresh on each +//! [`next_step`](LoopMachine::next_step) call via [`MachinePolicy`]. use serde::{Deserialize, Serialize}; +use super::lifecycle::{StopReason, ToolCall}; use crate::compact::types::CompactReason; -use crate::config::ParallelDispatchConfig; -use crate::engine::loop_core::{StopReason, ToolCall}; use crate::error::LoopError; use crate::message::{Message, MessagePart, Role, ToolContent}; -/// Per-run configuration. -/// -/// The slice of agent configuration that varies across `run()` calls on the -/// same agent (turn/token budgets, compaction policy, dispatch mode). It is -/// owned by the machine so that the machine's decisions are self-contained and -/// a serialized machine carries its configuration with it. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct RunConfig { - /// Maximum number of turns before forcing completion. - /// - /// A safety cap on runaway loops: once the machine has taken this many - /// turns without the model finishing, the run ends with - /// [`MachineOutcome::MaxTurnsExceeded`]. Defaults to `200`. Set it per-run - /// to bound long-running tool loops. - pub max_turns: usize, - - /// Maximum tokens for each API response. - /// - /// The per-response cap the driver sends to the provider, limiting the - /// length of a single model reply. Defaults to `16_384`. - pub max_tokens: u32, - - /// Context window size in tokens. - /// - /// The hard upper bound on tokens the model accepts in one request. The - /// machine uses it as the denominator for the compaction trigger: the - /// [`MachineStep::Compact`] step fires once estimated context size reaches - /// [`compact_threshold`](Self::compact_threshold) of this window. Defaults - /// to `200_000`; should match the configured model's real window. - pub context_window: u64, - - /// Threshold to trigger auto-compaction, in basis points (`0–10_000`; - /// `10_000` = 100% of the context window). - /// - /// The machine emits [`MachineStep::Compact`] once the current context size - /// exceeds this fraction of [`context_window`](Self::context_window). - /// Defaults to `8_000` (80%). Only consulted when - /// [`auto_compact`](Self::auto_compact) is `true`. - pub compact_threshold: u16, - - /// Whether auto-compaction is enabled. - /// - /// When `false`, the machine never emits [`MachineStep::Compact`] on its - /// own — the host must manage context size manually (useful in tests or - /// fixed-length runs). When `true` (the default), the machine gates - /// compaction with [`compact_threshold`](Self::compact_threshold). - pub auto_compact: bool, - - /// How independent tool calls within a single turn are dispatched. - /// - /// Controls whether the driver runs a turn's independent tool calls - /// sequentially or in parallel, up to a concurrency limit. See - /// [`ParallelDispatchConfig`]. - pub parallel_tool_dispatch: ParallelDispatchConfig, -} - -impl Default for RunConfig { - fn default() -> Self { - Self { - max_turns: 200, - max_tokens: 16_384, - context_window: 200_000, - compact_threshold: 8_000, - auto_compact: true, - parallel_tool_dispatch: ParallelDispatchConfig::default(), - } - } -} - /// One unit of work the driver must perform before feeding an outcome back. /// /// Returned by [`LoopMachine::next_step`]. The driver matches on this and, for @@ -104,7 +27,7 @@ pub enum MachineStep { /// /// The driver builds the feed (the messages actually sent to the LLM) from /// [`LoopMachine::history`], calls the provider, and feeds the completed - /// [`ModelTurn`] back via [`LoopMachine::model_response`]. `turn` is the + /// [`ModelResponse`] back via [`LoopMachine::model_response`]. `turn` is the /// 1-indexed number of the turn being requested. CallLLM { /// The 1-indexed turn number being requested. @@ -158,7 +81,7 @@ pub enum MachineStep { /// A tool call awaiting dispatch, with an optional preresolved result. /// /// The machine classifies each model-emitted tool call against the names the -/// driver reported as available ([`ModelTurn::available_tools`]). For a name the +/// driver reported as available ([`ModelResponse::available_tools`]). For a name the /// driver did not advertise, the machine sets /// [`preresolved_result`](Self::preresolved_result) to a synthetic error /// [`Message`]; the driver feeds that back as the tool result without @@ -175,7 +98,7 @@ pub struct PendingToolCall { /// A pre-resolved result for calls suppressed by invalid-tool-call recovery. /// /// `Some(message)` when the tool name was not in the call's - /// [`available_tools`](ModelTurn::available_tools): the machine has built a + /// [`available_tools`](ModelResponse::available_tools): the machine has built a /// synthetic error [`Message`] and the driver should feed it back as the /// tool result *without* dispatching, so the model learns the name is /// unknown and can correct itself. `None` for valid calls, which the driver @@ -188,14 +111,13 @@ pub struct PendingToolCall { /// The driver consumes the provider's stream, accumulates the final assistant /// [`Message`], and packages it with its usage and stop reason into this /// struct. Streaming stays driver-side; the machine only ever sees a completed -/// turn. +/// response. #[derive(Debug, Clone, Serialize, Deserialize)] #[non_exhaustive] -pub struct ModelTurn { +pub struct ModelResponse { /// The assistant message the model produced. /// - /// The fully-accumulated response, including any text and tool-call - /// [`MessagePart`] parts. + /// The fully-accumulated response, including any text and tool-call [`MessagePart`] parts. pub message: Message, /// Input tokens reported for this call. @@ -236,29 +158,15 @@ pub enum MachineOutcome { /// The model finished without requesting further tools. /// /// The normal end state: the last model response carried no tool calls, so - /// the machine treated it as the final answer. The captured totals let a - /// host record exactly what the completed run cost. + /// the machine treated it as the final answer. Token and turn totals are + /// not carried here — the owning `Run` is the + /// single source for accounting, derived from its `turns` list. Completed { /// The final text the model produced. /// /// Concatenation of the text parts of the last assistant message. May be /// empty if the model's final message carried only non-text parts. final_text: String, - - /// Total input tokens consumed across the run. - /// - /// Sum of every turn's [`ModelTurn::input_tokens`]. - total_input_tokens: u64, - - /// Total output tokens consumed across the run. - /// - /// Sum of every turn's [`ModelTurn::output_tokens`]. - total_output_tokens: u64, - - /// Number of turns executed. - /// - /// How many model calls the run made before finishing. - turns_taken: usize, }, /// The run was halted because `max_turns` was reached. @@ -266,13 +174,8 @@ pub enum MachineOutcome { /// The model kept requesting tools (or otherwise not finishing) until the /// turn budget was exhausted, so the machine stopped it rather than loop /// indefinitely. Usually a sign the agent isn't converging — raise - /// [`max_turns`](RunConfig::max_turns) or investigate the loop. - MaxTurnsExceeded { - /// Number of turns executed before the limit was hit. - /// - /// Equal to the configured [`RunConfig::max_turns`]. - turns_taken: usize, - }, + /// `max_turns` or investigate the loop. + MaxTurnsExceeded, /// The run was cancelled. /// @@ -300,22 +203,23 @@ pub enum MachineOutcome { /// Where the machine is in its request/respond cycle. /// /// One value is produced at each point in the loop and is also exposed by -/// [`LoopMachine::state`] for inspection. It is the machine's own step-protocol -/// bookkeeping — distinct from the public `LoopState` observation surface, which -/// a driver derives for observers. +/// [`LoopMachine::state`] and the trait method for inspection. +/// It is the machine's own step-protocol bookkeeping, surfaced to observers so +/// a host can see where a run currently is. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub enum MachineState { /// Ready to request the next model call. /// /// The starting state, and the state the machine returns to after tool /// results or a compaction result are fed back — i.e. whenever it is ready - /// to call the model again. + /// to call the model again. The driver also reports this as the idle + /// pre-run state before the first turn. Start, /// A model call has been requested; awaiting the response. /// /// Entered when the machine emits [`MachineStep::CallLLM`] and left when the - /// driver feeds the [`ModelTurn`] back via [`LoopMachine::model_response`]. + /// driver feeds the [`ModelResponse`] back via [`LoopMachine::model_response`]. AwaitingModel { /// The 1-indexed turn number in flight. /// @@ -355,27 +259,67 @@ pub enum MachineState { Terminal(MachineOutcome), } -/// The sans-IO agent-loop state machine. +/// Policy inputs passed to [`LoopMachine::next_step`] on each call. /// -/// Owns the append-only history and every decision the loop makes. Construct -/// one with [`LoopMachine::new`], then repeatedly call [`Self::next_step`] and -/// feed the result back. The machine performs no IO; a driver is required to -/// actually call the LLM, dispatch tools, and compact context. +/// The machine is policy-free: it tracks state and reports facts. The turn +/// budget and compaction thresholds live here, not on the machine, so the +/// machine serializes as pure state. Build one from your session/run config +/// and pass it fresh each `next_step` call. +#[derive(Debug, Clone, Copy)] +pub struct MachinePolicy { + /// Maximum turns before the run is capped. + /// + /// The turn budget for this run. The machine checks `turns_taken` against + /// this on every `next_step` call; reaching the cap ends the run with + /// [`MachineOutcome::MaxTurnsExceeded`]. + pub max_turns: usize, + + /// Model context window in tokens. + /// + /// The compaction trigger compares the current context size against this + /// value (via `compact_threshold`) and the 95% emergency line. + pub context_window: u64, + + /// Compaction threshold as a percentage of the context window (0–100). + /// + /// When `context_tokens` reaches this percentage of `context_window`, the + /// machine emits a [`MachineStep::Compact`] with a threshold-exceeded + /// reason (if `auto_compact` is `true`). + pub compact_threshold: u8, + + /// Whether automatic compaction is enabled. + /// + /// When `true` the machine emits `Compact` steps once the threshold (or the + /// 95% emergency line) is reached. When `false` only the emergency line + /// fires; threshold-compaction is left to the driver. + pub auto_compact: bool, +} + +/// The sans-IO agent-loop state machine — pure state, no policy. +/// +/// Owns the append-only history and tracks where the loop is in its +/// request/respond cycle. Construct one with +/// [`LoopMachine::from_history`], then repeatedly call +/// [`Self::next_step`] (passing a fresh [`MachinePolicy`]) and feed the +/// result back. The machine performs no IO and stores no policy — it +/// serializes as pure state. /// /// # Example /// -/// Driving a machine to completion (illustrative — a real driver performs IO in -/// each arm): +/// Driving a machine to completion (illustrative — a real driver performs IO +/// in each arm): /// /// ```no_run -/// # use loopctl::engine::machine::{LoopMachine, RunConfig, ModelTurn, MachineStep, MachineOutcome}; -/// # fn build_turn(_: &LoopMachine) -> ModelTurn { unimplemented!() } -/// let mut machine = LoopMachine::new(RunConfig::default(), "Hello"); +/// # use loopctl::engine::core::{LoopMachine, MachinePolicy, ModelResponse, MachineStep, MachineOutcome}; +/// # use loopctl::message::Message; +/// # fn build_response(_: &LoopMachine) -> ModelResponse { unimplemented!() } +/// let policy = MachinePolicy { max_turns: 10, context_window: 200_000, compact_threshold: 80, auto_compact: true }; +/// let mut machine = LoopMachine::from_history(vec![Message::user("Hello")]); /// loop { -/// match machine.next_step() { +/// match machine.next_step(policy) { /// MachineStep::CallLLM { .. } => { -/// let turn = build_turn(&machine); -/// machine.model_response(turn); +/// let response = build_response(&machine); +/// machine.model_response(response); /// } /// MachineStep::CallTools { .. } => { /// machine.tool_results(Vec::new()); @@ -383,7 +327,7 @@ pub enum MachineState { /// MachineStep::Compact { .. } => { /// machine.compaction_result(machine.history().to_vec(), 0); /// } -/// MachineStep::Done(MachineOutcome::Completed { final_text, .. }) => { +/// MachineStep::Done(MachineOutcome::Completed { final_text }) => { /// assert_eq!(final_text, "Hi there"); /// break; /// } @@ -393,12 +337,6 @@ pub enum MachineState { /// ``` #[derive(Debug, Clone, Serialize, Deserialize)] pub struct LoopMachine { - /// The per-run configuration the machine decides against. - /// - /// The turn budget, compaction policy, and dispatch mode that govern this - /// run. Owned by the machine so it travels with a serialized instance. - config: RunConfig, - /// The append-only conversation history. /// /// The complete record of the run: the opening user message, every @@ -415,32 +353,17 @@ pub struct LoopMachine { /// How many turns have completed. /// - /// The count of model responses accepted so far. Bounded by - /// [`RunConfig::max_turns`]; reaching that cap ends the run with + /// The count of model responses accepted so far. The driver checks this + /// against the turn budget passed to [`next_step`](Self::next_step); + /// reaching the cap ends the run with /// [`MachineOutcome::MaxTurnsExceeded`]. turns_taken: usize, - /// Cumulative input tokens across the run, for accounting. - /// - /// Total prompt-side token consumption from every accepted turn. Reported - /// in [`MachineOutcome::Completed`] when the run finishes. Distinct from - /// [`context_tokens`](Self::context_tokens), which tracks current size for - /// the compaction trigger rather than lifetime consumption. - total_input_tokens: u64, - - /// Cumulative output tokens across the run, for accounting. - /// - /// Total generation-side token consumption from every accepted turn. - /// Reported in [`MachineOutcome::Completed`] when the run finishes. - total_output_tokens: u64, - /// Current context-size estimate, in tokens. /// - /// The machine's view of how large the conversation currently is, used to - /// decide when to compact. Compared against - /// [`compact_threshold`](RunConfig::compact_threshold) of - /// [`context_window`](RunConfig::context_window); crossing it triggers a - /// [`MachineStep::Compact`]. + /// The machine's view of how large the conversation currently is. + /// [`next_step`](Self::next_step) compares it against the compaction + /// policy to decide whether to emit a [`MachineStep::Compact`]. context_tokens: u64, /// Whether [`Self::cancel`] has been called. @@ -458,37 +381,35 @@ pub struct LoopMachine { } impl LoopMachine { - /// Construct a machine for a new run with the given `user_input`. + /// Construct a machine seeded with an existing conversation history. + /// + /// The driver calls this at the top of every `run()`, passing the + /// session's accumulated messages so the model sees prior turns. For a + /// fresh single-message conversation, pass `vec![Message::user(prompt)]`. /// - /// The user's message is appended to the (initially empty) history as the - /// first entry, and the first [`Self::next_step`] call requests the initial - /// [`MachineStep::CallLLM`] at turn 1. + /// The machine starts in pure state: zero turns taken this run, zero + /// context tokens, not cancelled. The turn budget, compaction window, + /// and policy knobs are not stored — they are passed to + /// [`next_step`](Self::next_step) on each call so the machine stays + /// policy-free and serializes as pure state. #[must_use] - pub fn new(config: RunConfig, user_input: impl Into) -> Self { + pub fn from_history(history: Vec) -> Self { Self { - config, - history: vec![Message::user(user_input)], + history, state: MachineState::Start, turns_taken: 0, - total_input_tokens: 0, - total_output_tokens: 0, context_tokens: 0, cancelled: false, pending_tools: Vec::new(), } } - /// The configuration the machine makes its decisions against. - /// - /// The driver reads it to build provider requests (e.g. - /// [`RunConfig::max_tokens`], [`RunConfig::parallel_tool_dispatch`]). - #[must_use] - pub fn config(&self) -> &RunConfig { - &self.config - } - /// Return the next step the driver must perform. /// + /// `policy` supplies the turn budget and compaction thresholds the machine + /// consults to decide between `CallLLM`, `Compact`, and `MaxTurnsExceeded`. + /// It is not stored — pass it fresh each call from the session/run config. + /// /// This is pure and idempotent: calling it twice with no intervening feed /// method ([`Self::model_response`], [`Self::tool_results`], /// [`Self::compaction_result`], [`Self::cancel`]) returns an equal step. @@ -499,7 +420,7 @@ impl LoopMachine { /// /// If [`Self::cancel`] has been called, the next call returns /// [`MachineStep::Done`] with [`MachineOutcome::Cancelled`]. - pub fn next_step(&mut self) -> MachineStep { + pub fn next_step(&mut self, policy: MachinePolicy) -> MachineStep { if let MachineState::Terminal(outcome) = &self.state { let outcome = outcome.clone(); return MachineStep::Done(outcome); @@ -511,9 +432,14 @@ impl LoopMachine { return MachineStep::Done(outcome); } - let turn = self.turns_taken.saturating_add(1); + let turn = match &self.state { + MachineState::AwaitingModel { turn } => *turn, + _ => self.turns_taken.saturating_add(1), + }; match self.state.clone() { - MachineState::Start | MachineState::AwaitingModel { .. } => self.request_model(turn), + MachineState::Start | MachineState::AwaitingModel { .. } => { + self.request_model(turn, policy) + } MachineState::AwaitingTools { .. } => { let calls = std::mem::take(&mut self.pending_tools); MachineStep::CallTools { calls } @@ -529,14 +455,18 @@ impl LoopMachine { /// Returns [`MachineOutcome::MaxTurnsExceeded`] when the budget is spent, /// [`MachineStep::Compact`] when the context size has crossed the threshold /// and auto-compaction is on, or [`MachineStep::CallLLM`] otherwise. - fn request_model(&mut self, turn: usize) -> MachineStep { - if self.turns_taken >= self.config.max_turns { - let turns_taken = self.turns_taken; - let outcome = MachineOutcome::MaxTurnsExceeded { turns_taken }; + fn request_model(&mut self, turn: usize, policy: MachinePolicy) -> MachineStep { + if self.turns_taken >= policy.max_turns { + let outcome = MachineOutcome::MaxTurnsExceeded; self.state = MachineState::Terminal(outcome.clone()); return MachineStep::Done(outcome); } - if self.config.auto_compact && self.should_compact() { + if self.is_emergency(policy) { + let reason = CompactReason::Emergency; + self.state = MachineState::AwaitingCompaction { reason }; + return MachineStep::Compact { reason }; + } + if policy.auto_compact && self.should_compact(policy) { let reason = CompactReason::ThresholdExceeded; self.state = MachineState::AwaitingCompaction { reason }; return MachineStep::Compact { reason }; @@ -548,20 +478,31 @@ impl LoopMachine { /// Decide whether the current context size warrants a compaction pass. /// /// Compares the machine's current context-size estimate against the - /// configured [`compact_threshold`](RunConfig::compact_threshold) of the - /// [`context_window`](RunConfig::context_window): returns `true` when the - /// estimate has grown past that fraction of the window, signalling that the - /// next step should be [`MachineStep::Compact`] rather than another model - /// call. Returns `false` when compaction is disabled by a zero threshold or - /// window, or when the context still fits comfortably. - fn should_compact(&self) -> bool { - const BASIS: u128 = 10_000; - if self.config.context_window == 0 || self.config.compact_threshold == 0 { + /// configured `compact_threshold` of the model's context window: returns + /// `true` when the estimate has grown past that fraction of the window, + /// signalling that the next step should be [`MachineStep::Compact`] rather + /// than another model call. Returns `false` when compaction is disabled by + /// a zero threshold or window, or when the context still fits comfortably. + fn should_compact(&self, policy: MachinePolicy) -> bool { + if policy.context_window == 0 || policy.compact_threshold == 0 { + return false; + } + let limit = policy + .context_window + .saturating_mul(u64::from(policy.compact_threshold)) + / 100; + self.context_tokens > limit + } + + /// Whether the context is in the emergency zone (≥ 95% of the window). + /// + /// Fires regardless of `auto_compact` or the configured threshold — a + /// safety net against context overflow. + fn is_emergency(&self, policy: MachinePolicy) -> bool { + if policy.context_window == 0 { return false; } - let limit = u128::from(self.config.compact_threshold) - .saturating_mul(u128::from(self.config.context_window)); - u128::from(self.context_tokens).saturating_mul(BASIS) > limit + self.context_tokens >= policy.context_window.saturating_mul(95) / 100 } /// Feed a completed model call back into the machine. @@ -574,31 +515,32 @@ impl LoopMachine { /// [`MachineStep::CallTools`]. /// /// Has no effect once the machine is terminal. - pub fn model_response(&mut self, turn: ModelTurn) { + pub fn model_response(&mut self, response: ModelResponse) { if self.is_terminal() { return; } - let message = turn.message; - let tool_calls = Self::extract_tool_calls(&message); + let message = response.message; + let tool_calls: Vec = message + .tool_call_parts() + .into_iter() + .map(|(id, tool, input)| ToolCall { + id: id.to_string(), + tool: tool.to_string(), + input: input.clone(), + }) + .collect(); self.history.push(message); - self.total_input_tokens = self.total_input_tokens.saturating_add(turn.input_tokens); - self.total_output_tokens = self.total_output_tokens.saturating_add(turn.output_tokens); - self.context_tokens = turn.input_tokens; + self.context_tokens = response.input_tokens; self.turns_taken = self.turns_taken.saturating_add(1); if tool_calls.is_empty() { let final_text = self .history .last() - .map(Self::extract_text) + .map(Message::text_content) .unwrap_or_default(); - let outcome = MachineOutcome::Completed { - final_text, - total_input_tokens: self.total_input_tokens, - total_output_tokens: self.total_output_tokens, - turns_taken: self.turns_taken, - }; + let outcome = MachineOutcome::Completed { final_text }; self.state = MachineState::Terminal(outcome); return; } @@ -606,7 +548,7 @@ impl LoopMachine { let turn_number = self.turns_taken; self.pending_tools = tool_calls .into_iter() - .map(|call| Self::classify(call, &turn.available_tools)) + .map(|call| Self::classify(call, &response.available_tools)) .collect(); self.state = MachineState::AwaitingTools { turn: turn_number }; } @@ -626,6 +568,23 @@ impl LoopMachine { self.state = MachineState::Start; } + /// Inject an arbitrary message into the history. + /// + /// The driver calls this to add a message that did not come from a model + /// response, a tool result, or a compaction pass — for example a host + /// steering message, or the output of a [`ContextContributor`] that + /// re-emits the agent's goal before the next turn. The message becomes part + /// of the record the driver builds the feed from on the next + /// [`MachineStep::CallLLM`]. Has no effect once the machine is terminal. + /// + /// [`ContextContributor`]: crate::engine::ContextContributor + pub fn inject(&mut self, message: Message) { + if self.is_terminal() { + return; + } + self.history.push(message); + } + /// Feed compacted history back into the machine. /// /// The driver calls this after servicing a [`MachineStep::Compact`], passing @@ -653,6 +612,22 @@ impl LoopMachine { self.cancelled = true; } + /// Record a terminal failure in the machine's state. + /// + /// A driver that hits an unrecoverable error (stream failure, dispatch + /// error, compaction overflow) calls this so the machine transitions to + /// [`MachineState::Terminal`] with [`MachineOutcome::Failed`] carrying + /// `error`. Without this, a serialized machine that errored has no failure + /// record and would resume as if the error never happened. The driver + /// still returns the error from `run()`; this call ensures the machine's + /// own state agrees. Has no effect once the machine is already terminal. + pub fn fail(&mut self, error: LoopError) { + if self.is_terminal() { + return; + } + self.state = MachineState::Terminal(MachineOutcome::Failed { error }); + } + /// Whether [`Self::cancel`] has been called. /// /// When `true`, the next [`Self::next_step`] goes terminal with @@ -693,24 +668,6 @@ impl LoopMachine { self.turns_taken } - /// Total input tokens accumulated so far. - /// - /// The run's lifetime input consumption, for accounting. Also carried in - /// [`MachineOutcome::Completed`] once the run ends. - #[must_use] - pub fn total_input_tokens(&self) -> u64 { - self.total_input_tokens - } - - /// Total output tokens accumulated so far. - /// - /// The run's total generation cost, for accounting. Also carried in - /// [`MachineOutcome::Completed`] once the run ends. - #[must_use] - pub fn total_output_tokens(&self) -> u64 { - self.total_output_tokens - } - /// Whether the machine has reached a terminal outcome. /// /// `true` once the machine has emitted a [`MachineOutcome`] (completion, @@ -761,41 +718,6 @@ impl LoopMachine { )], ) } - - /// Concatenate every text part of a message into a single string. - /// - /// Walks the message's parts in order and joins all text content, skipping - /// non-text parts (tool calls, tool results, images) entirely. Used to - /// derive the final answer text from the model's last response. - fn extract_text(message: &Message) -> String { - message - .parts - .iter() - .filter_map(|part| part.as_text()) - .collect::>() - .join("") - } - - /// Collect the tool calls a message requests, in order. - /// - /// Scans the message's parts for tool-call parts and maps each to a - /// [`ToolCall`] (carrying its id, tool name, and JSON input). Returns an - /// empty vector when the message contains no tool calls — the signal that - /// the turn is complete rather than a request to dispatch tools. - fn extract_tool_calls(message: &Message) -> Vec { - message - .parts - .iter() - .filter_map(|part| match part { - MessagePart::ToolCall { id, name, input } => Some(ToolCall { - id: id.clone(), - tool: name.clone(), - input: input.clone(), - }), - _ => None, - }) - .collect() - } } #[cfg(test)] @@ -803,18 +725,21 @@ mod tests { use super::*; use serde_json::Value; - fn small_machine(max_turns: usize) -> LoopMachine { - LoopMachine::new( - RunConfig { - max_turns, - ..RunConfig::default() - }, - "hello", - ) + fn small_machine(_max_turns: usize) -> LoopMachine { + LoopMachine::from_history(vec![Message::user("hello")]) + } + + fn test_policy(max_turns: usize) -> MachinePolicy { + MachinePolicy { + max_turns, + context_window: 200_000, + compact_threshold: 80, + auto_compact: true, + } } - fn text_turn(text: &str, input_tokens: u64, output_tokens: u64) -> ModelTurn { - ModelTurn { + fn text_response(text: &str, input_tokens: u64, output_tokens: u64) -> ModelResponse { + ModelResponse { message: Message::assistant(text), input_tokens, output_tokens, @@ -823,9 +748,9 @@ mod tests { } } - fn tool_turn(tool: &str, available: &[&str], input_tokens: u64) -> ModelTurn { + fn tool_response(tool: &str, available: &[&str], input_tokens: u64) -> ModelResponse { let part = MessagePart::tool_call("call_1", tool, Value::Object(serde_json::Map::new())); - ModelTurn { + ModelResponse { message: Message::new(Role::Assistant, vec![part]), input_tokens, output_tokens: 10, @@ -840,8 +765,8 @@ mod tests { #[test] fn calling_llm_from_new_emits_call_llm_step() { - let mut machine = LoopMachine::new(RunConfig::default(), "hello"); - let step = machine.next_step(); + let mut machine = LoopMachine::from_history(vec![Message::user("hello")]); + let step = machine.next_step(test_policy(5)); let MachineStep::CallLLM { turn } = step else { panic!("expected CallLLM, got {step:?}"); }; @@ -851,23 +776,20 @@ mod tests { #[test] fn machine_api_has_no_async_no_tokio_no_apiclient() { - // If any driving method were async or needed a client/runtime, this - // body would not compile or would require a tokio context. It touches - // only the pure machine surface. - let mut machine = LoopMachine::new(RunConfig::default(), "hello"); - assert!(matches!(machine.next_step(), MachineStep::CallLLM { .. })); - machine.model_response(text_turn("hi", 5, 3)); + let mut machine = LoopMachine::from_history(vec![Message::user("hello")]); + assert!(matches!( + machine.next_step(test_policy(5)), + MachineStep::CallLLM { .. } + )); + machine.model_response(text_response("hi", 5, 3)); // A text-only turn completes the run, so the machine is terminal. assert!(machine.is_terminal()); assert_eq!(machine.turns_taken(), 1); - assert_eq!(machine.total_input_tokens(), 5); - assert_eq!(machine.total_output_tokens(), 3); assert_eq!(machine.history().len(), 2); assert!(matches!( machine.state(), MachineState::Terminal(MachineOutcome::Completed { .. }) )); - assert_eq!(machine.config().max_turns, RunConfig::default().max_turns); machine.cancel(); assert!(machine.is_cancelled()); } @@ -875,26 +797,24 @@ mod tests { #[test] fn resume_after_model_response_round_trips() { let mut machine = small_machine(5); - let _ = machine.next_step(); - machine.model_response(tool_turn("echo", &["echo"], 10)); - // Now AwaitingTools. + let _ = machine.next_step(test_policy(5)); + machine.model_response(tool_response("echo", &["echo"], 10)); let snapshot = serde_json::to_string(&machine).expect("serialize"); let mut restored: LoopMachine = serde_json::from_str(&snapshot).expect("deserialize"); - let a = machine.next_step(); - let b = restored.next_step(); + let a = machine.next_step(test_policy(5)); + let b = restored.next_step(test_policy(5)); assert!(same_step(&a, &b), "steps diverged after round-trip"); } #[test] fn resume_after_tool_results_round_trips() { let mut machine = small_machine(5); - let _ = machine.next_step(); - machine.model_response(tool_turn("echo", &["echo"], 10)); - let step = machine.next_step(); + let _ = machine.next_step(test_policy(5)); + machine.model_response(tool_response("echo", &["echo"], 10)); + let step = machine.next_step(test_policy(5)); let MachineStep::CallTools { calls } = &step else { panic!("expected CallTools, got {step:?}"); }; - // Synthesize the tool-result message the driver would build. let results: Vec = calls .iter() .map(|c| { @@ -911,35 +831,33 @@ mod tests { machine.tool_results(results); let snapshot = serde_json::to_string(&machine).expect("serialize"); let mut restored: LoopMachine = serde_json::from_str(&snapshot).expect("deserialize"); - let a = machine.next_step(); - let b = restored.next_step(); + let a = machine.next_step(test_policy(5)); + let b = restored.next_step(test_policy(5)); assert!(same_step(&a, &b), "steps diverged after round-trip"); } #[test] fn resume_after_compaction_result_round_trips() { - // window = 100, threshold = 0.5 ⇒ compact once tokens > 50. - let mut machine = LoopMachine::new( - RunConfig { - max_turns: 5, - context_window: 100, - compact_threshold: 5_000, - auto_compact: true, - ..RunConfig::default() - }, - "hello", - ); - // Drive a tool-call turn past the threshold so the next step compacts. - let _ = machine.next_step(); // CallLLM - machine.model_response(tool_turn("echo", &["echo"], 60)); // AwaitingTools - let _ = machine.next_step(); // CallTools + let policy = MachinePolicy { + max_turns: 5, + context_window: 100, + compact_threshold: 50, + auto_compact: true, + }; + let mut machine = LoopMachine::from_history(vec![Message::user("hello")]); + let _ = machine.next_step(policy); // CallLLM + machine.model_response(tool_response("echo", &["echo"], 60)); // AwaitingTools + let _ = machine.next_step(policy); // CallTools machine.tool_results(vec![Message::user("tool-out")]); // → Start - assert!(matches!(machine.next_step(), MachineStep::Compact { .. })); + assert!(matches!( + machine.next_step(policy), + MachineStep::Compact { .. } + )); machine.compaction_result(vec![Message::user("compacted")], 0); let snapshot = serde_json::to_string(&machine).expect("serialize"); let mut restored: LoopMachine = serde_json::from_str(&snapshot).expect("deserialize"); - let a = machine.next_step(); - let b = restored.next_step(); + let a = machine.next_step(policy); + let b = restored.next_step(policy); assert!(same_step(&a, &b), "steps diverged after round-trip"); } @@ -948,25 +866,23 @@ mod tests { let mut machine = small_machine(2); // Turn 1. assert!(matches!( - machine.next_step(), + machine.next_step(test_policy(2)), MachineStep::CallLLM { turn: 1 } )); - machine.model_response(tool_turn("echo", &["echo"], 1)); - let _ = machine.next_step(); + machine.model_response(tool_response("echo", &["echo"], 1)); + let _ = machine.next_step(test_policy(2)); machine.tool_results(vec![Message::user("r")]); // Turn 2. assert!(matches!( - machine.next_step(), + machine.next_step(test_policy(2)), MachineStep::CallLLM { turn: 2 } )); - machine.model_response(tool_turn("echo", &["echo"], 1)); - let _ = machine.next_step(); + machine.model_response(tool_response("echo", &["echo"], 1)); + let _ = machine.next_step(test_policy(2)); machine.tool_results(vec![Message::user("r")]); // Turn 3 must be denied. - match machine.next_step() { - MachineStep::Done(MachineOutcome::MaxTurnsExceeded { turns_taken }) => { - assert_eq!(turns_taken, 2); - } + match machine.next_step(test_policy(2)) { + MachineStep::Done(MachineOutcome::MaxTurnsExceeded) => {} other => panic!("expected MaxTurnsExceeded, got {other:?}"), } assert!(machine.is_terminal()); @@ -975,25 +891,56 @@ mod tests { #[test] fn cancel_returns_done_cancelled_at_next_step() { let mut machine = small_machine(5); - let _ = machine.next_step(); + let _ = machine.next_step(test_policy(5)); machine.cancel(); - match machine.next_step() { + match machine.next_step(test_policy(5)) { MachineStep::Done(MachineOutcome::Cancelled) => {} other => panic!("expected Done(Cancelled), got {other:?}"), } // Idempotent: stays terminal-cancelled. assert!(matches!( - machine.next_step(), + machine.next_step(test_policy(5)), MachineStep::Done(MachineOutcome::Cancelled) )); } + #[test] + fn fail_returns_done_failed_at_next_step() { + let mut machine = small_machine(5); + let _ = machine.next_step(test_policy(5)); + let err = LoopError::Api("stream failed".to_string()); + machine.fail(err.clone()); + match machine.next_step(test_policy(5)) { + MachineStep::Done(MachineOutcome::Failed { error }) => { + assert_eq!(error, err); + } + other => panic!("expected Done(Failed), got {other:?}"), + } + assert!(machine.is_terminal()); + } + + #[test] + fn failed_outcome_survives_round_trip() { + let mut machine = small_machine(5); + let _ = machine.next_step(test_policy(5)); + let err = LoopError::Api("stream failed".to_string()); + machine.fail(err.clone()); + let snapshot = serde_json::to_string(&machine).expect("serialize"); + let mut restored: LoopMachine = serde_json::from_str(&snapshot).expect("deserialize"); + match restored.next_step(test_policy(5)) { + MachineStep::Done(MachineOutcome::Failed { error }) => { + assert_eq!(error, err, "failure record survives round-trip"); + } + other => panic!("expected Done(Failed) after resume, got {other:?}"), + } + } + #[test] fn unknown_tool_call_gets_preresolved_result() { let mut machine = small_machine(5); - let _ = machine.next_step(); - machine.model_response(tool_turn("ghost", &["echo", "ls"], 3)); - let step = machine.next_step(); + let _ = machine.next_step(test_policy(5)); + machine.model_response(tool_response("ghost", &["echo", "ls"], 3)); + let step = machine.next_step(test_policy(5)); let MachineStep::CallTools { calls } = step else { panic!("expected CallTools, got {step:?}"); }; @@ -1012,9 +959,9 @@ mod tests { #[test] fn known_tool_call_emits_plain_pending_call() { let mut machine = small_machine(5); - let _ = machine.next_step(); - machine.model_response(tool_turn("echo", &["echo", "ls"], 3)); - let step = machine.next_step(); + let _ = machine.next_step(test_policy(5)); + machine.model_response(tool_response("echo", &["echo", "ls"], 3)); + let step = machine.next_step(test_policy(5)); let MachineStep::CallTools { calls } = step else { panic!("expected CallTools, got {step:?}"); }; @@ -1028,25 +975,24 @@ mod tests { #[test] fn compaction_triggered_when_tokens_exceed_threshold() { // window = 100, threshold = 0.5 ⇒ compact once tokens > 50. - let mut machine = LoopMachine::new( - RunConfig { - max_turns: 5, - context_window: 100, - compact_threshold: 5_000, - auto_compact: true, - ..RunConfig::default() - }, - "hello", - ); - // First step: no tokens yet (0 > 50 is false) ⇒ CallLLM, not Compact. - assert!(matches!(machine.next_step(), MachineStep::CallLLM { .. })); - // A tool-call turn keeps the run going and accumulates 60 input tokens - // (exceeds 50). A text-only turn would complete the run instead. - machine.model_response(tool_turn("echo", &["echo"], 60)); - assert!(matches!(machine.next_step(), MachineStep::CallTools { .. })); + let policy = MachinePolicy { + max_turns: 5, + context_window: 100, + compact_threshold: 50, + auto_compact: true, + }; + let mut machine = LoopMachine::from_history(vec![Message::user("hello")]); + assert!(matches!( + machine.next_step(policy), + MachineStep::CallLLM { .. } + )); + machine.model_response(tool_response("echo", &["echo"], 60)); + assert!(matches!( + machine.next_step(policy), + MachineStep::CallTools { .. } + )); machine.tool_results(vec![Message::user("tool-out")]); - // The next step must compact before calling the model again. - match machine.next_step() { + match machine.next_step(policy) { MachineStep::Compact { reason } => { assert_eq!(reason, CompactReason::ThresholdExceeded); } @@ -1054,25 +1000,52 @@ mod tests { } } + #[test] + fn emergency_compaction_fires_at_95_percent() { + let policy = MachinePolicy { + max_turns: 5, + context_window: 100, + compact_threshold: 50, + auto_compact: false, + }; + let mut machine = LoopMachine::from_history(vec![Message::user("hello")]); + assert!(matches!( + machine.next_step(policy), + MachineStep::CallLLM { .. } + )); + machine.model_response(tool_response("echo", &["echo"], 96)); + assert!(matches!( + machine.next_step(policy), + MachineStep::CallTools { .. } + )); + machine.tool_results(vec![Message::user("r")]); + match machine.next_step(policy) { + MachineStep::Compact { reason } => { + assert_eq!(reason, CompactReason::Emergency); + } + other => panic!("expected emergency Compact, got {other:?}"), + } + } + #[test] fn compaction_result_replaces_history() { - // window = 100, threshold = 0.5 ⇒ compact once tokens > 50. - let mut machine = LoopMachine::new( - RunConfig { - max_turns: 5, - context_window: 100, - compact_threshold: 5_000, - auto_compact: true, - ..RunConfig::default() - }, - "hello", - ); + // window = 100, threshold = 50 ⇒ compact once tokens > 50. + let policy = MachinePolicy { + max_turns: 5, + context_window: 100, + compact_threshold: 50, + auto_compact: true, + }; + let mut machine = LoopMachine::from_history(vec![Message::user("hello")]); // Drive a tool-call turn past the threshold so the next step compacts. - let _ = machine.next_step(); // CallLLM - machine.model_response(tool_turn("echo", &["echo"], 60)); // AwaitingTools - let _ = machine.next_step(); // CallTools + let _ = machine.next_step(policy); // CallLLM + machine.model_response(tool_response("echo", &["echo"], 60)); // AwaitingTools + let _ = machine.next_step(policy); // CallTools machine.tool_results(vec![Message::user("tool-out")]); // → Start - assert!(matches!(machine.next_step(), MachineStep::Compact { .. })); + assert!(matches!( + machine.next_step(policy), + MachineStep::Compact { .. } + )); let compacted = vec![Message::user("compacted-only")]; machine.compaction_result(compacted.clone(), 0); // Compare by serialized form: Message is not PartialEq. @@ -1080,15 +1053,18 @@ mod tests { let want = serde_json::to_string(&compacted).expect("serialize expected"); assert_eq!(got, want, "history must be replaced by the compacted slice"); // After compaction the next step is the deferred CallLLM. - assert!(matches!(machine.next_step(), MachineStep::CallLLM { .. })); + assert!(matches!( + machine.next_step(test_policy(5)), + MachineStep::CallLLM { .. } + )); } #[test] fn history_accumulates_user_assistant_tool_round() { let mut machine = small_machine(5); - let _ = machine.next_step(); - machine.model_response(tool_turn("echo", &["echo"], 1)); - let step = machine.next_step(); + let _ = machine.next_step(test_policy(5)); + machine.model_response(tool_response("echo", &["echo"], 1)); + let step = machine.next_step(test_policy(5)); let MachineStep::CallTools { calls } = step else { panic!("expected CallTools, got {step:?}"); }; @@ -1127,4 +1103,129 @@ mod tests { assert_eq!(back, reason); } } + + fn make_call() -> ToolCall { + ToolCall { + id: "test".to_string(), + tool: "Read".to_string(), + input: serde_json::json!({"path": "/tmp"}), + } + } + + #[test] + fn input_fix_accepts_json_object() { + use crate::reflection::{Correction, CorrectionResult, CorrectionType}; + let mut call = ToolCall { + id: "test".to_string(), + tool: "Read".to_string(), + input: serde_json::json!({"path": "/tmp"}), + }; + let correction = Correction { + correction_type: CorrectionType::InputFix, + description: "fix path".into(), + modified_input: Some(serde_json::json!({"path": "/tmp/fixed"})), + alternative_tool: None, + guidance: None, + }; + let result = call.apply_correction(&correction); + assert!(matches!(result, CorrectionResult::Applied)); + assert_eq!(call.input, serde_json::json!({"path": "/tmp/fixed"})); + } + + #[test] + fn input_fix_fails_when_modified_input_missing() { + use crate::reflection::{Correction, CorrectionResult, CorrectionType}; + let mut call = make_call(); + let correction = Correction { + correction_type: CorrectionType::InputFix, + description: "fix path".into(), + modified_input: None, + alternative_tool: None, + guidance: None, + }; + let result = call.apply_correction(&correction); + assert!(matches!(result, CorrectionResult::Failed(_))); + } + + #[test] + fn input_fix_fails_when_modified_input_not_object() { + use crate::reflection::{Correction, CorrectionResult, CorrectionType}; + let mut call = make_call(); + let correction = Correction { + correction_type: CorrectionType::InputFix, + description: "fix path".into(), + modified_input: Some(serde_json::json!("not an object")), + alternative_tool: None, + guidance: None, + }; + let result = call.apply_correction(&correction); + assert!(matches!(result, CorrectionResult::Failed(_))); + } + + #[test] + fn tool_change_swaps_tool_name() { + use crate::reflection::{Correction, CorrectionResult, CorrectionType}; + let mut call = make_call(); + let correction = Correction { + correction_type: CorrectionType::ToolChange, + description: "use alt tool".into(), + modified_input: None, + alternative_tool: Some("Write".into()), + guidance: None, + }; + let result = call.apply_correction(&correction); + assert!(matches!(result, CorrectionResult::Applied)); + assert_eq!(call.tool, "Write"); + } + + #[test] + fn tool_change_fails_without_alternative() { + use crate::reflection::{Correction, CorrectionResult, CorrectionType}; + let mut call = make_call(); + let correction = Correction { + correction_type: CorrectionType::ToolChange, + description: "use alt tool".into(), + modified_input: None, + alternative_tool: None, + guidance: None, + }; + let result = call.apply_correction(&correction); + assert!(matches!(result, CorrectionResult::Failed(_))); + } + + #[test] + fn prerequisite_fix_approach_change_escalate_all_skip() { + use crate::reflection::{Correction, CorrectionResult, CorrectionType}; + let mut call = make_call(); + for ct in [ + CorrectionType::PrerequisiteFix, + CorrectionType::ApproachChange, + CorrectionType::Escalate, + ] { + let correction = Correction { + correction_type: ct, + description: "n/a".into(), + modified_input: None, + alternative_tool: None, + guidance: None, + }; + let result = call.apply_correction(&correction); + assert!( + matches!(result, CorrectionResult::Skipped), + "{ct:?} should skip" + ); + } + } + + #[test] + fn configs_carry_no_model_field() { + use crate::config::SessionConfig; + let session = SessionConfig::default(); + let _: &Option = &session.system_prompt; + let _: u64 = session.context_window; + + let run = crate::engine::RunConfig::default(); + let _: usize = run.max_turns; + let _ = run.parallel_tool_dispatch; + } } diff --git a/src/engine/loop_core.rs b/src/engine/loop_core.rs deleted file mode 100644 index 9fbfec7..0000000 --- a/src/engine/loop_core.rs +++ /dev/null @@ -1,992 +0,0 @@ -//! 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)] -#[non_exhaustive] -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 configuration). - /// 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 { - /// A human-readable description of the error that caused the failure. - error: String, - }, - - /// The agent was cancelled by the user or a shutdown signal. - /// - /// Terminal state, distinct from [`Failed`](Self::Failed): cancellation is - /// a clean, cooperative termination ([`LoopError::Cancelled`]), not an - /// error condition. The agent may hold partial results in its - /// conversation history. - Cancelled, -} - -// ================================================== -// 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)] -#[non_exhaustive] -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)] -#[non_exhaustive] -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, -} - -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 { - 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(), - ) - } - } - 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 -// ================================================== - -/// 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)] -#[non_exhaustive] -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 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. - /// - /// 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) - } - - /// 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() - /// .with_session_id(Uuid::nil()) - /// .with_total_turns(5) - /// .with_tool_calls(3) - /// .with_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 with_session_id(mut self, id: Uuid) -> Self { - self.session_id = id; - self - } - - /// Set the total turns executed. - #[must_use] - pub fn with_total_turns(mut self, turns: usize) -> Self { - self.total_turns = turns; - self - } - - /// Set the total input tokens. - #[must_use] - pub fn with_input_tokens(mut self, tokens: u64) -> Self { - self.input_tokens = tokens; - self - } - - /// Set the total output tokens. - #[must_use] - pub fn with_output_tokens(mut self, tokens: u64) -> Self { - self.output_tokens = tokens; - self - } - - /// Set the total session duration. - #[must_use] - pub fn with_total_duration(mut self, duration: Duration) -> Self { - self.total_duration = duration; - self - } - - /// Set the total tool calls made. - #[must_use] - pub fn with_tool_calls(mut self, calls: usize) -> Self { - self.tool_calls = calls; - self - } - - /// Set whether the session succeeded. - #[must_use] - pub fn with_success(mut self, success: bool) -> Self { - self.success = success; - self - } - - /// Set the final output text. - #[must_use] - pub fn with_final_output(mut self, output: impl Into) -> Self { - self.final_output = Some(output.into()); - self - } - - /// Set the error message. - #[must_use] - pub fn with_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, - } - } -} - -// ================================================== -// 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 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, - ) -> 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); - - /// 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 { - // Clone is required: `initialize` takes `&mut self` which - // conflicts with the shared borrow from `config()`. - let config = self.config().clone(); - self.initialize(&config).await?; - - let mut is_first_turn = true; - loop { - if !self.should_continue() { - break; - } - - // 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; - } - 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 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 - /// 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"})); - } -} diff --git a/src/error.rs b/src/error.rs index 50922d7..b81fdd2 100644 --- a/src/error.rs +++ b/src/error.rs @@ -97,9 +97,18 @@ pub enum LoopError { /// [`tool_not_found`](LoopError::tool_not_found). #[error("Tool not found: {tool}. Available: {available}")] ToolNotFound { - /// The name of the requested tool (not normalised or lowercased). + /// The name of the requested tool. + /// + /// Preserved verbatim — no lowercasing or trimming is applied — + /// so the caller can display exactly what the agent asked for, + /// useful for case-sensitivity debugging. tool: String, - /// Available tool names, capped at ten with "… and N more" suffix. + + /// Human-readable list of registered tool names. + /// + /// Formatted as a comma-separated string, capped at ten names + /// with a trailing "… and N more" summary for large registries. + /// Built by [`tool_not_found`](LoopError::tool_not_found). available: String, }, @@ -111,9 +120,18 @@ pub enum LoopError { /// [`is_recoverable`](LoopError::is_recoverable). #[error("Tool execution error: {tool}: {message}")] ToolExecution { - /// Name of the tool that failed (matches the [`ToolRegistry`](crate::tool::ToolRegistry) key). + /// Name of the tool that failed. + /// + /// Matches the key used when registering the tool in the + /// [`ToolRegistry`](crate::tool::ToolRegistry), so callers can + /// correlate the failure with a specific implementation. tool: String, - /// Description of the execution failure returned by the tool. + + /// Description of the execution failure. + /// + /// The free-form message returned by the tool itself — typically + /// its `Err` rendering, or a wrapped IO/panic message. Use this + /// for logging and user-facing diagnostics. message: String, }, @@ -126,7 +144,15 @@ pub enum LoopError { /// [`is_recoverable`](LoopError::is_recoverable) because /// retrying the same input will produce the same result. #[error("Invalid input: {0}")] - InvalidInput(String), + InvalidInput( + /// Description of the invalid input. + /// + /// Names the field or argument that failed validation and + /// explains the constraint, for example `"parameter 'count' + /// must be non-negative"`. Rendering is stable enough for + /// log grep. + String, + ), /// An API call failed (e.g. LLM provider error). /// @@ -137,6 +163,11 @@ pub enum LoopError { #[error("API error: {0}")] Api( /// Upstream error message or status description. + /// + /// The provider's own error text, or a short rendering of the + /// HTTP status and body. Suitable for logging and for + /// surfacing to the caller when a retry strategy is being + /// chosen. String, ), @@ -147,7 +178,12 @@ pub enum LoopError { /// investigate why the agent is not converging. #[error("Max turns exceeded: {max}")] MaxTurnsExceeded { - /// The configured maximum turn count. See [`LoopConfig::max_turns`](crate::config::LoopConfig::max_turns). + /// The configured maximum turn count. + /// + /// Reflects the value supplied via + /// [`RunConfig::max_turns`](crate::engine::RunConfig::max_turns) + /// (or its default) when the run was started, so the caller can + /// decide whether to raise the limit and retry. max: usize, }, @@ -161,8 +197,18 @@ pub enum LoopError { #[error("Context window exceeded: used {used} of {limit} tokens")] ContextExceeded { /// Number of tokens consumed when the limit was exceeded. + /// + /// An estimate derived from the current message history; it + /// includes both the prompt and any retained tool results. Pair + /// with `limit` to report utilization to the caller. used: u64, - /// Maximum tokens allowed by the model. See [`LoopConfig::context_window`](crate::config::LoopConfig::context_window). + + /// Maximum tokens allowed by the model. + /// + /// Sourced from + /// [`SessionConfig::context_window`](crate::config::SessionConfig::context_window). + /// When `used` exceeds this after compaction, the run cannot + /// continue on the current model. limit: u64, }, @@ -175,9 +221,18 @@ pub enum LoopError { /// without parsing error messages. #[error("Phase '{phase}' failed: {message}")] PhaseFailed { - /// Name of the phase that failed (e.g. `"pre_process"`, `"reflection"`). + /// Name of the phase that failed. + /// + /// One of the framework's well-known phase identifiers, such as + /// `"pre_process"`, `"reflection"`, or `"post_process"`. Callers + /// can match on this to apply phase-specific recovery. phase: String, + /// Description of the phase failure. + /// + /// Free-form detail supplied by the failing phase — usually the + /// underlying error rendered to a string. Suitable for logs and + /// user-facing diagnostics. message: String, }, @@ -187,7 +242,14 @@ pub enum LoopError { /// connection failure, a serialization error, or a capacity limit /// reached in the backing store. #[error("Memory error: {0}")] - Memory(String), + Memory( + /// Description of the memory-backend failure. + /// + /// Captures the underlying store error — a database connection + /// failure, a serialization fault, or a capacity limit reached + /// in the backing store — rendered as a short diagnostic. + String, + ), /// Reflection / self-correction cycle failed. /// @@ -196,7 +258,14 @@ pub enum LoopError { /// variant is *recoverable* — the framework may retry the /// reflection or fall back to a simpler strategy. #[error("Reflection error: {0}")] - Reflection(String), + Reflection( + /// Description of the reflection failure. + /// + /// Free-form detail from the reflection phase — typically the + /// underlying error message or a note on which correction + /// strategy was abandoned. Use this for logging. + String, + ), /// The agent was cancelled by the user or a shutdown signal. /// @@ -216,7 +285,12 @@ pub enum LoopError { /// `message` field describes the detected pattern. #[error("Loop detected: {message}")] LoopDetected { - /// Description of the detected loop (e.g. `"Tool Read called 5 times with identical arguments"`). + /// Description of the detected loop. + /// + /// A human-readable summary of the repeated pattern, for + /// example `"Tool Read called 5 times with identical + /// arguments"`. Use this to explain to the caller why the agent + /// was halted. message: String, }, @@ -229,7 +303,12 @@ pub enum LoopError { /// to increase the limit or accept the partial result. #[error("Tool limit reached: {message}")] ToolLimitReached { - /// Description of the limit reached (e.g. `"Session tool limit of 100 reached"`). + /// Description of the limit reached. + /// + /// Names the scope (session or turn) and the cap that was hit, + /// for example `"Session tool limit of 100 reached"`. The caller + /// can use this to decide between raising the limit and + /// accepting the partial result. message: String, }, @@ -241,7 +320,12 @@ pub enum LoopError { /// whether to retry from the last complete message. #[error("Stream error: {0}")] StreamError( - /// Description of the streaming failure (network drop, malformed SSE, timeout). + /// Description of the streaming failure. + /// + /// Covers the common mid-stream failure modes — network drops, + /// malformed server-sent events, provider timeouts — rendered + /// as a short diagnostic string. The caller may use it to + /// decide whether to resume from the last complete message. String, ), @@ -254,7 +338,12 @@ pub enum LoopError { /// configuration and retry. #[error("Configuration error: {0}")] Config( - /// Description of the configuration problem (e.g. `"max_turns must be > 0"`). + /// Description of the configuration problem. + /// + /// Names the offending field and the constraint it violated, + /// for example `"max_turns must be > 0"` or `"model must be + /// non-empty"`. Use this to guide the caller toward the exact + /// setting to fix. String, ), @@ -265,7 +354,15 @@ pub enum LoopError { /// this only for truly unexpected conditions (e.g. a poisoned /// mutex, an allocation failure). #[error("{0}")] - Internal(String), + Internal( + /// Description of the internal error. + /// + /// A catch-all diagnostic for conditions that do not fit a more + /// specific variant — poisoned mutexes, allocation failures, + /// or invariant violations. Prefer a more specific variant + /// wherever possible. + String, + ), /// Rate-limit retries on the current model were exhausted. /// @@ -277,9 +374,17 @@ pub enum LoopError { #[error("Rate-limit escalation after {attempts} retries (retry-after {retry_after:?})")] RateLimitEscalation { /// Number of rate-limit retries honored before escalating. + /// + /// Counts the 429/529 responses the stream handler retried + /// (honoring the provider's `Retry-After`) before giving up on + /// the current model and handing control to the circuit breaker. attempts: u32, - /// Last server-advised `Retry-After` hint, after clamping. Used for - /// diagnostics; `None` when the provider sent no header. + + /// Last server-advised `Retry-After` hint, after clamping. + /// + /// Preserved for diagnostics and back-off tuning. `None` when + /// the provider sent no header on the final rate-limited + /// response. retry_after: Option, }, } @@ -375,10 +480,6 @@ impl LoopError { } } -// =================================================== -// Mutex poison recovery helper -// =================================================== - /// Recover a mutex guard from a `LockResult`, ignoring poison. /// /// Use this at lock sites protecting data where a panic mid-hold diff --git a/src/fallback.rs b/src/fallback.rs index 382b1b4..9d1b43a 100644 --- a/src/fallback.rs +++ b/src/fallback.rs @@ -546,29 +546,37 @@ impl FallbackEntry { #[derive(Debug, Clone)] pub struct FallbackConfig { /// Consecutive API failures required to trip the circuit open. Defaults to `3`. + /// + /// Once [`FallbackManager`] records this many failures in a row on the primary + /// model, it transitions from [`FallbackState::Primary`] to + /// [`FallbackState::Fallback`]. A higher value tolerates transient blips; a + /// lower value reacts faster to a degrading provider. pub trip_threshold: usize, + /// Minimum time in fallback before probing the primary model again. Defaults to 60 s. + /// + /// Cooldown enforced by [`FallbackManager::should_try_resume_primary`]: the + /// manager stays on the fallback model for at least this long before it is + /// willing to probe the primary again via [`FallbackState::Recovering`]. pub recovery_timeout: Duration, + /// Consecutive successes during recovering before the circuit fully closes. Defaults to `2`. + /// + /// Number of healthy responses the primary model must produce while in + /// [`FallbackState::Recovering`] before the manager transitions back to + /// [`FallbackState::Primary`]. Requiring more than one guards against a + /// single lucky success masking an ongoing outage. pub recovery_successes_needed: usize, + /// Per-model failure threshold before a fallback model is skipped. Defaults to `2`. + /// + /// Applied to each [`FallbackEntry`] in the fallback chain: once a single + /// model accumulates this many recorded attempts, it is taken out of + /// rotation and [`FallbackManager::active_model`] advances to the next + /// candidate. pub max_fail_count: usize, } -/// Produces a [`FallbackConfig`] with sensible production defaults. -/// -/// Defaults: `trip_threshold = 3`, `recovery_timeout = 60 s`, -/// `recovery_successes_needed = 2`, `max_fail_count = 2`. -/// -/// # Example -/// -/// ```rust -/// use loopctl::fallback::FallbackConfig; -/// -/// let config = FallbackConfig::default(); -/// assert_eq!(config.trip_threshold, 3); -/// assert_eq!(config.max_fail_count, 2); -/// ``` impl Default for FallbackConfig { fn default() -> Self { Self { @@ -630,13 +638,36 @@ impl Default for FallbackConfig { /// reads that could occur when acquiring separate locks sequentially. #[derive(Default)] struct FallbackInner { - /// Original model name (before fallback). + /// Original model name, before any fallback was activated. + /// + /// Captured when the manager first trips so the recovering state + /// knows which model to resume to. `None` until the circuit breaker + /// has switched away from the primary at least once. original_model: Option, - /// Ordered fallback models with failure status. + + /// Ordered list of fallback models with their per-model failure + /// status. + /// + /// Each entry pairs a model name with whether it has already been + /// tried-and-failed. The manager walks this list in order when + /// selecting the next fallback, skipping entries marked failed. fallback_models: Vec, - /// Cached first non-failed fallback model name. + + /// Cached name of the first non-failed fallback model. + /// + /// Computed once and reused while the manager is in the fallback + /// state, so callers asking "which model am I on now?" get an + /// `O(1)` answer without re-scanning `fallback_models`. Recomputed + /// whenever the set of failed entries changes. active_fallback: Option, - /// Time when fallback was activated. + + /// Instant at which fallback was activated. + /// + /// Recorded when the manager transitions from `Primary` to + /// `Fallback`, and compared against the current time to compute + /// cooldown eligibility for + /// [`should_try_resume_primary`](FallbackManager::should_try_resume_primary). + /// `None` while the manager has never left the primary. fallback_switched_at: Option, } @@ -646,27 +677,79 @@ struct FallbackInner { /// threads (e.g. via `Arc`). No `&mut self` is needed /// for any operation. pub struct FallbackManager { - /// Failures before switching to fallback. + /// Number of consecutive failures on the primary that trips the + /// circuit. + /// + /// Once [`consecutive_failures`](Self::consecutive_failures) reaches + /// this value the manager transitions from `Primary` to `Fallback`. + /// Set at construction and immutable thereafter. fallback_threshold: usize, - /// Successes needed on primary before resuming. + + /// Number of consecutive successes required on the primary during + /// recovery before transitioning back. + /// + /// Counts towards this via + /// [`record_model_success`](FallbackManager::record_model_success) + /// while in the `Recovering` state; reaching it transitions to + /// `Primary`. Set at construction and immutable thereafter. primary_resume_threshold: usize, - /// Per-model max failure count for new [`FallbackEntry`] instances. + + /// Per-model max failure count seeded into new + /// [`FallbackEntry`] instances. + /// + /// When a fallback model is added without an explicit cap, it + /// inherits this value as its `max_fail_count`. Set at construction + /// and immutable thereafter. default_max_fail_count: usize, - /// Consecutive API failure counter. + + /// Consecutive API failure counter on the current model. + /// + /// Incremented by + /// [`record_model_failure`](FallbackManager::record_model_failure), + /// reset to zero on any success. Hitting `fallback_threshold` trips + /// the circuit. Atomic so it updates lock-free on the hot path. consecutive_failures: AtomicUsize, - /// Whether fallback has been activated (sticky flag). + + /// Sticky flag recording whether fallback has ever been activated. + /// + /// `true` once the circuit has tripped at least once, and never + /// reset for the life of the manager. Lets observers tell "still + /// on primary, never failed" from "recovered back to primary". fallback_activated: AtomicBool, - /// Circuit breaker state (0=Primary, 1=Fallback, 2=Recovering). + + /// Current circuit-breaker state, encoded as an atomic. + /// + /// `0` = `Primary`, `1` = `Fallback`, `2` = `Recovering` (see + /// [`FallbackState`]). Stored as `AtomicU8` for lock-free reads on + /// every request; mutated only on state transitions. fallback_state: AtomicU8, - /// Consecutive successes on primary during recovery. + + /// Consecutive successes on the primary accumulated during the + /// `Recovering` state. + /// + /// Incremented by + /// [`record_model_success`](FallbackManager::record_model_success); + /// reaching `primary_resume_threshold` transitions back to + /// `Primary` and resets this to zero. Atomic so it updates + /// lock-free. primary_success_count: AtomicUsize, + /// 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. + + /// How long to remain in fallback before attempting primary + /// recovery. + /// + /// Compared against the elapsed time since + /// [`fallback_switched_at`](FallbackInner::fallback_switched_at) + /// by + /// [`should_try_resume_primary`](FallbackManager::should_try_resume_primary) + /// to decide when a cooldown has elapsed. Set at construction and + /// immutable thereafter. recovery_timeout: Duration, } @@ -833,10 +916,6 @@ impl FallbackManager { mgr } - // ================================================== - // Accessors - // ================================================== - /// Get the current circuit breaker state. /// /// # Example @@ -1232,10 +1311,6 @@ impl FallbackManager { self.recompute_active_fallback(); } - // ================================================== - // Fallback failure tracking - // ================================================== - /// Mark a fallback model as failed by name. /// /// Call this when the fallback model itself returns errors. The model @@ -1490,10 +1565,6 @@ impl FallbackManager { found } - // ================================================== - // Recording - // ================================================== - /// 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. @@ -1520,20 +1591,31 @@ impl FallbackManager { /// assert!(!mgr.record_api_failure()); // 4 — already activated, no re-trip /// ``` pub fn record_api_failure(&self) -> bool { - let failures = self - .consecutive_failures - .fetch_add(1, Ordering::Relaxed) - .saturating_add(1); - if failures >= self.fallback_threshold && !self.fallback_activated.load(Ordering::Relaxed) { - warn!( - consecutive_failures = failures, - threshold = self.fallback_threshold, - "Fallback threshold reached" - ); - self.fallback_activated.store(true, Ordering::Relaxed); - true - } else { - false + match self.state() { + FallbackState::Primary => { + let failures = self + .consecutive_failures + .fetch_add(1, Ordering::Relaxed) + .saturating_add(1); + if failures >= self.fallback_threshold { + warn!( + consecutive_failures = failures, + threshold = self.fallback_threshold, + "Fallback threshold reached" + ); + self.transition_to_fallback(); + return true; + } + false + } + FallbackState::Fallback | FallbackState::Recovering => { + let failures = self.consecutive_failures.load(Ordering::Relaxed); + warn!( + consecutive_failures = failures, + "API failure recorded while not in Primary state; counter unchanged" + ); + false + } } } @@ -1686,10 +1768,6 @@ impl FallbackManager { } } - // ================================================== - // State transitions - // ================================================== - /// Check if we should try resuming the primary model. /// /// Returns `true` if **both** conditions hold: @@ -1862,10 +1940,6 @@ impl FallbackManager { self.clear_all_fallback_failed(); } - // ================================================== - // Private helpers - // ================================================== - /// Recompute the cached [`active_fallback`](Self::active_fallback) from /// the current fallback chain. /// @@ -1884,22 +1958,6 @@ impl FallbackManager { } } -/// Produces a [`FallbackManager`] with production defaults. -/// -/// Equivalent to `FallbackManager::new(3, 2)` — trips after 3 -/// consecutive failures and resumes after 2 consecutive successes. -/// No model name is stored; use [`FallbackManager::for_model`] or -/// [`FallbackManager::set_original_model`] to configure one. -/// -/// # Example -/// -/// ```rust -/// use loopctl::fallback::FallbackManager; -/// -/// let mgr = FallbackManager::default(); -/// assert_eq!(mgr.consecutive_failures(), 0); -/// assert!(!mgr.is_using_fallback()); -/// ``` impl Default for FallbackManager { fn default() -> Self { Self::new(3, 2) @@ -2104,8 +2162,9 @@ mod tests { fn test_consolidated_mutex_reset_clears_all() { let mgr = FallbackManager::for_model("primary-model"); mgr.add_fallback_model("fallback-model"); - mgr.transition_to_fallback(); + // Record a failure while in Primary to dirty the counter, then trip. mgr.record_failure(); + mgr.transition_to_fallback(); // State is dirty. assert!(mgr.consecutive_failures() > 0); diff --git a/src/hooks.rs b/src/hooks.rs index f1b9433..f45a2f6 100644 --- a/src/hooks.rs +++ b/src/hooks.rs @@ -52,10 +52,6 @@ use context::{ SessionEndContext, SessionStartContext, }; -// =================================================== -// Hook trait -// =================================================== - /// Hook trait for bidirectional lifecycle control. /// /// Hooks differ from observers ([`crate::observer::LoopObserver`]) in two key ways: @@ -110,10 +106,6 @@ pub trait Hook: Send + Sync { /// produced a given result. fn name(&self) -> &str; - // ================================================== - // Pre-hooks (return value controls flow) - // ================================================== - /// Called before a tool is executed. /// /// Return `Some(HookAction::Block{...})` to prevent execution. @@ -133,10 +125,6 @@ pub trait Hook: Send + Sync { None } - // ================================================== - // Post-hooks (notification only, no flow control) - // ================================================== - /// Called after a tool completes execution. /// /// Cannot block. Use this for audit logging, metric recording, @@ -148,10 +136,6 @@ pub trait Hook: Send + Sync { /// Notification-only; cannot alter the compaction result. fn on_post_compact(&self, _ctx: &PostCompactContext) {} - // ================================================== - // Session lifecycle - // ================================================== - /// Called when an agent session starts. /// /// Fires once at session start. Use for initialization, @@ -165,10 +149,6 @@ pub trait Hook: Send + Sync { fn on_session_end(&self, _ctx: &SessionEndContext) {} } -// =================================================== -// Interactivity -// =================================================== - /// Whether the session can interact with a human operator. /// /// This controls how [`HookAction::Ask`] is handled by the @@ -189,10 +169,6 @@ pub enum Interactivity { Interactive, } -// =================================================== -// HookAction -// =================================================== - /// Action returned by a pre-hook to control flow. /// /// Pre-hooks return `Option`: @@ -248,6 +224,11 @@ impl fmt::Display for HookAction { impl HookAction { /// Convenience constructor for [`HookAction::Block`]. + /// + /// Accepts any `Into` so string literals work directly. + /// The carried `reason` is returned to the model as a tool error, + /// giving it a chance to adjust its approach. Prefer this over + /// struct-literal construction at call sites that block. pub fn block(reason: impl Into) -> Self { Self::Block { reason: reason.into(), @@ -255,6 +236,13 @@ impl HookAction { } /// Convenience constructor for [`HookAction::Ask`]. + /// + /// Accepts any `Into` so string literals work directly. + /// The carried `message` is presented to the user for confirmation; + /// the executor delegates to a + /// [`ConfirmationHandler`](crate::hooks::builtin::ConfirmationHandler) + /// to resolve it. In headless mode `Ask` is downgraded to `Block` + /// by the [`HookExecutor`]. pub fn ask(message: impl Into) -> Self { Self::Ask { message: message.into(), @@ -262,24 +250,43 @@ impl HookAction { } /// Returns `true` if this action allows the operation to proceed. + /// + /// `true` only for [`HookAction::Allow`]. Use this in match arms + /// where the default path is to proceed and only the block/ask + /// branches need explicit handling. #[must_use] pub fn is_allow(&self) -> bool { matches!(self, Self::Allow) } /// Returns `true` if this action blocks the operation. + /// + /// `true` only for [`HookAction::Block`]. Note that in headless + /// mode an [`HookAction::Ask`] is downgraded to `Block`, so a + /// post-executor action that reports `is_block` may have originated + /// as an `Ask`. #[must_use] pub fn is_block(&self) -> bool { matches!(self, Self::Block { .. }) } /// Returns `true` if this action requests interactive confirmation. + /// + /// `true` only for [`HookAction::Ask`]. Note this reflects the + /// action as produced by the hook, before the executor applies its + /// interactivity policy — an `Ask` returned by a hook may surface + /// as `Block` after the executor downgrades it in headless mode. #[must_use] pub fn is_ask(&self) -> bool { matches!(self, Self::Ask { .. }) } /// Returns the block reason, if this is a [`HookAction::Block`]. + /// + /// `Some(reason)` when this action is a `Block`, `None` for `Allow` + /// and `Ask`. Useful for logging why an action was refused without + /// a full match arm — pair with [`is_block`](Self::is_block) when + /// you need to distinguish the variants. #[must_use] pub fn block_reason(&self) -> Option<&str> { match self { diff --git a/src/hooks/builtin/auto_commit.rs b/src/hooks/builtin/auto_commit.rs index 1840814..36ef6ab 100644 --- a/src/hooks/builtin/auto_commit.rs +++ b/src/hooks/builtin/auto_commit.rs @@ -2,6 +2,16 @@ //! //! Tracks file modifications from tool calls (Write, Edit) and commits //! them at session end. Useful for keeping a git trail of agent actions. +//! +//! # How it fits together +//! +//! 1. [`AutoCommitHook`] observes post-tool-use events, recording each +//! `file_path` touched by the configured tracking tools. +//! 2. At session end the hook hands the recorded paths (plus the +//! [`AutoCommitConfig`]) to [`GitExecutor`], which shells out to +//! `git` to stage, commit, and optionally push. +//! 3. [`AutoCommitResult`] describes the outcome so callers can log +//! failures without inspecting git's stderr. use std::process::{Command, Output, Stdio}; use std::sync::Mutex; @@ -11,71 +21,147 @@ use crate::hooks::Hook; use crate::hooks::context::{PostToolUseContext, SessionEndContext, SessionStartContext}; /// Default timeout for git subprocess invocations. +/// +/// Applied to every `git` call made by [`GitExecutor`] so a stuck git +/// process (for example a hung GPG-signing prompt or an unreachable +/// remote during a fetch-adjacent operation) cannot block the agent +/// loop indefinitely. After this elapses the child is killed and the +/// call surfaces as [`GitExecutorError::Timeout`]. const GIT_TIMEOUT: Duration = Duration::from_secs(30); -// =================================================== -// GitExecutor -// =================================================== - /// Errors that can occur during git operations. +/// +/// Every `git` invocation made by [`GitExecutor`] returns `Result<_, +/// GitExecutorError>`. The variants distinguish the failure modes that +/// callers are likely to want to react to differently — a missing +/// repository, an empty commit, or a hung process — while folding the +/// remainder into [`GitExecutorError::GitError`] with git's own stderr +/// attached for diagnostics. #[derive(Debug, thiserror::Error)] pub enum GitExecutorError { - /// Failed to execute git command. + /// The `git` binary could not be spawned or the worker thread panicked. + /// + /// Typically means `git` is not on `PATH`, the working directory is + /// inaccessible, or the OS refused to create the subprocess. The + /// carried string is the underlying `std::io::Error` rendering. #[error("Failed to execute git: {0}")] ExecutionFailed(String), - /// Git command returned non-zero exit code. + + /// Git exited with a non-zero status code. + /// + /// The carried string is git's own `stderr` output, which usually + /// identifies the failure (merge conflict, lock file held, bad + /// ref). Inspect it before deciding whether a retry is worthwhile. #[error("Git error: {0}")] GitError(String), - /// No changes to commit. + + /// There is nothing to commit. + /// + /// Returned when `git commit` refuses to create a commit because + /// the index is empty or the staged changes are identical to + /// `HEAD`. [`AutoCommitConfig::skip_if_clean`] governs whether this + /// is treated as a silent no-op or surfaced to the caller. #[error("No changes to commit")] NoChanges, - /// Invalid repository state. + + /// The repository is in a state that prevents the operation. + /// + /// For example a rebase in progress, a detached `HEAD`, or a + /// missing upstream. The carried string describes the specific + /// condition git reported. #[error("Invalid repository state: {0}")] InvalidState(String), - /// Git command timed out. + + /// The git subprocess exceeded the configured timeout deadline. + /// + /// The carried [`Duration`] echoes the elapsed timeout so callers + /// can correlate it with their own deadline tracking. The child + /// process has been killed by the time this is returned. #[error("Git command timed out after {0:?}")] Timeout(Duration), } /// Result of an auto-commit operation. +/// +/// Returned by [`GitExecutor::auto_commit`] and +/// [`GitExecutor::auto_commit_with_files`] to describe the outcome +/// without surfacing raw git output. Callers typically match on this +/// enum to decide whether to log the run as informational (a clean +/// working tree), as a silent success, or as a warning. #[derive(Debug, Clone, PartialEq, Eq)] pub enum AutoCommitResult { - /// Commit was created successfully. + /// A commit was created successfully. + /// + /// The carried SHA lets the caller record the commit (for example + /// to attribute it to the session) without re-running `git`. Committed { - /// Short SHA of the new commit. + /// SHA of the newly created commit. + /// + /// The full `HEAD` object name as returned by + /// `git rev-parse HEAD`, trimmed of surrounding whitespace. + /// Suitable for linking from session logs or PR descriptions. sha: String, }, - /// No changes to commit. + + /// There was nothing to commit. + /// + /// The working tree was clean (or the staged set matched `HEAD`). + /// This is a successful no-op, not an error — + /// [`AutoCommitConfig::skip_if_clean`] decides whether the clean + /// state was a deliberate skip or an explicit no-op. NoChanges, - /// Commit was skipped due to configuration. + + /// The commit was skipped due to configuration. + /// + /// Emitted when auto-commit is disabled + /// ([`AutoCommitConfig::enabled`] is `false`) or when + /// `skip_if_clean` is `false` and no changes were found. The + /// carried reason explains which condition applied. Skipped { - /// Why the commit was skipped. + /// Human-readable reason the commit was skipped. + /// + /// Short diagnostic string such as `"Auto-commit disabled"` or + /// `"No changes detected"`. Intended for log output rather than + /// programmatic matching. reason: String, }, - /// Commit failed. + + /// The commit could not be created. + /// + /// One of the underlying git operations failed; the carried + /// message is the [`GitExecutorError`] rendering. The working tree + /// is left in whatever partial state git reached (staged but not + /// committed, committed but not pushed, etc.). Failed { - /// Error message. + /// Description of the failure. + /// + /// The full [`GitExecutorError`] message — typically git's own + /// stderr or an IO failure description. Use it for logging and + /// user-facing diagnostics. error: String, }, } /// Executes git operations for auto-commit. /// -/// Each method runs a git subprocess and returns the result. -/// All methods are static — no state is held. +/// A zero-sized type whose methods each shell out to `git` and return +/// the result. No state is held between calls, so every invocation +/// starts a fresh subprocess; this keeps the executor trivially +/// thread-safe at the cost of one `git` process per operation. pub struct GitExecutor; impl GitExecutor { /// Run a git subprocess with a bounded timeout. /// - /// Spawns the child, waits up to `GIT_TIMEOUT`, and kills it if the - /// deadline expires. Returns `(stdout, stderr)` on success. + /// Spawns the child, waits up to the configured timeout, and kills it if + /// the deadline expires. Returns `(stdout, stderr)` on success. /// /// # Errors /// - /// Returns [`GitExecutorError::ExecutionFailed`] if the child cannot be spawned - /// or the worker thread panics, [`GitExecutorError::GitError`] if git exits - /// non-zero, or [`GitExecutorError::Timeout`] if the deadline expires. + /// Returns [`GitExecutorError::ExecutionFailed`] if the child + /// cannot be spawned or the worker thread panics, + /// [`GitExecutorError::GitError`] if git exits non-zero, or + /// [`GitExecutorError::Timeout`] if the deadline expires. fn run_git(args: &[&str]) -> Result<(Vec, Vec), GitExecutorError> { let child = Command::new("git") .args(args) @@ -109,29 +195,41 @@ impl GitExecutor { } } - /// Check if there are uncommitted changes in the working tree. + /// Check whether the working tree has uncommitted changes. + /// + /// Runs `git status --porcelain` and returns `true` when any output + /// is produced. Use this to short-circuit a commit attempt before + /// staging, avoiding the overhead of a `git commit` invocation that + /// would refuse to run anyway. /// /// # Errors /// - /// Returns [`GitExecutorError::ExecutionFailed`] if git is not available, - /// or [`GitExecutorError::GitError`] if the command fails. + /// Returns [`GitExecutorError::ExecutionFailed`] if git is not + /// available, or [`GitExecutorError::GitError`] if the command + /// fails (for example, outside a repository). 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()) } - /// Stage files for commit. Empty slice stages all changes. + /// Stage files for commit. + /// + /// Each path in `files` is passed to `git add` individually. Passing + /// an empty slice stages every change in the working tree + /// (`git add -A`), which is convenient for "commit everything" + /// semantics but will include unrelated modifications. /// /// # Errors /// - /// Returns [`GitExecutorError`] if the git command fails. + /// Returns [`GitExecutorError`] if any `git add` invocation fails; + /// earlier files in the list may already be staged. pub fn stage_files(files: &[String]) -> Result<(), GitExecutorError> { if files.is_empty() { Self::run_git(&["add", "-A"])?; } else { for file in files { - Self::run_git(&["add", file])?; + Self::run_git(&["add", "--", file])?; } } Ok(()) @@ -139,9 +237,16 @@ impl GitExecutor { /// Create a commit with the given message. /// + /// Runs `git commit` (optionally with `--amend` when `amend` is + /// `true`) and returns the SHA of the resulting commit. The + /// "nothing to commit" case is mapped to + /// [`GitExecutorError::NoChanges`] so callers can distinguish it + /// from a genuine failure without parsing git's stderr. + /// /// # Errors /// - /// Returns [`GitExecutorError`] if the git command fails. + /// Returns [`GitExecutorError::NoChanges`] if there is nothing to + /// commit, or any other [`GitExecutorError`] variant on failure. pub fn commit(message: &str, amend: bool) -> Result { let mut args: Vec<&str> = vec!["commit"]; if amend { @@ -165,11 +270,18 @@ impl GitExecutor { } } - /// Push to remote. + /// Push the current branch to the configured remote. + /// + /// Resolves the branch from `branch` when given, otherwise from the + /// currently checked-out branch via [`current_branch`](Self::current_branch). + /// The push targets `origin` and is non-forceful — it will fail if + /// the remote has diverged. /// /// # Errors /// - /// Returns [`GitExecutorError`] if the git command fails. + /// Returns [`GitExecutorError`] if branch resolution or the push + /// itself fails (for example, no upstream configured or + /// non-fast-forward). pub fn push(branch: Option<&str>) -> Result<(), GitExecutorError> { let current_branch = Self::current_branch()?; let branch = branch.unwrap_or(¤t_branch); @@ -178,21 +290,30 @@ impl GitExecutor { Ok(()) } - /// Get the current branch name. + /// Get the currently checked-out branch name. + /// + /// Runs `git rev-parse --abbrev-ref HEAD` and trims whitespace. + /// Returns `"HEAD"` when in a detached-HEAD state (git's own + /// convention). /// /// # Errors /// - /// Returns [`GitExecutorError`] if the git command fails. + /// Returns [`GitExecutorError`] if git is unavailable or not inside + /// a repository. pub fn current_branch() -> Result { let (stdout, _stderr) = Self::run_git(&["rev-parse", "--abbrev-ref", "HEAD"])?; Ok(String::from_utf8_lossy(&stdout).trim().to_string()) } - /// Get the HEAD commit SHA. + /// Get the full SHA of the current `HEAD` commit. + /// + /// Runs `git rev-parse HEAD` and trims whitespace. Returns the full + /// object name (40 characters for SHA-1) rather than the short form. /// /// # Errors /// - /// Returns [`GitExecutorError`] if the git command fails. + /// Returns [`GitExecutorError`] if git is unavailable or not inside + /// a repository. pub fn get_head_sha() -> Result { let (stdout, _stderr) = Self::run_git(&["rev-parse", "HEAD"])?; Ok(String::from_utf8_lossy(&stdout).trim().to_string()) @@ -258,60 +379,114 @@ impl GitExecutor { } } - /// Perform auto-commit with the given configuration. + /// Perform auto-commit using the file list from the config. /// - /// Convenience wrapper around [`auto_commit_with_files`](Self::auto_commit_with_files) - /// that uses the file list from `config.files`. + /// Convenience wrapper around + /// [`auto_commit_with_files`](Self::auto_commit_with_files) that + /// passes `None` for the session-files override, so + /// [`AutoCommitConfig::files`] controls what gets staged. #[must_use] pub fn auto_commit(config: &AutoCommitConfig) -> AutoCommitResult { Self::auto_commit_with_files(config, None) } } -// =================================================== -// CommitMode -// =================================================== - /// How the auto-commit hook creates commits. +/// +/// Selected via [`AutoCommitConfig::commit_mode`] to control whether a +/// session's edits produce a fresh commit or fold into the previous +/// one. The default is [`CommitMode::Create`]. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum CommitMode { - /// Create a new commit. + /// Create a new commit on top of `HEAD`. + /// + /// The standard git behavior: each auto-commit produces a distinct + /// history entry. Choose this when you want a per-session audit + /// trail in the commit log. #[default] Create, + /// Amend the previous commit instead of creating a new one. + /// + /// Runs `git commit --amend`, folding the new changes into `HEAD`. + /// Useful for squashing a session's worth of edits into a single + /// history entry, but rewrites the previous commit's SHA — avoid + /// on branches that may already be shared. Amend, } impl CommitMode { /// Returns `true` if this mode amends the previous commit. + /// + /// Convenience predicate used by [`GitExecutor::commit`] to decide + /// whether to pass `--amend` to `git commit`. Equivalent to + /// comparing against [`CommitMode::Amend`]. #[must_use] pub fn is_amend(self) -> bool { self == Self::Amend } } -// =================================================== -// AutoCommitConfig -// =================================================== - /// Configuration for auto-commit behavior. +/// +/// Bundles every tunable that [`GitExecutor`] and [`AutoCommitHook`] +/// consult: whether to run at all, which tools to track, what to stage, +/// how to phrase the message, and whether to push afterwards. Build +/// with [`AutoCommitConfigBuilder`] or construct directly — +/// [`Default`] supplies sensible values for a local-only workflow. #[derive(Debug, Clone)] pub struct AutoCommitConfig { /// Enable auto-commit functionality. + /// + /// Master switch for the whole hook. When `false`, no file tracking or + /// git operations occur regardless of the other fields. pub enabled: bool, + /// Commit message template (supports `{{tool}}`, `{{session}}` placeholders). + /// + /// Template string used to build the commit message; `{{tool}}` and + /// `{{session}}` are expanded at session end into the triggering tool + /// name and the session identifier respectively. pub message_template: String, + /// Automatically push after commit. + /// + /// When `true`, a successful commit is immediately followed by a `git + /// push` to the configured (or current) branch. When `false` commits + /// stay local. pub auto_push: bool, + /// Branch to push to (`None` = current branch). + /// + /// Explicit remote branch name passed to `git push`. `None` pushes to the + /// currently checked-out branch, which is usually what you want. pub push_branch: Option, + /// Skip commit if no changes detected. + /// + /// When `true` a clean working tree is treated as a no-op + /// ([`AutoCommitResult::NoChanges`]) rather than a skipped run, so the + /// hook stays quiet during sessions that made no file changes. pub skip_if_clean: bool, + /// Whether to amend the previous commit or create a new one. + /// + /// [`CommitMode::Amend`] folds new changes into the previous commit + /// instead of producing a separate one, useful for squashing a session's + /// edits into a single history entry. pub commit_mode: CommitMode, + /// Tools that trigger file tracking (e.g., `"Write"`, `"Edit"`). + /// + /// Tool names whose output should be watched for a `file_path`; when one + /// of these tools runs, the touched file is recorded and staged at + /// session end. pub commit_on_tools: Vec, + /// Files to add before commit (empty = all files). + /// + /// Explicit allow-list of paths to `git add`. An empty vector stages every + /// change in the working tree (`git add -A`). pub files: Vec, } @@ -331,12 +506,21 @@ impl Default for AutoCommitConfig { } /// Builder for [`AutoCommitConfig`]. +/// +/// Provides a fluent API for constructing a config without naming every +/// field. Each `with_*` method overrides one setting; call +/// [`build`](Self::build) to finish. Omitted fields keep their +/// [`AutoCommitConfig::default`] values. pub struct AutoCommitConfigBuilder { + /// The in-progress configuration, mutated by each `with_*` call. config: AutoCommitConfig, } impl AutoCommitConfigBuilder { - /// Create a new builder with default configuration. + /// Create a new builder seeded with default configuration. + /// + /// Equivalent to [`AutoCommitConfigBuilder::default`]; every field + /// starts at its [`AutoCommitConfig::default`] value. #[must_use] pub fn new() -> Self { Self { @@ -345,6 +529,9 @@ impl AutoCommitConfigBuilder { } /// Enable or disable auto-commit. + /// + /// Mirrors [`AutoCommitConfig::enabled`]; `false` suppresses all + /// file tracking and git operations. #[must_use] pub fn with_enabled(mut self, enabled: bool) -> Self { self.config.enabled = enabled; @@ -352,27 +539,40 @@ impl AutoCommitConfigBuilder { } /// Set the commit message template. + /// + /// Accepts any `Into` so literals work directly. The + /// template may include `{{tool}}` and `{{session}}` placeholders + /// expanded at session end. #[must_use] pub fn with_message_template(mut self, template: impl Into) -> Self { self.config.message_template = template.into(); self } - /// Set whether to auto-push. + /// Set whether commits are followed by a `git push`. + /// + /// Mirrors [`AutoCommitConfig::auto_push`]; `true` enables + /// automatic pushing after each successful commit. #[must_use] pub fn with_auto_push(mut self, value: bool) -> Self { self.config.auto_push = value; self } - /// Set files to add before commit. + /// Set the explicit file allow-list to stage before commit. + /// + /// Mirrors [`AutoCommitConfig::files`]; an empty vector stages every + /// change in the working tree. #[must_use] pub fn with_files(mut self, files: Vec) -> Self { self.config.files = files; self } - /// Build the configuration. + /// Finalize and return the built configuration. + /// + /// Consumes the builder. The returned [`AutoCommitConfig`] is ready + /// to pass to [`GitExecutor`] or [`AutoCommitHook::with_config`]. #[must_use] pub fn build(self) -> AutoCommitConfig { self.config @@ -385,10 +585,6 @@ impl Default for AutoCommitConfigBuilder { } } -// =================================================== -// AutoCommitHook -// =================================================== - /// A hook that tracks file modifications from tool calls and commits /// them at session end. /// @@ -409,12 +605,29 @@ impl Default for AutoCommitConfigBuilder { /// assert_eq!(executor.hook_count(), 1); /// ``` pub struct AutoCommitHook { + /// Configuration consulted at each lifecycle callback. + /// + /// Set via [`new`](Self::new) (defaults) or + /// [`with_config`](Self::with_config) and held immutably for the + /// hook's lifetime. Controls which tools trigger tracking, the + /// commit message, and whether to push. config: AutoCommitConfig, + + /// File paths recorded by tracked tools this session. + /// + /// A deduplicated, insertion-ordered list updated by + /// [`track_modification`](Self::track_modification) and cleared on + /// [`clear_modifications`](Self::clear_modifications). Guarded by a + /// [`Mutex`] so the hook can be shared across threads via + /// [`HookExecutor`](crate::hooks::HookExecutor). modified_files: Mutex>, } impl AutoCommitHook { /// Create a new auto-commit hook with default configuration. + /// + /// Equivalent to [`AutoCommitHook::default`]; the hook starts with + /// an empty modification list and [`AutoCommitConfig::default`]. #[must_use] pub fn new() -> Self { Self { @@ -423,13 +636,22 @@ impl AutoCommitHook { } } - /// Set custom configuration. + /// Replace this hook's configuration. + /// + /// Builder-style setter intended for use during hook assembly. The + /// existing modification list is preserved so reconfiguring + /// mid-session does not lose already-tracked files. #[must_use] pub fn with_config(mut self, config: AutoCommitConfig) -> Self { self.config = config; self } + /// Record a file path touched by a tracked tool. + /// + /// Deduplicates against the existing list — calling twice with the + /// same path records it only once. Lock failures are silently + /// ignored so a poisoned mutex cannot block the agent loop. fn track_modification(&self, file: &str) { if let Ok(mut files) = self.modified_files.lock() && !files.contains(&file.to_string()) @@ -438,6 +660,11 @@ impl AutoCommitHook { } } + /// Clear all recorded file modifications. + /// + /// Called at session start so a fresh session does not inherit + /// modifications from the previous one. Lock failures are silently + /// ignored. fn clear_modifications(&self) { if let Ok(mut files) = self.modified_files.lock() { files.clear(); diff --git a/src/hooks/context.rs b/src/hooks/context.rs index e9c7386..e47745e 100644 --- a/src/hooks/context.rs +++ b/src/hooks/context.rs @@ -2,77 +2,146 @@ //! //! Each context type captures the relevant state for its trigger point. //! Contexts are read-only snapshots; hooks cannot mutate agent state -//! directly (they return `HookAction` or `CompactResult` to influence flow). +//! directly (they return `HookAction` or [`CompactResult`] to influence +//! flow). +//! +//! # Lifecycle +//! +//! Every session opens with [`SessionStartContext`], then interleaves +//! any number of tool calls (`PreToolUse` → tool → `PostToolUse`) and +//! compaction passes (`PreCompact` → compact → `PostCompact`), and +//! finally closes with [`SessionEndContext`]. Every callback receives a +//! context whose fields mirror the moment it fires: pre-contexts carry +//! the *about-to-happen* inputs, post-contexts carry the +//! *what-happened* results, and the session contexts bracket the whole +//! run. use serde::{Deserialize, Serialize}; use serde_json::Value; -// =================================================== -// Tool-use contexts -// =================================================== - /// Context provided to `on_pre_tool_use` hooks. /// -/// Contains everything a hook needs to decide whether to allow, -/// block, or ask about a tool invocation. +/// Carries everything a hook needs to decide whether to allow, block, or +/// ask about a tool invocation: the tool name, its JSON input, and the +/// surrounding session/turn coordinates. The fields are the same values +/// that will reach the tool, so a pre-hook acts as a gatekeeper with +/// full visibility into the upcoming call. #[derive(Debug, Clone)] pub struct PreToolUseContext { /// Name of the tool about to be invoked. + /// + /// Matches the identifier the tool was registered under, so hooks can + /// dispatch on it (e.g. treat `"Edit"` differently from `"Read"`). pub tool_name: String, + /// JSON input that will be passed to the tool. + /// + /// The full arguments object the model generated, exactly as it will be + /// delivered to the tool's execution. Hooks may inspect it to validate, + /// redact, or block based on field values. pub input: Value, + /// Session ID of the running agent. + /// + /// Correlates this tool call with the surrounding agent session, letting + /// hooks share per-session state across lifecycle callbacks. pub session_id: uuid::Uuid, + /// Current turn number within the session (0-indexed). + /// + /// How many turns have elapsed since the session started, so hooks can + /// reason about timing (e.g. only enforce a policy after turn *N*). pub turn_number: usize, } /// Context provided to `on_post_tool_use` hooks. /// -/// Captures the outcome of a tool invocation. Post-hooks are -/// notification-only and cannot alter the result. +/// Captures the outcome of a tool invocation: the input that was sent, +/// the output (or error) it produced, how long it took, and the +/// surrounding session/turn coordinates. Post-hooks are +/// notification-only and cannot alter the result — use them for +/// logging, metrics, and side-effects like +/// [`AutoCommitHook`](crate::hooks::builtin::AutoCommitHook)'s file +/// tracking. #[derive(Debug, Clone)] pub struct PostToolUseContext { /// Name of the tool that was invoked. + /// + /// The same identifier the tool was registered under, mirrored from the + /// matching [`PreToolUseContext::tool_name`]. pub tool_name: String, + /// JSON input that was passed to the tool. + /// + /// The arguments object the tool actually received, useful for logging + /// or correlating an error with its triggering input. pub input: Value, + /// Output produced by the tool (stringified). + /// + /// The tool's return value rendered as a string, whether the call + /// succeeded or failed. For errors this typically holds the error message. pub output: String, + /// Whether the tool reported an error. + /// + /// `true` when the tool returned an error result rather than a normal + /// output, so hooks can branch on failure without parsing `output`. pub is_error: bool, + /// Wall-clock execution time in milliseconds. + /// + /// Measured from tool dispatch to completion, for latency tracking and + /// slow-tool detection. pub duration_ms: u64, + /// Session ID of the running agent. + /// + /// Correlates this result with the surrounding agent session, mirroring + /// the value carried by [`PreToolUseContext::session_id`]. pub session_id: uuid::Uuid, + /// Turn number within the session. + /// + /// The 0-indexed turn in which this tool call executed, mirroring the + /// value from the corresponding [`PreToolUseContext::turn_number`]. pub turn_number: usize, } -// =================================================== -// Compaction contexts -// =================================================== - /// Why compaction was triggered. /// -/// Distinguishes between automatic threshold-based compaction -/// and manual compaction requested by the agent or user. +/// Distinguishes between automatic threshold-based compaction and +/// manual compaction requested by the agent or user. Hooks consult this +/// to apply different policies per source — for example, allowing a +/// manual compaction to proceed while aborting an automatic one during a +/// sensitive operation. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum CompactTrigger { - /// Automatic compaction triggered by token threshold. + /// Automatic compaction triggered by a token threshold. + /// + /// Fired when token usage crosses + /// [`SessionConfig::compact_threshold`](crate::config::SessionConfig::compact_threshold) + /// or the 95% emergency line. The compactor decided to run on its + /// own; no caller asked for it. Auto, + /// Manual compaction requested by the agent or user. + /// + /// Fired by an explicit compaction request — for example the agent + /// invoking a "compact now" tool, or a host application reacting to + /// a user gesture. The caller asked for it, independent of token + /// usage. Manual, } impl From for CompactTrigger { - /// Map the compaction pipeline's [`CompactReason`](crate::compact::types::CompactReason) into the hook-level + /// Map the compaction pipeline's `CompactReason` into the hook-level /// [`CompactTrigger`]. /// /// [`ThresholdExceeded`](crate::compact::types::CompactReason::ThresholdExceeded) - /// and - /// [`Emergency`](crate::compact::types::CompactReason::Emergency) - /// are both automatic triggers, while + /// and [`Emergency`](crate::compact::types::CompactReason::Emergency) + /// are both automatic triggers (they collapse to + /// [`Auto`](CompactTrigger::Auto)), while /// [`Manual`](crate::compact::types::CompactReason::Manual) maps to /// [`Manual`](CompactTrigger::Manual). fn from(reason: crate::compact::types::CompactReason) -> Self { @@ -86,68 +155,155 @@ impl From for CompactTrigger { /// Context provided to `on_pre_compact` hooks. /// -/// Hooks can inspect the current state and decide to abort -/// compaction or inject additional instructions. +/// Carries the state at the moment compaction was about to run — why it +/// was triggered, how big the conversation is, and how close token usage +/// is to the limit. Hooks use this to decide whether to abort compaction +/// or to inject extra instructions/context the summarizer should +/// preserve. The returned [`CompactResult`] accumulates across hooks. #[derive(Debug, Clone)] pub struct PreCompactContext { /// Why compaction was triggered. + /// + /// [`CompactTrigger::Auto`] when a token threshold was exceeded, + /// [`CompactTrigger::Manual`] when the agent or user requested it. Hooks + /// can use this to apply different policies per trigger source. pub trigger: CompactTrigger, + /// Custom instructions from prior hooks (accumulated). + /// + /// Instructions that earlier `on_pre_compact` hooks have already appended, + /// so a later hook can build on them rather than overwrite. `None` when no + /// hook has contributed custom instructions yet. pub custom_instructions: Option, + /// Number of messages in the conversation before compaction. + /// + /// The size of the message history at the moment compaction was + /// triggered, useful for deciding how aggressively to summarize. pub message_count: usize, + /// Estimated token count before compaction. + /// + /// The conversation's approximate token usage right before the compactor + /// runs, against which `context_window` is compared to trigger compaction. pub tokens_before: u64, + /// Model context window size in tokens. + /// + /// The configured window for the active model, providing the denominator + /// against which `tokens_before` is compared. pub context_window: u64, + /// Session ID of the running agent. + /// + /// Correlates this compaction event with the surrounding agent session. pub session_id: uuid::Uuid, } /// Context provided to `on_post_compact` hooks. /// -/// Notification-only; cannot alter the compaction result. +/// Reports how compaction turned out — how many messages it removed, +/// how many tokens it saved, how long it took — plus the trigger that +/// started it. Notification-only; cannot alter the compaction result. +/// Use it for metrics, budget tracking, and post-compaction logging. #[derive(Debug, Clone)] pub struct PostCompactContext { /// Why compaction was triggered. + /// + /// The same trigger value carried by the matching + /// [`PreCompactContext::trigger`], so post-hooks can correlate. pub trigger: CompactTrigger, + /// Number of messages removed by compaction. + /// + /// How many messages were dropped or summarized away, indicating how + /// aggressively the compactor pruned the history. pub messages_compacted: usize, + /// Tokens saved by compaction. + /// + /// The net token reduction achieved, computed as `tokens_before - + /// tokens_after`, surfaced for metrics and budget tracking. pub tokens_saved: u64, + /// Estimated tokens after compaction. + /// + /// The conversation's approximate token usage once the compactor has + /// finished, which will be used for subsequent model calls. pub tokens_after: u64, + /// Compaction duration in milliseconds. + /// + /// Wall-clock time spent running the compactor, useful for spotting + /// expensive summarization strategies. pub duration_ms: u64, + /// Session ID of the running agent. + /// + /// Correlates this compaction result with the surrounding agent session. pub session_id: uuid::Uuid, } /// Result of a pre-compact hook check. /// -/// Unlike tool hooks (Allow/Block/Ask), compact hooks have richer semantics: -/// multiple hooks may want to inject context or instructions. The executor -/// accumulates results, with abort taking priority. +/// Unlike tool hooks (Allow/Block/Ask), compact hooks have richer +/// semantics: multiple hooks may want to inject context or instructions, +/// and any one of them may abort. The executor accumulates results +/// across hooks, with [`abort`](Self::abort) taking priority over every +/// other field. Construct with [`allow`](Self::allow) for the common +/// "proceed" case, or use the builder methods +/// ([`with_context`](Self::with_context), +/// [`with_instructions`](Self::with_instructions)) to shape the +/// summary. #[derive(Debug, Clone, Default)] pub struct CompactResult { /// Whether to abort compaction. + /// + /// When `true` the executor skips compaction entirely for this trigger; + /// abort takes priority over every other field. Set via + /// [`CompactResult::abort`]. pub abort: bool, + /// Reason for abort (shown to the user/agent). + /// + /// Human-readable explanation surfaced when [`abort`](Self::abort) is + /// `true`, so callers can report why compaction was suppressed. `None` + /// unless compaction was aborted. pub abort_reason: Option, + /// Override custom instructions for the summarizer. + /// + /// When `Some`, replaces the summarizer's default instructions entirely. + /// Use sparingly; prefer appending to + /// [`additional_context`](Self::additional_context) when you only want to + /// add guidance. pub new_instructions: Option, + /// Additional context strings to include in the summary. + /// + /// Extra facts the hook wants preserved during summarization, accumulated + /// across hooks. The executor feeds these to the summarizer alongside the + /// conversation history. pub additional_context: Vec, } impl CompactResult { - /// Allow compaction to proceed (default). + /// Allow compaction to proceed with no modifications. + /// + /// Returns the default result — no abort, no new instructions, no + /// extra context. This is the value a hook returns when it has no + /// opinion about the compaction. #[must_use] pub fn allow() -> Self { Self::default() } - /// Abort compaction with a reason. + /// Abort compaction and surface `reason` to the user/agent. + /// + /// Sets [`abort`](Self::abort) to `true` and records `reason` in + /// [`abort_reason`](Self::abort_reason). The executor treats this as + /// a veto: compaction is skipped for this trigger, regardless of + /// what other hooks returned. pub fn abort(reason: impl Into) -> Self { Self { abort: true, @@ -156,14 +312,26 @@ impl CompactResult { } } - /// Inject additional context into the summary. + /// Append an additional context string for the summarizer to preserve. + /// + /// Builder-style helper that pushes `ctx` onto + /// [`additional_context`](Self::additional_context). Multiple calls + /// accumulate; the summarizer receives every entry alongside the + /// conversation history. Prefer this over + /// [`with_instructions`](Self::with_instructions) when you only want + /// to add a fact, not rewrite the whole prompt. #[must_use] pub fn with_context(mut self, ctx: impl Into) -> Self { self.additional_context.push(ctx.into()); self } - /// Override the summarization instructions. + /// Replace the summarizer's default instructions entirely. + /// + /// Builder-style helper that sets + /// [`new_instructions`](Self::new_instructions). Use sparingly — it + /// overrides the summarizer's prompt rather than augmenting it. + /// Prefer [`with_context`](Self::with_context) for additive guidance. #[must_use] pub fn with_instructions(mut self, instructions: impl Into) -> Self { self.new_instructions = Some(instructions.into()); @@ -171,52 +339,121 @@ impl CompactResult { } } -// =================================================== -// Session lifecycle contexts -// =================================================== - /// Why the session ended. +/// +/// Carried by [`SessionEndContext::reason`] so end-hooks can branch on +/// the terminal condition. The variants cover the four ways a session +/// can stop — normal completion, cancellation, an unrecoverable error, +/// hitting a turn cap — plus context overflow when compaction could not +/// recover. Use this for outcome-specific logging and cleanup. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum SessionEndReason { - /// Normal completion (agent converged). + /// The agent finished normally. + /// + /// The agent declared its task complete (for example the model + /// emitted a terminal stop reason) and the loop wound down cleanly. + /// This is the "happy path" outcome. Complete, - /// User interrupted (Ctrl+C / cancel signal). + + /// The user interrupted the session. + /// + /// A cancel signal — Ctrl+C in an interactive host, or a programmatic + /// cancellation — stopped the loop. The session may have partial + /// results available in its state at the time of interruption. Cancelled, - /// Agent error. + + /// The agent hit an unrecoverable error. + /// + /// A [`LoopError`](crate::error::LoopError) (other than + /// cancellation) terminated the loop — for example an API failure + /// after retries were exhausted, or a misconfiguration. The carried + /// error is available in the agent state. Error, - /// Maximum turns reached. + + /// The configured turn cap was reached. + /// + /// The agent ran for [`RunConfig::max_turns`](crate::engine::RunConfig::max_turns) + /// turns without converging. Not necessarily a failure — long tasks + /// can legitimately hit the cap — but worth surfacing so the caller + /// can decide whether to raise the limit or accept partial results. MaxTurns, - /// Context limit exceeded and compaction failed. + + /// Token usage exceeded the context window and compaction failed. + /// + /// The conversation grew past [`SessionConfig::context_window`](crate::config::SessionConfig::context_window), + /// compaction was attempted but could not bring it back under the + /// limit, and the loop terminated. This usually means the task + /// genuinely needs a larger model context than configured. ContextOverflow, } /// Context provided to `on_session_start` hooks. /// -/// Notification-only; fires once at session start. +/// Fires exactly once per session, before any turns run. Carries the +/// session identifier, the model that will handle the first request, +/// and the working directory the agent treats as its root. +/// Notification-only — use it to initialize per-session state (open +/// resources, set up tracking, log the start). #[derive(Debug, Clone)] pub struct SessionStartContext { /// Session ID. + /// + /// Unique identifier for the agent session that is starting, letting hooks + /// initialize any per-session state they track. pub session_id: uuid::Uuid, + /// Model identifier. + /// + /// The model the session will use for its first request (e.g. a provider + /// name like `"claude-3-opus"`). May change later if fallback or an + /// explicit model switch occurs. pub model: String, + /// Working directory the agent was started in. + /// + /// The filesystem path the agent treats as its root for relative tool + /// calls, useful for hooks that log or restrict operations by location. pub working_directory: String, } /// Context provided to `on_session_end` hooks. /// -/// Notification-only; fires once at session end. +/// Fires exactly once per session, after the loop has terminated. +/// Carries the session identifier, the terminal +/// [`SessionEndReason`], and aggregate counters (turns, tokens, +/// duration) for bookkeeping and cost accounting. Notification-only — +/// use it to flush resources, finalize tracking, and emit a summary log +/// line. #[derive(Debug, Clone)] pub struct SessionEndContext { /// Session ID. + /// + /// Unique identifier for the session that is ending, matching the value + /// carried by [`SessionStartContext::session_id`]. pub session_id: uuid::Uuid, + /// Why the session ended. + /// + /// The terminal condition — normal completion, cancellation, error, turn + /// cap, or context overflow — so hooks can branch on the outcome. pub reason: SessionEndReason, + /// Total turns executed. + /// + /// How many turns ran to completion during the session, for bookkeeping + /// and budget reporting. pub total_turns: usize, + /// Total tokens consumed (input + output across all turns). + /// + /// Aggregate token usage over the whole session, combining input and + /// output tokens from every turn for cost accounting. pub total_tokens: u64, + /// Wall-clock session duration in seconds. + /// + /// Elapsed time from session start to session end, rounded to whole + /// seconds, for latency and uptime reporting. pub duration_secs: u64, } diff --git a/src/hooks/executor.rs b/src/hooks/executor.rs index 3221cd2..33f1986 100644 --- a/src/hooks/executor.rs +++ b/src/hooks/executor.rs @@ -55,7 +55,21 @@ use crate::hooks::context::{ /// - [`Interactivity::Interactive`] — `Ask` passes through unchanged, /// allowing the agent to present a prompt to the user. pub struct HookExecutor { + /// Registered hooks in the order they will be invoked. + /// + /// Stored as `Arc` so a single hook instance can be + /// shared across executors cheaply, and so the executor itself is + /// `Send + Sync`. The vector is append-only after construction and + /// never reordered. hooks: Vec>, + + /// How [`HookAction::Ask`] results are handled. + /// + /// Set at construction (default + /// [`Interactivity::Headless`]) and applied by + /// [`apply_interactivity`](Self::apply_interactivity) to every + /// pre-tool-use result. Headless downgrades `Ask` to `Block`; + /// interactive passes it through unchanged. interactivity: Interactivity, } @@ -78,30 +92,38 @@ impl HookExecutor { } } - /// Set the interactivity mode. + /// Set the interactivity mode (builder style). /// - /// Use this builder method to change the mode after construction, or - /// [`with_hook`](Self::with_hook) to add hooks via the builder pattern. + /// Overrides the default [`Interactivity::Headless`] set by + /// [`new`](Self::new). Switch to [`Interactivity::Interactive`] when + /// a human is available to confirm [`HookAction::Ask`] prompts, so + /// they pass through unchanged instead of being downgraded to + /// [`HookAction::Block`]. #[must_use] pub fn with_interactivity(mut self, interactivity: Interactivity) -> Self { self.interactivity = interactivity; self } - /// Register a hook (builder pattern). + /// Register a hook via the builder pattern. /// - /// Hooks are called in registration order. Returns `self` - /// for chaining: `HookExecutor::new().with_hook(a).with_hook(b)`. + /// Appends `hook` to the end of the execution list, so hooks fire + /// in registration order. Returns `self` for chaining — for example + /// `HookExecutor::new().with_hook(a).with_hook(b)`. Use + /// [`register`](Self::register) instead when you need to mutate an + /// existing executor. #[must_use] pub fn with_hook(mut self, hook: Arc) -> Self { self.hooks.push(hook); self } - /// Register a hook (mutating). + /// Register a hook by mutating the executor in place. /// - /// Appends a hook to the end of the execution list. - /// Unlike [`with_hook`](Self::with_hook), this takes `&mut self`. + /// Appends `hook` to the end of the execution list, mirroring + /// [`with_hook`](Self::with_hook) but taking `&mut self` instead of + /// consuming the executor. Useful when hooks are registered + /// conditionally after construction. pub fn register(&mut self, hook: Arc) { self.hooks.push(hook); } @@ -120,18 +142,16 @@ impl HookExecutor { } } - /// Number of registered hooks. + /// Number of hooks currently registered. /// - /// Returns 0 for a freshly constructed executor. + /// Returns `0` for a freshly constructed executor. Cheap (`O(1)`) + /// since it reads the vector length directly; safe to call from + /// hot paths. #[must_use] pub fn hook_count(&self) -> usize { self.hooks.len() } - // ================================================== - // Pre-hook checks (short-circuit on first non-None) - // ================================================== - /// Check pre-tool-use hooks. /// /// Returns the first non-Allow action (e.g., `Block`, `Ask`), continuing @@ -167,10 +187,16 @@ impl HookExecutor { Box::pin(async move { action }) } - /// Check pre-compact hooks. Merges results from all hooks: - /// - If any hook aborts, returns immediately with abort. - /// - Instructions from later hooks override earlier ones. - /// - Additional contexts accumulate. + /// Check pre-compact hooks, merging every hook's [`CompactResult`]. + /// + /// Unlike [`check_pre_tool_use`](Self::check_pre_tool_use), there is + /// no short-circuit on the first result — all hooks run, and their + /// results are combined: if any hook returns + /// [`abort: true`](CompactResult::abort), the merged result is that + /// abort (the first one wins, returned immediately); otherwise the + /// last hook's `new_instructions` overrides earlier ones, and every + /// hook's `additional_context` entries accumulate in registration + /// order. #[must_use] pub fn check_pre_compact(&self, ctx: &PreCompactContext) -> CompactResult { let mut result = CompactResult::allow(); @@ -203,13 +229,13 @@ impl HookExecutor { Box::pin(async move { result }) } - // ================================================== - // Post-hook notifications (all hooks run) - // ================================================== - - /// Notify all post-tool-use hooks. + /// Notify every registered post-tool-use hook. /// - /// All registered hooks are called regardless of return value. + /// Fires [`Hook::on_post_tool_use`] on each hook in registration + /// order. There is no short-circuit and no return value — + /// post-hooks are notification-only, so every hook always runs. + /// Use this for logging, metrics, and side-effects like file + /// tracking. pub fn notify_post_tool_use(&self, ctx: &PostToolUseContext) { for hook in &self.hooks { hook.on_post_tool_use(ctx); @@ -229,27 +255,38 @@ impl HookExecutor { Box::pin(async {}) } - /// Notify all post-compact hooks. + /// Notify every registered post-compact hook. /// - /// All registered hooks are called regardless of return value. + /// Fires [`Hook::on_post_compact`] on each hook in registration + /// order with the compaction outcome (messages removed, tokens + /// saved, duration). Notification-only — every hook always runs, + /// regardless of any prior hook's behaviour. Use it for budget + /// tracking and post-compaction logging. pub fn notify_post_compact(&self, ctx: &PostCompactContext) { for hook in &self.hooks { hook.on_post_compact(ctx); } } - /// Notify all session-start hooks. + /// Notify every registered session-start hook. /// - /// All registered hooks are called regardless of return value. + /// Fires [`Hook::on_session_start`] on each hook in registration + /// order, once, at the beginning of a session. Notification-only — + /// every hook always runs. Use it to initialize per-session state + /// (open resources, reset counters, emit a start log line). pub fn notify_session_start(&self, ctx: &SessionStartContext) { for hook in &self.hooks { hook.on_session_start(ctx); } } - /// Notify all session-end hooks. + /// Notify every registered session-end hook. /// - /// All registered hooks are called regardless of return value. + /// Fires [`Hook::on_session_end`] on each hook in registration + /// order, once, after the loop has terminated. Notification-only — + /// every hook always runs. Use it to flush resources, finalize + /// tracking, and emit a summary log line keyed off + /// [`SessionEndContext::reason`]. pub fn notify_session_end(&self, ctx: &SessionEndContext) { for hook in &self.hooks { hook.on_session_end(ctx); diff --git a/src/lib.rs b/src/lib.rs index 0a8b0f7..e682758 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,7 +7,7 @@ //! //! - **[`message`]** — Core conversation types: messages, parts, tool results. //! - **[`error`]** — Central error enum ([`LoopError`](error::LoopError)) for all framework operations. -//! - **[`config`]** — Session configuration ([`LoopConfig`](config::LoopConfig)). +//! - **[`config`]** — Session configuration ([`SessionConfig`](config::SessionConfig)). //! - **[`cancel`]** — Cooperative cancellation signal (`CancelSignal`). //! //! ## Subsystems @@ -30,7 +30,7 @@ //! ## Runtime & Capabilities //! //! - **[`capabilities`]** — Capability traits ([`Observable`](capabilities::Observable), [`Detectable`](capabilities::Detectable), etc.). -//! - **[`runtime`]** — [`LoopRuntime`](runtime::LoopRuntime) — the default infrastructure bundle. +//! - **[`managers`]** — [`LoopManagers`](managers::LoopManagers) — the default infrastructure bundle. //! //! ## Engine //! @@ -74,6 +74,7 @@ pub mod error; pub mod fallback; #[cfg(feature = "hooks")] pub mod hooks; +pub mod managers; pub mod memory; pub mod message; pub mod middleware; @@ -82,7 +83,6 @@ pub mod presets; #[cfg(feature = "providers")] pub mod provider; pub mod reflection; -pub mod runtime; pub mod stream; pub mod structured; #[cfg(feature = "testing")] diff --git a/src/runtime.rs b/src/managers.rs similarity index 52% rename from src/runtime.rs rename to src/managers.rs index 9ff1cbd..7504395 100644 --- a/src/runtime.rs +++ b/src/managers.rs @@ -1,4 +1,4 @@ -//! Loop runtime — capability traits and runtime infrastructure for agent loops. +//! Manager bundle — capability traits and runtime infrastructure for agent loops. //! //! Two categories of types govern how agent loops //! interact with their infrastructure: @@ -8,7 +8,7 @@ //! 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, +//! models, etc. The [`LoopManagers`] struct implements all capability traits, //! providing a concrete, all-in-one infrastructure bundle. //! //! | Trait | Purpose | @@ -18,18 +18,18 @@ //! | [`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 | +//! | [`Hookable`] | Bidirectional hooks that can block actions | //! | [`PipelineAware`] | Dispatch tools through a middleware pipeline | -//! | `HealthTrackable` | Per-tool health tracking with circuit breakers | +//! | [`HealthTrackable`] | Per-tool health tracking with circuit breakers | //! -//! # `LoopRuntime` +//! # `LoopManagers` //! -//! [`LoopRuntime`] is the framework's default infrastructure bundle. +//! [`LoopManagers`] 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 +//! LoopManagers //! ├── ObserverHost → Observable //! ├── DetectionManager → Detectable //! ├── FallbackManager → FallbackCapable @@ -42,25 +42,24 @@ //! //! # Design Philosophy //! -//! The trait hierarchy separates **what the runtime can do** (capability -//! traits) from **how it's composed** (the `LoopRuntime` struct). This +//! The trait hierarchy separates **what the managers can do** (capability +//! traits) from **how it's composed** (the `LoopManagers` 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 +//! - **Testing** — swap `LoopManagers` for a stub that implements only the //! traits under test. //! ```rust,ignore -//! let runtime = LoopRuntime::builder() +//! let managers = LoopManagers::new() //! .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 +//! Every capability is optional. A bundle 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. @@ -69,7 +68,7 @@ use std::sync::Arc; use crate::compact::ContextManager; use crate::detection::DetectionManager; use crate::detection::{ConvergenceAction, DetectedPattern}; -use crate::error::LoopError; + use crate::fallback::FallbackManager; #[cfg(feature = "hooks")] use crate::hooks::HookExecutor; @@ -82,38 +81,31 @@ 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 +/// `LoopManagers` 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::builder()`] to compose only the capabilities you need, -/// or [`LoopRuntime::new()`] for a runtime with default managers and no -/// optional components: +/// Use [`LoopManagers::new()`] then chain `.with_*()` methods to compose +/// only the capabilities you need: /// /// ``` -/// # use loopctl::runtime::LoopRuntime; -/// # use loopctl::runtime::FallbackCapable; +/// # use loopctl::managers::LoopManagers; +/// # use loopctl::managers::FallbackCapable; /// use loopctl::fallback::FallbackManager; /// -/// // 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")); +/// let managers = LoopManagers::new() +/// .with_fallback(FallbackManager::for_model("llm-70b")); +/// assert_eq!(managers.fallback().active_model().as_deref(), Some("llm-70b")); /// ``` /// /// # Capability Traits /// -/// `LoopRuntime` implements all capability traits defined in this module: +/// `LoopManagers` implements all capability traits defined in this module: /// /// - [`Observable`] — via the internal [`ObserverHost`] /// - [`Detectable`] — via the internal [`DetectionManager`] @@ -126,123 +118,78 @@ pub use crate::capabilities::*; /// /// # Reset /// -/// Call [`reset_all`](LoopRuntime::reset_all) at the start of a new +/// Call [`reset_all`](LoopManagers::reset_all) at the start of a new /// session to reinitialise every manager and observer to its default state. -pub struct LoopRuntime { +pub struct LoopManagers { /// Circuit breaker for API model fallback. - pub fallback: FallbackManager, + /// + /// Tracks consecutive API failures and, when the threshold is exceeded, + /// switches to a configured backup model. Reset at the start of each run + /// via [`reset_all`](Self::reset_all). + fallback: FallbackManager, + /// Loop and convergence detection orchestrator. - pub detection: DetectionManager, + /// + /// Monitors the agent's responses for repetitive patterns (identical tool + /// calls, text convergence) and fires observer events. When thresholds + /// are exceeded, returns an error that aborts the run. + detection: DetectionManager, + /// Observer host for lifecycle event fan-out. + /// + /// Holds the list of registered [`LoopObserver`](crate::observer::LoopObserver) + /// instances and dispatches every lifecycle event (turn start/end, stream + /// deltas, tool dispatch, compaction) to all of them. observer_host: ObserverHost, + /// Optional middleware pipeline wrapping tool dispatch. + /// + /// When set, every tool call passes through the pipeline before reaching + /// the tool itself. Middleware can modify inputs, cache results, enforce + /// permissions, or cap output size. tool_pipeline: Option, + /// Optional context manager for automatic compaction. + /// + /// When set, the driver runs the context manager's compactor after the + /// machine's compaction trigger fires, reducing the history to fit within + /// the context window budget. context_manager: Option>, - /// Optional stream handler for resilient streaming with retries. + + /// Optional stream handler for resilient streaming. + /// + /// When set, wraps streaming calls with retry logic, timeout enforcement, + /// and automatic fallback to non-streaming when the provider drops the + /// connection mid-stream. stream_handler: Option, + /// Optional hook executor for bidirectional lifecycle interception. + /// + /// When set, hooks can inspect and approve or reject actions before they + /// execute (pre-tool, pre-compact) and observe outcomes after (post-tool, + /// post-compact). Unlike observers, hooks can control flow. #[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>, -} - -/// 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. + /// Optional per-tool health tracker with circuit breakers. + /// + /// 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. #[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 - } + health_registry: Option>, } -impl LoopRuntime { +impl LoopManagers { /// 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()); + /// # use loopctl::managers::LoopManagers; + /// # use loopctl::managers::FallbackCapable; + /// let managers = LoopManagers::new(); + /// assert!(managers.fallback().active_model().is_none()); /// ``` #[must_use] pub fn new() -> Self { @@ -260,45 +207,17 @@ impl LoopRuntime { } } - // ================================================== - // 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 /// /// ``` - /// # use loopctl::runtime::LoopRuntime; - /// # use loopctl::runtime::FallbackCapable; + /// # use loopctl::managers::LoopManagers; + /// # use loopctl::managers::FallbackCapable; /// # use loopctl::fallback::FallbackManager; - /// let runtime = LoopRuntime::new() + /// let managers = LoopManagers::new() /// .with_fallback(FallbackManager::for_model("llm-70b")); - /// assert_eq!(runtime.fallback().active_model().as_deref(), Some("llm-70b")); + /// assert_eq!(managers.fallback().active_model().as_deref(), Some("llm-70b")); /// ``` #[must_use] pub fn with_fallback(mut self, fallback: FallbackManager) -> Self { @@ -311,16 +230,16 @@ impl LoopRuntime { /// # Example /// /// ``` - /// # use loopctl::runtime::LoopRuntime; - /// # use loopctl::runtime::Detectable; + /// # use loopctl::managers::LoopManagers; + /// # use loopctl::managers::Detectable; /// # use loopctl::detection::{DetectionManager, DetectionConfig}; /// let config = DetectionConfig { /// loop_threshold: 5, /// ..Default::default() /// }; - /// let runtime = LoopRuntime::new() + /// let managers = LoopManagers::new() /// .with_detection(DetectionManager::new_with_config(config).unwrap()); - /// assert_eq!(runtime.detection().config().loop_threshold, 5); + /// assert_eq!(managers.detection().config().loop_threshold, 5); /// ``` #[must_use] pub fn with_detection(mut self, detection: DetectionManager) -> Self { @@ -328,25 +247,6 @@ impl LoopRuntime { 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); - } - /// Register an observer and return `self` for chaining. /// /// Builder-style alias for [`register_observer`](Self::register_observer). @@ -354,10 +254,9 @@ impl LoopRuntime { /// # Example /// /// ```rust,ignore - /// let runtime = LoopRuntime::builder() + /// let managers = LoopManagers::new() /// .with_observer(Arc::new(logging_observer)) /// .with_observer(Arc::new(metrics_observer)) - /// .build(); /// ``` #[must_use] pub fn with_observer(mut self, observer: Arc) -> Self { @@ -366,6 +265,9 @@ impl LoopRuntime { } /// Access the observer host directly. + /// + /// The observer host manages fan-out to all registered observers. + /// Use this to fire events or inspect the registered observer list. pub fn observers(&self) -> &ObserverHost { &self.observer_host } @@ -375,9 +277,8 @@ impl LoopRuntime { /// # Example /// /// ```rust,ignore - /// let runtime = LoopRuntime::builder() + /// let managers = LoopManagers::new() /// .with_pipeline(builder.build()?) - /// .build(); /// ``` #[must_use] pub fn with_pipeline(mut self, pipeline: ToolPipeline) -> Self { @@ -385,39 +286,27 @@ impl LoopRuntime { 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 (builder-style). + /// Set the tool health registry for per-tool health tracking (builder-style). /// - /// # Example + /// 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. /// - /// ```rust,ignore - /// let runtime = LoopRuntime::builder() - /// .with_context_manager(Arc::new(manager)) - /// .build(); - /// ``` + /// *Requires `tool_health` feature.* #[must_use] - pub fn with_context_manager(mut self, manager: Arc) -> Self { - self.context_manager = Some(manager); + #[cfg(feature = "tool_health")] + pub fn with_health_registry(mut self, registry: Arc) -> Self { + self.health_registry = Some(registry); 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 (builder-style). /// /// # Example /// /// ```rust,ignore - /// let runtime = LoopRuntime::builder() + /// let managers = LoopManagers::new() /// .with_stream_handler(handler) - /// .build(); /// ``` #[must_use] pub fn with_stream_handler(mut self, handler: StreamHandler) -> Self { @@ -425,11 +314,6 @@ impl LoopRuntime { 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 (builder-style). /// /// *Requires `hooks` feature.* @@ -440,29 +324,72 @@ impl LoopRuntime { self } - /// Set the hook executor (`&mut self` variant). + /// Register an observer with the observer host. /// - /// *Requires `hooks` feature.* - #[cfg(feature = "hooks")] - pub fn set_hook_executor(&mut self, executor: Arc) { - self.hook_executor = Some(executor); + /// # Example + /// + /// ```rust,ignore + /// use loopctl::observer::LoopObserver; + /// use std::sync::Arc; + /// + /// let mut managers = LoopManagers::new(); + /// managers.register_observer(Arc::new(MyObserver)); + /// ``` + pub fn register_observer(&mut self, observer: Arc) { + self.observer_host.register(observer); } - /// Set the tool health registry for per-tool health tracking (builder-style). + /// Set the middleware pipeline for tool dispatch. /// - /// 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. + /// Non-consuming variant of [`with_pipeline`](Self::with_pipeline) for + /// cases where the managers is already constructed. + pub fn set_pipeline(&mut self, pipeline: ToolPipeline) { + self.tool_pipeline = Some(pipeline); + } + + /// Set the context manager for automatic compaction (builder-style). /// - /// *Requires `tool_health` feature.* + /// # Example + /// + /// ```rust,ignore + /// let managers = LoopManagers::new() + /// .with_context_manager(Arc::new(manager)) + /// ``` #[must_use] - #[cfg(feature = "tool_health")] - pub fn with_health_registry(mut self, registry: Arc) -> Self { - self.health_registry = Some(registry); + pub fn with_context_manager(mut self, manager: Arc) -> Self { + self.context_manager = Some(manager); self } - /// Set the tool health registry (`&mut self` variant). + /// Set the context manager for automatic compaction. + /// + /// Non-consuming variant of [`with_context_manager`](Self::with_context_manager). + /// The context manager's context window is synced from the session + /// config when attached via [`BareLoop::set_context_manager`](crate::engine::BareLoop::set_context_manager). + pub fn set_context_manager(&mut self, manager: Arc) { + self.context_manager = Some(manager); + } + + /// Set the stream handler for resilient streaming. + /// + /// Non-consuming variant of [`with_stream_handler`](Self::with_stream_handler). + pub fn set_stream_handler(&mut self, handler: StreamHandler) { + self.stream_handler = Some(handler); + } + + /// Set the hook executor for bidirectional lifecycle interception. + /// + /// Non-consuming variant of [`with_hook_executor`](Self::with_hook_executor). + /// + /// *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. + /// + /// Non-consuming variant of [`with_health_registry`](Self::with_health_registry). /// /// *Requires `tool_health` feature.* #[cfg(feature = "tool_health")] @@ -470,10 +397,6 @@ impl LoopRuntime { self.health_registry = Some(registry); } - // ================================================== - // Lifecycle - // ================================================== - /// Reset all managers and observers to their initial state. /// /// Delegates to each manager's `reset()` method and calls @@ -483,10 +406,10 @@ impl LoopRuntime { /// # Example /// /// ``` - /// # use loopctl::runtime::LoopRuntime; - /// let runtime = LoopRuntime::new(); + /// # use loopctl::managers::LoopManagers; + /// let managers = LoopManagers::new(); /// // ... after a session ... - /// runtime.reset_all(); + /// managers.reset_all(); /// // All managers are back to their initial state /// ``` pub fn reset_all(&self) { @@ -495,25 +418,18 @@ impl LoopRuntime { self.observer_host.reset_all(); } - // ================================================== - // Detection interpretation - // ================================================== - - /// Interpret a [`DetectedPattern`] and decide whether to abort. + /// Fire tracing log lines and observer events for a detected pattern. /// - /// Generic framework logic: checks loop-detection thresholds, - /// maps convergence actions, and notifies observers. Returns - /// `None` to continue the session, or `Some(error)` to abort. - /// - /// Called by the agent loop after each response is recorded with the - /// [`DetectionManager`]. - pub fn handle_detected_pattern( + /// Called by the agent loop after each response or tool operation is + /// recorded with the [`DetectionManager`]. Does **not** decide whether + /// to abort — that is the engine's job. + pub fn notify_detected_pattern( &self, pattern: &crate::detection::DetectedPattern, turn: usize, - ) -> Option { + ) { match pattern { - DetectedPattern::NoPattern => None, + DetectedPattern::NoPattern => {} DetectedPattern::LoopDetected { repetitions, @@ -530,20 +446,6 @@ impl LoopRuntime { 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(LoopError::LoopDetected { - message: format!("{pattern_description} repeated {repetitions} times"), - }) - } else { - None - } } DetectedPattern::ConvergenceDetected { @@ -564,134 +466,45 @@ impl LoopRuntime { .on_convergence_detected(&ConvergenceDetectedContext { action: action_str.to_string(), }); - - match action { - ConvergenceAction::Stop => Some(LoopError::LoopDetected { - message: "agent stopped: convergence detected".into(), - }), - ConvergenceAction::AskUser => Some(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 { - /// Produce a [`LoopRuntime`] with default managers and no optional components. +impl Default for LoopManagers { + /// Produce a [`LoopManagers`] with default managers and no optional components. /// - /// Equivalent to [`LoopRuntime::new`]. + /// Equivalent to [`LoopManagers::new`]. fn default() -> Self { Self::new() } } -// ================================================== -// Capability trait implementations -// ================================================== - -impl crate::capabilities::Observable for LoopRuntime { +impl crate::capabilities::Observable for LoopManagers { fn observers(&self) -> &ObserverHost { &self.observer_host } } -impl crate::capabilities::Detectable for LoopRuntime { +impl crate::capabilities::Detectable for LoopManagers { fn detection(&self) -> &DetectionManager { &self.detection } } -impl crate::capabilities::FallbackCapable for LoopRuntime { +impl crate::capabilities::FallbackCapable for LoopManagers { fn fallback(&self) -> &FallbackManager { &self.fallback } } -impl crate::capabilities::Compactable for LoopRuntime { +impl crate::capabilities::Compactable for LoopManagers { fn context_manager(&self) -> Option<&Arc> { self.context_manager.as_ref() } } -impl crate::capabilities::StreamCapable for LoopRuntime { +impl crate::capabilities::StreamCapable for LoopManagers { fn stream_handler(&self) -> &StreamHandler { self.stream_handler .as_ref() @@ -700,20 +513,20 @@ impl crate::capabilities::StreamCapable for LoopRuntime { } #[cfg(feature = "hooks")] -impl crate::capabilities::Hookable for LoopRuntime { +impl crate::capabilities::Hookable for LoopManagers { fn hook_executor(&self) -> Option<&HookExecutor> { self.hook_executor.as_deref() } } -impl crate::capabilities::PipelineAware for LoopRuntime { +impl crate::capabilities::PipelineAware for LoopManagers { fn pipeline(&self) -> Option<&ToolPipeline> { self.tool_pipeline.as_ref() } } #[cfg(feature = "tool_health")] -impl crate::capabilities::HealthTrackable for LoopRuntime { +impl crate::capabilities::HealthTrackable for LoopManagers { fn health_registry(&self) -> Option<&ToolHealthRegistry> { self.health_registry.as_deref() } @@ -726,31 +539,31 @@ mod tests { #[test] fn test_runtime_default() { - let runtime = LoopRuntime::default(); - assert!(runtime.fallback().active_model().is_none()); + let managers = LoopManagers::default(); + assert!(managers.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); + let managers = LoopManagers::new().with_fallback(fallback); assert_eq!( - runtime.fallback().active_model().as_deref(), + managers.fallback().active_model().as_deref(), Some("my-model") ); } #[test] fn test_reset_all() { - let runtime = LoopRuntime::new(); - runtime.reset_all(); + let managers = LoopManagers::new(); + managers.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); + let managers = LoopManagers::new(); + assert_eq!(managers.detection().config().loop_threshold, 3); + assert_eq!(managers.detection().config().stop_threshold, 10); } #[test] @@ -763,24 +576,23 @@ mod tests { ..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()); + let managers = LoopManagers::new().with_detection(detection); + assert_eq!(managers.detection().config().loop_threshold, 7); + assert_eq!(managers.detection().config().stop_threshold, 20); + assert!(managers.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); + let managers = LoopManagers::new(); + let _ = managers.detection().record_tool_call("Read", 12345); + managers.reset_all(); + let pattern = managers.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) {} @@ -788,37 +600,37 @@ mod tests { 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); + let managers = LoopManagers::new(); + _assert_observable(&managers); + _assert_detectable(&managers); + _assert_fallback(&managers); + _assert_pipeline(&managers); + _assert_compactable(&managers); + _assert_stream_capable(&managers); } #[test] fn test_pipeline_defaults_to_none() { - let runtime = LoopRuntime::new(); - assert!(runtime.pipeline().is_none()); + let managers = LoopManagers::new(); + assert!(managers.pipeline().is_none()); } #[test] fn test_observers_accessible_via_trait() { - let runtime = LoopRuntime::new(); - let _: &ObserverHost = runtime.observers(); + let managers = LoopManagers::new(); + let _: &ObserverHost = managers.observers(); } #[test] fn test_context_manager_defaults_to_none() { - let runtime = LoopRuntime::new(); - assert!(runtime.context_manager().is_none()); + let managers = LoopManagers::new(); + assert!(managers.context_manager().is_none()); } #[test] fn test_stream_handler_defaults_to_passthrough() { - let runtime = LoopRuntime::new(); - let handler = runtime.stream_handler(); + let managers = LoopManagers::new(); + let handler = managers.stream_handler(); assert_eq!( handler.timeout_config().total_stream_timeout, std::time::Duration::MAX @@ -833,10 +645,10 @@ mod tests { 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()); + let mut managers = LoopManagers::new(); + assert!(managers.pipeline().is_none()); + managers.set_pipeline(pipeline); + assert!(managers.pipeline().is_some()); } #[test] @@ -845,10 +657,10 @@ mod tests { 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()); + let mut managers = LoopManagers::new(); + assert!(managers.context_manager().is_none()); + managers.set_context_manager(Arc::new(manager)); + assert!(managers.context_manager().is_some()); } #[test] @@ -856,18 +668,18 @@ mod tests { use crate::stream::handler::StreamHandler; let handler = StreamHandler::new(); - let mut runtime = LoopRuntime::new(); + let mut managers = LoopManagers::new(); assert_eq!( - runtime + managers .stream_handler() .timeout_config() .total_stream_timeout, std::time::Duration::MAX ); - runtime.set_stream_handler(handler); + managers.set_stream_handler(handler); // After set_stream_handler, the production defaults apply (15 min total). assert_eq!( - runtime + managers .stream_handler() .timeout_config() .total_stream_timeout, @@ -887,20 +699,20 @@ mod tests { } } - 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); + let mut managers = LoopManagers::new(); + assert!(managers.observers().is_empty()); + managers.register_observer(Arc::new(NopObserver)); + assert_eq!(managers.observers().len(), 1); + managers.register_observer(Arc::new(NopObserver)); + assert_eq!(managers.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()); + let managers = LoopManagers::new(); + let _ = managers.fallback().record_api_failure(); + managers.reset_all(); + assert!(managers.fallback().active_model().is_none()); } #[test] @@ -915,13 +727,13 @@ mod tests { } } - let mut runtime = LoopRuntime::new(); - runtime.register_observer(Arc::new(NopObserver)); - assert_eq!(runtime.observers().len(), 1); - runtime.reset_all(); + let mut managers = LoopManagers::new(); + managers.register_observer(Arc::new(NopObserver)); + assert_eq!(managers.observers().len(), 1); + managers.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); + assert_eq!(managers.observers().len(), 1); } #[test] @@ -934,13 +746,13 @@ mod tests { 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); + let managers = LoopManagers::new(); + accepts_observable(&managers); + accepts_detectable(&managers); + accepts_fallback(&managers); + accepts_compactable(&managers); + accepts_stream_capable(&managers); + accepts_pipeline(&managers); + accepts_multi_bound(&managers); } } diff --git a/src/memory/builtin.rs b/src/memory/builtin.rs index 949cc56..ea91754 100644 --- a/src/memory/builtin.rs +++ b/src/memory/builtin.rs @@ -122,10 +122,6 @@ pub struct InMemoryStore { entries: RwLock>, } -// =================================================== -// Construction -// =================================================== - impl InMemoryStore { /// Create a new empty store. /// @@ -177,10 +173,6 @@ 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. diff --git a/src/memory/entry.rs b/src/memory/entry.rs index 4962968..6e9c8a7 100644 --- a/src/memory/entry.rs +++ b/src/memory/entry.rs @@ -27,20 +27,58 @@ use uuid::Uuid; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct MemoryEntry { /// UUID v4 for deduplication and stable reference during consolidation. + /// + /// Generated at construction by [`MemoryEntry::new`] and never reused, so + /// two entries are distinct even if their text is identical. Used as the + /// stable key when merging or pruning the store. pub id: Uuid, + /// Entry category, influencing ranking and consolidation rules. + /// + /// The [`MemoryCategory`] this entry belongs to (insight, error pattern, + /// strategy, etc.), which retrieval and consolidation strategies use to + /// weight or group entries differently. pub category: MemoryCategory, + /// What the agent learned — free-form text. + /// + /// The actual knowledge payload, in whatever natural-language form the + /// agent captured it. This is the primary content returned by retrieval. pub memory: String, + /// Arbitrary labels for categorization and retrieval (e.g. `"performance"`, `"security"`). + /// + /// User- or agent-supplied tags that broaden retrieval beyond the formal + /// [`category`](Self::category), letting queries filter by topic without + /// parsing the free-form `memory` text. pub tags: Vec, + /// Timestamp set by [`MemoryEntry::new`]. Recency-based strategies use this. + /// + /// Wall-clock time the entry was created, captured once at construction. + /// Recency-aware ranking and consolidation use the age derived from it to + /// favour fresher knowledge. pub created_at: SystemTime, + /// Relevance score (0.0–1.0). Starts at 1.0; implementations may decay over time. + /// + /// Numeric prominence used to rank entries during retrieval and decide + /// which to prune. New entries begin at `1.0`; stores may decay it over + /// time or boost it on repeated access. pub relevance: f32, + /// Number of times this entry has been retrieved. + /// + /// Popularity counter incremented each time the entry is surfaced, + /// feeding into ranking (frequently retrieved entries are deemed more + /// useful) and protecting high-traffic entries from pruning. pub access_count: usize, + /// Whether this entry has been validated. Consolidation prefers keeping validated entries. + /// + /// Confidence flag set via the [`validated`](MemoryEntry::validated) + /// builder. Consolidation treats validated entries as higher-trust and is + /// less likely to prune them during a cleanup pass. pub validated: bool, } @@ -202,13 +240,34 @@ pub enum MemoryCategory { #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct ConsolidationStats { /// Number of entries before consolidation. + /// + /// Size of the store at the start of the consolidation pass, giving the + /// baseline against which `entries_after`, `pruned`, and `merged` are + /// compared to measure the pass's effect. pub entries_before: usize, + /// Number of entries after consolidation. + /// + /// Resulting store size once pruning and merging are complete, so callers + /// can see the net footprint change at a glance. pub entries_after: usize, + /// Entries removed (low relevance, stale, superseded). + /// + /// Count of entries discarded outright during the pass — typically those + /// whose relevance decayed below a threshold or that were superseded by + /// newer entries. pub pruned: usize, + /// Entries merged (duplicates combined). + /// + /// Count of duplicate or near-duplicate entries folded into single + /// consolidated entries, rather than removed entirely. pub merged: usize, + /// Estimated storage reclaimed in bytes. + /// + /// Rough byte savings from pruning and merging, useful for reporting how + /// much memory or disk the consolidation pass freed up. pub bytes_saved: usize, } diff --git a/src/message.rs b/src/message.rs index 520773b..87ed2d5 100644 --- a/src/message.rs +++ b/src/message.rs @@ -23,21 +23,13 @@ //! //! # Message Flow //! -//! ```text -//! User sends text ──▶ Message::user("query") -//! │ -//! API processes -//! │ -//! ◀── Message::assistant("thinking...") -//! ◀── MessagePart::ToolCall { id, name, input } -//! │ -//! Framework executes tool -//! │ -//! ──▶ MessagePart::ToolResult { call_id, output } -//! │ -//! API continues -//! ◀── Message::assistant("final answer") -//! ``` +//! A turn starts with a user message (`Message::user`), which is sent to the +//! LLM API. The API responds with an assistant message that may contain plain +//! text plus one or more `MessagePart::ToolCall` parts. For each tool call, +//! the framework executes the tool and appends a user-role message carrying a +//! `MessagePart::ToolResult` part back into the history. The API then +//! continues, typically producing a final assistant text message that +//! completes the turn. //! //! # Quick Start //! @@ -69,10 +61,6 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use std::fmt; -// ================================================== -// Message -// ================================================== - /// A message in the conversation. /// /// Messages are the fundamental unit of communication between users, @@ -197,25 +185,39 @@ impl Message { pub const fn new(role: Role, parts: Vec) -> Self { Self { role, parts } } + + /// Concatenate every text part into a single string. + /// + /// Walks the message's parts in order and joins all text content, + /// skipping non-text parts (tool calls, tool results, images). + #[must_use] + pub fn text_content(&self) -> String { + self.parts + .iter() + .filter_map(MessagePart::as_text) + .collect::>() + .join("") + } + + /// Collect the tool calls this message requests, in order. + /// + /// Scans the message's parts for tool-call parts and maps each to a + /// `(id, name, input)` tuple. Returns an empty vector when the message + /// contains no tool calls. + #[must_use] + pub fn tool_call_parts(&self) -> Vec<(&str, &str, &serde_json::Value)> { + self.parts + .iter() + .filter_map(|part| match part { + MessagePart::ToolCall { id, name, input } => { + Some((id.as_str(), name.as_str(), input)) + } + _ => None, + }) + .collect() + } } -/// 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 -/// `[Tool: {name} with input: {input}]`; tool-result parts render -/// as `[Tool Result: {content}]`; image parts render as -/// `[Image: {media_type}]`. -/// -/// For structured serialization, use [`serde_json::to_string`] instead. -/// -/// # Example -/// -/// ```rust -/// use loopctl::message::{Message}; -/// let msg = Message::user("Hello"); -/// println!("{msg}"); // prints "Hello" -/// ``` impl fmt::Display for Message { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let mut chunks = Vec::new(); @@ -241,10 +243,6 @@ impl fmt::Display for Message { } } -// ================================================== -// Role -// ================================================== - /// The role of a message sender in the conversation. /// /// Each [`Message`] has exactly one role that identifies who produced @@ -290,19 +288,6 @@ pub enum Role { System, } -/// Formats a [`Role`] as its lowercase API string. -/// -/// This implementation is used by [`Display`](fmt::Display) to produce -/// the string that LLM APIs expect: `"user"`, `"assistant"`, or `"system"`. -/// It is also used for logging and building request payloads -/// without pulling in the serde serializer. -/// -/// # Example -/// -/// ```rust -/// use loopctl::message::{Role}; -/// assert_eq!(format!("Role: {}", Role::User), "Role: user"); -/// ``` impl fmt::Display for Role { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { @@ -313,10 +298,6 @@ impl fmt::Display for Role { } } -// ================================================== -// MessagePart -// ================================================== - /// A part of message content within a [`Message`]. /// /// Messages can contain multiple types of content interleaved in a @@ -621,10 +602,6 @@ impl MessagePart { } } -// ================================================== -// ImageSource -// ================================================== - /// Source data for an image part in loopctl API format. /// /// Represents a base64-encoded image with its MIME type. Serialized @@ -702,10 +679,6 @@ impl ImageSource { } } -// ================================================== -// ToolContent -// ================================================== - /// Content that can be returned from a tool invocation. /// /// Supports both simple string results and multipart results with @@ -844,66 +817,18 @@ impl Default for ToolContent { } } -/// Converts a [`String`] into a [`ToolContent::Text`]. -/// -/// This blanket conversion allows passing owned strings directly -/// into APIs that accept [`Into`], for example -/// [`MessagePart::tool_result`]. -/// -/// # Example -/// -/// ```rust -/// use loopctl::message::{ToolContent}; -/// let content: ToolContent = "File found".to_string().into(); -/// assert!(content.is_string()); -/// ``` impl From for ToolContent { fn from(s: String) -> Self { Self::Text(s) } } -/// Converts a `&str` into a [`ToolContent::Text`]. -/// -/// This conversion allocates a new [`String`] from the borrowed slice. -/// It enables ergonomic usage with string literals: -/// -/// # Example -/// -/// ```rust -/// use loopctl::message::{ToolContent}; -/// let content: ToolContent = "ok".into(); -/// assert!(content.is_string()); -/// ``` impl From<&str> for ToolContent { fn from(s: &str) -> Self { Self::Text(s.to_owned()) } } -/// Formats a [`ToolContent`] for display. -/// -/// For the [`String`](ToolContent::Text) variant, writes the -/// text directly. For the [`Multipart`](ToolContent::Multipart) -/// variant, concatenates all [`ToolContentPart::Text`] parts with -/// newlines, silently skipping any non-text parts (e.g. images). -/// -/// For quick debugging and logging, use [`Display`](std::fmt::Display). -/// For full structured serialization, use [`serde_json::to_string`]. -/// -/// # Example -/// -/// ```rust -/// use loopctl::message::{ToolContent, ToolContentPart}; -/// let content = ToolContent::from_string("Done!"); -/// assert_eq!(content.to_string(), "Done!"); -/// -/// let multipart = ToolContent::from_multipart(vec![ -/// ToolContentPart::text("line 1"), -/// ToolContentPart::text("line 2"), -/// ]); -/// assert_eq!(multipart.to_string(), "line 1\nline 2"); -/// ``` impl fmt::Display for ToolContent { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { @@ -925,10 +850,6 @@ impl fmt::Display for ToolContent { } } -// ================================================== -// ToolContentPart -// ================================================== - /// A part of a Multipart [`ToolContent`]. /// /// Represents a single piece of content within a multi-part tool result. @@ -1028,10 +949,6 @@ impl ToolContentPart { } } -// ================================================== -// Tests -// ================================================== - #[cfg(test)] mod tests { use super::*; diff --git a/src/middleware.rs b/src/middleware.rs index 411a8fa..3b8c510 100644 --- a/src/middleware.rs +++ b/src/middleware.rs @@ -533,10 +533,6 @@ impl Default for ToolPipelineBuilder { } } -// =================================================== -// Tests -// =================================================== - #[cfg(test)] mod tests { use super::*; diff --git a/src/middleware/memoize.rs b/src/middleware/memoize.rs index b809f32..2801ddd 100644 --- a/src/middleware/memoize.rs +++ b/src/middleware/memoize.rs @@ -48,10 +48,6 @@ use crate::message::{ToolContent, ToolContentPart}; use super::{ToolDispatchContext, ToolDispatchResult, ToolMiddleware, ToolPipeline}; -// =================================================== -// PathExtractor trait -// =================================================== - /// Extracts the filesystem paths a tool call touches, for cache invalidation. /// /// Implement this per tool (or per tool family) to tell @@ -116,10 +112,6 @@ impl PathExtractor for NoopPathExtractor { } } -// =================================================== -// Cache types (private) -// =================================================== - /// Cache key: the tool name plus the hash of its canonical-serialized input. /// /// Both components are required: the same input hash could theoretically @@ -172,10 +164,6 @@ struct CacheEntry { paths: Vec, } -// =================================================== -// MemoizingMiddleware -// =================================================== - /// Deduplicates repeat tool calls by `(tool_name, hash(canonical input))`. /// /// On a cache hit (memoized tool + key present + not TTL-expired), diff --git a/src/middleware/timeout.rs b/src/middleware/timeout.rs index cdcefd9..c25e17f 100644 --- a/src/middleware/timeout.rs +++ b/src/middleware/timeout.rs @@ -10,10 +10,24 @@ use std::time::Duration; #[derive(Debug, Clone)] pub struct TimeoutConfig { /// Per-tool execution deadline. + /// + /// Each tool invocation is given this long to complete before it is + /// cancelled and (optionally) retried. Set to [`Duration::MAX`] via + /// [`TimeoutMiddleware::none`] to disable the timeout entirely. pub timeout: Duration, + /// Retry with exponential backoff up to [`max_retries`](TimeoutConfig::max_retries) times. + /// + /// When `true`, a timed-out call is retried with a doubled deadline + /// rather than failing immediately; when `false` the first timeout + /// produces the final error. pub retry_on_timeout: bool, + /// Maximum number of retries (0 = no retry). Total attempts = `1 + max_retries`. + /// + /// Upper bound on additional attempts after the initial one. Each retry + /// doubles the deadline, so this also bounds how long the middleware will + /// keep trying before giving up. pub max_retries: u32, } diff --git a/src/middleware/unknown_tool.rs b/src/middleware/unknown_tool.rs index cf79e3d..9429e2f 100644 --- a/src/middleware/unknown_tool.rs +++ b/src/middleware/unknown_tool.rs @@ -112,7 +112,8 @@ impl UnknownToolMiddleware { 0.0 }; - f64::from(lcs_u32) / f64::from(max_len) + prefix_bonus + let score = f64::from(lcs_u32) / f64::from(max_len) + prefix_bonus; + score.clamp(0.0, 1.0) } /// Compute the length of the longest common subsequence of two character slices. diff --git a/src/middleware/verify.rs b/src/middleware/verify.rs index 6caa209..711c012 100644 --- a/src/middleware/verify.rs +++ b/src/middleware/verify.rs @@ -29,10 +29,6 @@ use crate::tool::ToolContext; use super::{ToolDispatchContext, ToolDispatchResult, ToolMiddleware, ToolPipeline}; -// =================================================== -// Verifier trait + VerifyResult -// =================================================== - /// A pluggable verifier run after write-class tools. /// /// Implement this trait to supply the verification logic for @@ -146,10 +142,6 @@ impl Verifier for NoopVerifier { } } -// =================================================== -// VerifyMiddleware -// =================================================== - /// Post-execution middleware that runs a [`Verifier`] after configured /// write-class tools. /// diff --git a/src/observer.rs b/src/observer.rs index e95a1cc..1cb0f50 100644 --- a/src/observer.rs +++ b/src/observer.rs @@ -47,13 +47,10 @@ pub use context::{ StreamFailureContext, TextDeltaContext, ThinkingDeltaContext, ToolCallReceivedContext, ToolPostContext, ToolPreContext, TurnEndContext, TurnStartContext, }; -// ================================================== -// LoopObserver Trait -// ================================================== /// A notification observer that receives typed callbacks at agent loop lifecycle points. /// -/// Observers are registered via [`LoopRuntime`](crate::runtime::LoopRuntime) and called at each +/// Observers are registered via [`LoopManagers`](crate::managers::LoopManagers) and called at each /// lifecycle point in registration order. All methods are **notification-only** — they /// return `()`. Use the hook system (requires `hooks` feature) if you need to control /// flow (block/allow actions). @@ -246,10 +243,6 @@ pub trait LoopObserver: Send + Sync { fn reset(&self) {} } -// ================================================== -// ObserverHost -// ================================================== - /// Holds registered observers and dispatches notifications to each. /// /// Observers run in registration order. All observers are always notified — @@ -279,13 +272,44 @@ impl ObserverHost { self.observers.push(observer); } - /// Number of registered observers. + /// Dispatch a closure to each observer, isolating panics. + /// + /// If an observer panics, it is logged and the remaining observers + /// still receive the event. This matches the panic-isolation applied + /// to tool dispatch. + fn dispatch(&self, f: F) + where + F: Fn(&dyn LoopObserver), + { + use std::panic::{AssertUnwindSafe, catch_unwind}; + + for obs in &self.observers { + let obs: &dyn LoopObserver = obs.as_ref(); + if let Err(payload) = catch_unwind(AssertUnwindSafe(|| f(obs))) { + tracing::error!( + observer = obs.name(), + payload_downcast = payload.is::<&str>(), + "observer panicked; continuing with remaining observers" + ); + } + } + } + + /// Number of observers currently registered. + /// + /// Cheap (`O(1)`) read of the observer vector length. Returns `0` + /// for a freshly constructed host, in which case every dispatch + /// method is effectively a no-op (it iterates an empty `Vec`). #[must_use] pub fn len(&self) -> usize { self.observers.len() } /// Whether no observers are registered. + /// + /// `true` when [`len`](Self::len) is `0`. When `true`, every + /// dispatch call short-circuits through an empty loop, so an + /// observerless host adds negligible overhead to the agent loop. #[must_use] pub fn is_empty(&self) -> bool { self.observers.is_empty() @@ -296,72 +320,72 @@ impl ObserverHost { /// 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(); - } + self.dispatch(|obs| obs.reset()); } /// Dispatch [`LoopObserver::on_session_start`] to all observers. /// - /// Iterates registered observers in registration order. + /// Fired once per session, before the first turn begins. Iterates + /// registered observers in registration order; use it to initialize + /// per-session observer state. pub fn on_session_start(&self, ctx: &SessionStartContext) { - for obs in &self.observers { - obs.on_session_start(ctx); - } + self.dispatch(|obs| obs.on_session_start(ctx)); } /// Dispatch [`LoopObserver::on_session_end`] to all observers. /// - /// Iterates registered observers in registration order. + /// Fired once per session, after the loop has terminated. Iterates + /// registered observers in registration order; check + /// [`SessionEndContext::success`] in each observer to distinguish a + /// normal exit from a failure. pub fn on_session_end(&self, ctx: &SessionEndContext) { - for obs in &self.observers { - obs.on_session_end(ctx); - } + self.dispatch(|obs| obs.on_session_end(ctx)); } /// Dispatch [`LoopObserver::on_turn_start`] to all observers. /// - /// Iterates registered observers in registration order. + /// Fired before the model is called for the turn. Iterates + /// registered observers in registration order. pub fn on_turn_start(&self, ctx: &TurnStartContext) { - for obs in &self.observers { - obs.on_turn_start(ctx); - } + self.dispatch(|obs| obs.on_turn_start(ctx)); } /// Dispatch [`LoopObserver::on_turn_end`] to all observers. /// - /// Iterates registered observers in registration order. + /// Fired after the turn's model call and any tool dispatch have + /// completed. Iterates registered observers in registration order; + /// check [`TurnEndContext::success`] to detect turn-level failures. pub fn on_turn_end(&self, ctx: &TurnEndContext) { - for obs in &self.observers { - obs.on_turn_end(ctx); - } + self.dispatch(|obs| obs.on_turn_end(ctx)); } /// Dispatch [`LoopObserver::on_stream_success`] to all observers. /// - /// Iterates registered observers in registration order. + /// Fired when a streaming model call completes successfully. + /// Iterates registered observers in registration order. Not fired + /// when the stream fails — see + /// [`on_stream_failure`](Self::on_stream_failure). pub fn on_stream_success(&self, ctx: &StreamContext) { - for obs in &self.observers { - obs.on_stream_success(ctx); - } + self.dispatch(|obs| obs.on_stream_success(ctx)); } /// Dispatch [`LoopObserver::on_stream_failure`] to all observers. /// - /// Iterates registered observers in registration order. + /// Fired when a streaming model call fails (timeout, transport + /// error, rate limit). Iterates registered observers in registration + /// order; the loop may retry or fall back to another model after + /// this notification. pub fn on_stream_failure(&self, ctx: &StreamFailureContext) { - for obs in &self.observers { - obs.on_stream_failure(ctx); - } + self.dispatch(|obs| obs.on_stream_failure(ctx)); } /// Dispatch [`LoopObserver::on_response`] to all observers. /// - /// Iterates registered observers in registration order. + /// Fired once the full model response is assembled — the committed + /// assistant text plus any tool calls, before the framework resolves + /// tool results. Iterates registered observers in registration order. pub fn on_response(&self, ctx: &ResponseContext) { - for obs in &self.observers { - obs.on_response(ctx); - } + self.dispatch(|obs| obs.on_response(ctx)); } /// Dispatch [`LoopObserver::on_text_delta`] to all observers. @@ -369,9 +393,7 @@ impl ObserverHost { /// Iterates registered observers in registration order. Called once per /// streamed text delta, so each observer's `on_text_delta` must be cheap. pub fn on_text_delta(&self, ctx: &TextDeltaContext) { - for obs in &self.observers { - obs.on_text_delta(ctx); - } + self.dispatch(|obs| obs.on_text_delta(ctx)); } /// Dispatch [`LoopObserver::on_thinking_delta`] to all observers. @@ -380,88 +402,87 @@ impl ObserverHost { /// streamed reasoning delta, so each observer's `on_thinking_delta` must /// be cheap. pub fn on_thinking_delta(&self, ctx: &ThinkingDeltaContext) { - for obs in &self.observers { - obs.on_thinking_delta(ctx); - } + self.dispatch(|obs| obs.on_thinking_delta(ctx)); } /// Dispatch [`LoopObserver::on_tool_pre`] to all observers. /// - /// Iterates registered observers in registration order. + /// Fired before a tool is dispatched, carrying the call's name and + /// input. Iterates registered observers in registration order. + /// Notification-only — use the hook system (requires the `hooks` + /// feature) for flow control. pub fn on_tool_pre(&self, ctx: &ToolPreContext) { - for obs in &self.observers { - obs.on_tool_pre(ctx); - } + self.dispatch(|obs| obs.on_tool_pre(ctx)); } /// Dispatch [`LoopObserver::on_tool_call_received`] to all observers. /// - /// Iterates registered observers in registration order. + /// Fired when the framework receives a tool call from the model + /// response, before it is dispatched. Iterates registered observers + /// in registration order; useful for pending-call tracking. pub fn on_tool_call_received(&self, ctx: &ToolCallReceivedContext) { - for obs in &self.observers { - obs.on_tool_call_received(ctx); - } + self.dispatch(|obs| obs.on_tool_call_received(ctx)); } /// Dispatch [`LoopObserver::on_tool_post`] to all observers. /// - /// Iterates registered observers in registration order. + /// Fired after a tool completes, carrying its output and timing. + /// Iterates registered observers in registration order; useful for + /// loop-detection correlation. pub fn on_tool_post(&self, ctx: &ToolPostContext) { - for obs in &self.observers { - obs.on_tool_post(ctx); - } + self.dispatch(|obs| obs.on_tool_post(ctx)); } /// Dispatch [`LoopObserver::on_compaction`] to all observers. /// - /// Iterates registered observers in registration order. + /// Fired after the context manager reduces the history, carrying the + /// before/after token counts. Iterates registered observers in + /// registration order; fired only when compaction actually occurred, + /// not on no-action passes. pub fn on_compaction(&self, ctx: &CompactedContext) { - for obs in &self.observers { - obs.on_compaction(ctx); - } + self.dispatch(|obs| obs.on_compaction(ctx)); } /// Dispatch [`LoopObserver::on_fallback`] to all observers. /// - /// Iterates registered observers in registration order. + /// Fired when the fallback manager activates a fallback model, + /// carrying the reason and the selected model. Iterates registered + /// observers in registration order; the fallback model will be used + /// for subsequent requests. pub fn on_fallback(&self, ctx: &FallbackContext) { - for obs in &self.observers { - obs.on_fallback(ctx); - } + self.dispatch(|obs| obs.on_fallback(ctx)); } /// Dispatch [`LoopObserver::on_model_switched`] to all observers. /// - /// Iterates registered observers in registration order. + /// Fired after an explicit model switch is accepted by the client, + /// carrying the old and new model names. Iterates registered + /// observers in registration order. pub fn on_model_switched(&self, ctx: &ModelSwitchedContext) { - for obs in &self.observers { - obs.on_model_switched(ctx); - } + self.dispatch(|obs| obs.on_model_switched(ctx)); } /// Dispatch [`LoopObserver::on_loop_detected`] to all observers. /// - /// Iterates registered observers in registration order. + /// Fired when the loop detector observes the same operation + /// repeatedly, exceeding the configured threshold. Iterates + /// registered observers in registration order. pub fn on_loop_detected(&self, ctx: &LoopDetectedContext) { - for obs in &self.observers { - obs.on_loop_detected(ctx); - } + self.dispatch(|obs| obs.on_loop_detected(ctx)); } - /// Dispatch [`LoopObserver::on_convergence_detected`] to all observers. + /// Dispatch [`LoopObserver::on_convergence_detected`] to all + /// observers. /// - /// Iterates registered observers in registration order. + /// Fired when the convergence detector observes semantically + /// similar assistant responses, as determined by the convergence + /// detection policy. Iterates registered observers in registration + /// order. pub fn on_convergence_detected(&self, ctx: &ConvergenceDetectedContext) { - for obs in &self.observers { - obs.on_convergence_detected(ctx); - } + self.dispatch(|obs| obs.on_convergence_detected(ctx)); } } -// ================================================== -// Tests -// ================================================== - #[cfg(test)] mod tests { use super::*; diff --git a/src/observer/context.rs b/src/observer/context.rs index 560e941..7034789 100644 --- a/src/observer/context.rs +++ b/src/observer/context.rs @@ -4,10 +4,6 @@ //! 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 @@ -431,8 +427,15 @@ pub struct FallbackContext { #[derive(Debug, Clone)] pub struct ModelSwitchedContext { /// Model identifier before the switch. + /// + /// The model the session was using up to the point of the hot-swap, so + /// observers can log or revert the transition. pub from: String, + /// Model identifier after the switch. + /// + /// The model that will handle subsequent requests, which may differ in + /// capability or cost from the previous one. pub to: String, } diff --git a/src/presets.rs b/src/presets.rs index 486c8fb..f587821 100644 --- a/src/presets.rs +++ b/src/presets.rs @@ -8,7 +8,8 @@ //! of the small-model machinery installed. //! //! Apply a profile with [`ConstrainedProfile::apply`] (pipeline + contributor) -//! or compose the individual pieces ([`ConstrainedProfile::loop_config`], +//! or compose the individual pieces ([`ConstrainedProfile::session_config`], +//! [`ConstrainedProfile::run_config`], //! [`ConstrainedProfile::pipeline_builder`], //! [`ConstrainedProfile::request_options`]) by hand. //! @@ -21,14 +22,15 @@ //! use std::sync::Arc; //! //! // `client` is any ApiClient impl; `registry` is your tool set. -//! let mut agent = BareLoop::new(client, registry, ConstrainedProfile::loop_config()); +//! let mut agent = BareLoop::new(client, registry, ConstrainedProfile::session_config()); //! agent.set_request_options(ConstrainedProfile::request_options()); //! ConstrainedProfile::apply(&mut agent).unwrap(); //! ``` use std::sync::Arc; -use crate::config::LoopConfig; +use crate::config::SessionConfig; +use crate::engine::RunConfig; use crate::engine::{BareLoop, ContextContributor, ContributorContext}; use crate::error::LoopError; use crate::message::{Message, MessagePart, Role}; @@ -49,10 +51,6 @@ const WRITE_TOOLS: &[&str] = &["Write", "Edit", "MultiEdit"]; /// Default memoized tool names the preset wires into its middleware. Advisory. const MEMOIZED_TOOLS: &[&str] = &["Read", "Glob", "Grep", "LS"]; -// =================================================== -// ConstrainedProfile -// =================================================== - /// The small-model-tuned runtime profile. /// /// Bundles aggressive context budgeting (smaller window, fewer turns), @@ -76,20 +74,32 @@ const MEMOIZED_TOOLS: &[&str] = &["Read", "Glob", "Grep", "LS"]; pub struct ConstrainedProfile; impl ConstrainedProfile { - /// A [`LoopConfig`] tuned for a small model. + /// A [`SessionConfig`] tuned for a small model. + /// + /// Sets a smaller context window than the default — small models degrade + /// faster as context fills, so the window is tightened. Compaction + /// threshold is unchanged (tightness comes from the window, not from + /// compacting earlier). + /// + /// Value: `context_window = 32_768`. + #[must_use] + pub fn session_config() -> SessionConfig { + SessionConfig::default().with_context_window(32_768) + } + + /// A [`RunConfig`] tuned for a small model. /// - /// Sets a smaller context window and fewer max turns than the v0.1.0 - /// defaults — small models degrade faster as context fills, so the - /// budget is tighter. Compaction threshold is unchanged (tightness comes - /// from the window, not from compacting earlier). + /// Fewer max turns than the default, since small models are more prone to + /// non-converging tool loops. All other run-scoped knobs stay at their + /// defaults. /// - /// Values: `context_window = 120_000`, `max_turns = 100`, - /// `max_tokens = 16_384`, `compact_threshold = 0.80`, `auto_compact = true`. + /// Value: `max_turns = 100`. #[must_use] - pub fn loop_config() -> LoopConfig { - LoopConfig::default() - .with_context_window(120_000) - .with_max_turns(100) + pub fn run_config() -> RunConfig { + RunConfig { + max_turns: 100, + ..RunConfig::default() + } } /// A [`ToolPipeline`] builder pre-loaded with the small-model middleware stack. @@ -148,7 +158,7 @@ impl ConstrainedProfile { /// use std::sync::Arc; /// /// // `client` is any ApiClient impl. - /// let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), ConstrainedProfile::loop_config()); + /// let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), ConstrainedProfile::session_config()); /// ConstrainedProfile::apply(&mut agent).unwrap(); /// ``` pub fn apply(loop_: &mut BareLoop) -> Result<(), LoopError> { @@ -158,28 +168,33 @@ impl ConstrainedProfile { } } -// =================================================== -// FrontierProfile -// =================================================== - /// The frontier-opt-out profile. /// -/// Produces default [`LoopConfig`], no small-model +/// Produces default session/run configuration, no small-model /// middleware, no tool-call constraint. Named (rather than just "don't use /// [`ConstrainedProfile`]") so the opt-out is explicit and discoverable. pub struct FrontierProfile; impl FrontierProfile { - /// A default [`LoopConfig`] — no small-model tightening. + /// A default [`SessionConfig`] — no small-model tightening. /// - /// Literally [`LoopConfig::default`]: the full v0.1.0 context budget - /// (200k window, 200 turns), default compaction threshold, no tightening. - /// The opt-out counterpart to - /// [`ConstrainedProfile::loop_config`](ConstrainedProfile::loop_config); + /// Literally [`SessionConfig::default`]: the full context budget + /// (200k window). The opt-out counterpart to + /// [`ConstrainedProfile::session_config`](ConstrainedProfile::session_config); /// the two are swap-in replacements at the call site. #[must_use] - pub fn loop_config() -> LoopConfig { - LoopConfig::default() + pub fn session_config() -> SessionConfig { + SessionConfig::default() + } + + /// A default [`RunConfig`] — no small-model tightening. + /// + /// Literally [`RunConfig::default`]: the full turn/token budget + /// (200 turns). The opt-out counterpart to + /// [`ConstrainedProfile::run_config`](ConstrainedProfile::run_config). + #[must_use] + pub fn run_config() -> RunConfig { + RunConfig::default() } /// An empty middleware pipeline — no small-model middleware installed. @@ -206,10 +221,6 @@ impl FrontierProfile { } } -// =================================================== -// GoalReminder -// =================================================== - /// A [`ContextContributor`] that re-injects the first user message as a /// [`Role::System`] reminder every `n` turns. /// @@ -283,13 +294,12 @@ mod tests { use super::*; #[test] - fn test_constrained_loop_config_values() { - let cfg = ConstrainedProfile::loop_config(); - assert_eq!(cfg.context_window, 120_000); - assert_eq!(cfg.max_turns, 100); - assert_eq!(cfg.max_tokens, 16_384); - assert_eq!(cfg.compact_threshold, 8_000); - assert!(cfg.auto_compact); + fn test_constrained_config_values() { + let session = ConstrainedProfile::session_config(); + assert_eq!(session.context_window, 32_768); + + let run = ConstrainedProfile::run_config(); + assert_eq!(run.max_turns, 100); } #[test] @@ -300,14 +310,16 @@ mod tests { } #[test] - fn test_frontier_loop_config_matches_default() { - let preset = FrontierProfile::loop_config(); - let default = LoopConfig::default(); - assert_eq!(preset.context_window, default.context_window); - assert_eq!(preset.max_turns, default.max_turns); - assert_eq!(preset.max_tokens, default.max_tokens); - assert_eq!(preset.compact_threshold, default.compact_threshold); - assert_eq!(preset.auto_compact, default.auto_compact); + fn test_frontier_config_matches_default() { + let session = FrontierProfile::session_config(); + assert_eq!( + session.context_window, + SessionConfig::default().context_window + ); + + let run = FrontierProfile::run_config(); + let default = RunConfig::default(); + assert_eq!(run.max_turns, default.max_turns); } #[test] diff --git a/src/provider.rs b/src/provider.rs index 01fce98..c848387 100644 --- a/src/provider.rs +++ b/src/provider.rs @@ -26,7 +26,8 @@ //! ```rust,ignore //! use loopctl::provider; //! use loopctl::engine::BareLoop; -//! use loopctl::engine::loop_core::Loop; +//! use loopctl::engine::RunConfig; +//! use loopctl::engine::core::Loop; //! //! // OpenAI: //! let client = provider::OpenAiClient::from_env()?; @@ -51,15 +52,44 @@ //! tool_registry, //! config, //! ); -//! let result = agent.run("Hello!").await?; +//! let result = agent.run("Hello!", &RunConfig::default()).await?; //! ``` -use crate::{ - api::error::ApiError, - message::{MessagePart, Role}, -}; +use crate::api::error::ApiError; +#[cfg(any(feature = "anthropic", feature = "gemini"))] +use crate::message::{MessagePart, Role}; use std::time::Duration; +// SSE line-framing shared by every streaming provider. Each provider keeps +// its own event-extraction logic (`next_data` / `next_event`); the struct, +// `from_response`, `take_line`, and the buffer-overflow guard live here. +#[cfg(any(feature = "openai", feature = "anthropic", feature = "gemini"))] +mod sse; + +/// Maximum accepted response body size (10 MB). +/// +/// Guards against unbounded memory growth from a misbehaving or hostile +/// provider that returns a very large non-streaming response. +#[cfg(any(feature = "openai", feature = "anthropic", feature = "gemini"))] +pub(super) const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; + +/// Reject a response body that exceeds [`MAX_RESPONSE_BODY`]. +/// +/// Shared guard used by every provider's non-streaming path. +/// +/// # Errors +/// +/// Returns [`ApiError`] when `len` exceeds [`MAX_RESPONSE_BODY`]. +#[cfg(any(feature = "openai", feature = "anthropic", feature = "gemini"))] +pub(super) fn check_response_body(len: usize) -> Result<(), ApiError> { + if len > MAX_RESPONSE_BODY { + return Err(ApiError::http(format!( + "response body too large: {len} bytes (max {MAX_RESPONSE_BODY})" + ))); + } + Ok(()) +} + /// Shared HTTP-client configuration embedded by every provider builder. /// /// Holds the timeout, connection-pool, and TCP knobs that are identical @@ -295,10 +325,6 @@ pub use gemini::GeminiClient; #[cfg(feature = "grammar")] pub use grammar::{JsonSchemaGrammar, ToolGrammarProvider}; -// ======================================================= -// Default endpoints / models for convenience constructors -// ======================================================= - #[cfg(feature = "ollama")] const OLLAMA_BASE_URL: &str = "http://localhost:11434/v1"; @@ -320,34 +346,38 @@ 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`). +/// Read an environment variable or return a default. +#[cfg(any( + feature = "ollama", + feature = "deepseek", + feature = "grok", + feature = "zai", + feature = "openai" +))] +fn env_or_default(name: &str, default: &str) -> String { + std::env::var(name).unwrap_or_else(|_| default.into()) +} + +/// Read a primary env var, falling back to a secondary if unset. +#[cfg(any(feature = "deepseek", feature = "grok", feature = "zai"))] 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. +/// Look up a required API key from the environment. /// /// # Errors /// -/// Returns [`ApiError::auth_invalid_key`] if neither environment -/// variable is set. +/// Returns [`ApiError::auth_invalid_key`] if neither environment variable is set. +#[cfg(any(feature = "deepseek", feature = "grok", feature = "zai"))] fn require_api_key(primary: &str, fallback: Option<&str>) -> Result { if let Some(fb) = fallback { if let Some(val) = env_or_fallback(primary, fb) { @@ -362,15 +392,11 @@ fn require_api_key(primary: &str, fallback: Option<&str>) -> Result( messages: &'a [crate::message::Message], system: Option<&str>, @@ -488,7 +514,9 @@ pub fn deepseek() -> Result { #[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); + let model = std::env::var("XAI_MODEL") + .or_else(|_| std::env::var("GROK_MODEL")) + .unwrap_or_else(|_| GROK_DEFAULT_MODEL.into()); OpenAiClient::builder() .with_api_key(api_key) @@ -559,14 +587,17 @@ pub fn self_hosted(base_url: &str, model: &str) -> Result {{ // SAFETY: This is only used in single-threaded test code where @@ -575,6 +606,13 @@ mod tests { }}; } + #[cfg(any( + feature = "ollama", + feature = "deepseek", + feature = "grok", + feature = "zai", + feature = "openai" + ))] macro_rules! env_remove { ($($arg:tt)*) => {{ // SAFETY: This is only used in single-threaded test code where @@ -583,6 +621,7 @@ mod tests { }}; } + #[cfg(any(feature = "deepseek", feature = "grok", feature = "zai"))] #[test] fn env_or_fallback_primary_set() { env_set!("LOOPCTL_TEST_PRIMARY", "primary-val"); @@ -594,6 +633,7 @@ mod tests { env_remove!("LOOPCTL_TEST_PRIMARY"); } + #[cfg(any(feature = "deepseek", feature = "grok", feature = "zai"))] #[test] fn env_or_fallback_fallback_used_when_primary_missing() { env_remove!("LOOPCTL_TEST_PRIMARY2"); @@ -605,6 +645,7 @@ mod tests { env_remove!("LOOPCTL_TEST_FALLBACK2"); } + #[cfg(any(feature = "deepseek", feature = "grok", feature = "zai"))] #[test] fn env_or_fallback_none_when_both_missing() { env_remove!("LOOPCTL_TEST_NEITHER_A"); @@ -615,6 +656,13 @@ mod tests { ); } + #[cfg(any( + feature = "ollama", + feature = "deepseek", + feature = "grok", + feature = "zai", + feature = "openai" + ))] #[test] fn env_or_default_uses_env_when_set() { env_set!("LOOPCTL_TEST_DEFAULT", "from-env"); @@ -625,6 +673,13 @@ mod tests { env_remove!("LOOPCTL_TEST_DEFAULT"); } + #[cfg(any( + feature = "ollama", + feature = "deepseek", + feature = "grok", + feature = "zai", + feature = "openai" + ))] #[test] fn env_or_default_uses_default_when_unset() { env_remove!("LOOPCTL_TEST_DEFAULT2"); @@ -634,6 +689,7 @@ mod tests { ); } + #[cfg(any(feature = "deepseek", feature = "grok", feature = "zai"))] #[test] fn require_api_key_primary_set() { env_set!("LOOPCTL_TEST_KEY_PRIMARY", "secret"); @@ -647,6 +703,7 @@ mod tests { env_remove!("LOOPCTL_TEST_KEY_PRIMARY"); } + #[cfg(any(feature = "deepseek", feature = "grok", feature = "zai"))] #[test] fn require_api_key_fallback_used() { env_remove!("LOOPCTL_TEST_KEY_PRIMARY2"); @@ -660,6 +717,7 @@ mod tests { env_remove!("LOOPCTL_TEST_KEY_FALLBACK2"); } + #[cfg(any(feature = "deepseek", feature = "grok", feature = "zai"))] #[test] fn require_api_key_no_fallback_set() { env_set!("LOOPCTL_TEST_KEY_ONLY", "only-val"); @@ -668,6 +726,7 @@ mod tests { env_remove!("LOOPCTL_TEST_KEY_ONLY"); } + #[cfg(any(feature = "deepseek", feature = "grok", feature = "zai"))] #[test] fn require_api_key_errors_when_missing() { env_remove!("LOOPCTL_TEST_MISSING_KEY"); @@ -675,6 +734,7 @@ mod tests { assert!(err.to_string().contains("LOOPCTL_TEST_MISSING_KEY")); } + #[cfg(any(feature = "deepseek", feature = "grok", feature = "zai"))] #[test] fn require_api_key_errors_when_both_missing() { env_remove!("LOOPCTL_TEST_MISSING_A"); @@ -735,12 +795,14 @@ mod tests { } /// Build a `Role::System` message carrying the given text parts. + #[cfg(any(feature = "anthropic", feature = "gemini"))] fn sys_msg(texts: &[&str]) -> crate::message::Message { use crate::message::{MessagePart, Role}; let parts: Vec = texts.iter().map(|t| MessagePart::text(*t)).collect(); crate::message::Message::new(Role::System, parts) } + #[cfg(any(feature = "anthropic", feature = "gemini"))] #[test] fn fold_system_no_system_messages_no_caller_returns_none() { let msgs = [crate::message::Message::user("hi")]; @@ -749,6 +811,7 @@ mod tests { assert!(system.is_none(), "no system content → None"); } + #[cfg(any(feature = "anthropic", feature = "gemini"))] #[test] fn fold_system_caller_only_passes_through() { let msgs = [crate::message::Message::user("hi")]; @@ -757,6 +820,7 @@ mod tests { assert_eq!(system.as_deref(), Some("be brief")); } + #[cfg(any(feature = "anthropic", feature = "gemini"))] #[test] fn fold_system_single_system_message_removed_and_folded() { let msgs = [ @@ -769,6 +833,7 @@ mod tests { assert_eq!(system.as_deref(), Some("stay on task")); } + #[cfg(any(feature = "anthropic", feature = "gemini"))] #[test] fn fold_system_caller_prompt_prepended_to_folded() { let msgs = [crate::message::Message::user("hi"), sys_msg(&["reminder"])]; @@ -788,6 +853,7 @@ mod tests { ); } + #[cfg(any(feature = "anthropic", feature = "gemini"))] #[test] fn fold_system_multiple_system_messages_joined_with_newlines() { let msgs = [ @@ -801,6 +867,7 @@ mod tests { assert_eq!(system, "first reminder\nsecond reminder"); } + #[cfg(any(feature = "anthropic", feature = "gemini"))] #[test] fn fold_system_only_text_parts_are_folded() { // A System message carrying a tool-call part (unusual, but defensive): @@ -818,6 +885,7 @@ mod tests { assert_eq!(system.as_deref(), Some("keep this")); } + #[cfg(any(feature = "anthropic", feature = "gemini"))] #[test] fn fold_system_preserves_relative_order_of_non_system_messages() { let msgs = [ @@ -839,6 +907,7 @@ mod tests { assert_eq!(texts, vec!["first", "second", "third"]); } + #[cfg(any(feature = "anthropic", feature = "gemini"))] #[test] fn fold_system_empty_text_part_contributes_nothing() { // A System message whose text is empty: folded string stays empty, so @@ -852,6 +921,7 @@ mod tests { ); } + #[cfg(any(feature = "anthropic", feature = "gemini"))] #[test] fn fold_system_multiple_text_parts_in_one_message_joined() { let msgs = [sys_msg(&["part one", "part two"])]; diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs index 57a524c..a167e75 100644 --- a/src/provider/anthropic.rs +++ b/src/provider/anthropic.rs @@ -21,8 +21,7 @@ use std::future::Future; use std::pin::Pin; -use futures::stream::{Stream, StreamExt}; -use reqwest::Response; +use futures::stream::Stream; use serde_json::Value; use std::time::Duration; @@ -37,10 +36,6 @@ use crate::structured::ToolConstraint; use crate::structured::tighten_json_schema; 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"; @@ -48,14 +43,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; -const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb -const SSE_MAX_BUFFER: usize = 1024 * 1024; // 1 Mb const MAX_ERROR_BODY: usize = 8 * 1024; // 8 Kb -// ================================================== -// Client -// ================================================== - /// An Anthropic Claude chat client with streaming support. /// /// Implements [`ApiClient`] by translating between the framework's @@ -312,13 +301,7 @@ impl ApiClient for AnthropicClient { .bytes() .await .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 - ))); - } + super::check_response_body(resp.len())?; serde_json::from_slice::(&resp).map_err(|e| ApiError::http(e.to_string())) }) } @@ -339,10 +322,6 @@ impl ApiClient for AnthropicClient { } } -// ================================================== -// Builder -// ================================================== - /// Builder for [`AnthropicClient`]. /// /// Created via [`AnthropicClientBuilder::default`] or @@ -524,10 +503,6 @@ impl AnthropicClientBuilder { } } -// ================================================== -// Request body construction -// ================================================== - /// The per-request inputs to [`build_request_body`]. /// /// Carries the request-shape knobs: the model, conversation, tools, and @@ -775,33 +750,9 @@ fn convert_tools(tools: &[ToolSchema], strict: bool) -> Vec { .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, -} +use super::sse::SseReader; 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. @@ -817,24 +768,8 @@ impl SseReader { 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 { - 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 - } - } - }; + let parsed = Self::parse_event_data(&data, &event_type); return Ok(Some((event_type, parsed))); } continue; @@ -844,10 +779,6 @@ impl SseReader { event_type = ev.into(); have_event = true; } else if let Some(d) = line.strip_prefix(SSE_DATA_PREFIX) { - // 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 { @@ -858,57 +789,36 @@ impl SseReader { } } - match self.bytes.next().await { - 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. - if have_event { - let parsed = if data.is_empty() { - None - } else { - 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))); - } - return Ok(None); + if self.next_chunk().await?.is_none() { + if have_event { + let parsed = Self::parse_event_data(&data, &event_type); + 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) + /// Parse the accumulated `data` string into JSON, logging on failure. + fn parse_event_data(data: &str, event_type: &str) -> Option { + if data.is_empty() { + return None; + } + 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 + } + } } } -// ================================================== -// Stream event emitter -// ================================================== - /// Stateful translator that converts Anthropic SSE events into /// [`StreamEvent`]s. /// @@ -1324,10 +1234,6 @@ impl StreamEmitter { } } -// ================================================== -// Tests -// ================================================== - #[cfg(test)] mod tests { use super::*; @@ -1878,7 +1784,9 @@ mod tests { 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(), + buf: "event: message_start\ndata: {}\n\n" + .to_string() + .into_bytes(), }; assert_eq!(reader.take_line(), Some("event: message_start".to_string())); assert_eq!(reader.take_line(), Some("data: {}".to_string())); @@ -1888,10 +1796,11 @@ mod tests { #[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 stream = + futures::stream::iter(vec![Ok::(chunk.to_string().into())]); let mut reader = SseReader { bytes: Box::pin(stream), - buf: String::new(), + buf: Vec::new(), }; let result = reader.next_event().await.unwrap(); assert!(result.is_some()); @@ -1903,10 +1812,11 @@ mod tests { #[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 stream = + futures::stream::iter(vec![Ok::(chunk.to_string().into())]); let mut reader = SseReader { bytes: Box::pin(stream), - buf: String::new(), + buf: Vec::new(), }; let result = reader.next_event().await.unwrap(); assert!(result.is_some()); @@ -1926,10 +1836,11 @@ mod tests { // 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 stream = + futures::stream::iter(vec![Ok::(chunk.to_string().into())]); let mut reader = SseReader { bytes: Box::pin(stream), - buf: String::new(), + buf: Vec::new(), }; let result = reader.next_event().await.unwrap(); assert!(result.is_some()); @@ -1940,11 +1851,11 @@ mod tests { #[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 { + let huge = "x".repeat(2 * 1024 * 1024); + let stream = futures::stream::iter(vec![Ok::(huge.into())]); + let mut reader = super::SseReader { bytes: Box::pin(stream), - buf: String::new(), + buf: Vec::new(), }; let result = reader.next_event().await; assert!(result.is_err(), "should error on buffer overflow"); @@ -1957,7 +1868,7 @@ mod tests { #[test] fn max_response_body_is_ten_mb() { - assert_eq!(MAX_RESPONSE_BODY, 10 * 1024 * 1024); + assert_eq!(super::super::MAX_RESPONSE_BODY, 10 * 1024 * 1024); } #[test] diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs index 27c4b98..2552b0a 100644 --- a/src/provider/gemini.rs +++ b/src/provider/gemini.rs @@ -22,8 +22,7 @@ use std::future::Future; use std::pin::Pin; -use futures::stream::{Stream, StreamExt}; -use reqwest::Response; +use futures::stream::Stream; use serde_json::Value; use std::time::Duration; @@ -38,23 +37,13 @@ use crate::structured::ToolConstraint; use crate::structured::tighten_json_schema; 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; const THINKING_PART_INDEX: usize = 1; -const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb -const SSE_MAX_BUFFER: usize = 1024 * 1024; // 1 Mb const MAX_ERROR_BODY: usize = 8 * 1024; // 8 Kb -// ================================================== -// Client -// ================================================== - /// A Google Gemini API client with streaming support. /// /// Implements [`ApiClient`] by translating between the framework's @@ -257,7 +246,7 @@ impl ApiClient for GeminiClient { let mut sse = SseReader::from_response(resp); let mut emitter = StreamEmitter::default(); - while let Some(data) = sse.next_data().await? { + while let Some(data) = sse.next_gemini_data().await? { emitter.process_chunk(&data); for ev in emitter.drain() { yield ev; @@ -295,13 +284,7 @@ impl ApiClient for GeminiClient { .bytes() .await .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 - ))); - } + super::check_response_body(resp.len())?; serde_json::from_slice::(&resp).map_err(|e| ApiError::http(e.to_string())) }) } @@ -334,7 +317,7 @@ impl ApiClient for GeminiClient { let mut sse = SseReader::from_response(resp); let mut emitter = StreamEmitter::default(); - while let Some(data) = sse.next_data().await? { + while let Some(data) = sse.next_gemini_data().await? { emitter.process_chunk(&data); for ev in emitter.drain() { yield ev; @@ -373,13 +356,7 @@ impl ApiClient for GeminiClient { .bytes() .await .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 - ))); - } + super::check_response_body(resp.len())?; serde_json::from_slice::(&resp).map_err(|e| ApiError::http(e.to_string())) }) } @@ -395,9 +372,6 @@ impl ApiClient for GeminiClient { .unwrap_or_else(|| Value::String(text.to_string())) } } -// ================================================== -// Builder -// ================================================== /// Builder for [`GeminiClient`]. /// @@ -583,10 +557,6 @@ impl GeminiClientBuilder { } } -// ================================================== -// 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 @@ -726,40 +696,9 @@ fn convert_tools(tools: &[ToolSchema], strict: bool) -> Vec { .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, -} +use super::sse::SseReader; impl SseReader { - /// Wrap a streaming HTTP response into an SSE reader. - /// - /// Takes the response body's byte stream and converts it into a - /// line-oriented reader that yields SSE event data. Used by - /// [`stream_messages`](crate::api::ApiClient::stream_messages) and its - /// `*_with_options` variant to parse Gemini's streaming - /// `streamGenerateContent` responses. - 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. @@ -767,13 +706,12 @@ impl SseReader { /// # Errors /// /// Returns [`ApiError`] if the underlying HTTP stream fails. - async fn next_data(&mut self) -> Result, ApiError> { + async fn next_gemini_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) { match serde_json::from_str::(data) { Ok(json) => return Ok(Some(json)), @@ -787,42 +725,13 @@ impl SseReader { } } } - - match self.bytes.next().await { - 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), + if self.next_chunk().await?.is_none() { + return Ok(None); } } } - - /// Pop the first `\n`-terminated line from the internal buffer. - /// - /// Returns the line (trimmed) if a newline is present, and removes it - /// (plus the newline) from the buffer. Returns `None` if the buffer - /// does not yet contain a complete line — the caller should wait for - /// more bytes from the HTTP stream. - 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. /// @@ -865,11 +774,14 @@ struct StreamEmitter { /// Set by [`finish`](Self::finish) when it appends the synthetic /// [`StreamEvent::MessageStop`]. Guards against emitting a second /// `MessageStop` if `finish` is called again after the stream ends. - /// (Note: unlike the OpenAI/Anthropic emitters, Gemini's finish - /// reason arrives inside a regular data chunk and is handled by - /// [`extract_finish_reason`](Self::extract_finish_reason); this flag - /// only governs the final `MessageStop` synthesis.) - finished: bool, + message_stop_emitted: bool, + + /// Whether `finishReason` has already been processed. + /// + /// Guards against a second `finishReason` chunk (e.g. from proxies that + /// re-emit). Distinct from `message_stop_emitted` — the finish-reason + /// processing and the terminal `MessageStop` synthesis are independent. + finish_reason_processed: bool, /// Buffered [`StreamEvent`]s waiting to be yielded to the consumer. /// @@ -1050,9 +962,9 @@ impl StreamEmitter { /// [`PartStop`](StreamEvent::PartStop), then emits a /// [`MessageDelta`](StreamEvent::MessageDelta) carrying the mapped /// [`StreamStopReason`]. No-op on a second `finishReason` chunk (the - /// `finished` guard defends against proxies/gateways that re-emit). + /// `finish_reason_processed` guard defends against proxies/gateways that re-emit). fn extract_finish_reason(&mut self, json: &Value) { - if self.finished { + if self.finish_reason_processed { return; } let Some(reason) = json @@ -1061,7 +973,7 @@ impl StreamEmitter { else { return; }; - self.finished = true; + self.finish_reason_processed = true; let stop = match reason { "MAX_TOKENS" => StreamStopReason::MaxTokens, _ => StreamStopReason::EndTurn, @@ -1100,11 +1012,11 @@ impl StreamEmitter { /// /// Drains any remaining pending events and appends the stop event. /// Safe to call exactly once at the end of the stream; subsequent calls - /// return an empty vec (the `finished` flag guards against double-stop). + /// return an empty vec (the `message_stop_emitted` flag guards against double-stop). fn finish(&mut self) -> Vec { let mut out = self.drain(); - if self.started && !self.finished { - self.finished = true; + if self.started && !self.message_stop_emitted { + self.message_stop_emitted = true; out.push(StreamEvent::MessageStop); } out @@ -1121,10 +1033,6 @@ impl StreamEmitter { } } -// ================================================== -// Tests -// ================================================== - #[cfg(test)] mod tests { use super::*; @@ -1932,7 +1840,7 @@ mod tests { fn emitter_finish_emits_message_stop_if_needed() { let mut em = StreamEmitter::default(); em.started = true; - em.finished = false; + em.message_stop_emitted = false; let events = em.finish(); assert!(events.iter().any(|e| matches!(e, StreamEvent::MessageStop))); @@ -1942,12 +1850,32 @@ mod tests { fn emitter_finish_noop_if_already_stopped() { let mut em = StreamEmitter::default(); em.started = true; - em.finished = true; + em.message_stop_emitted = true; let events = em.finish(); assert!(events.is_empty()); } + #[test] + fn emitter_finish_reason_does_not_suppress_message_stop() { + // Regression: extract_finish_reason sets finish_reason_processed, + // but finish() must still emit MessageStop. Previously both used the + // same `finished` flag, causing finish() to skip MessageStop after + // a finishReason chunk. + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{"finishReason": "STOP"}] + })); + em.drain(); + assert!(em.finish_reason_processed); + let events = em.finish(); + assert!( + events.iter().any(|e| matches!(e, StreamEvent::MessageStop)), + "MessageStop must be emitted even after finishReason was processed" + ); + } + #[test] fn sse_reader_take_line_extracts_newline() { let mut reader = SseReader { @@ -2025,7 +1953,7 @@ mod tests { 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(), + buf: "data: {\"candidates\":[]}\n\n".to_string().into_bytes(), }; assert_eq!( reader.take_line(), @@ -2038,12 +1966,13 @@ mod tests { #[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 stream = + futures::stream::iter(vec![Ok::(chunk.to_string().into())]); let mut reader = SseReader { bytes: Box::pin(stream), - buf: String::new(), + buf: Vec::new(), }; - let result = reader.next_data().await.unwrap(); + let result = reader.next_gemini_data().await.unwrap(); assert!(result.is_some()); let json = result.unwrap(); assert!(json["candidates"].is_array()); @@ -2052,13 +1981,14 @@ mod tests { #[tokio::test] async fn sse_reader_next_data_malformed_returns_none() { let chunk = "data: not valid json\n\ndata: {\"ok\":true}\n\n"; - let stream = futures::stream::iter(vec![Ok::(chunk.to_string())]); + let stream = + futures::stream::iter(vec![Ok::(chunk.to_string().into())]); let mut reader = SseReader { bytes: Box::pin(stream), - buf: String::new(), + buf: Vec::new(), }; // First call should skip malformed and return the valid one. - let result = reader.next_data().await.unwrap(); + let result = reader.next_gemini_data().await.unwrap(); assert!(result.is_some()); let json = result.unwrap(); assert_eq!(json["ok"], true); @@ -2066,13 +1996,13 @@ mod tests { #[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 { + let huge = "x".repeat(2 * 1024 * 1024); + let stream = futures::stream::iter(vec![Ok::(huge.into())]); + let mut reader = super::SseReader { bytes: Box::pin(stream), - buf: String::new(), + buf: Vec::new(), }; - let result = reader.next_data().await; + let result = reader.next_gemini_data().await; assert!(result.is_err(), "should error on buffer overflow"); let err_msg = result.unwrap_err().to_string(); assert!( @@ -2083,7 +2013,7 @@ mod tests { #[test] fn max_response_body_is_ten_mb() { - assert_eq!(MAX_RESPONSE_BODY, 10 * 1024 * 1024); + assert_eq!(super::super::MAX_RESPONSE_BODY, 10 * 1024 * 1024); } #[test] diff --git a/src/provider/openai.rs b/src/provider/openai.rs index b8a66a7..2bfd931 100644 --- a/src/provider/openai.rs +++ b/src/provider/openai.rs @@ -24,8 +24,7 @@ use std::future::Future; use std::pin::Pin; use std::time::Duration; -use futures::stream::{Stream, StreamExt}; -use reqwest::Response; +use futures::stream::Stream; use serde::Deserialize; use serde_json::Value; @@ -40,24 +39,14 @@ use crate::structured::ToolConstraint; use crate::structured::tighten_json_schema; 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; const THINKING_PART_INDEX: usize = 1; -const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb -const SSE_MAX_BUFFER: usize = 1024 * 1024; // 1 Mb const MAX_ERROR_BODY: usize = 8 * 1024; // 8 Kb -// ================================================== -// Client -// ================================================== - /// An OpenAI-compatible chat completions client with streaming support. /// /// Implements [`ApiClient`] by translating between the framework's @@ -246,7 +235,7 @@ impl ApiClient for OpenAiClient { let mut sse = SseReader::from_response(resp); let mut emitter = StreamEmitter::default(); - while let Some(data) = sse.next_data().await? { + while let Some(data) = sse.next_openai_data().await? { let Some(chunk) = OpenAiChunk::parse(&data) else { continue; }; @@ -290,13 +279,7 @@ impl ApiClient for OpenAiClient { .bytes() .await .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 - ))); - } + super::check_response_body(resp.len())?; serde_json::from_slice::(&resp).map_err(|e| ApiError::http(e.to_string())) }) } @@ -330,7 +313,7 @@ impl ApiClient for OpenAiClient { let mut sse = SseReader::from_response(resp); let mut emitter = StreamEmitter::default(); - while let Some(data) = sse.next_data().await? { + while let Some(data) = sse.next_openai_data().await? { let Some(chunk) = OpenAiChunk::parse(&data) else { continue; }; @@ -376,13 +359,7 @@ impl ApiClient for OpenAiClient { .bytes() .await .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 - ))); - } + super::check_response_body(resp.len())?; serde_json::from_slice::(&resp).map_err(|e| ApiError::http(e.to_string())) }) } @@ -404,9 +381,6 @@ impl ApiClient for OpenAiClient { } } } -// ================================================== -// Builder -// ================================================== /// Builder for [`OpenAiClient`]. /// @@ -565,10 +539,6 @@ impl OpenAiClientBuilder { } } -// ================================================== -// Request body construction -// ================================================== - /// A built OpenAI Chat Completions request body. /// /// Separating construction from serialization lets us reuse the same @@ -640,7 +610,7 @@ impl RequestBody { } for m in messages { - msgs.push(convert_message(m)); + msgs.extend(convert_message(m)); } let (tools, guided_json) = if response_format.is_some() { @@ -711,8 +681,10 @@ impl RequestBody { /// /// 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 { +/// text to use a simple `{role, content}` pair. A single loopctl message +/// with multiple tool-result parts expands to one OpenAI `tool` message per +/// result, so the return is a vector. +fn convert_message(m: &Message) -> Vec { let role = match m.role { Role::User => "user", Role::Assistant => "assistant", @@ -749,11 +721,11 @@ fn convert_message(m: &Message) -> Value { } if !tool_calls.is_empty() { - build_assistant_message(role, &tool_calls, &text_parts) + vec![build_assistant_message(role, &tool_calls, &text_parts)] } else if !tool_results.is_empty() { - merge_tool_results(&tool_results) + tool_results } else { - serde_json::json!({ "role": role, "content": text_parts.join("") }) + vec![serde_json::json!({ "role": role, "content": text_parts.join("") })] } } @@ -778,18 +750,6 @@ fn build_assistant_message(role: &str, tool_calls: &[Value], text_parts: &[&str] }) } -/// 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 framework tool schemas into the OpenAI `tools` array shape. /// /// Each [`ToolSchema`] becomes a JSON object with `type: "function"` and a @@ -852,36 +812,9 @@ fn input_to_string(v: &Value) -> 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, -} +use super::sse::SseReader; 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. /// @@ -890,9 +823,8 @@ impl SseReader { /// # Errors /// /// Returns [`ApiError`] if the underlying HTTP stream fails. - async fn next_data(&mut self) -> Result, ApiError> { + async fn next_openai_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; @@ -902,47 +834,40 @@ impl SseReader { } 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); - 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), + if self.next_chunk().await?.is_none() { + return Ok(None); } } } - - /// Pop the first `\n`-terminated line from the internal buffer. - /// - /// Returns the trimmed line if a newline is present, and removes it - /// (plus the newline) from the buffer. Returns `None` if the buffer - /// does not yet contain a complete line — the caller should wait for - /// more bytes from the HTTP stream. - 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. +/// +/// Each `data:` line in a streamed Chat Completions response deserializes +/// into one of these. The chunk carries the message identity, the model +/// that produced it, and a list of [`OpenAiChoice`] deltas that the +/// [`StreamEmitter`] assembles into [`StreamEvent`]s. #[derive(Deserialize)] struct OpenAiChunk { + /// Server-assigned identifier for the overall completion. + /// + /// Stable across every chunk in a single streamed response (for + /// example `chatcmpl-abc123`); emitted once as the message id in + /// the initial [`MessageStart`]. id: String, + + /// Name of the model that produced this chunk. + /// + /// Echoed back by the server on the first chunk; forwarded in the + /// initial [`MessageStart`] so downstream consumers know which model + /// answered, even after a fallback switch. model: String, + + /// One entry per alternative the model is generating. + /// + /// In practice OpenAI streams a single choice (`n=1`), so this + /// vector usually holds exactly one [`OpenAiChoice`] carrying the + /// incremental [`OpenAiDelta`] for this chunk. choices: Vec, } @@ -966,36 +891,106 @@ impl OpenAiChunk { } } +/// A single alternative within an [`OpenAiChunk`]. +/// +/// Carries the incremental content for this turn (in `delta`) and, on +/// the final chunk for the choice, the reason the model stopped (in +/// `finish_reason`). #[derive(Deserialize)] struct OpenAiChoice { + /// Incremental content for this chunk, or `None` on the terminal + /// chunk that carries only a `finish_reason`. delta: Option, + + /// Why the model stopped generating, present only on the last chunk. + /// + /// Common values are `"stop"`, `"tool_calls"`, and `"length"`; the + /// [`StreamEmitter`] maps it to a [`StreamEvent`] stop reason. finish_reason: Option, } +/// Incremental content delivered by one chunk. +/// +/// Mirrors the `delta` object in OpenAI's streaming protocol. Every +/// field is optional because a single chunk typically populates only +/// the field it is extending (text content, reasoning, or a tool call). #[derive(Deserialize)] struct OpenAiDelta { + /// Incremental assistant text for this chunk. + /// + /// Concatenated across chunks to reconstruct the full message body. content: Option, + + /// Incremental chain-of-thought / reasoning text. + /// + /// Some models (e.g. o1-style reasoning models) emit their private + /// reasoning here under `reasoning_content`; the `reasoning` alias + /// covers providers that use the shorter key. #[serde(alias = "reasoning")] reasoning_content: Option, + + /// Incremental tool-call fragments for this chunk. + /// + /// Tool calls arrive across multiple chunks keyed by `index`; the + /// [`StreamEmitter`] accumulates them per index until each call's + /// arguments are complete. tool_calls: Option>, } +/// Incremental fragment of a single tool call within a chunk. +/// +/// OpenAI streams tool calls in pieces: the first chunk for a given +/// `index` carries the call `id` and function `name`, subsequent chunks +/// append to `arguments`. The [`StreamEmitter`] reassembles these per +/// index. #[derive(Deserialize)] struct OpenAiToolCallDelta { - index: usize, + /// Server-assigned identifier for the tool call. + /// + /// Present on the first chunk for this `index`; the emitter forwards + /// it as the call id so the host can match the result later. id: String, + + /// Position of this tool call in the request's tool list. + /// + /// Used to correlate fragments across chunks — chunks with the same + /// `index` belong to the same tool call. + index: usize, + + /// Function name and accumulated arguments for this tool call. + /// + /// `None` on chunks that carry no function update (for example a + /// chunk that only extends a different tool call's `arguments`). + /// When present, the inner [`OpenAiToolCallFunction`] holds either + /// the function `name` (first chunk for this `index`) or a fragment + /// of the JSON `arguments` (subsequent chunks) — the + /// [`StreamEmitter`] reassembles both per index. function: Option, } +/// Name and arguments of a tool call, as carried by an [`OpenAiToolCallDelta`]. +/// +/// The `name` arrives on the first chunk for a tool call; `arguments` +/// is a JSON string that may itself arrive in fragments across several +/// chunks and must be concatenated before parsing. #[derive(Deserialize)] struct OpenAiToolCallFunction { - name: String, + /// Accumulated JSON arguments for the tool call. + /// + /// A partial JSON string that grows across chunks; the emitter + /// buffers it per `index` and hands the complete string to the + /// caller once the part stops. arguments: String, -} -// ================================================== -// Stream event emitter -// ================================================== + /// Fully-qualified name of the tool to invoke. + /// + /// Matches the `name` the tool was registered under in the request's + /// `tools` array. Arrives on the first chunk for a given `index` + /// only; subsequent chunks for the same call carry just argument + /// fragments, so the emitter latches this value and reuses it for + /// every later chunk with the same index. + name: String, +} /// Stateful translator that converts a sequence of [`OpenAiChunk`]s /// into [`StreamEvent`]s. @@ -1271,10 +1266,6 @@ impl StreamEmitter { } } -// ================================================== -// Tests -// ================================================== - #[cfg(test)] mod tests { use super::*; @@ -1360,7 +1351,7 @@ mod tests { #[test] fn convert_message_user_text() { let m = Message::user("hello world"); - let v = convert_message(&m); + let v = convert_message(&m).remove(0); assert_eq!(v["role"], "user"); assert_eq!(v["content"], "hello world"); } @@ -1368,7 +1359,7 @@ mod tests { #[test] fn convert_message_assistant_text() { let m = Message::new(Role::Assistant, vec![MessagePart::text("hi there")]); - let v = convert_message(&m); + let v = convert_message(&m).remove(0); assert_eq!(v["role"], "assistant"); assert_eq!(v["content"], "hi there"); } @@ -1383,7 +1374,7 @@ mod tests { input: serde_json::json!({"message": "hi"}), }], ); - let v = convert_message(&m); + let v = convert_message(&m).remove(0); assert_eq!(v["role"], "assistant"); assert!(v["content"].is_null()); let calls = v["tool_calls"].as_array().unwrap(); @@ -1408,12 +1399,40 @@ mod tests { is_error: None, }], ); - let v = convert_message(&m); + let v = convert_message(&m).remove(0); assert_eq!(v["role"], "tool"); assert_eq!(v["tool_call_id"], "call_1"); assert!(v["content"].is_string()); } + #[test] + fn convert_message_multiple_tool_results_expand() { + // A single loopctl message with two tool-result parts must expand to + // two separate OpenAI tool messages (not one array-valued element, + // which Ollama/OpenAI reject). + let m = Message::new( + Role::User, + vec![ + MessagePart::ToolResult { + call_id: "call_1".into(), + output: ToolContent::from_string("a"), + is_error: None, + }, + MessagePart::ToolResult { + call_id: "call_2".into(), + output: ToolContent::from_string("b"), + is_error: None, + }, + ], + ); + let vs = convert_message(&m); + assert_eq!(vs.len(), 2, "two tool results expand to two messages"); + assert_eq!(vs[0]["role"], "tool"); + assert_eq!(vs[0]["tool_call_id"], "call_1"); + assert_eq!(vs[1]["role"], "tool"); + assert_eq!(vs[1]["tool_call_id"], "call_2"); + } + #[test] fn convert_tools_shape() { let tools = vec![ @@ -1725,24 +1744,6 @@ mod tests { assert!(events.is_empty()); } - #[test] - fn merge_tool_results_single() { - let r = serde_json::json!({"role": "tool", "content": "ok"}); - let merged = merge_tool_results(std::slice::from_ref(&r)); - 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 { @@ -1797,7 +1798,7 @@ mod tests { 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(), + buf: "data: hello\ndata: world\n".to_string().into_bytes(), }; assert_eq!(reader.take_line(), Some("data: hello".to_string())); assert_eq!(reader.take_line(), Some("data: world".to_string())); @@ -1807,39 +1808,39 @@ mod tests { #[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 stream = + futures::stream::iter(vec![Ok::(data.to_string().into())]); let mut reader = SseReader { bytes: Box::pin(stream), - buf: String::new(), + buf: Vec::new(), }; - let result = reader.next_data().await.unwrap(); + let result = reader.next_openai_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 stream = futures::stream::iter(vec![Ok::( + "data: [DONE]\n\n".into(), + )]); let mut reader = SseReader { bytes: Box::pin(stream), - buf: String::new(), + buf: Vec::new(), }; - let result = reader.next_data().await.unwrap(); + let result = reader.next_openai_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 { + let huge = "x".repeat(2 * 1024 * 1024); + let stream = futures::stream::iter(vec![Ok::(huge.into())]); + let mut reader = super::SseReader { bytes: Box::pin(stream), - buf: String::new(), + buf: Vec::new(), }; - let result = reader.next_data().await; + let result = reader.next_openai_data().await; assert!(result.is_err(), "should error on buffer overflow"); let err_msg = result.unwrap_err().to_string(); assert!( @@ -1850,20 +1851,20 @@ mod tests { #[test] fn max_response_body_is_ten_mb() { - assert_eq!(MAX_RESPONSE_BODY, 10 * 1024 * 1024); + assert_eq!(super::super::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); + let oversized = super::super::MAX_RESPONSE_BODY + 1; + assert!(oversized > super::super::MAX_RESPONSE_BODY); } #[test] fn body_size_check_accepts_within_limit() { - let within = MAX_RESPONSE_BODY; - assert!(within <= MAX_RESPONSE_BODY); + let within = super::super::MAX_RESPONSE_BODY; + assert!(within <= super::super::MAX_RESPONSE_BODY); } #[test] @@ -2151,7 +2152,7 @@ mod tests { // Role::System message is emitted verbatim (not folded into a // top-level field). let msg = Message::new(Role::System, vec![MessagePart::text("stay on task")]); - let value = convert_message(&msg); + let value = convert_message(&msg).remove(0); assert_eq!(value["role"], "system"); assert_eq!(value["content"], "stay on task"); } diff --git a/src/provider/sse.rs b/src/provider/sse.rs new file mode 100644 index 0000000..ab8b69a --- /dev/null +++ b/src/provider/sse.rs @@ -0,0 +1,100 @@ +//! Shared SSE line-framing for streaming providers. +//! +//! Each provider that streams responses (OpenAI, Anthropic, Gemini) uses +//! the same line-oriented framing over an HTTP byte stream: buffer raw +//! bytes, split on newlines, decode each complete line as UTF-8, and guard +//! against unbounded buffer growth. The provider-specific event extraction +//! (`next_data` / `next_event`) lives in each provider file as an +//! `impl super::sse::SseReader` block. + +use std::pin::Pin; + +use futures::stream::{Stream, StreamExt}; +use reqwest::Response; + +use crate::api::error::ApiError; + +/// Upper bound on the internal line buffer (`1 MiB`). +/// +/// A single SSE event that exceeds this without a newline is treated as a +/// protocol error, preventing unbounded memory growth from a malformed or +/// hostile stream. +const SSE_MAX_BUFFER: usize = 1024 * 1024; + +/// Minimal SSE line reader over an HTTP byte stream. +/// +/// Buffers **raw bytes** from the response (not pre-decoded strings — +/// HTTP/TCP chunk boundaries are not UTF-8 aligned, so per-chunk decoding +/// corrupts multi-byte sequences split across chunks). Splits on `\n` in +/// the byte buffer, then decodes only complete lines. +pub(super) struct SseReader { + /// Underlying HTTP byte stream yielding raw `Bytes`. + /// + /// Pinned because it is polled as an async stream. + pub(super) bytes: Pin> + Send>>, + + /// Byte accumulator for data that has not yet formed a complete line. + /// + /// Newly arrived chunks are appended here as raw bytes; complete + /// `\n`-terminated lines are drained by [`take_line`](SseReader::take_line) + /// and decoded as UTF-8. + pub(super) buf: Vec, +} + +impl SseReader { + /// Wrap a streaming HTTP response into an SSE reader. + /// + /// The byte stream yields raw `Bytes`; decoding happens per-line in + /// [`take_line`](Self::take_line), after splitting on `\n`, so + /// multi-byte UTF-8 sequences split across chunks are reassembled + /// correctly. + pub(super) fn from_response(resp: Response) -> Self { + let bytes = resp + .bytes_stream() + .map(|res| res.map_err(|e| ApiError::http(e.to_string()))); + Self { + bytes: Box::pin(bytes), + buf: Vec::new(), + } + } + + /// Pop the first `\n`-terminated line from the byte buffer and decode it. + /// + /// Returns the trimmed line as a `String` if a newline is present, + /// removing it (plus the newline) from the buffer. Returns `None` if + /// the buffer does not yet contain a complete line. Uses + /// `String::from_utf8_lossy` on the complete line only — by this point + /// any split multi-byte sequence has been reassembled. + pub(super) fn take_line(&mut self) -> Option { + let pos = self.buf.iter().position(|&b| b == b'\n')?; + let rest_start = pos.saturating_add(1); + let line_bytes: Vec = self.buf.drain(..rest_start).collect(); + let line = String::from_utf8_lossy(&line_bytes); + Some(line.trim().to_string()) + } + + /// Fetch the next chunk from the HTTP stream and append it to the buffer. + /// + /// Returns `Ok(Some(()))` when a chunk arrived, `Ok(None)` at + /// end-of-stream. + /// + /// # Errors + /// + /// Returns [`ApiError`] if the transport failed or the buffer + /// exceeded [`SSE_MAX_BUFFER`]. + pub(super) async fn next_chunk(&mut self) -> Result, ApiError> { + match self.bytes.next().await { + Some(Ok(chunk)) => { + self.buf.extend_from_slice(&chunk); + if self.buf.len() > SSE_MAX_BUFFER { + return Err(ApiError::http(format!( + "SSE buffer exceeded {SSE_MAX_BUFFER} bytes" + ))); + } + Ok(Some(())) + } + Some(Err(e)) => Err(e), + None => Ok(None), + } + } +} diff --git a/src/reflection.rs b/src/reflection.rs index 567f066..a442bc7 100644 --- a/src/reflection.rs +++ b/src/reflection.rs @@ -46,10 +46,6 @@ use std::future::Future; use std::pin::Pin; use std::time::Duration; -// =================================================== -// FailureSeverity -// =================================================== - /// How severe a failure is. /// /// Ordered from least to most severe. The [`RecoveryStrategy`] may use @@ -140,10 +136,6 @@ impl fmt::Display for FailureSeverity { } } -// =================================================== -// ReflectionContext -// =================================================== - /// Context provided to [`Reflector::analyze()`] describing the retry state. /// /// Built by the framework before invoking the reflector. Carries what the @@ -210,10 +202,6 @@ pub struct ReflectionContext { pub max_attempts: u32, } -// =================================================== -// Correction -// =================================================== - /// Type of correction to apply. /// /// Categorizes the fix strategy that the reflection system has determined @@ -368,10 +356,6 @@ pub enum CorrectionResult { Skipped, } -// =================================================== -// FailureAnalysis -// =================================================== - /// Result of analyzing a failure via [`Reflector::analyze()`]. /// /// Describes what went wrong, how severe it is, whether it's worth @@ -536,10 +520,6 @@ impl crate::structured::StructuredOutput for FailureAnalysis { } } -// =================================================== -// ReflectionError -// =================================================== - /// Errors produced by [`Reflector::analyze()`]. /// /// The reflector can either produce a valid analysis, skip analysis @@ -562,10 +542,6 @@ pub enum ReflectionError { Internal(String), } -// =================================================== -// RecoveryAction -// =================================================== - /// What the framework should do after a failure. /// /// Produced by [`RecoveryStrategy::decide()`] after the @@ -735,10 +711,6 @@ impl fmt::Display for RecoveryAction { } } -// =================================================== -// Reflector trait -// =================================================== - /// Analyses failed tool calls to determine recoverability and corrections. /// /// Implement this trait to provide custom failure analysis — for example, @@ -821,10 +793,6 @@ pub trait Reflector: Send + Sync { ) -> Pin> + Send + '_>>; } -// =================================================== -// RecoveryStrategy trait -// =================================================== - /// Decides what to do after a failure has been analyzed. /// /// Takes a [`FailureAnalysis`] and the current retry state, returns a @@ -884,10 +852,6 @@ pub trait RecoveryStrategy: Send + Sync { ) -> Pin + Send + '_>>; } -// =================================================== -// NoopReflector -// =================================================== - /// A no-op reflector that marks everything as non-recoverable. /// /// Useful as a safe default when reflection is disabled, or as a base diff --git a/src/reflection/backoff.rs b/src/reflection/backoff.rs index b329e74..3cff9c8 100644 --- a/src/reflection/backoff.rs +++ b/src/reflection/backoff.rs @@ -18,10 +18,6 @@ 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 @@ -48,10 +44,28 @@ use std::time::Duration; #[derive(Debug, Clone)] pub struct ExponentialBackoffRecovery { /// Maximum number of retry attempts before giving up. + /// + /// Counts only retries issued by this strategy; the framework's own + /// `max_attempts` budget is honoured separately and whichever ceiling + /// is reached first wins. Exposed via + /// [`max_retries`](Self::max_retries). max_retries: u32, - /// Initial delay applied to the first retry, doubled on each subsequent attempt. + + /// Initial delay applied to the first retry, doubled on each + /// subsequent attempt. + /// + /// The delay for attempt `n` is `base_delay * 2^n` (computed in + /// [`delay_for_attempt`](Self::delay_for_attempt)). Exposed via + /// [`base_delay`](Self::base_delay). base_delay: Duration, - /// Upper bound on the computed delay, regardless of exponential growth. + + /// Upper bound on the computed delay, regardless of exponential + /// growth. + /// + /// Once `base_delay * 2^n` would exceed this value, every subsequent + /// attempt returns `max_delay` instead, preventing unbounded wait + /// times on long retry chains. Exposed via + /// [`max_delay`](Self::max_delay). max_delay: Duration, } @@ -70,35 +84,50 @@ impl ExponentialBackoffRecovery { /// Set a custom base delay for the exponential backoff. /// - /// The delay for attempt `n` is `base_delay * 2^n`, capped at `max_delay`. + /// Builder-style override of the default 100 ms base. The delay for + /// attempt `n` is `base_delay * 2^n`, capped at `max_delay`; raising + /// the base scales every subsequent attempt proportionally. #[must_use] pub fn with_base_delay(mut self, delay: Duration) -> Self { self.base_delay = delay; self } - /// Set a maximum delay cap. + /// Set the maximum delay cap. /// - /// Even as the exponential grows, the delay never exceeds this value. + /// Builder-style override of the default 30 s cap. Even as the + /// exponential grows, the computed delay never exceeds this value, + /// bounding the worst-case wait between retries. #[must_use] pub fn with_max_delay(mut self, delay: Duration) -> Self { self.max_delay = delay; self } - /// Maximum retry attempts before giving up. + /// Maximum retry attempts before giving up, as configured at + /// construction. + /// + /// Note the framework's own `max_attempts` budget is honoured + /// separately; whichever ceiling is reached first terminates the + /// retry loop. #[must_use] pub fn max_retries(&self) -> u32 { self.max_retries } - /// Base delay for exponential backoff calculation. + /// Base delay used as the multiplier in the exponential backoff. + /// + /// The first retry (attempt 0) waits exactly this long; each later + /// retry doubles it until `max_delay` is reached. #[must_use] pub fn base_delay(&self) -> Duration { self.base_delay } - /// Maximum delay cap. + /// Maximum delay cap that the exponential growth will not exceed. + /// + /// Once `base_delay * 2^n` surpasses this value, every subsequent + /// retry waits exactly `max_delay`. #[must_use] pub fn max_delay(&self) -> Duration { self.max_delay diff --git a/src/reflection/llm.rs b/src/reflection/llm.rs index ffc88f2..1f3bc9e 100644 --- a/src/reflection/llm.rs +++ b/src/reflection/llm.rs @@ -237,28 +237,26 @@ fn validate_modified_input( return Ok(()); } - let Some(correction) = &analysis.correction else { - return Ok(()); - }; - let Some(modified_input) = &correction.modified_input else { - return Ok(()); - }; - let Some(schema) = tool_schema else { - // No schema available — the engine couldn't resolve the tool. - // Skip validation rather than reject a possibly-correct fix. - return Ok(()); - }; - #[cfg(feature = "schema_validation")] { + let Some(correction) = &analysis.correction else { + return Ok(()); + }; + let Some(modified_input) = &correction.modified_input else { + return Ok(()); + }; + let Some(schema) = tool_schema else { + return Ok(()); + }; + if !jsonschema::is_valid(schema, modified_input) { return Err(ReflectionError::Internal( "corrected input does not match the tool's schema".to_string(), )); } - } - Ok(()) + Ok(()) + } } #[cfg(test)] diff --git a/src/stream.rs b/src/stream.rs index 3818ec5..e39ab60 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -35,8 +35,6 @@ //! //! - **[`handler`]** — [`handler::StreamHandler`] with retry, timeout, //! and fallback for resilient streaming. -//! - **[`heartbeat`]** — [`heartbeat::HeartbeatStream`] composable heartbeat -//! and timeout wrapper for any stream. //! //! # Quick Start //! @@ -61,16 +59,11 @@ use serde_json::Value; use std::fmt; pub mod handler; -pub mod heartbeat; pub mod rate_limit; pub use handler::{DetectedRateLimit, RateLimitConfig, RateLimitKind}; pub use rate_limit::{RateLimiter, TokenBucket}; -// ================================================== -// StreamError -// ================================================== - /// Errors that can occur during stream event processing. /// /// Returned by [`StreamAccumulator::process`] when an event cannot be @@ -104,10 +97,6 @@ impl std::error::Error for StreamError { } } -// ================================================== -// StreamEvent -// ================================================== - /// An event from a streaming LLM API response. /// /// Events follow the SSE (Server-Sent Events) protocol used by @@ -218,10 +207,6 @@ pub enum StreamEvent { Ping, } -// ================================================== -// MessageStart -// ================================================== - /// The start of a new message from the API. /// /// Wraps [`MessageMetadata`] and is always the first event in a @@ -255,10 +240,6 @@ pub struct MessageStart { pub message: MessageMetadata, } -// ================================================== -// MessageMetadata -// ================================================== - /// Metadata about a message from the API. /// /// Carries identifying information for the streaming response: the @@ -308,10 +289,6 @@ pub struct MessageMetadata { pub model: String, } -// ================================================== -// PartStart -// ================================================== - /// The start of a new part within the response. /// /// Emitted once per part, before any [`IndexedDelta`] @@ -354,10 +331,6 @@ pub struct PartStart { pub part: Option, } -// ================================================== -// IndexedDelta -// ================================================== - /// A delta (incremental update) for the current part. /// /// Carries a [`DeltaPart`] payload that should be appended to the @@ -392,10 +365,6 @@ pub struct IndexedDelta { pub delta: DeltaPart, } -// ================================================== -// DeltaPart -// ================================================== - /// A delta (incremental update) for content within a streaming response. /// /// Each variant represents a different kind of incremental data that @@ -517,10 +486,6 @@ pub enum DeltaPart { }, } -// ================================================== -// StreamStopReason -// ================================================== - /// Reason why the model stopped generating tokens. /// /// Streaming / API-level stop reason returned by the LLM @@ -562,10 +527,9 @@ pub enum StreamStopReason { /// The model reached the configured maximum token limit. /// - /// The response was truncated. The caller may want to request - /// continuation or increase the token budget. - /// - /// *Note: `LoopConfig` will be available once the builder module is complete.* + /// The response was truncated because the model produced `max_tokens` of + /// output before finishing. The caller may want to request continuation or + /// raise the provider's max-tokens limit. MaxTokens, /// The model hit a configured stop sequence. @@ -672,10 +636,6 @@ impl StreamStopReason { } } -// ================================================== -// MessageDelta -// ================================================== - /// A delta update for the message, typically emitted at the end of the stream. /// /// Carries the final [`StreamStopReason`] (why the model stopped) and @@ -716,10 +676,6 @@ pub struct MessageDelta { pub usage: Option, } -// ================================================== -// MessageDeltaPayload -// ================================================== - /// The delta details within a [`MessageDelta`] event. /// /// Contains the stop reason string that explains why the model @@ -753,10 +709,6 @@ pub struct MessageDeltaPayload { pub stop_reason: Option, } -// ================================================== -// Usage -// ================================================== - /// Token usage statistics from an API response. /// /// Tracks input and output token counts for a single streaming @@ -853,10 +805,6 @@ impl Usage { } } -// ================================================== -// Stream accumulator -// ================================================== - /// Accumulates streaming events into a complete [`Message`]. /// /// Stateful builder that tracks the progress of a streaming @@ -1170,10 +1118,6 @@ impl StreamAccumulator { } } -// ================================================== -// Tests -// ================================================== - #[cfg(test)] mod tests { use super::*; diff --git a/src/stream/handler.rs b/src/stream/handler.rs index 0544823..a58e1b2 100644 --- a/src/stream/handler.rs +++ b/src/stream/handler.rs @@ -46,10 +46,6 @@ use std::pin::Pin; use std::sync::Arc; use std::time::{Duration, Instant}; -// =================================================== -// StreamTimeoutConfig -// =================================================== - /// Configuration for the [`StreamHandler`]'s timeout behaviour. /// /// Controls how long the handler waits for events at each phase of the @@ -183,10 +179,6 @@ impl StreamTimeoutConfig { } } -// =================================================== -// StreamRetryConfig -// =================================================== - /// Configuration for retry behaviour when stream initialization fails. /// /// Uses exponential backoff with jitter to avoid thundering herd when @@ -321,10 +313,6 @@ impl StreamRetryConfig { } } -// =================================================== -// RateLimitConfig + detection -// =================================================== - /// Policy for handling 429 / 503 rate-limit responses from LLM providers. /// /// Governs backoff and retry behaviour when the server signals that the client @@ -443,6 +431,12 @@ impl RateLimitConfig { if self.max_retries == 0 { return Err("max_retries must be >= 1".into()); } + if self.fallback_after_retries > self.max_retries { + return Err(format!( + "fallback_after_retries ({}) must be <= max_retries ({})", + self.fallback_after_retries, self.max_retries + )); + } Ok(()) } @@ -831,10 +825,6 @@ impl EventDiagnostics { } } -// =================================================== -// StreamOutcome -// =================================================== - /// Why the stream ended. /// /// Each variant captures the relevant context for how streaming @@ -856,9 +846,19 @@ pub enum StreamOutcome { /// Happy path. The [`StreamAccumulator`] /// contains the full response. Completed { - /// Number of SSE events processed. + /// Number of SSE events processed before `MessageStop`. + /// + /// Counts every event the accumulator accepted, including + /// keep-alive heartbeats and metadata events. Useful as a + /// throughput signal alongside `duration`. events_processed: u64, - /// Wall-clock duration of the stream. + + /// Wall-clock duration of the stream from first byte to + /// `MessageStop`. + /// + /// Measured inside the handler; divide `events_processed` by + /// this to get the average events-per-second rate for + /// telemetry. duration: Duration, }, @@ -868,11 +868,26 @@ pub enum StreamOutcome { /// [`total_stream_timeout`](StreamTimeoutConfig::total_stream_timeout). /// Partial data may be available in the accumulator. TotalTimeout { - /// Whether partial content was accumulated before timeout. + /// Whether partial content was accumulated before the timeout. + /// + /// `true` when at least one content event arrived before the + /// deadline; the caller may inspect the accumulator to decide + /// whether the partial result is usable. has_partial_data: bool, - /// Events processed before timeout. + + /// Events processed before the timeout fired. + /// + /// How many SSE events the accumulator accepted before the + /// overall stream deadline elapsed — zero implies the stream + /// stalled immediately. events_processed: u64, - /// Duration before timeout was triggered. + + /// Elapsed time from stream start to the timeout trigger. + /// + /// Approximately equal to + /// [`total_stream_timeout`](StreamTimeoutConfig::total_stream_timeout), + /// reported for diagnostics so callers can correlate the + /// observed wait with the configured ceiling. duration: Duration, }, @@ -882,9 +897,21 @@ pub enum StreamOutcome { /// [`per_event_timeout`](StreamTimeoutConfig::per_event_timeout). /// Partial data may be available. EventTimeout { - /// Whether partial content was accumulated. + /// Whether partial content was accumulated before the failure. + /// + /// `true` when at least one content event arrived before the + /// consecutive-timeout threshold was reached; the caller may + /// inspect the accumulator to decide whether the partial result + /// is usable. has_partial_data: bool, - /// Consecutive timeouts that triggered the failure. + + /// Consecutive per-event timeouts that triggered the failure. + /// + /// Reaches + /// [`max_consecutive_timeouts`](StreamTimeoutConfig::max_consecutive_timeouts) + /// when the handler gives up. Each individual timeout equals + /// [`per_event_timeout`](StreamTimeoutConfig::per_event_timeout); + /// this count tells the caller how many gaps were observed. consecutive_timeouts: u32, }, @@ -895,11 +922,27 @@ pub enum StreamOutcome { /// [`DetectedRateLimit`] carries the parsed `Retry-After`, if the server /// provided one. RateLimited { - /// Decoded rate-limit detail (kind + parsed `Retry-After`). + /// Decoded rate-limit detail. + /// + /// Carries the rate-limit kind (429 / 503 / 529) and the parsed + /// `Retry-After` hint, if the server provided one. Downstream + /// layers (the fallback manager, the retry loop) read this to + /// honour the provider's back-off guidance without re-parsing + /// the response. detail: DetectedRateLimit, + /// Whether partial content was accumulated before the rate limit. + /// + /// `true` when the stream produced content events before the + /// provider started rate-limiting; the caller may inspect the + /// accumulator to decide whether the partial result is usable. has_partial_data: bool, + /// Events processed before the rate limit fired. + /// + /// How many SSE events the accumulator accepted before the + /// 429/503/529 response arrived — zero implies the provider + /// rejected the stream early. events_processed: u64, }, @@ -909,8 +952,19 @@ pub enum StreamOutcome { /// No data was accumulated. InitFailed { /// The last error from the final retry attempt. + /// + /// Rendered to a string for diagnostics; typically the + /// underlying transport or HTTP error that prevented the + /// handler from receiving a first event. Surface this in logs + /// so the caller can see why the stream never started. last_error: String, - /// Number of retry attempts made. + + /// Number of retry attempts made before giving up. + /// + /// Counts the (re-)connection attempts up to + /// [`StreamRetryConfig`]'s ceiling; reaching this count without + /// a first event means the provider was unreachable or + /// rejecting the request outright. attempts: u32, }, @@ -1007,10 +1061,6 @@ impl fmt::Display for StreamOutcome { } } -// =================================================== -// StreamHandlerError -// =================================================== - /// Errors produced by [`StreamHandler`]. /// /// Each variant captures the specific failure mode, allowing callers @@ -1036,9 +1086,21 @@ pub enum StreamHandlerError { /// The handler attempted a fallback `create_message()` call after /// streaming failed, but the fallback also produced an error. FallbackFailed { - /// The streaming failure that triggered the fallback. + /// The streaming failure that triggered the fallback attempt. + /// + /// Preserved verbatim so the caller can see both halves of the + /// double failure — why streaming gave up (timeout, mid-stream + /// error, rate limit) and why the non-streaming retry then + /// failed — without losing the original context. stream_outcome: StreamOutcome, - /// The error from the fallback request. + + /// The error returned by the non-streaming fallback request. + /// + /// Rendered to a string for diagnostics; typically the + /// underlying transport or HTTP error from the + /// `create_message()` call. Surface both this and + /// `stream_outcome` in logs so the caller can see the full + /// chain of failures. fallback_error: String, }, @@ -1057,8 +1119,21 @@ pub enum StreamHandlerError { /// same-model non-streaming fallback. RateLimitEscalation { /// Number of rate-limit retries honored before escalating. + /// + /// Counts the 429/503/529 responses the handler retried + /// (honoring the provider's `Retry-After`) before giving up on + /// the current model and handing control to the model circuit + /// breaker. Reaching + /// [`RateLimitConfig::fallback_after_retries`] triggers this + /// variant. attempts: u32, - /// Last server-advised `Retry-After` hint, after clamping. `None` when no header sent. + + /// Last server-advised `Retry-After` hint, after clamping. + /// + /// Preserved for diagnostics and back-off tuning so the caller + /// can correlate the escalation with the provider's last + /// guidance. `None` when the provider sent no `Retry-After` + /// header on the final rate-limited response. retry_after: Option, }, } @@ -1091,10 +1166,6 @@ impl fmt::Display for StreamHandlerError { impl std::error::Error for StreamHandlerError {} -// =================================================== -// StreamProgress -// =================================================== - /// A snapshot of stream progress for external reporting. /// /// Plain data struct carrying the two progress signals a consumer is likely @@ -1133,10 +1204,6 @@ pub struct StreamProgress { pub events_processed: u64, } -// =================================================== -// StreamHandler -// =================================================== - /// Holds configuration for the streaming resilience layer. /// /// `StreamHandler` wraps an [`ApiClient`]'s streaming path with timeout, @@ -1218,6 +1285,13 @@ pub struct StreamHandler { /// than risking a 429. `None` (the default) means reactive-only handling: /// no pre-throttling, server 429s still handled via `rate_limit_config`. rate_limiter: Option>, + + /// Upper bound on how long `gate_on_rate_limit` blocks for a token. + /// + /// When the limiter's `acquire` returns a wait exceeding this, the + /// handler proceeds anyway rather than hanging the agent. Defaults + /// to 30 seconds. + rate_limit_max_wait: Duration, } impl fmt::Debug for StreamHandler { @@ -1227,6 +1301,7 @@ impl fmt::Debug for StreamHandler { .field("retry_config", &self.retry_config) .field("rate_limit_config", &self.rate_limit_config) .field("rate_limiter", &self.rate_limiter) + .field("rate_limit_max_wait", &self.rate_limit_max_wait) .finish() } } @@ -1290,6 +1365,7 @@ impl StreamHandler { ..Default::default() }, rate_limiter: None, + rate_limit_max_wait: Duration::from_secs(30), } } @@ -1328,6 +1404,7 @@ impl StreamHandler { retry_config: StreamRetryConfig::default(), rate_limit_config: RateLimitConfig::default(), rate_limiter: None, + rate_limit_max_wait: Duration::from_secs(30), } } @@ -1432,9 +1509,16 @@ impl StreamHandler { self } - // ================================================== - // Runtime methods - // ================================================== + /// Set the max-wait ceiling for `gate_on_rate_limit` (builder style). + /// + /// Defaults to 30 seconds. When the limiter's `acquire` returns a + /// wait exceeding this, the handler proceeds anyway rather than + /// hanging the agent. + #[must_use] + pub fn with_rate_limit_max_wait(mut self, max_wait: Duration) -> Self { + self.rate_limit_max_wait = max_wait; + self + } /// Drive one turn as a stream of [`HandlerEvent`]s. /// @@ -1476,12 +1560,13 @@ impl StreamHandler { // "no total deadline" (`None`) rather than wrapping to `Instant::now()` // (which would immediately fire a spurious TotalTimeout). let total_deadline = Instant::now().checked_add(self.timeout_config.total_stream_timeout); + let stream_start = Instant::now(); let max_attempts = self.retry_config.max_retries.saturating_add(1); Box::pin(async_stream::try_stream! { let mut rate_limit_retries: u32 = 0; let mut transport_attempts: u32 = 0; - let mut attempt: u32 = 0; + let mut first_attempt = true; // Shadow accumulator tracks partial-data presence for diagnostics. let mut shadow = StreamAccumulator::new(); @@ -1493,16 +1578,16 @@ impl StreamHandler { StreamOutcome::TotalTimeout { has_partial_data: !shadow.peek_parts().is_empty(), events_processed: 0, - duration: self.timeout_config.total_stream_timeout, + duration: stream_start.elapsed(), }, ))?; } - attempt = attempt.saturating_add(1); - if attempt > 1 { + if !first_attempt { shadow = StreamAccumulator::new(); yield HandlerEvent::AttemptReset; } + first_attempt = false; self.gate_on_rate_limit(client, cancel, total_deadline).await?; let request = crate::api::StreamRequest { @@ -1536,7 +1621,7 @@ impl StreamHandler { if let Err(e) = shadow.process(&event) { Err(StreamHandlerError::StreamFailed( StreamOutcome::InitFailed { - attempts: 1, + attempts: transport_attempts.saturating_add(1), last_error: e.to_string(), }, ))?; @@ -1678,7 +1763,7 @@ impl StreamHandler { return Ok(()); }; let key = client.base_url(); - let max_wait = limiter.max_wait(); + let max_wait = self.rate_limit_max_wait; let mut waited = Duration::ZERO; loop { match limiter.acquire(&key) { @@ -1930,12 +2015,23 @@ pub enum HandlerEvent { /// Always the last event before the stream ends (when the fallback path /// is taken). Fallback { - /// The fallback assistant message (text extracted from the JSON - /// response's first text part). + /// The fallback assistant message produced by the non-streaming + /// request. + /// + /// Built from the first text part of the + /// [`create_message`](crate::api::ApiClient::create_message) JSON + /// response. The engine should treat this as the authoritative + /// turn output — the streaming accumulator's partial state from + /// failed attempts is discarded on this path. message: Message, - /// Stop reason parsed from the JSON response’s `stop_reason` field, - /// defaulting to [`EndTurn`](StreamStopReason::EndTurn) when absent - /// or unrecognized. + + /// Stop reason parsed from the JSON response's `stop_reason` + /// field. + /// + /// Defaults to [`EndTurn`](StreamStopReason::EndTurn) when the + /// field is absent or holds an unrecognized value, so the engine + /// always has a concrete reason to act on. Drives the same + /// downstream behaviour as a streaming `MessageStop`. stop_reason: StreamStopReason, }, } @@ -3367,8 +3463,10 @@ mod tests { // process_events deadline checks report the actual TotalTimeout — the // gate's job is just to not overrun the budget. use crate::stream::rate_limit::RateLimiter; - let limiter = Arc::new(RateLimiter::new(1).with_max_wait(Duration::from_mins(2))); - let handler = StreamHandler::new().with_rate_limiter(Arc::clone(&limiter)); + let limiter = Arc::new(RateLimiter::new(1)); + let handler = StreamHandler::new() + .with_rate_limiter(Arc::clone(&limiter)) + .with_rate_limit_max_wait(Duration::from_mins(2)); let client = GateMock { url: "openai" }; let cancel = Arc::new(CancelSignal::new()); @@ -3404,8 +3502,10 @@ mod tests { // after ~80ms — proving it honors the turn ceiling rather than the // 60s refill wait or the 120s max_wait. use crate::stream::rate_limit::RateLimiter; - let limiter = Arc::new(RateLimiter::new(1).with_max_wait(Duration::from_mins(2))); - let handler = StreamHandler::new().with_rate_limiter(Arc::clone(&limiter)); + let limiter = Arc::new(RateLimiter::new(1)); + let handler = StreamHandler::new() + .with_rate_limiter(Arc::clone(&limiter)) + .with_rate_limit_max_wait(Duration::from_mins(2)); let client = GateMock { url: "openai" }; let cancel = Arc::new(CancelSignal::new()); @@ -3439,7 +3539,9 @@ mod tests { // back-to-back through stream_turn (end-to-end wiring). use crate::stream::rate_limit::RateLimiter; let limiter = Arc::new(RateLimiter::new(60)); - let handler = StreamHandler::new().with_rate_limiter(Arc::clone(&limiter)); + let handler = StreamHandler::new() + .with_rate_limiter(Arc::clone(&limiter)) + .with_rate_limit_max_wait(Duration::from_mins(2)); let client = GateMock { url: "openai" }; let cancel = Arc::new(CancelSignal::new()); @@ -3464,8 +3566,10 @@ mod tests { // second turn must wait ~60s. Cancelling during that wait should return // promptly. use crate::stream::rate_limit::RateLimiter; - let limiter = Arc::new(RateLimiter::new(1).with_max_wait(Duration::from_mins(2))); - let handler = StreamHandler::new().with_rate_limiter(limiter); + let limiter = Arc::new(RateLimiter::new(1)); + let handler = StreamHandler::new() + .with_rate_limiter(limiter) + .with_rate_limit_max_wait(Duration::from_millis(50)); let client = GateMock { url: "openai" }; // First turn consumes the only token. @@ -3500,8 +3604,10 @@ mod tests { // 1 RPM but max_wait = 50ms. The second turn must wait ~60s for a token, // but the clamp caps the cumulative wait at 50ms, so the turn proceeds. use crate::stream::rate_limit::RateLimiter; - let limiter = Arc::new(RateLimiter::new(1).with_max_wait(Duration::from_millis(50))); - let handler = StreamHandler::new().with_rate_limiter(limiter); + let limiter = Arc::new(RateLimiter::new(1)); + let handler = StreamHandler::new() + .with_rate_limiter(limiter) + .with_rate_limit_max_wait(Duration::from_millis(50)); let client = GateMock { url: "openai" }; let cancel = Arc::new(CancelSignal::new()); diff --git a/src/stream/heartbeat.rs b/src/stream/heartbeat.rs deleted file mode 100644 index 2b75c9a..0000000 --- a/src/stream/heartbeat.rs +++ /dev/null @@ -1,556 +0,0 @@ -//! Composable heartbeat and timeout wrapper for any stream. -//! -//! [`HeartbeatStream`] wraps any `Stream>` -//! and adds two behaviours: -//! -//! 1. **Heartbeat callbacks** — Fires a callback at regular intervals to report -//! elapsed time and timeout status. -//! 2. **Hard timeout** — Returns an [`ApiError`] if the stream exceeds a -//! configured maximum duration. -//! -//! It does **not** retry or fall back — that's [`StreamHandler`](super::handler::StreamHandler)'s -//! job. Use this when you need heartbeat/timeout on a stream you've already opened. -//! -//! # Architecture -//! -//! On each `poll_next`, `HeartbeatStream` runs three checks in order before -//! delegating to the inner stream: -//! -//! 1. **Heartbeat interval** — if the configured interval has elapsed -//! since the last beat, fire the heartbeat callback with the elapsed -//! time and current timeout status. -//! 2. **Hard timeout** — if the total elapsed time has exceeded the -//! configured maximum, return an [`ApiError`] without consulting the -//! inner stream. -//! 3. **Delegate** — otherwise, forward to the inner stream's `poll_next` -//! and pass through its result. -//! -//! [`ApiError`]: crate::api::error::ApiError -//! -//! # Quick Start -//! -//! ```rust -//! use loopctl::stream::heartbeat::{HeartbeatStream, HeartbeatConfig, HeartbeatData}; -//! use std::time::Duration; -//! use std::sync::{Arc, Mutex}; -//! -//! let callbacks = Arc::new(Mutex::new(Vec::new())); -//! let cb = callbacks.clone(); -//! -//! let config = HeartbeatConfig::new( -//! Duration::from_secs(30), // heartbeat_interval -//! Duration::from_secs(600), // timeout -//! Box::new(move |data: HeartbeatData| { -//! cb.lock().unwrap().push(data.elapsed); -//! }), -//! ); -//! ``` - -use crate::api::error::ApiError; -use crate::stream::StreamEvent; -use futures::Stream; -use std::pin::Pin; -use std::task::{Context, Poll}; -use std::time::{Duration, Instant}; - -// =================================================== -// HeartbeatData -// =================================================== - -/// Data emitted on each heartbeat callback. -/// -/// Passed to the callback registered in [`HeartbeatConfig`] at each -/// heartbeat interval. Carries a snapshot of how long the stream has -/// been running and whether it has crossed its configured hard-timeout -/// deadline, so a UI or metrics collector can render progress without -/// owning a clock itself. -/// -/// # Example -/// -/// ```rust -/// use loopctl::stream::heartbeat::HeartbeatData; -/// use std::time::Duration; -/// -/// let data = HeartbeatData { -/// elapsed: Duration::from_secs(45), -/// is_timeout: false, -/// }; -/// assert!(!data.is_timeout); -/// ``` -#[derive(Debug, Clone)] -pub struct HeartbeatData { - /// Time elapsed since the wrapped stream was constructed. - /// - /// Monotonic — measured from the [`Instant`] captured in - /// [`HeartbeatStream::new`], not wall-clock time. Useful for - /// progress UIs ("streaming for 45s"), metrics, and detecting - /// stalls without each consumer holding its own start instant. - /// - /// [`Instant`]: std::time::Instant - pub elapsed: Duration, - - /// Whether the stream has exceeded its configured hard timeout. - /// - /// Set to `elapsed > config.timeout` at the moment the heartbeat - /// fires. Note this is *advisory*: the heartbeat may observe - /// `is_timeout == true` a beat after the deadline actually crossed, - /// because callbacks only fire on `poll_next`. The next - /// [`HeartbeatStream::poll_next`] after the deadline will return - /// the hard-timeout error regardless of when the callback last - /// fired. - pub is_timeout: bool, -} - -// =================================================== -// HeartbeatCallback -// =================================================== - -/// Callback type for heartbeat events. -/// -/// A `Box` that is called at each -/// heartbeat interval with the current stream status. -pub type HeartbeatCallback = Box; - -// =================================================== -// HeartbeatConfig -// =================================================== - -/// Configuration for a [`HeartbeatStream`]. -/// -/// Holds the heartbeat interval, hard timeout, and callback function. -/// Created via [`HeartbeatConfig::new()`]. -/// -/// # Example -/// -/// ```rust -/// use loopctl::stream::heartbeat::{HeartbeatConfig, HeartbeatData}; -/// use std::time::Duration; -/// -/// let config = HeartbeatConfig::new( -/// Duration::from_secs(30), -/// Duration::from_secs(600), -/// Box::new(|_data: HeartbeatData| {}), -/// ); -/// assert_eq!(config.heartbeat_interval(), Duration::from_secs(30)); -/// assert_eq!(config.timeout(), Duration::from_secs(600)); -/// ``` -pub struct HeartbeatConfig { - /// Interval between heartbeat callbacks. - /// - /// Checked on every `poll_next`: when `last_heartbeat.elapsed()` - /// reaches this value, the callback fires and `last_heartbeat` is - /// reset. There is no background timer — heartbeats only fire while - /// the stream is actively being polled. - heartbeat_interval: Duration, - - /// Maximum total stream duration before triggering a hard timeout. - /// - /// Once `start.elapsed()` exceeds this value, the next `poll_next` - /// returns an [`ApiError`] instead of delegating to the inner - /// stream. Choose a value comfortably above the expected p99 - /// response time so transient slowness doesn't trip it. - /// - /// [`ApiError`]: crate::api::error::ApiError - timeout: Duration, - - /// Callback invoked at each heartbeat interval. - /// - /// A `Box` so it can be shared - /// across the runtime and mutate captured state (typically an - /// `Arc>`-protected metrics struct or channel sender). - /// Called synchronously from `poll_next` — keep the body cheap to - /// avoid stalling the stream's task. - on_heartbeat: HeartbeatCallback, -} - -impl std::fmt::Debug for HeartbeatConfig { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("HeartbeatConfig") - .field("heartbeat_interval", &self.heartbeat_interval) - .field("timeout", &self.timeout) - .finish_non_exhaustive() - } -} - -impl HeartbeatConfig { - /// Create a new heartbeat configuration. - /// - /// # Arguments - /// - /// - `heartbeat_interval` — How often to fire the callback. - /// - `timeout` — Maximum stream duration before returning an error. - /// - `on_heartbeat` — Callback invoked at each interval. - /// - /// # Example - /// - /// ```rust - /// use loopctl::stream::heartbeat::{HeartbeatConfig, HeartbeatData}; - /// use std::time::Duration; - /// - /// let config = HeartbeatConfig::new( - /// Duration::from_secs(15), - /// Duration::from_secs(300), - /// Box::new(|data: HeartbeatData| { - /// println!("heartbeat: {:.1}s elapsed", data.elapsed.as_secs_f64()); - /// }), - /// ); - /// ``` - #[must_use] - pub fn new( - heartbeat_interval: Duration, - timeout: Duration, - on_heartbeat: HeartbeatCallback, - ) -> Self { - Self { - heartbeat_interval, - timeout, - on_heartbeat, - } - } - - /// Returns the configured heartbeat interval. - /// - /// Exposed so callers (e.g. a metrics reporter that wants to align - /// its own cadence with the heartbeat) can read the value the - /// stream was constructed with. - #[must_use] - pub fn heartbeat_interval(&self) -> Duration { - self.heartbeat_interval - } - - /// Returns the configured hard timeout. - /// - /// Exposed so callers can render the deadline ("stream will time out - /// after 600s") or compute remaining budget from the elapsed time. - #[must_use] - pub fn timeout(&self) -> Duration { - self.timeout - } -} - -// =================================================== -// HeartbeatStream -// =================================================== - -/// A stream wrapper that emits heartbeat callbacks and enforces a hard timeout. -/// -/// Wraps any `Stream>` and adds: -/// - Periodic heartbeat callbacks via [`HeartbeatConfig`]. -/// - A hard timeout that returns an [`ApiError`] when exceeded. -/// -/// Does **not** retry or fallback — use [`StreamHandler`](super::handler::StreamHandler) instead. -/// -/// # Composability -/// -/// `HeartbeatStream` implements `Stream` directly, so it composes with -/// any other stream wrapper. Use it on any stream you've already opened -/// when you need heartbeat/timeout without the full handler lifecycle. -/// -/// # Example -/// -/// ```rust -/// use loopctl::stream::heartbeat::{HeartbeatStream, HeartbeatConfig, HeartbeatData}; -/// use std::time::Duration; -/// -/// let config = HeartbeatConfig::new( -/// Duration::from_secs(30), -/// Duration::from_secs(600), -/// Box::new(|_data: HeartbeatData| {}), -/// ); -/// -/// // Wrap any stream: -/// // let heartbeat_stream = HeartbeatStream::new(inner_stream, config); -/// // while let Some(result) = futures::StreamExt::next(&mut heartbeat_stream).await { -/// // // ... -/// // } -/// ``` -pub struct HeartbeatStream { - /// The inner stream being wrapped. - /// - /// All `poll_next` calls that survive the heartbeat check and the - /// hard-timeout check are delegated to this stream. The wrapper - /// does not buffer or transform items — it passes them through - /// verbatim, including errors from the inner stream. - inner: S, - - /// Heartbeat and timeout configuration. - /// - /// Holds the interval, the timeout, and the callback. Owned by - /// the stream (not shared) so the wrapper can call the callback - /// without synchronization. - config: HeartbeatConfig, - - /// Time of the last heartbeat callback. - /// - /// Compared against `Instant::now()` on every `poll_next` to - /// decide whether the interval has elapsed. Reset to `Instant::now()` - /// (not `start + n*interval`) after each fire so drift from - /// irregular polling doesn't accumulate. - last_heartbeat: Instant, - - /// Time the stream was created. - /// - /// The reference for all `elapsed` calculations — heartbeat - /// `elapsed`, and the hard-timeout comparison. Captured once in - /// [`new`](Self::new) and never mutated for the life of the stream. - start: Instant, - - /// A `Sleep` future that fires at the hard-timeout deadline. - /// - /// Ensures the runtime wakes this task when the timeout expires, - /// even if the inner stream is `Pending` and nobody re-polls. - /// Polled proactively in `poll_next` so a ready `Sleep` short- - /// circuits to the timeout error without delegating to the inner - /// stream first. - timeout_sleep: std::pin::Pin>, -} - -impl HeartbeatStream { - /// Create a new heartbeat stream wrapping the given inner stream. - /// - /// The heartbeat timer starts immediately upon construction: `start` - /// and `last_heartbeat` are captured at this moment, and the - /// hard-timeout `Sleep` is armed against `start + config.timeout`. - /// The first heartbeat callback fires after `heartbeat_interval` - /// elapses (checked on each `poll_next` — there is no background - /// timer). - /// - /// # Timeout overflow - /// - /// `start + config.timeout` is computed via `checked_add`. For any - /// realistic `Duration` this always succeeds; the fallback (a - /// deadline 30 years in the future) only triggers for extreme values - /// near `Duration::MAX`, where the deadline is effectively "never" - /// either way. - /// - /// # Example - /// - /// ```rust - /// use loopctl::stream::heartbeat::{HeartbeatStream, HeartbeatConfig, HeartbeatData}; - /// use std::time::Duration; - /// - /// let config = HeartbeatConfig::new( - /// Duration::from_secs(30), - /// Duration::from_secs(600), - /// Box::new(|_data: HeartbeatData| {}), - /// ); - /// // let stream = HeartbeatStream::new(inner_stream, config); - /// ``` - 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: 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, - ))); - Self { - inner, - config, - last_heartbeat: now, - start: now, - timeout_sleep, - } - } -} - -impl Stream for HeartbeatStream -where - S: Stream> + Unpin, -{ - type Item = Result; - - fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let this = self.get_mut(); - - // Check heartbeat interval — fire callback if elapsed. - if this.last_heartbeat.elapsed() >= this.config.heartbeat_interval { - let elapsed = this.start.elapsed(); - let data = HeartbeatData { - elapsed, - is_timeout: elapsed > this.config.timeout, - }; - (this.config.on_heartbeat)(data); - this.last_heartbeat = Instant::now(); - } - - // Hard timeout — check the Sleep first (proactive wake-up), - // then fall back to elapsed() for the sync case. - if this.timeout_sleep.as_mut().poll(cx).is_ready() - || this.start.elapsed() > this.config.timeout - { - return Poll::Ready(Some(Err(ApiError::Api(format!( - "Stream timeout after {}s", - this.config.timeout.as_secs() - ))))); - } - - // Delegate to inner stream. - Pin::new(&mut this.inner).poll_next(cx) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use futures::StreamExt; - - struct VecStream { - items: Vec>, - } - - impl Stream for VecStream { - type Item = Result; - - fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(self.get_mut().items.pop()) - } - } - - fn make_config( - callbacks: &std::sync::Arc>>, - ) -> HeartbeatConfig { - let cb = callbacks.clone(); - HeartbeatConfig::new( - Duration::from_millis(10), - Duration::from_mins(1), - Box::new(move |data: HeartbeatData| { - cb.lock().unwrap().push(data); - }), - ) - } - - #[test] - fn heartbeat_data_fields() { - let data = HeartbeatData { - elapsed: Duration::from_secs(30), - is_timeout: true, - }; - assert_eq!(data.elapsed, Duration::from_secs(30)); - assert!(data.is_timeout); - } - - #[test] - fn config_accessors() { - let config = HeartbeatConfig::new( - Duration::from_secs(15), - Duration::from_mins(5), - Box::new(|_| {}), - ); - assert_eq!(config.heartbeat_interval(), Duration::from_secs(15)); - assert_eq!(config.timeout(), Duration::from_mins(5)); - } - - #[test] - fn config_debug() { - let config = HeartbeatConfig::new( - Duration::from_secs(30), - Duration::from_mins(10), - Box::new(|_| {}), - ); - let debug = format!("{config:?}"); - assert!(debug.contains("HeartbeatConfig")); - assert!(debug.contains("heartbeat_interval")); - assert!(debug.contains("timeout")); - } - - #[tokio::test] - async fn passes_through_events() { - let callbacks: std::sync::Arc>> = - std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); - let config = make_config(&callbacks); - - let inner = VecStream { - items: vec![Ok(StreamEvent::Ping), Ok(StreamEvent::Ping)], - }; - - let mut stream = HeartbeatStream::new(inner, config); - let first = stream.next().await; - assert!(first.is_some()); - - let second = stream.next().await; - assert!(second.is_some()); - - let third = stream.next().await; - assert!(third.is_none()); - } - - #[tokio::test] - async fn passes_through_errors() { - let callbacks: std::sync::Arc>> = - std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); - let config = make_config(&callbacks); - - let inner = VecStream { - items: vec![Err(ApiError::Api("test error".to_string()))], - }; - - let mut stream = HeartbeatStream::new(inner, config); - let result = stream.next().await; - assert!(matches!(result, Some(Err(ApiError::Api(_))))); - } - - #[tokio::test] - async fn fires_heartbeat_on_interval_sync() { - let callbacks: std::sync::Arc>> = - std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); - let config = make_config(&callbacks); - let inner = VecStream { - items: vec![Ok(StreamEvent::Ping)], - }; - let mut stream = HeartbeatStream::new(inner, config); - - 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); - let result = Pin::new(&mut stream).poll_next(&mut cx); - - assert!(matches!(result, Poll::Ready(Some(Ok(StreamEvent::Ping))))); - - let cbs = callbacks.lock().unwrap(); - assert_eq!(cbs.len(), 1); - assert!(cbs[0].elapsed > Duration::ZERO); - } - - #[tokio::test] - async fn timeout_returns_error_sync() { - // Verify that poll_next returns a timeout error when the timeout - // has elapsed. We construct the stream, manually advance time by - // setting start into the past, then poll. - let callbacks: std::sync::Arc>> = - std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); - let config = HeartbeatConfig::new( - Duration::from_millis(10), - Duration::from_millis(1), - Box::new(move |data: HeartbeatData| { - callbacks.lock().unwrap().push(data); - }), - ); - - // VecStream that returns Pending on first poll (simulates waiting). - let inner = VecStream { items: vec![] }; - let mut stream = HeartbeatStream::new(inner, config); - - // Manually set start into the past so timeout has elapsed. - 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(); - let mut cx = Context::from_waker(&waker); - let result = Pin::new(&mut stream).poll_next(&mut cx); - - assert!( - matches!(result, Poll::Ready(Some(Err(ApiError::Api(msg)))) if msg.contains("timeout")) - ); - } -} diff --git a/src/stream/rate_limit.rs b/src/stream/rate_limit.rs index 831751e..2f4465a 100644 --- a/src/stream/rate_limit.rs +++ b/src/stream/rate_limit.rs @@ -26,9 +26,6 @@ use std::sync::Arc; use std::sync::Mutex; use std::time::{Duration, Instant}; -/// Default upper bound on how long a single turn waits for a token (30s). -const DEFAULT_MAX_WAIT: Duration = Duration::from_secs(30); - /// A continuous-fill token bucket: the standard client-side rate-limiting algorithm. /// /// - **Capacity** = max burst size. @@ -41,16 +38,41 @@ const DEFAULT_MAX_WAIT: Duration = Duration::from_secs(30); /// `Mutex` critical section (a few float ops + one `Instant::now()`). #[derive(Debug)] pub struct TokenBucket { + /// Maximum burst size, in tokens. + /// + /// The bucket starts full at this value and never refills above it, + /// so a long idle period always restores the full burst budget. capacity: f64, + + /// Continuous refill rate in tokens per second. + /// + /// Computed once at construction as `capacity / 60.0`, so a + /// 60-token bucket refills at one token per second. Float-valued + /// to preserve sub-second refill precision on each lazy update. refill_per_sec: f64, + + /// Mutable bucket state behind a short-lived mutex. + /// + /// Held only for the few float ops and one `Instant::now()` of each + /// [`take`](Self::take) / [`available`](Self::available) call, so + /// contention is negligible under normal turn rates. state: Mutex, } #[derive(Debug, Clone, Copy)] struct BucketState { - /// Current token count (float for sub-second refill precision). + /// Current token count. + /// + /// Float-valued for sub-second refill precision: a partial token + /// accrues between calls and is banked on the next one. Capped at + /// the bucket's `capacity` by [`elapsed_refill`]. tokens: f64, - /// Last instant at which `tokens` was updated. + + /// Instant at which `tokens` was last reconciled with elapsed time. + /// + /// Each call to [`elapsed_refill`] advances this to the current + /// `at`, so the next caller only accounts for the gap since the + /// previous call rather than recomputing from construction. last_refill: Instant, } @@ -141,7 +163,11 @@ impl TokenBucket { elapsed_refill(&mut state, at, self.capacity, self.refill_per_sec) } - /// Tokens available right now. See [`available_at`](Self::available_at). + /// Tokens available right now (after a lazy refill). Non-consuming. + /// + /// Thin wrapper around [`available_at`](Self::available_at) that + /// plugs in `Instant::now()`. Useful for observability — reporting + /// the current budget without consuming a token. #[must_use] pub fn available(&self) -> f64 { self.available_at(Instant::now()) @@ -175,14 +201,26 @@ fn elapsed_refill(state: &mut BucketState, at: Instant, capacity: f64, refill_pe /// This is the type [`StreamHandler`](super::handler::StreamHandler) holds. #[derive(Debug)] pub struct RateLimiter { + /// Per-provider token buckets, keyed by base URL. + /// + /// Lazily populated: the first [`acquire`](Self::acquire) for a + /// given `base_url` creates its bucket at the configured + /// `requests_per_minute`; later acquires for the same URL share it + /// via [`Arc`]. Guarded by a `Mutex` so concurrent turns acquire + /// safely. buckets: Mutex>>, + + /// Configured request budget per provider, in requests per minute. + /// + /// `0` disables the limiter entirely — [`acquire`](Self::acquire) + /// short-circuits to `Ok(())` and never allocates a bucket. Stored + /// as `u32` because it doubles as each bucket's capacity. requests_per_minute: u32, - max_wait: Duration, } impl RateLimiter { /// Build a limiter allowing `requests_per_minute` requests per minute per - /// provider. `0` disables. The `max_wait` ceiling defaults to 30s. + /// provider. `0` disables. /// /// # Example /// @@ -197,19 +235,9 @@ impl RateLimiter { Self { buckets: Mutex::new(HashMap::new()), requests_per_minute, - max_wait: DEFAULT_MAX_WAIT, } } - /// Override the max-wait ceiling: the longest a single turn will block - /// waiting for a token before proceeding anyway ("better to risk a 429 than - /// hang the agent"). - #[must_use] - pub fn with_max_wait(mut self, max_wait: Duration) -> Self { - self.max_wait = max_wait; - self - } - /// Get (or lazily create) the bucket for `base_url`, then take one token. /// /// When the limiter is disabled (`requests_per_minute == 0`) this returns @@ -236,17 +264,15 @@ impl RateLimiter { bucket.take() } - /// Whether this limiter is active (`requests_per_minute > 0`). + /// Whether this limiter is active. + /// + /// Returns `true` when `requests_per_minute > 0`. When `false`, + /// [`acquire`](Self::acquire) short-circuits to `Ok(())` and never + /// allocates a bucket, so the limiter is effectively a no-op. #[must_use] pub fn is_enabled(&self) -> bool { self.requests_per_minute > 0 } - - /// The configured max-wait ceiling. - #[must_use] - pub fn max_wait(&self) -> Duration { - self.max_wait - } } #[cfg(test)] diff --git a/src/structured.rs b/src/structured.rs index eed7b4d..17c9dbc 100644 --- a/src/structured.rs +++ b/src/structured.rs @@ -340,6 +340,7 @@ pub enum StructuredError { /// /// Returns `None` if the content cannot be parsed as JSON (even after the /// lenient rescue). +#[cfg(any(feature = "openai", feature = "gemini"))] pub(crate) fn parse_json_lenient(text: &str) -> Option { if let Ok(v) = serde_json::from_str(text) { return Some(v); @@ -354,6 +355,7 @@ pub(crate) fn parse_json_lenient(text: &str) -> Option { /// the substring up to the matching close. String-aware: braces/brackets /// inside JSON string literals (`"..."`) do not affect depth, and `\"` /// escapes are honored. +#[cfg(any(feature = "openai", feature = "gemini"))] pub(crate) fn extract_json_substring(text: &str) -> Option { let bytes = text.as_bytes(); let mut start = None; @@ -451,6 +453,12 @@ pub(crate) fn extract_json_substring(text: &str) -> Option { /// assert_eq!(tightened["properties"]["filter"]["additionalProperties"], false); /// assert_eq!(tightened["properties"]["filter"]["required"], json!(["lang"])); /// ``` +#[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" +))] pub(crate) fn tighten_json_schema(schema: &serde_json::Value) -> serde_json::Value { let mut out = schema.clone(); tighten_in_place(&mut out); @@ -489,6 +497,12 @@ pub(crate) fn tighten_json_schema(schema: &serde_json::Value) -> serde_json::Val /// are returned unchanged. /// /// Idempotent: re-running on an already-tight schema leaves it unchanged. +#[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" +))] fn tighten_in_place(schema: &mut serde_json::Value) { let Some(obj) = schema.as_object_mut() else { return; diff --git a/src/testing.rs b/src/testing.rs index 851e172..b658b75 100644 --- a/src/testing.rs +++ b/src/testing.rs @@ -40,9 +40,7 @@ //! - [`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 [`LoopConfig`] with sensible defaults. -//! - [`test_config_with_id`] — Create a test [`LoopConfig`] with a specific -//! session ID. +//! - [`test_config`] — Create a test [`SessionConfig`] with sensible defaults. //! //! # Quick Start //! @@ -91,7 +89,7 @@ use crate::api::ApiClient; use crate::api::error::ApiError; -use crate::config::LoopConfig; +use crate::config::SessionConfig; use crate::message::{Message, MessagePart, Role}; use crate::stream::{ DeltaPart, IndexedDelta, MessageDelta, MessageDeltaPayload, MessageMetadata, MessageStart, @@ -105,11 +103,6 @@ use std::pin::Pin; use std::sync::Arc; use std::sync::Mutex; -use uuid::Uuid; - -// ================================================== -// MockApiClient -// ================================================== /// A mock [`ApiClient`] that returns preconfigured streaming responses. /// @@ -305,10 +298,6 @@ pub struct MockToolCall { pub input: Value, } -// =================================================== -// MockApiClient — construction & builder methods -// =================================================== - /// Construction and builder methods for [`MockApiClient`]. /// /// The builder pattern lets you configure mock responses fluently. @@ -547,10 +536,6 @@ impl MockApiClient { } } -// =================================================== -// MockApiClient — ApiClient trait implementation -// =================================================== - /// Trait implementation that turns canned [`MockResponse`] values into /// real [`StreamEvent`] sequences and JSON payloads. /// @@ -598,6 +583,14 @@ impl ApiClient for MockApiClient { crate::error::recover_guard(self.model_name.lock()).clone() } + /// Hot-swap the mock's model name at runtime. + /// + /// Unlike the trait default (which returns `false`), the mock stores + /// its model behind a mutex and updates it in place so tests can + /// exercise [`BareLoop::switch_model`](crate::engine::BareLoop::switch_model) + /// and verify the new name is observed by subsequent [`model`](Self::model) + /// calls. Returns `false` (no-op) when `model` is empty or whitespace, + /// matching the trait contract that an empty model is not a valid switch. fn set_model(&self, model: &str) -> bool { if model.trim().is_empty() { return false; @@ -745,10 +738,6 @@ impl ApiClient for MockApiClient { } } -// ================================================== -// MockTool -// ================================================== - /// A mock [`Tool`] that returns a fixed result or error when called. /// /// Useful for testing tool dispatch, registries, and agent-loop tool @@ -877,10 +866,6 @@ pub struct MockTool { system_prompt: Option, } -// =================================================== -// MockTool — construction & builder methods -// =================================================== - /// Construction and builder methods for [`MockTool`]. /// /// The builder pattern lets you configure the mock's behaviour fluently. @@ -1104,10 +1089,6 @@ impl MockTool { } } -// =================================================== -// MockTool — Tool trait implementation -// =================================================== - /// Trait implementation that returns canned tool metadata and results. /// /// Every method delegates to the fields configured via the builder @@ -1243,10 +1224,6 @@ impl Tool for MockTool { } } -// ================================================== -// Fixture Factories -// ================================================== - /// Create a test user [`Message`] with the given text. /// /// Shorthand for `Message::user(text)`. Useful when building message @@ -1340,26 +1317,15 @@ pub fn test_tool_use_message(calls: &[(&str, &str, Value)]) -> Message { Message::new(Role::Assistant, blocks) } -/// Create a test [`LoopConfig`] with sensible defaults. +/// Create a test [`SessionConfig`] with sensible defaults. /// /// The returned config has: /// -/// - A random [`Uuid`] session ID. -/// - `max_turns` set to `10`. /// - `system_prompt` set to `"You are a test assistant."`. /// - All other fields at their [`Default`] values. /// /// Useful as a starting point when you don't care about specific -/// configuration values. If you need a deterministic session ID for -/// log correlation or assertion, use [`test_config_with_id`] instead. -/// -/// # Fields -/// -/// ```text -/// session_id → Uuid::new_v4() -/// max_turns → 10 -/// system_prompt → Some("You are a test assistant.") -/// ``` +/// configuration values. /// /// # Example /// @@ -1367,57 +1333,16 @@ pub fn test_tool_use_message(calls: &[(&str, &str, Value)]) -> Message { /// use loopctl::testing::test_config; /// /// let config = test_config(); -/// assert_eq!(config.max_turns, 10); -/// ``` -#[must_use] -pub fn test_config() -> LoopConfig { - LoopConfig { - session_id: Uuid::new_v4(), - max_turns: 10, - system_prompt: Some("You are a test assistant.".to_string()), - ..Default::default() - } -} - -/// 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 -/// log correlation, session persistence, or that a session resumes -/// with the correct identity after a restart. -/// -/// # Fields -/// -/// ```text -/// session_id → -/// max_turns → 10 -/// system_prompt → Some("You are a test assistant.") -/// ``` -/// -/// # Example -/// -/// ```rust -/// use loopctl::testing::test_config_with_id; -/// use uuid::Uuid; -/// -/// let id = Uuid::new_v4(); -/// let config = test_config_with_id(id); -/// assert_eq!(config.session_id, id); +/// assert_eq!(config.system_prompt.as_deref(), Some("You are a test assistant.")); /// ``` #[must_use] -pub fn test_config_with_id(id: Uuid) -> LoopConfig { - LoopConfig { - session_id: id, - max_turns: 10, +pub fn test_config() -> SessionConfig { + SessionConfig { system_prompt: Some("You are a test assistant.".to_string()), - ..Default::default() + ..SessionConfig::default() } } -// ================================================== -// Tests -// ================================================== - /// Unit tests for the testing module itself. /// /// These tests verify that the mock implementations and fixture @@ -1644,8 +1569,10 @@ mod tests { #[test] fn test_fixture_test_config() { let config = test_config(); - assert_eq!(config.max_turns, 10); - assert!(config.system_prompt.is_some()); + assert_eq!( + config.system_prompt.as_deref(), + Some("You are a test assistant.") + ); } #[test] diff --git a/src/tool.rs b/src/tool.rs index f9fb13c..3e111a2 100644 --- a/src/tool.rs +++ b/src/tool.rs @@ -82,10 +82,6 @@ pub mod registry; pub use permission::PermissionCheck; pub use registry::{FnTool, ToolRegistry}; -// =================================================== -// ToolSchema -// =================================================== - /// Schema descriptor for a tool, used for LLM API tool definitions. /// /// When the agent loop sends a list of available tools to the LLM, each @@ -139,10 +135,6 @@ pub struct ToolSchema { pub input_schema: Value, } -// =================================================== -// ToolOutput -// =================================================== - /// Advisory rendering hint attached to a [`ToolOutput`]. /// /// Tells presentation layers (TUI, headless console, loggers) *how* the tool @@ -562,50 +554,17 @@ impl ToolOutput { } impl From for ToolOutput { - /// Convert a [`String`] into a successful text [`ToolOutput`]. - /// - /// Enables `let result: ToolOutput = s.into()` where `s: String`. - /// Equivalent to calling [`ToolOutput::text`]. Sets - /// [`is_error`](ToolOutput::is_error) to `false`. - /// - /// # Example - /// - /// ```rust - /// use loopctl::tool::{ToolOutput, ToolError, ToolSchema, ToolContext, PermissionCheck, ToolRegistry}; - /// - /// let s = String::from("hello"); - /// let result: ToolOutput = s.into(); - /// assert_eq!(result.text_content(), "hello"); - /// ``` fn from(s: String) -> Self { Self::text(s) } } impl From<&str> for ToolOutput { - /// Convert a `&str` into a successful text [`ToolOutput`]. - /// - /// Enables `let result: ToolOutput = "ok".into()`. Equivalent to - /// calling [`ToolOutput::text`]. Sets - /// [`is_error`](ToolOutput::is_error) to `false`. - /// - /// # Example - /// - /// ```rust - /// use loopctl::tool::{ToolOutput, ToolError, ToolSchema, ToolContext, PermissionCheck, ToolRegistry}; - /// - /// let result: ToolOutput = "ok".into(); - /// assert_eq!(result.text_content(), "ok"); - /// ``` fn from(s: &str) -> Self { Self::text(s) } } -// =================================================== -// ToolDispatchResult -// =================================================== - /// The outcome of a single tool invocation. /// /// Produced after the framework dispatches a tool call and collects @@ -837,10 +796,6 @@ impl From for ToolDispatchResult { } } -// =================================================== -// ToolError -// =================================================== - /// Error type for tool invocations. /// /// Covers the full range of failure modes a tool can encounter — from @@ -1038,10 +993,6 @@ impl ToolError { } } -// =================================================== -// ToolContext -// =================================================== - /// Session-level context provided to every tool invocation. /// /// The agent loop constructs a [`ToolContext`] at session start and passes @@ -1162,18 +1113,6 @@ impl ToolContext { } impl fmt::Debug for ToolContext { - /// Format the context for debug output, summarizing extensions. - /// - /// Produces a human-readable struct dump. The `extensions` field is - /// rendered as a count (e.g., `"3 entries"`) rather than printing every - /// entry, because extensions can be arbitrarily large and their types - /// may not implement [`Debug`](std::fmt::Debug). - /// - /// # Example output - /// - /// ```text - /// ToolContext { cwd: "/tmp/ws", session_id: 0123-4567-..., temp_dir: "/tmp", is_non_interactive: false, user_context: {}, extensions: 2 entries } - /// ``` fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("ToolContext") .field("cwd", &self.cwd) @@ -1186,30 +1125,30 @@ impl fmt::Debug for ToolContext { } } +/// Produce a context with sensible defaults. +/// +/// Creates a ready-to-use [`ToolContext`] suitable for most agent +/// sessions. The defaults are: +/// +/// | Field | Default | +/// |---------------------|--------------------------------------| +/// | `cwd` | `"."` (process current directory) | +/// | `session_id` | new UUID v4 | +/// | `temp_dir` | [`std::env::temp_dir`] | +/// | `is_non_interactive`| `false` | +/// | `user_context` | empty [`HashMap`] | +/// | `extensions` | empty [`HashMap`] | +/// +/// # Example +/// +/// ```rust +/// use loopctl::tool::{ToolOutput, ToolError, ToolSchema, ToolContext, PermissionCheck, ToolRegistry}; +/// +/// let ctx = ToolContext::default(); +/// assert_eq!(ctx.cwd, "."); +/// assert!(!ctx.is_non_interactive); +/// ``` impl Default for ToolContext { - /// Produce a context with sensible defaults. - /// - /// Creates a ready-to-use [`ToolContext`] suitable for most agent - /// sessions. The defaults are: - /// - /// | Field | Default | - /// |---------------------|--------------------------------------| - /// | `cwd` | `"."` (process current directory) | - /// | `session_id` | new UUID v4 | - /// | `temp_dir` | [`std::env::temp_dir`] | - /// | `is_non_interactive`| `false` | - /// | `user_context` | empty [`HashMap`] | - /// | `extensions` | empty [`HashMap`] | - /// - /// # Example - /// - /// ```rust - /// use loopctl::tool::{ToolOutput, ToolError, ToolSchema, ToolContext, PermissionCheck, ToolRegistry}; - /// - /// let ctx = ToolContext::default(); - /// assert_eq!(ctx.cwd, "."); - /// assert!(!ctx.is_non_interactive); - /// ``` fn default() -> Self { Self { cwd: ".".to_string(), @@ -1221,9 +1160,6 @@ impl Default for ToolContext { } } } -// =================================================== -// Tool trait -// =================================================== /// The trait that all agent tools must implement. /// @@ -1520,15 +1456,13 @@ pub trait Tool: Send + Sync { } } -// =================================================== -// Tests -// =================================================== - #[cfg(test)] mod tests { use super::*; #[cfg(feature = "testing")] - use crate::engine::loop_core::Loop; + use crate::engine::RunConfig; + #[cfg(feature = "testing")] + use crate::engine::core::Loop; use serde_json::json; struct EchoTool; @@ -2105,7 +2039,7 @@ mod tests { let mut agent = crate::engine::BareLoop::new( std::sync::Arc::new(client), registry, - crate::config::LoopConfig::default(), + crate::config::SessionConfig::default(), ); let capture = std::sync::Arc::new(PostCapture::default()); agent.register_observer(capture.clone()); @@ -2120,7 +2054,10 @@ mod tests { "@@ -1 +1 @@\n-old\n+new", Some(DisplayHint::Diff), ); - agent.run("edit the file").await.unwrap(); + agent + .run("edit the file", &RunConfig::default()) + .await + .unwrap(); let posts = crate::error::recover_guard(capture.posts.lock()).clone(); assert_eq!(posts.len(), 1, "exactly one tool call this turn"); @@ -2136,7 +2073,7 @@ mod tests { #[tokio::test] async fn no_hint_tool_yields_none_at_observer() { let (mut agent, capture) = hinted_dispatch_setup("plain", "just text", None); - agent.run("go").await.unwrap(); + agent.run("go", &RunConfig::default()).await.unwrap(); let posts = crate::error::recover_guard(capture.posts.lock()).clone(); assert_eq!(posts.len(), 1); @@ -2151,7 +2088,7 @@ mod tests { async fn suppress_hint_keeps_full_payload_into_conversation() { let (mut agent, capture) = hinted_dispatch_setup("reader", "the quick brown fox", Some(DisplayHint::Suppress)); - agent.run("read it").await.unwrap(); + agent.run("read it", &RunConfig::default()).await.unwrap(); let posts = crate::error::recover_guard(capture.posts.lock()).clone(); assert_eq!(posts[0].display_hint, Some(DisplayHint::Suppress)); diff --git a/src/tool/health.rs b/src/tool/health.rs index ab35418..a0c7980 100644 --- a/src/tool/health.rs +++ b/src/tool/health.rs @@ -63,10 +63,6 @@ use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; -// =================================================== -// HealthStatus -// =================================================== - /// Classified health status for a tool. /// /// Determined by combining the tool's composite [`ToolStats::health_score`] @@ -83,11 +79,25 @@ use std::time::{Duration, Instant}; /// | Any | Open | [`Unhealthy`](HealthStatus::Unhealthy) | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum HealthStatus { - /// Tool is operating normally (health score ≥ 0.8, breaker closed). + /// Tool is operating normally. + /// + /// Assigned when the composite health score is ≥ 0.8 and the + /// circuit breaker is `Closed`. The router routes calls to this + /// tool without hesitation. Healthy, - /// Tool is experiencing elevated errors or latency (score 0.5–0.8). + + /// Tool is experiencing elevated errors or latency. + /// + /// Assigned when the health score is in the 0.5–0.8 band. The tool + /// is still called (the breaker has not tripped), but the router + /// may prefer a healthy alternative when one is available. Degraded, - /// Tool is failing frequently or circuit breaker is open (score < 0.5). + + /// Tool is failing frequently or its circuit breaker is open. + /// + /// Assigned when the health score drops below 0.5, or whenever the + /// breaker is `Open` regardless of score. The router treats the + /// tool as unavailable and redirects to a fallback. Unhealthy, } @@ -101,10 +111,6 @@ impl fmt::Display for HealthStatus { } } -// =================================================== -// ToolStats — Lock-free Atomic Counters -// =================================================== - /// Fixed-point scale for the EWMA success rate. /// /// Stored as `u64` where `1.0` = `1_000_000`. This avoids floating-point @@ -150,11 +156,46 @@ const EWMA_SCALE: u64 = 1_000_000; /// assert!(stats.health_score() > 0.0); /// ``` pub struct ToolStats { + /// Total number of calls recorded (successes + failures). + /// + /// Incremented once per `record_success` / `record_failure` call; + /// the denominator for [`success_rate`](Self::success_rate). total_calls: AtomicU64, + + /// Number of calls that completed successfully. + /// + /// Incremented by [`record_success`](Self::record_success); paired + /// with `total_calls` to compute the all-time success rate. success_count: AtomicU64, + + /// Number of calls that failed. + /// + /// Incremented by [`record_failure`](Self::record_failure). Kept + /// separately from `success_count` so both rates are available + /// without re-deriving from the total. failure_count: AtomicU64, + + /// Sum of per-call durations in nanoseconds, across all calls. + /// + /// Accumulated by both record methods; divided by `total_calls` in + /// [`avg_duration`](Self::avg_duration). Saturates at `u64::MAX` on + /// overflow rather than wrapping. total_duration_ns: AtomicU64, + + /// High-water mark for the longest single-call duration, in + /// nanoseconds. + /// + /// Updated via `fetch_max` so it only ever grows; exposed by + /// [`max_duration`](Self::max_duration). Useful for spotting + /// tail-latency outliers. max_duration_ns: AtomicU64, + + /// Exponentially-weighted moving average of success, in fixed-point. + /// + /// Stored as `u64` on a `1.0 = EWMA_SCALE` scale so it can be + /// updated atomically without floating-point atomics. Decays with a + /// 0.7 factor on every call, making the health score responsive to + /// recent degradation. ewma_success: AtomicU64, } @@ -216,18 +257,28 @@ impl ToolStats { } /// Total number of calls recorded (successes + failures). + /// + /// Lock-free relaxed load of the counter incremented on every + /// record call. Use as the denominator when computing custom rates. #[must_use] pub fn total_calls(&self) -> u64 { self.total_calls.load(Ordering::Relaxed) } - /// Number of successful calls. + /// Number of calls that completed successfully. + /// + /// Lock-free relaxed load. Pair with [`total_calls`](Self::total_calls) + /// to derive the all-time success rate, or use + /// [`success_rate`](Self::success_rate) directly. #[must_use] pub fn success_count(&self) -> u64 { self.success_count.load(Ordering::Relaxed) } - /// Number of failed calls. + /// Number of calls that failed. + /// + /// Lock-free relaxed load. Pair with [`total_calls`](Self::total_calls) + /// to derive the all-time failure rate. #[must_use] pub fn failure_count(&self) -> u64 { self.failure_count.load(Ordering::Relaxed) @@ -314,31 +365,36 @@ fn update_ewma(prev: u64, is_success: bool) -> u64 { next.min(EWMA_SCALE) } -// =================================================== -// CircuitState + ToolCircuitBreaker -// =================================================== - /// Circuit breaker state. /// -/// The breaker follows the standard three-state pattern: -/// -/// ```text -/// Closed ──(threshold failures)──► Open -/// ▲ │ -/// │ │ (recovery_duration elapsed) -/// │ ▼ -/// └──(probe succeeds)──────── HalfOpen -/// │ -/// (probe fails) ────►│──► Open -/// ``` +/// The breaker follows the standard three-state pattern. Starting from +/// `Closed`, `failure_threshold` consecutive failures transition it to +/// `Open`, where requests are blocked. Once `recovery_duration` +/// elapses the next `allow_request` moves it to `HalfOpen` and lets a +/// single probe call through. A successful probe closes the breaker; a +/// failed probe reopens it. #[derive(Debug, Clone, Copy, PartialEq, Eq)] #[repr(u32)] enum CircuitState { /// Normal operation — requests are allowed. + /// + /// The resting state. Consecutive failures are counted toward + /// `failure_threshold`; reaching it transitions to [`Open`](Self::Open). Closed = 0, + /// Too many failures — requests are blocked. + /// + /// Entered after `failure_threshold` consecutive failures (or when + /// a [`HalfOpen`](Self::HalfOpen) probe fails). Requests are + /// refused until `recovery_duration` elapses, after which the next + /// `allow_request` transitions to `HalfOpen`. Open = 1, + /// Recovery probe — one request is allowed to test recovery. + /// + /// A single probe call runs; a success closes the breaker, a + /// failure reopens it. Additional probes are refused to avoid a + /// thundering herd. HalfOpen = 2, } @@ -372,6 +428,7 @@ pub struct CircuitBreakerConfig { /// /// Defaults to 3. pub failure_threshold: u64, + /// How long to wait in the `Open` state before transitioning to `HalfOpen`. /// /// Defaults to 30 seconds. @@ -419,10 +476,37 @@ impl Default for CircuitBreakerConfig { /// assert!(!breaker.allow_request()); /// ``` pub struct ToolCircuitBreaker { + /// Current circuit-breaker state, encoded as a [`CircuitState`] in + /// an `AtomicU32`. + /// + /// Read on every `allow_request` check and mutated only on state + /// transitions; atomic so the hot path is lock-free. state: AtomicU32, + + /// Consecutive failures recorded since the last success. + /// + /// Reset to zero by [`record_success`](Self::record_success); + /// reaching `failure_threshold` while in `Closed` transitions the + /// breaker to `Open`. consecutive_failures: AtomicU64, + + /// Number of consecutive failures that trips the breaker. + /// + /// Set at construction and immutable thereafter; compared against + /// `consecutive_failures` on each `record_failure`. failure_threshold: u64, + + /// How long to remain `Open` before allowing a `HalfOpen` probe. + /// + /// Set at construction and immutable thereafter; compared against + /// the elapsed time since `last_failure_time` in `allow_request`. recovery_duration: Duration, + + /// Instant of the most recent failure, or `None` if none yet. + /// + /// Guarded by a `Mutex` because `Instant` is not atomic and this is + /// only written on the failure cold path; the read-only + /// `allow_request` check acquires the lock briefly. last_failure_time: Mutex>, } @@ -443,6 +527,10 @@ impl ToolCircuitBreaker { } /// Create a circuit breaker from a [`CircuitBreakerConfig`]. + /// + /// Convenience constructor that unpacks the threshold and recovery + /// duration from a config struct, delegating to [`new`](Self::new). + /// Useful when many breakers share a single config. #[must_use] pub fn from_config(config: &CircuitBreakerConfig) -> Self { Self::new(config.recovery_duration, config.failure_threshold) @@ -523,9 +611,11 @@ impl ToolCircuitBreaker { } } - /// Current state of the breaker as a string. + /// Current state of the breaker as a human-readable string. /// - /// Returns `"closed"`, `"open"`, or `"half-open"`. + /// Returns `"closed"`, `"open"`, or `"half-open"`. Intended for + /// logs and metrics where a string label is preferable to the + /// numeric encoding. #[must_use] pub fn state_label(&self) -> &'static str { match self.state.load(Ordering::Acquire).into() { @@ -536,12 +626,19 @@ impl ToolCircuitBreaker { } /// Number of consecutive failures recorded since the last success. + /// + /// Lock-free relaxed load of the counter reset to zero on every + /// success; reaching `failure_threshold` trips the breaker. #[must_use] pub fn consecutive_failures(&self) -> u64 { self.consecutive_failures.load(Ordering::Relaxed) } /// Whether the breaker is currently in the Closed (healthy) state. + /// + /// `true` when requests are allowed unconditionally. Note a breaker + /// that just closed may still have a low health score, so callers + /// that want the composite picture should consult the registry. #[must_use] pub fn is_closed(&self) -> bool { matches!( @@ -551,6 +648,10 @@ impl ToolCircuitBreaker { } /// Whether the breaker is currently in the Open (blocking) state. + /// + /// `true` when requests are refused outright (subject to the + /// recovery-duration transition handled inside + /// [`allow_request`](Self::allow_request)). #[must_use] pub fn is_open(&self) -> bool { matches!( @@ -559,7 +660,11 @@ impl ToolCircuitBreaker { ) } - /// Whether the breaker is currently in the `HalfOpen` (probing) state. + /// Whether the breaker is currently in the `HalfOpen` (probing) + /// state. + /// + /// `true` when a single probe call is in flight and additional + /// probes are refused to avoid a thundering herd. #[must_use] pub fn is_half_open(&self) -> bool { matches!( @@ -569,10 +674,6 @@ impl ToolCircuitBreaker { } } -// =================================================== -// ToolHealthRegistry -// =================================================== - /// Global health registry for all tools. /// /// A concrete struct (not a trait) because every agent uses the same health @@ -617,8 +718,26 @@ impl ToolCircuitBreaker { /// assert!(*score > 0.9); /// ``` pub struct ToolHealthRegistry { + /// Per-tool statistics, keyed by tool name. + /// + /// Lazily populated on first sighting of a tool; entries hold + /// `Arc` so callers can read counters without holding the + /// lock. The `Mutex` is only acquired on the cold path of inserting + /// a previously-unseen tool. stats: Mutex>>, + + /// Per-tool circuit breakers, keyed by tool name. + /// + /// Lazily populated alongside `stats`; each entry is configured from + /// `breaker_config` at insertion time. Same cold-path locking + /// strategy as `stats`. breakers: Mutex>>, + + /// Configuration applied to every newly-created circuit breaker. + /// + /// Set at construction (default 3 failures / 30 s recovery) and + /// cloned into each breaker on first sight of a tool; changing it + /// after the fact does not retroactively update existing breakers. breaker_config: CircuitBreakerConfig, } @@ -641,7 +760,11 @@ impl ToolHealthRegistry { } } - /// Set custom circuit-breaker configuration. + /// Set custom circuit-breaker configuration (builder style). + /// + /// Overrides the default (3 failures / 30 s recovery). Applied only + /// to breakers created *after* this call — existing per-tool + /// breakers keep their original thresholds. #[must_use] pub fn with_config(mut self, config: CircuitBreakerConfig) -> Self { self.breaker_config = config; @@ -769,16 +892,16 @@ impl ToolHealthRegistry { } /// Number of distinct tools currently tracked. + /// + /// Lock-free in the sense that it briefly acquires the stats map's + /// `Mutex` to read its length. Returns the count of tools seen at + /// least once via `record_*` / `get_stats` / `get_circuit_breaker`. #[must_use] pub fn tool_count(&self) -> usize { crate::error::recover_guard(self.stats.lock()).len() } } -// =================================================== -// HealthRouterMiddleware -// =================================================== - /// A mapping from primary tool names to alternative tool names. /// /// When the primary tool is [`Unhealthy`](HealthStatus::Unhealthy) or @@ -802,11 +925,20 @@ impl ToolHealthRegistry { /// ``` #[derive(Debug, Clone, Default)] pub struct HealthRouter { + /// Primary tool name → ordered list of fallback tool names. + /// + /// Populated via [`HealthRouterBuilder`]; the router consults this + /// map (in order) when the primary tool is unavailable, returning + /// the first healthy alternative. fallbacks: HashMap>, } impl HealthRouter { /// Create an empty router with no fallback mappings. + /// + /// Equivalent to [`HealthRouter::default`]. An empty router always + /// resolves to the primary tool name since no fallbacks are + /// configured. #[must_use] pub fn new() -> Self { Self { @@ -817,7 +949,8 @@ impl HealthRouter { /// Get the ordered list of fallback tool names for a primary tool. /// /// Returns an empty slice if no fallbacks are configured for the - /// given tool name. + /// given tool name. The order matches the order supplied to + /// [`HealthRouterBuilder::add_fallback`]. #[must_use] pub fn fallbacks_for(&self, tool_name: &str) -> &[String] { self.fallbacks.get(tool_name).map_or(&[], Vec::as_slice) @@ -864,11 +997,18 @@ impl HealthRouter { /// ``` #[derive(Debug, Clone, Default)] pub struct HealthRouterBuilder { + /// In-progress fallback map, mutated by each `add_fallback` call. + /// + /// Consumed by [`build`](Self::build) to produce the immutable + /// [`HealthRouter`]. fallbacks: HashMap>, } impl HealthRouterBuilder { /// Create an empty builder. + /// + /// Equivalent to [`HealthRouterBuilder::default`]; every primary + /// starts with no fallbacks until `add_fallback` is called. #[must_use] pub fn new() -> Self { Self { @@ -878,7 +1018,9 @@ impl HealthRouterBuilder { /// Register a list of fallback tools for a primary tool name. /// - /// Fallbacks are tried in the order provided. + /// Replaces any previously-registered fallbacks for `primary`. + /// Fallbacks are tried in the order provided when + /// [`HealthRouter::resolve_tool`] walks the list. #[must_use] pub fn add_fallback(mut self, primary: &str, alternatives: Vec) -> Self { self.fallbacks.insert(primary.to_string(), alternatives); @@ -886,6 +1028,10 @@ impl HealthRouterBuilder { } /// Build the [`HealthRouter`]. + /// + /// Consumes the builder and freezes the fallback map into an + /// immutable router. The returned router is cheap to clone for + /// sharing across dispatch paths. #[must_use] pub fn build(self) -> HealthRouter { HealthRouter { @@ -894,10 +1040,6 @@ impl HealthRouterBuilder { } } -// =================================================== -// Tests -// =================================================== - #[cfg(test)] mod tests { use super::*; diff --git a/src/tool/permission.rs b/src/tool/permission.rs index 4d93a8b..cf2ebc5 100644 --- a/src/tool/permission.rs +++ b/src/tool/permission.rs @@ -6,10 +6,6 @@ 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 @@ -38,7 +34,13 @@ pub enum PermissionCheck { /// [`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. + /// Explanation forwarded to the LLM as part of the error. + /// + /// Surfaced inside [`ToolError::Permission`](super::ToolError::Permission) + /// so the model can see *why* its call was rejected and adjust + /// its next action. Keep it concrete and actionable — for + /// example `"shell execution is disabled"` rather than just + /// `"denied"`. reason: String, }, @@ -48,7 +50,13 @@ pub enum PermissionCheck { /// 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 text to present to the user for approval. + /// + /// Should clearly describe the action and its potential side + /// effects so the user can make an informed decision — for + /// example `"Allow write to /etc/config.yaml?"`. In + /// non-interactive sessions the loop treats an `Ask` as a + /// `Deny`, so this string is primarily for interactive hosts. prompt: String, }, @@ -58,7 +66,14 @@ pub enum PermissionCheck { /// `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). + /// Rewritten input to pass to [`Tool::call`](super::Tool::call) + /// in place of the original. + /// + /// Must conform to the tool's + /// [`ToolSchema::input_schema`](super::ToolSchema::input_schema) + /// — the loop does not re-validate it. Use this to sanitise + /// paths, redact secrets, or inject default values before the + /// tool sees the input. modified_input: Value, }, } diff --git a/src/tool/registry.rs b/src/tool/registry.rs index 46d5927..a9b877e 100644 --- a/src/tool/registry.rs +++ b/src/tool/registry.rs @@ -12,10 +12,6 @@ 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 @@ -26,9 +22,11 @@ use super::{Tool, ToolContext, ToolError, ToolOutput, ToolSchema}; /// /// # 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>`. +/// `ToolRegistry` is `Send + Sync` because [`Tool`] requires `Send + Sync`. +/// The registry is created once during session setup and then shared via +/// `Arc` for read-only access during dispatch, including +/// concurrent reads in parallel tool dispatch. No wrapping in `RwLock` is +/// needed. /// /// # Example /// @@ -45,7 +43,14 @@ use super::{Tool, ToolContext, ToolError, ToolOutput, ToolSchema}; /// let schemas = registry.all_schemas(); /// ``` pub struct ToolRegistry { - /// `Box` keyed by [`Tool::name`]. + /// Registered tools, keyed by their [`Tool::name`]. + /// + /// Each value is a boxed trait object so the registry can hold + /// tools of different concrete types uniformly. Insertion + /// ([`register`](ToolRegistry::register)) overwrites any existing + /// entry under the same name; lookup + /// ([`get`](ToolRegistry::get)) and schema enumeration + /// ([`all_schemas`](ToolRegistry::all_schemas)) iterate this map. tools: HashMap>, } @@ -206,29 +211,25 @@ impl 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()); +/// ``` 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: @@ -530,24 +531,15 @@ impl FnTool { } } -/// [`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(), @@ -556,8 +548,6 @@ impl Tool for FnTool { } } - /// Delegates execution to the stored [`tool_fn`](FnTool::tool_fn) - /// function pointer, forwarding the input and context unchanged. fn call( &self, input: Value, @@ -566,34 +556,23 @@ impl Tool for FnTool { (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)) } - /// Delegates to [`resource_key_fn`](FnTool::resource_key_fn) when set, - /// otherwise returns `None` (no declarable resource). fn resource_key(&self, input: &Value) -> Option { self.resource_key_fn.and_then(|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 21563d7..34cf438 100644 --- a/src/tool/shield.rs +++ b/src/tool/shield.rs @@ -66,10 +66,6 @@ use std::sync::Mutex; use serde_json::Value; -// =================================================== -// RiskLevel -// =================================================== - /// Risk classification produced by shield evaluation. /// /// Each dimension of the shield (single-turn, multi-turn, combination) @@ -144,10 +140,6 @@ impl std::fmt::Display for RiskLevel { } } -// =================================================== -// SafetyAction -// =================================================== - /// The action the shield recommends for a tool call. /// /// Produced as the [`SafetyDecision::action`] field. The middleware is @@ -194,10 +186,6 @@ pub enum SafetyAction { Block, } -// =================================================== -// SafetyDecision -// =================================================== - /// The shield's decision for a single tool invocation. /// /// Produced by [`ToolSafetyShield::evaluate`]. Carries the action the @@ -312,10 +300,6 @@ impl SafetyDecision { } } -// =================================================== -// ShieldContext -// =================================================== - /// Context provided to the shield for evaluation. /// /// Constructed by the middleware before each tool call. Carries what the @@ -368,10 +352,6 @@ pub struct ShieldContext { pub recent_calls: Vec<(String, usize)>, } -// =================================================== -// RiskPattern -// =================================================== - /// A named risk pattern matched against a tool's input. /// /// Patterns are stored per tool name. When a tool call is evaluated, @@ -414,10 +394,6 @@ pub struct RiskPattern { pub pattern: &'static str, } -// =================================================== -// CombinationRule -// =================================================== - /// A rule that scores a dangerous *sequence* of tool calls. /// /// A combination rule triggers when **all** of its [`triggers`](Self::triggers) @@ -466,10 +442,6 @@ pub struct CombinationRule { pub triggers: &'static [(&'static str, Option<&'static str>)], } -// =================================================== -// ToolSafetyShield trait -// =================================================== - /// Trait for evaluating tool call safety. /// /// Implementations assess whether a tool invocation is safe to execute, @@ -531,10 +503,6 @@ pub trait ToolSafetyShield: Send + Sync { fn watched_tools(&self) -> HashSet; } -// =================================================== -// UnixShield -// =================================================== - /// Reference shield implementation with Unix shell pattern matching. /// /// Detects dangerous patterns in tools commonly found in Unix-based @@ -930,10 +898,6 @@ impl ToolSafetyShield for UnixShield { } } -// =================================================== -// UnixShieldBuilder -// =================================================== - /// Builder for constructing a [`UnixShield`] with custom configuration. /// /// # Example @@ -1084,10 +1048,6 @@ impl Default for UnixShieldBuilder { } } -// =================================================== -// NullShield -// =================================================== - /// A no-op [`ToolSafetyShield`] that allows every call and watches no /// tools. /// diff --git a/tests/constrained_decode.rs b/tests/constrained_decode.rs new file mode 100644 index 0000000..b96bfc6 --- /dev/null +++ b/tests/constrained_decode.rs @@ -0,0 +1,111 @@ +//! Live test: tool-call constrained decoding. +//! +//! Run: `LOOPCTL_E2E=1 OLLAMA_MODEL=qwen2.5:7b cargo test --features ollama,grammar --test constrained_decode -- --nocapture` +//! +//! Override iteration count via `LOOPCTL_N=1000`. +//! Re-run only specific steps via `LOOPCTL_ONLY=3,7,12`. + +#![allow( + clippy::pedantic, + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::arithmetic_side_effects, + clippy::indexing_slicing, + clippy::redundant_clone +)] + +mod helpers; + +#[cfg(feature = "ollama")] +use loopctl::api::ApiClient; +#[cfg(feature = "ollama")] +use loopctl::structured::{RequestOptions, ToolConstraint}; + +#[cfg(feature = "ollama")] +const NOOP_INDICES: &[usize] = &[21, 22]; + +#[cfg(feature = "ollama")] +#[tokio::test] +async fn strict_mode_produces_valid_tool_calls() { + if std::env::var("LOOPCTL_E2E").as_deref() != Ok("1") { + eprintln!("skipping (set LOOPCTL_E2E=1)"); + return; + } + + let model = std::env::var("OLLAMA_MODEL").expect("set OLLAMA_MODEL"); + let client = loopctl::provider::ollama(&model).unwrap(); + let tools = helpers::test_tools(); + let prompts = helpers::tool_call_prompts(); + let n: u32 = std::env::var("LOOPCTL_N") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(50); + + let only: Option> = std::env::var("LOOPCTL_ONLY") + .ok() + .map(|s| s.split(',').filter_map(|n| n.trim().parse().ok()).collect()); + + let mut valid: u32 = 0; + let mut total: u32 = 0; + let mut errors: u32 = 0; + + for i in 0..n { + let prompt_idx = (i as usize) % prompts.len(); + + if let Some(ref indices) = only + && !indices.contains(&prompt_idx) + { + continue; + } + + let prompt = prompts[prompt_idx]; + total += 1; + let req = loopctl::api::StreamRequest::new(vec![loopctl::message::Message::user(prompt)]) + .with_tools(Some(tools.clone())); + let opts = RequestOptions::default().with_tool_constraint(ToolConstraint::Strict); + let stream = client.stream_messages_with_options(req, opts); + + let events = match helpers::collect_events(Box::pin(stream)).await { + Some(events) => events, + None => { + errors += 1; + println!( + "[{}/{}] \x1b[31mERROR:\x1b[0m (#{} {prompt})", + i + 1, + n, + prompt_idx + ); + continue; + } + }; + + let has_tool_call = helpers::extract_tool_call(&events) + .is_some_and(|(name, input)| tools.iter().any(|t| t.tool == name) && input.is_object()); + + let is_noop = NOOP_INDICES.contains(&prompt_idx); + let pass = if is_noop { + !has_tool_call + } else { + has_tool_call + }; + + if pass { + valid += 1; + } + + let status = if pass { + "\x1b[32mPASS:\x1b[0m" + } else { + "\x1b[31mFAIL:\x1b[0m" + }; + println!("[{}/{}] {status} (#{} {prompt})", i + 1, n, prompt_idx); + } + + let rate = f64::from(valid) / f64::from(total) * 100.0; + println!("Strict: {valid}/{total} ({rate:.0}%)"); + if errors > 0 { + println!("\x1b[33m{errors} stream errors (not counted in rate)\x1b[0m"); + } + assert!(rate >= 90.0, "got {rate:.0}%"); +} diff --git a/tests/examples_e2e.rs b/tests/examples_e2e.rs new file mode 100644 index 0000000..e05b535 --- /dev/null +++ b/tests/examples_e2e.rs @@ -0,0 +1,46 @@ +//! Live test: basic end-to-end run (Ollama only). +//! +//! Run: `LOOPCTL_E2E=1 OLLAMA_MODEL=qwen2.5:7b cargo test --features ollama --test examples_e2e -- --nocapture` + +#![allow( + clippy::pedantic, + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::arithmetic_side_effects, + clippy::indexing_slicing, + clippy::redundant_clone +)] + +mod helpers; + +#[cfg(feature = "ollama")] +use loopctl::engine::core::Loop; + +#[cfg(feature = "ollama")] +#[tokio::test] +async fn chat_example_works_end_to_end() { + if std::env::var("LOOPCTL_E2E").as_deref() != Ok("1") { + eprintln!("skipping (set LOOPCTL_E2E=1)"); + return; + } + + let model = std::env::var("OLLAMA_MODEL").expect("set OLLAMA_MODEL"); + let client = std::sync::Arc::new(loopctl::provider::ollama(&model).unwrap()); + let mut agent = loopctl::engine::BareLoop::new( + client, + loopctl::tool::ToolRegistry::new(), + loopctl::config::SessionConfig::default(), + ); + let result = agent + .run( + "Say hello in one sentence.", + &loopctl::engine::RunConfig::default(), + ) + .await + .expect("run should succeed"); + + let output = result.output.expect("should have output"); + assert!(!output.is_empty()); + println!("Response: {output}"); +} diff --git a/tests/helpers.rs b/tests/helpers.rs new file mode 100644 index 0000000..0efa2ac --- /dev/null +++ b/tests/helpers.rs @@ -0,0 +1,197 @@ +//! Shared helpers for live integration tests. + +#![allow( + dead_code, + clippy::pedantic, + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::arithmetic_side_effects, + clippy::indexing_slicing, + clippy::redundant_clone +)] + +pub fn test_tools() -> Vec { + vec![ + loopctl::tool::ToolSchema { + tool: "read_file".into(), + description: "Read a file from disk".into(), + input_schema: serde_json::json!({ + "type": "object", + "properties": { + "path": {"type": "string", "description": "Absolute file path"}, + "encoding": {"type": "string", "enum": ["utf-8", "binary"]} + }, + "required": ["path"] + }), + }, + loopctl::tool::ToolSchema { + tool: "search".into(), + description: "Search the web".into(), + input_schema: serde_json::json!({ + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query"}, + "max_results": {"type": "integer", "minimum": 1, "maximum": 50} + }, + "required": ["query"] + }), + }, + loopctl::tool::ToolSchema { + tool: "write_file".into(), + description: "Write content to a file".into(), + input_schema: serde_json::json!({ + "type": "object", + "properties": { + "path": {"type": "string"}, + "content": {"type": "string"}, + "create_dirs": {"type": "boolean"} + }, + "required": ["path", "content"] + }), + }, + loopctl::tool::ToolSchema { + tool: "run_command".into(), + description: "Execute a shell command".into(), + input_schema: serde_json::json!({ + "type": "object", + "properties": { + "command": {"type": "string", "description": "The command to execute"}, + "working_dir": {"type": "string"}, + "timeout_seconds": {"type": "integer", "minimum": 1}, + "env": { + "type": "object", + "description": "Environment variables as key-value pairs", + "additionalProperties": {"type": "string"} + } + }, + "required": ["command"] + }), + }, + loopctl::tool::ToolSchema { + tool: "git_commit".into(), + description: "Create a git commit with staged changes".into(), + input_schema: serde_json::json!({ + "type": "object", + "properties": { + "message": {"type": "string", "description": "Commit message"}, + "amend": {"type": "boolean", "description": "Amend the previous commit"}, + "co_authors": { + "type": "array", + "items": {"type": "string"}, + "description": "Co-author emails" + } + }, + "required": ["message"] + }), + }, + ] +} + +pub fn tool_call_prompts() -> Vec<&'static str> { + vec![ + // read_file + "Read the file /etc/hostname", + "Show me the contents of /var/log/syslog using binary encoding", + "I need to see what is in /home/user/.bashrc", + "Open and display the configuration file at /etc/nginx/nginx.conf", + // search + "Search the web for rust async patterns", + "Find information about best coding agent 2026, limit 10 results", + "Look up how to use tokio channels, max 5 results", + "Can you search for recent papers on small language model tool use?", + // write_file + "Write hello world to /tmp/test.txt", + "Save the text 'print(42)' to /home/user/script.py", + "Write 'test content' to /tmp/another.txt and create parent directories", + "Write a file at /tmp/config.json with content: {\"name\": \"test\", \"value\": 42}", + "Create an empty file at /tmp/placeholder.txt", + // run_command + "Run 'cargo build' in /home/user/project with a 120 second timeout", + "Execute 'npm test' in the current directory", + "Run 'echo $GREETING' with GREETING set to hello world in the environment", + // git_commit + "Commit with message 'fix: update README' and amend", + "Create a commit saying 'feat: add streaming support' and credit co-author jane@example.com", + "Make a git commit with the message 'wip'", + // ambiguous + "I need to see the git log, run git log in /home/user/repo", + "Create a file called notes.txt with the content 'remember to deploy'", + // no-op (should NOT call a tool) + "What is 2 plus 2?", + "Hello, how are you today?", + ] +} + +pub async fn collect_events( + stream: std::pin::Pin< + Box< + dyn futures::Stream< + Item = Result, + > + Send, + >, + >, +) -> Option> { + use futures::StreamExt; + let mut stream = stream; + let mut events = Vec::new(); + while let Some(result) = stream.next().await { + match result { + Ok(ev) => events.push(ev), + Err(e) => { + eprintln!("stream error: {e}"); + return None; + } + } + } + Some(events) +} + +pub fn extract_tool_call( + events: &[loopctl::stream::StreamEvent], +) -> Option<(String, serde_json::Value)> { + let mut name: Option = None; + let mut json_buf = String::new(); + + for ev in events { + match ev { + loopctl::stream::StreamEvent::PartStart(ps) => { + if let Some(loopctl::message::MessagePart::ToolCall { name: n, .. }) = &ps.part { + name = Some(n.clone()); + } + } + loopctl::stream::StreamEvent::IndexedDelta(d) => { + if name.is_some() + && let loopctl::stream::DeltaPart::InputJson { partial_json } = &d.delta + { + json_buf.push_str(partial_json); + } + } + loopctl::stream::StreamEvent::PartStop => { + if let Some(n) = name.take() { + let input = if json_buf.is_empty() { + serde_json::Value::Object(serde_json::Map::new()) + } else { + serde_json::from_str(&json_buf).unwrap_or(serde_json::Value::Null) + }; + json_buf.clear(); + return Some((n, input)); + } + } + _ => {} + } + } + None +} + +pub fn extract_text(events: &[loopctl::stream::StreamEvent]) -> String { + let mut text = String::new(); + for ev in events { + if let loopctl::stream::StreamEvent::IndexedDelta(d) = ev + && let loopctl::stream::DeltaPart::Text { text: delta } = &d.delta + { + text.push_str(delta); + } + } + text +} diff --git a/tests/provider_e2e.rs b/tests/provider_e2e.rs new file mode 100644 index 0000000..6389ce7 --- /dev/null +++ b/tests/provider_e2e.rs @@ -0,0 +1,160 @@ +//! Live test: provider end-to-end smoke test. +//! +//! Works with any provider. Set the API key for the ones you want to test. +//! +//! Run: +//! `set -a; source .env; set +a; LOOPCTL_E2E=1 cargo test --features ollama,openai,anthropic,gemini,grok,deepseek,zai --test provider_e2e -- --nocapture --test-threads=1` + +#![allow( + clippy::pedantic, + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::arithmetic_side_effects, + clippy::indexing_slicing, + clippy::redundant_clone +)] + +use futures::StreamExt; +use loopctl::api::ApiClient; + +const GREEN: &str = "\x1b[32m"; +const RED: &str = "\x1b[31m"; +const CYAN: &str = "\x1b[36m"; +const DIM: &str = "\x1b[2m"; +const RESET: &str = "\x1b[0m"; + +fn extract_text(events: &[loopctl::stream::StreamEvent]) -> String { + let mut text = String::new(); + for ev in events { + if let loopctl::stream::StreamEvent::IndexedDelta(d) = ev + && let loopctl::stream::DeltaPart::Text { text: delta } = &d.delta + { + text.push_str(delta); + } + } + text +} + +async fn run_provider_test(client: &dyn ApiClient, name: &str) { + let model = client.model(); + print!("{GREEN}PASS{RESET} {name} {DIM}({model}){RESET} → "); + + let req = loopctl::api::StreamRequest::new(vec![loopctl::message::Message::user( + "Say hello in exactly 3 words.", + )]); + let stream = client.stream_messages(req); + let mut stream = std::pin::pin!(stream); + let mut events = Vec::new(); + + while let Some(result) = stream.next().await { + match result { + Ok(ev) => events.push(ev), + Err(e) => { + println!("{RED}FAIL{RESET} {name} {DIM}({model}){RESET} → error: {e}"); + panic!("{name} stream error: {e}"); + } + } + } + + let text = extract_text(&events); + let has_stop = events + .iter() + .any(|e| matches!(e, loopctl::stream::StreamEvent::MessageStop)); + + println!("{CYAN}\"{text}\"{RESET}"); + + assert!(!text.is_empty(), "{name} should produce non-empty text"); + assert!(has_stop, "{name} stream should end with MessageStop"); +} + +#[cfg(feature = "ollama")] +#[tokio::test] +async fn ollama_test() { + if std::env::var("LOOPCTL_E2E").as_deref() != Ok("1") || std::env::var("OLLAMA_MODEL").is_err() + { + eprintln!("{DIM}skip{RESET} Ollama"); + return; + } + let model = std::env::var("OLLAMA_MODEL").unwrap(); + let client = loopctl::provider::ollama(&model).unwrap(); + run_provider_test(&client, "Ollama").await; +} + +#[cfg(feature = "openai")] +#[tokio::test] +async fn openai_test() { + if std::env::var("LOOPCTL_E2E").as_deref() != Ok("1") + || std::env::var("OPENAI_API_KEY").is_err() + { + eprintln!("{DIM}skip{RESET} OpenAI"); + return; + } + let client = loopctl::provider::OpenAiClient::from_env().unwrap(); + run_provider_test(&client, "OpenAI").await; +} + +#[cfg(feature = "anthropic")] +#[tokio::test] +async fn anthropic_test() { + if std::env::var("LOOPCTL_E2E").as_deref() != Ok("1") + || std::env::var("ANTHROPIC_API_KEY").is_err() + { + eprintln!("{DIM}skip{RESET} Anthropic"); + return; + } + let client = loopctl::provider::AnthropicClient::from_env().unwrap(); + run_provider_test(&client, "Anthropic").await; +} + +#[cfg(feature = "gemini")] +#[tokio::test] +async fn gemini_test() { + if std::env::var("LOOPCTL_E2E").as_deref() != Ok("1") + || (std::env::var("GEMINI_API_KEY").is_err() && std::env::var("GOOGLE_API_KEY").is_err()) + { + eprintln!("{DIM}skip{RESET} Gemini"); + return; + } + let client = loopctl::provider::GeminiClient::from_env().unwrap(); + run_provider_test(&client, "Gemini").await; +} + +#[cfg(feature = "grok")] +#[tokio::test] +async fn grok_test() { + if std::env::var("LOOPCTL_E2E").as_deref() != Ok("1") + || (std::env::var("XAI_API_KEY").is_err() && std::env::var("GROK_API_KEY").is_err()) + { + eprintln!("{DIM}skip{RESET} Grok"); + return; + } + let client = loopctl::provider::grok().unwrap(); + run_provider_test(&client, "Grok").await; +} + +#[cfg(feature = "deepseek")] +#[tokio::test] +async fn deepseek_test() { + if std::env::var("LOOPCTL_E2E").as_deref() != Ok("1") + || std::env::var("DEEPSEEK_API_KEY").is_err() + { + eprintln!("{DIM}skip{RESET} DeepSeek"); + return; + } + let client = loopctl::provider::deepseek().unwrap(); + run_provider_test(&client, "DeepSeek").await; +} + +#[cfg(feature = "zai")] +#[tokio::test] +async fn zai_test() { + if std::env::var("LOOPCTL_E2E").as_deref() != Ok("1") + || (std::env::var("ZAI_API_KEY").is_err() && std::env::var("ZHIPUAI_API_KEY").is_err()) + { + eprintln!("{DIM}skip{RESET} Z.ai"); + return; + } + let client = loopctl::provider::zai().unwrap(); + run_provider_test(&client, "Z.ai").await; +} diff --git a/tests/provider_survival.rs b/tests/provider_survival.rs new file mode 100644 index 0000000..f91221b --- /dev/null +++ b/tests/provider_survival.rs @@ -0,0 +1,60 @@ +//! Live test: provider survives consecutive requests (Ollama only). +//! +//! Run: `LOOPCTL_E2E=1 OLLAMA_MODEL=qwen2.5:7b cargo test --features ollama --test provider_survival -- --nocapture` + +#![allow( + clippy::pedantic, + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::arithmetic_side_effects, + clippy::indexing_slicing, + clippy::redundant_clone +)] + +mod helpers; + +#[cfg(feature = "ollama")] +use loopctl::engine::core::Loop; + +#[cfg(feature = "ollama")] +#[tokio::test] +async fn provider_survives_consecutive_requests() { + if std::env::var("LOOPCTL_E2E").as_deref() != Ok("1") { + eprintln!("skipping (set LOOPCTL_E2E=1)"); + return; + } + + let model = std::env::var("OLLAMA_MODEL").expect("set OLLAMA_MODEL"); + let client = std::sync::Arc::new(loopctl::provider::ollama(&model).unwrap()); + let session = loopctl::config::SessionConfig::default(); + let run = loopctl::engine::RunConfig::default(); + let n: u32 = std::env::var("LOOPCTL_N") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(100); + let mut ok: u32 = 0; + + for i in 0..n { + let mut agent = loopctl::engine::BareLoop::new( + std::sync::Arc::clone(&client), + loopctl::tool::ToolRegistry::new(), + session.clone(), + ); + let prompt = format!("Say exactly: response-{i}"); + let pass = agent.run(&prompt, &run).await.is_ok(); + if pass { + ok += 1; + } + let status = if pass { + "\x1b[32mPASS:\x1b[0m" + } else { + "\x1b[31mFAIL:\x1b[0m" + }; + println!("[{}/{}] {status} {prompt}", i + 1, n); + } + + let rate = f64::from(ok) / f64::from(n) * 100.0; + println!("{ok}/{n} ({rate:.0}%) succeeded"); + assert!(rate >= 95.0, "got {rate:.0}%"); +} diff --git a/tests/structured_output.rs b/tests/structured_output.rs new file mode 100644 index 0000000..3a7a20e --- /dev/null +++ b/tests/structured_output.rs @@ -0,0 +1,105 @@ +//! Live test: structured output schema adherence (Ollama only). +//! +//! Run: `LOOPCTL_E2E=1 OLLAMA_MODEL=qwen2.5:7b cargo test --features ollama --test structured_output -- --nocapture` + +#![allow( + clippy::pedantic, + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::arithmetic_side_effects, + clippy::indexing_slicing, + clippy::redundant_clone +)] + +mod helpers; + +#[cfg(feature = "ollama")] +use loopctl::structured::StructuredOutput; +#[cfg(feature = "ollama")] +use serde::{Deserialize, Serialize}; + +#[cfg(feature = "ollama")] +#[derive(Debug, Serialize, Deserialize)] +struct PersonInfo { + name: String, + age: u32, + city: String, +} + +#[cfg(feature = "ollama")] +impl StructuredOutput for PersonInfo { + fn name() -> &'static str { + "person_info" + } + fn schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer", "minimum": 0}, + "city": {"type": "string"} + }, + "required": ["name", "age", "city"] + }) + } + fn from_value(v: serde_json::Value) -> Result { + serde_json::from_value(v).map_err(loopctl::structured::StructuredError::from) + } +} + +#[cfg(feature = "ollama")] +#[tokio::test] +async fn structured_outputs_match_schema() { + if std::env::var("LOOPCTL_E2E").as_deref() != Ok("1") { + eprintln!("skipping (set LOOPCTL_E2E=1)"); + return; + } + + let model = std::env::var("OLLAMA_MODEL").expect("set OLLAMA_MODEL"); + let client = loopctl::provider::ollama(&model).unwrap(); + let n: u32 = std::env::var("LOOPCTL_N") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(20); + let prompts = [ + "Tell me about Ada Lovelace as JSON with fields: name, age, city", + "Tell me about Alan Turing as JSON with fields: name, age, city", + "Tell me about Grace Hopper as JSON with fields: name, age, city", + "Tell me about Linus Torvalds as JSON with fields: name, age, city", + "Tell me about Margaret Hamilton as JSON with fields: name, age, city", + ]; + let mut valid: u32 = 0; + + for i in 0..n { + let prompt = prompts[(i as usize) % prompts.len()]; + let messages = vec![loopctl::message::Message::user(prompt)]; + let pass = + match loopctl::structured::request_structured::(&client, messages, None) + .await + { + Ok(p) if !p.name.is_empty() && !p.city.is_empty() => true, + Ok(p) => { + eprintln!("req {i}: parsed but empty fields: {p:?}"); + false + } + Err(e) => { + eprintln!("req {i}: failed: {e}"); + false + } + }; + if pass { + valid += 1; + } + let status = if pass { + "\x1b[32mPASS:\x1b[0m" + } else { + "\x1b[31mFAIL:\x1b[0m" + }; + println!("[{}/{}] {status} {prompt}", i + 1, n); + } + + let rate = f64::from(valid) / f64::from(n) * 100.0; + println!("{valid}/{n} ({rate:.0}%) valid"); + assert!(rate >= 95.0, "got {rate:.0}%"); +} From 4807079a228753aaf71abe78e15279cd99994195 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 27 Jul 2026 14:49:22 +1200 Subject: [PATCH 05/21] fix: openai multi-chunk streaming --- CHANGELOG.md | 22 +++ src/provider/openai.rs | 164 ++++++++++++++++++-- src/stream.rs | 344 +++++++++++++++++++++++++++++------------ 3 files changed, 411 insertions(+), 119 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 65d76b3..a7f8f7c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -381,6 +381,28 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. - `StreamHandlerError::RateLimitEscalation.prior: StreamOutcome` field (never read by any consumer). +### Fixed + +- OpenAI streaming corrupted tool-call arguments that arrived across more + than one chunk. The `StreamEmitter` opened a `PartStart` for the tool + call on every chunk carrying a `function` field (which is all of them, + including argument-fragment chunks), so each fragment after the first + re-opened the part and wiped the previously buffered JSON in the + downstream `StreamAccumulator`. The emitter now tracks each tool call's + `index` and emits `PartStart` exactly once per call; argument fragments + append to the buffer as intended and the complete JSON parses at + `PartStop`. +- `StreamAccumulator` dropped parallel tool calls whose argument fragments + arrived interleaved (the shape OpenAI streams for + `parallel_tool_calls`). The accumulator tracked a single in-progress + part, so a second `PartStart` arriving before the first `PartStop` + overwrote the first call's buffer, and `IndexedDelta` fragments whose + `index` did not match the single current slot were silently dropped. + It now holds a `Vec` of open slots keyed by `index`, routes each delta + to the matching slot, and flushes slots in `PartStart` arrival order + (FIFO) on `PartStop`. Anthropic (strictly sequential) and Gemini + (atomic per-chunk tool calls) are unaffected. + ## [0.1.0] - 2025-07-01 Initial crates.io release. diff --git a/src/provider/openai.rs b/src/provider/openai.rs index 2bfd931..d126bd7 100644 --- a/src/provider/openai.rs +++ b/src/provider/openai.rs @@ -1031,12 +1031,24 @@ struct StreamEmitter { /// to the text lane. thinking_part_open: bool, + /// Tool-call indices that have already had their + /// [`StreamEvent::PartStart`] emitted. + /// + /// OpenAI streams a single tool call across many chunks, all sharing + /// the same `index`: the first carries the call `id` and function + /// `name`, and every later chunk carries only an `arguments` + /// fragment (still under `function`). Gating `PartStart` on + /// `function.is_some()` would re-emit it on every fragment and wipe + /// the accumulator's buffered arguments. Tracking seen indices lets + /// the emitter open each part exactly once. + seen_tool_indices: Vec, + /// Number of tool-call parts currently open. /// - /// Each `delta.tool_calls` entry with a `function` field opens a new - /// tool part via [`StreamEvent::PartStart`]. The counter drives the - /// matching batch of `PartStop` emissions on finish (one per open - /// tool) so callers see balanced part lifecycles. + /// Each distinct tool `index` opens one tool part via + /// [`StreamEvent::PartStart`]. The counter drives the matching batch + /// of `PartStop` emissions on finish (one per open tool) so callers + /// see balanced part lifecycles. open_tool_count: usize, /// Whether the terminal stop signal has been processed. @@ -1153,15 +1165,18 @@ impl StreamEmitter { /// Handle a single tool-call delta from the stream. /// - /// On the first delta for a tool call (when `function` is present), - /// emits a [`PartStart`](StreamEvent::PartStart) with a - /// [`ToolCall`](crate::message::MessagePart::ToolCall) part carrying the - /// tool ID and name. Subsequent deltas carrying `function.arguments` - /// fragments emit [`InputJson`](crate::stream::DeltaPart::InputJson) - /// events so the caller can accumulate the full JSON input. + /// OpenAI streams one tool call across many chunks that share an + /// `index`: the first carries the call `id` and function `name` + /// under `function`; every later chunk carries only an `arguments` + /// fragment (still under `function`). The emitter opens the part + /// with [`StreamEvent::PartStart`] exactly once per `index` (on the + /// first chunk it sees for that index), then forwards every + /// non-empty `arguments` fragment as an + /// [`InputJson`](crate::stream::DeltaPart::InputJson) delta so the + /// caller can concatenate them into the full JSON input. fn process_tool_call(&mut self, tc: &OpenAiToolCallDelta) { - if tc.function.is_some() { - // New tool call — emit PartStart. + if tc.function.is_some() && !self.seen_tool_indices.contains(&tc.index) { + self.seen_tool_indices.push(tc.index); self.push(StreamEvent::PartStart(PartStart { index: tc.index, part: Some(MessagePart::ToolCall { @@ -1177,7 +1192,6 @@ impl StreamEmitter { self.open_tool_count = self.open_tool_count.saturating_add(1); } - // Stream argument fragments. if let Some(func) = &tc.function && !func.arguments.is_empty() { @@ -1205,17 +1219,14 @@ impl StreamEmitter { } self.finished = true; - // Close any open text part. if self.text_part_open { self.push(StreamEvent::PartStop); } - // Close any open reasoning (thinking) part. if self.thinking_part_open { self.push(StreamEvent::PartStop); } - // Close each open tool-call part. for _ in 0..self.open_tool_count { self.push(StreamEvent::PartStop); } @@ -1583,6 +1594,127 @@ mod tests { )); } + #[test] + fn emitter_multi_chunk_tool_call_emits_part_start_once() { + let mut em = StreamEmitter::default(); + let chunk0 = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":null},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk0); + em.drain(); + + let header = 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(&header); + em.drain(); + + let fragment = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":1,"id":"call_1","function":{"name":"echo","arguments":"\"hi\"}"}}]},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&fragment); + let events = em.drain(); + let deltas: Vec<_> = events + .iter() + .filter(|e| matches!(e, StreamEvent::IndexedDelta(_))) + .collect(); + assert_eq!( + deltas.len(), + 1, + "follow-up chunk must emit only an argument delta, not a PartStart" + ); + assert!( + events + .iter() + .all(|e| !matches!(e, StreamEvent::PartStart(_))) + ); + assert_eq!(em.open_tool_count, 1); + } + + #[test] + fn emitter_multi_chunk_tool_call_accumulates_through_accumulator() { + use crate::stream::StreamAccumulator; + let mut em = StreamEmitter::default(); + let mut acc = StreamAccumulator::new(); + let chunks = [ + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":null},"finish_reason":null}]}"#, + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":1,"id":"call_1","function":{"name":"echo","arguments":"{\"msg\":"}}]},"finish_reason":null}]}"#, + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":1,"id":"call_1","function":{"name":"echo","arguments":"\"hi\"}"}}]},"finish_reason":null}]}"#, + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":null,"finish_reason":"tool_calls"}]}"#, + ]; + for raw in chunks { + let chunk = OpenAiChunk::parse(raw).unwrap(); + em.process_chunk(&chunk); + for ev in em.drain() { + acc.process(&ev).unwrap(); + } + } + for ev in em.finish() { + acc.process(&ev).unwrap(); + } + + let msg = acc.build(); + assert_eq!(msg.parts.len(), 1); + match &msg.parts[0] { + MessagePart::ToolCall { name, input, .. } => { + assert_eq!(name, "echo"); + assert_eq!(input, &serde_json::json!({"msg": "hi"})); + } + other => panic!("expected ToolCall, got {other:?}"), + } + } + + #[test] + fn emitter_two_interleaved_multi_chunk_tool_calls_accumulate() { + use crate::stream::StreamAccumulator; + + let mut em = StreamEmitter::default(); + let mut acc = StreamAccumulator::new(); + let chunks = [ + // Message start. + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":null},"finish_reason":null}]}"#, + // Call A (index 0): header + first fragment. + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_a","function":{"name":"echo","arguments":"{\"msg\":"}}]},"finish_reason":null}]}"#, + // Call B (index 1): header + first fragment. + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":1,"id":"call_b","function":{"name":"search","arguments":"{\"q\":"}}]},"finish_reason":null}]}"#, + // Call A: remaining fragment. + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_a","function":{"name":"echo","arguments":"\"a\"}"}}]},"finish_reason":null}]}"#, + // Call B: remaining fragment. + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":1,"id":"call_b","function":{"name":"search","arguments":"\"b\"}"}}]},"finish_reason":null}]}"#, + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":null,"finish_reason":"tool_calls"}]}"#, + ]; + for raw in chunks { + let chunk = OpenAiChunk::parse(raw).unwrap(); + em.process_chunk(&chunk); + for ev in em.drain() { + acc.process(&ev).unwrap(); + } + } + for ev in em.finish() { + acc.process(&ev).unwrap(); + } + + let msg = acc.build(); + assert_eq!(msg.parts.len(), 2, "two tool calls expected"); + match &msg.parts[0] { + MessagePart::ToolCall { name, input, .. } => { + assert_eq!(name, "echo"); + assert_eq!(input, &serde_json::json!({"msg": "a"})); + } + other => panic!("expected ToolCall, got {other:?}"), + } + match &msg.parts[1] { + MessagePart::ToolCall { name, input, .. } => { + assert_eq!(name, "search"); + assert_eq!(input, &serde_json::json!({"q": "b"})); + } + other => panic!("expected ToolCall, got {other:?}"), + } + } + #[test] fn emitter_finish_emits_part_stops_and_message_delta() { let mut em = StreamEmitter::default(); diff --git a/src/stream.rs b/src/stream.rs index e39ab60..6ad3742 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -838,67 +838,122 @@ impl Usage { /// calling [`new`](Self::new). #[derive(Debug, Default)] pub struct StreamAccumulator { - /// Accumulated parts from completed part sequences. + /// Fully assembled parts flushed by [`PartStop`](StreamEvent::PartStop). + /// + /// Each entry is a finished [`MessagePart`] produced when a + /// [`PartStop`](StreamEvent::PartStop) closes one of the entries in + /// [`open`](Self::open). These are the parts returned by + /// [`build`](Self::build). + completed: Vec, + + /// Parts currently receiving deltas, in [`PartStart`] arrival order. + /// + /// A provider may keep several parts open at once — for example, + /// OpenAI streams parallel tool calls by interleaving argument + /// fragments across distinct `index` values and only closing them + /// all at the terminal `finish_reason`. Each + /// [`PartStart`](StreamEvent::PartStart) pushes a new slot; + /// [`IndexedDelta`](StreamEvent::IndexedDelta) routes to the slot + /// whose `index` matches; [`PartStop`](StreamEvent::PartStop) + /// flushes the oldest slot still open. + open: Vec, + + /// The model name that produced this response. /// - /// Each entry is a fully assembled [`MessagePart`] produced when - /// a [`PartStop`](StreamEvent::PartStop) event - /// finalizes the current part. These are the parts that will - /// appear in the final [`Message`] returned by [`build`](Self::build). - parts: Vec, + /// Extracted from the [`MessageStart`](StreamEvent::MessageStart) + /// event. `None` until that event is processed. Can be used for + /// logging or routing after the stream completes. + model: Option, - /// Current text being accumulated from [`DeltaPart::Text`] deltas. + /// Token usage statistics from the response. /// - /// Grows as text deltas arrive. Cleared and flushed into [`parts`](Self::parts) - /// as a [`MessagePart::Text`] when - /// [`PartStop`](StreamEvent::PartStop) is received. - current_text: String, + /// Populated from the [`MessageDelta`](StreamEvent::MessageDelta) + /// event. `None` until that event is processed. Access via + /// [`usage`](Self::usage) after processing. + usage: Option, +} + +/// Which lane an in-progress [`OpenPart`] is accumulating. +/// +/// Distinguishes plain assistant text from a tool-call invocation so +/// [`PartStop`](StreamEvent::PartStop) knows which buffer to flush and +/// which [`MessagePart`] shape to build. Decided once, at +/// [`PartStart`](StreamEvent::PartStart) time, from the carried part. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +enum OpenPartKind { + /// Assistant text, reconstructed from `delta.content` fragments. + /// + /// The default lane: a [`PartStart`](StreamEvent::PartStart) that + /// carries no tool-call part opens a text slot, and its buffered + /// string becomes a [`MessagePart::Text`] on close. + #[default] + Text, + + /// A tool-call invocation, reconstructed from `function.arguments` + /// fragments. + /// + /// The slot latches the tool `id` and `name` from + /// [`PartStart::part`](crate::stream::PartStart::part) at open time + /// and accumulates the raw JSON arguments string; on close it parses + /// that string into the tool-call [`MessagePart::ToolCall`] input. + Tool, +} - /// The tool-call ID being accumulated for the current tool part. +/// A single in-progress part being accumulated between open and close. +/// +/// Holds the buffers for one lane — assistant text or a tool call — +/// from the [`PartStart`](StreamEvent::PartStart) that opens it to the +/// [`PartStop`](StreamEvent::PartStop) that flushes it. The +/// [`StreamAccumulator`] keeps a [`Vec`] of these so that providers +/// which leave several parts open at once (OpenAI's interleaved parallel +/// tool calls) accumulate each one independently. +#[derive(Debug, Default)] +struct OpenPart { + /// The `index` this slot was opened with. /// - /// Populated from [`PartStart`] when the part - /// is a tool-call invocation. Cleared on the next - /// [`PartStart`](StreamEvent::PartStart). - /// Used to correlate the tool-call part with its result. - current_tool_id: String, + /// Copied from [`PartStart::index`](crate::stream::PartStart::index) + /// at open time and used to route each + /// [`IndexedDelta`](StreamEvent::IndexedDelta) fragment to the slot + /// it belongs to. Stable for the lifetime of the slot. + index: usize, - /// The tool name being accumulated for the current tool part. + /// Whether this slot accumulates assistant text or a tool call. /// - /// Populated from [`PartStart`] when the part - /// is a tool-call invocation. Cleared on the next - /// [`PartStart`](StreamEvent::PartStart). - /// Determines whether a part is text or tool-call when - /// [`PartStop`](StreamEvent::PartStop) arrives. - current_tool_name: String, + /// Set once when the slot opens and never mutated; it selects which + /// buffer the deltas append to and which [`MessagePart`] variant + /// [`PartStop`](StreamEvent::PartStop) builds from the slot. + kind: OpenPartKind, - /// The raw JSON string being accumulated for the current tool input. + /// Buffered assistant text for a text-lane slot. /// - /// Built up from [`DeltaPart::InputJson`] and - /// [`DeltaPart::InputJson`] deltas, then parsed into a - /// [`Value`](serde_json::Value) when - /// [`PartStop`](StreamEvent::PartStop) is received. - /// If parsing fails, defaults to an empty JSON object `{}`. - current_tool_input: String, + /// Grown one fragment at a time by [`DeltaPart::Text`] deltas and + /// flushed into a [`MessagePart::Text`] on close. Empty and unused + /// for tool-call slots. + text: String, - /// Index of the part currently being accumulated. + /// Server-assigned identifier of the tool call. /// - /// `None` before the first [`PartStart`] event, then - /// `Some(index)` for the remainder of the stream. Used to - /// track which part is in progress. - current_index: Option, + /// Latched from [`PartStart::part`](crate::stream::PartStart::part) + /// when the slot opens as a tool call, then carried onto the built + /// [`MessagePart::ToolCall`] so the host can match the eventual + /// tool result back to this call. Empty for text slots. + tool_id: String, - /// The model name that produced this response. + /// Name of the tool being invoked. /// - /// Extracted from the [`MessageStart`](StreamEvent::MessageStart) - /// event. `None` until that event is processed. Can be used for - /// logging or routing after the stream completes. - model: Option, + /// Latched from [`PartStart::part`](crate::stream::PartStart::part) + /// at open time and used both to detect that this is a tool-call + /// slot (non-empty) and to label the built + /// [`MessagePart::ToolCall`]. Empty for text slots. + tool_name: String, - /// Token usage statistics from the response. + /// Raw JSON arguments string for a tool-call slot. /// - /// Populated from the [`MessageDelta`](StreamEvent::MessageDelta) - /// event. `None` until that event is processed. Access via - /// [`usage`](Self::usage) after processing. - usage: Option, + /// Grown fragment by fragment by [`DeltaPart::InputJson`] and + /// [`DeltaPart::ToolCall`] deltas and parsed into the + /// [`MessagePart::ToolCall`] input when the slot closes. Empty for + /// text slots. + tool_input: String, } impl StreamAccumulator { @@ -924,10 +979,16 @@ impl StreamAccumulator { /// Process a single stream event and update internal state. /// /// Called for each [`StreamEvent`] as it arrives from the SSE - /// connection. The accumulator tracks the current part - /// being built and flushes completed parts into the internal - /// list when [`PartStop`](StreamEvent::PartStop) - /// is received. + /// connection. The accumulator keeps a [`Vec`] of in-progress parts + /// so that providers which leave several parts open at once — OpenAI + /// interleaves parallel tool calls — accumulate each one + /// independently. Each [`PartStart`](StreamEvent::PartStart) opens a + /// new slot keyed by its `index`; each + /// [`IndexedDelta`](StreamEvent::IndexedDelta) routes to the matching + /// slot; each [`PartStop`](StreamEvent::PartStop) flushes the oldest + /// open slot (FIFO by `PartStart` arrival order). Providers that only + /// ever hold one slot open (Anthropic, Gemini) behave exactly as on a + /// single-slot accumulator. /// /// # Arguments /// @@ -936,12 +997,15 @@ impl StreamAccumulator { /// # Event Handling /// /// - [`MessageStart`](StreamEvent::MessageStart) — Captures the model name. - /// - [`PartStart`](StreamEvent::PartStart) — Resets buffers, - /// captures tool ID/name if present. - /// - [`IndexedDelta`](StreamEvent::IndexedDelta) — Appends text or - /// JSON fragments to the appropriate buffer. - /// - [`PartStop`](StreamEvent::PartStop) — Flushes the current - /// buffer into an accumulated [`MessagePart`]. + /// - [`PartStart`](StreamEvent::PartStart) — Opens a new in-progress + /// slot for the given `index` (text or tool call, decided by the + /// carried part). Several slots may be open at once. + /// - [`IndexedDelta`](StreamEvent::IndexedDelta) — Routes the + /// fragment to the open slot whose `index` matches and appends it + /// to that slot's buffer. + /// - [`PartStop`](StreamEvent::PartStop) — Flushes the oldest + /// still-open slot into [`completed`](Self::completed) as a + /// [`MessagePart`] and drops it. /// - [`MessageDelta`](StreamEvent::MessageDelta) — Captures usage statistics. /// - [`MessageStop`](StreamEvent::MessageStop) / [`Ping`](StreamEvent::Ping) — Ignored. /// @@ -981,65 +1045,73 @@ impl StreamAccumulator { Ok(()) } StreamEvent::PartStart(part_start) => { - self.current_index = Some(part_start.index); - self.current_text.clear(); - self.current_tool_id.clear(); - self.current_tool_name.clear(); - self.current_tool_input.clear(); - + let kind = match &part_start.part { + Some(MessagePart::ToolCall { .. }) => OpenPartKind::Tool, + _ => OpenPartKind::Text, + }; + let mut slot = OpenPart { + index: part_start.index, + kind, + ..Default::default() + }; if let Some(MessagePart::ToolCall { id, name, .. }) = &part_start.part { - self.current_tool_id.clone_from(id); - self.current_tool_name.clone_from(name); + slot.tool_id.clone_from(id); + slot.tool_name.clone_from(name); } + self.open.push(slot); Ok(()) } StreamEvent::IndexedDelta(delta) => { - // 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) { + let Some(slot) = self.open.iter_mut().find(|s| s.index == delta.index) else { return Ok(()); - } + }; match &delta.delta { DeltaPart::Text { text } => { - self.current_text.push_str(text); + slot.text.push_str(text); } DeltaPart::InputJson { partial_json } => { - self.current_tool_input.push_str(partial_json); + slot.tool_input.push_str(partial_json); } DeltaPart::ToolCall { partial_json } => { if let Some(s) = partial_json.as_str() { - self.current_tool_input.push_str(s); + slot.tool_input.push_str(s); } } - // Reasoning is not part of the assistant Message. Drop it - // here; observers already received it via on_thinking_delta - // upstream (engine/bare/stream.rs). DeltaPart::Thinking { .. } => {} } Ok(()) } StreamEvent::PartStop => { - if !self.current_text.is_empty() { - self.parts.push(MessagePart::text(&self.current_text)); - } else if !self.current_tool_name.is_empty() { - let input: Value = if self.current_tool_input.is_empty() { - Value::Object(serde_json::Map::new()) - } else { - serde_json::from_str(&self.current_tool_input).map_err(|e| { - StreamError::InvalidToolInputJson( - e, - std::mem::take(&mut self.current_tool_input), - ) - })? - }; - self.parts.push(MessagePart::tool_call( - &self.current_tool_id, - &self.current_tool_name, - input, - )); + let Some(slot) = self.open.first_mut() else { + return Ok(()); + }; + let flushed = match slot.kind { + OpenPartKind::Text if !slot.text.is_empty() => { + Some(MessagePart::text(std::mem::take(&mut slot.text))) + } + OpenPartKind::Tool if !slot.tool_name.is_empty() => { + let input: Value = if slot.tool_input.is_empty() { + Value::Object(serde_json::Map::new()) + } else { + serde_json::from_str(&slot.tool_input).map_err(|e| { + StreamError::InvalidToolInputJson( + e, + std::mem::take(&mut slot.tool_input), + ) + })? + }; + Some(MessagePart::tool_call( + std::mem::take(&mut slot.tool_id), + std::mem::take(&mut slot.tool_name), + input, + )) + } + _ => None, + }; + self.open.remove(0); + if let Some(part) = flushed { + self.completed.push(part); } - self.current_text.clear(); Ok(()) } StreamEvent::MessageDelta(delta) => { @@ -1057,7 +1129,7 @@ impl StreamAccumulator { /// received before the stream times out. #[must_use] pub fn peek_parts(&self) -> &[MessagePart] { - &self.parts + &self.completed } /// Consume the accumulator and produce the final [`Message`]. @@ -1087,7 +1159,7 @@ impl StreamAccumulator { pub fn build(self) -> Message { Message { role: crate::message::Role::Assistant, - parts: self.parts, + parts: self.completed, } } @@ -1244,6 +1316,79 @@ mod tests { assert!(msg.parts[0].is_tool_call()); } + #[test] + fn test_accumulator_interleaved_tool_calls() { + // Reproduces OpenAI's parallel-tool-call wire shape: two tool + // calls opened up front, argument fragments interleaved across + // their indices, both closed by bare PartStops at the end. + let mut acc = StreamAccumulator::new(); + acc.process(&StreamEvent::PartStart(PartStart { + index: 0, + part: Some(MessagePart::tool_call( + "call_a", + "echo", + Value::Object(serde_json::Map::new()), + )), + })) + .unwrap(); + acc.process(&StreamEvent::PartStart(PartStart { + index: 1, + part: Some(MessagePart::tool_call( + "call_b", + "search", + Value::Object(serde_json::Map::new()), + )), + })) + .unwrap(); + acc.process(&StreamEvent::IndexedDelta(IndexedDelta { + index: 0, + delta: DeltaPart::InputJson { + partial_json: r#"{"msg":"#.to_string(), + }, + })) + .unwrap(); + acc.process(&StreamEvent::IndexedDelta(IndexedDelta { + index: 1, + delta: DeltaPart::InputJson { + partial_json: r#"{"q":"#.to_string(), + }, + })) + .unwrap(); + acc.process(&StreamEvent::IndexedDelta(IndexedDelta { + index: 0, + delta: DeltaPart::InputJson { + partial_json: r#""a"}"#.to_string(), + }, + })) + .unwrap(); + acc.process(&StreamEvent::IndexedDelta(IndexedDelta { + index: 1, + delta: DeltaPart::InputJson { + partial_json: r#""b"}"#.to_string(), + }, + })) + .unwrap(); + acc.process(&StreamEvent::PartStop).unwrap(); + acc.process(&StreamEvent::PartStop).unwrap(); + + let msg = acc.build(); + assert_eq!(msg.parts.len(), 2); + match &msg.parts[0] { + MessagePart::ToolCall { name, input, .. } => { + assert_eq!(name, "echo"); + assert_eq!(input, &serde_json::json!({"msg": "a"})); + } + other => panic!("expected ToolCall, got {other:?}"), + } + match &msg.parts[1] { + MessagePart::ToolCall { name, input, .. } => { + assert_eq!(name, "search"); + assert_eq!(input, &serde_json::json!({"q": "b"})); + } + other => panic!("expected ToolCall, got {other:?}"), + } + } + #[test] fn test_accumulator_empty() { let acc = StreamAccumulator::new(); @@ -1264,7 +1409,6 @@ mod tests { #[test] fn test_stream_event_variants() { - // Just verify all variants can be constructed let _ = StreamEvent::Ping; let _ = StreamEvent::MessageStop; let _ = StreamEvent::PartStop; @@ -1337,7 +1481,6 @@ mod tests { })) .unwrap(); - // Delta with correct index 0 — should be applied. acc.process(&StreamEvent::IndexedDelta(IndexedDelta { index: 0, delta: DeltaPart::Text { @@ -1374,7 +1517,6 @@ mod tests { acc.process(&StreamEvent::PartStop).unwrap(); let msg = acc.build(); - // Empty text → no parts. assert!(msg.parts.is_empty()); } @@ -1387,7 +1529,6 @@ mod tests { })) .unwrap(); - // DeltaPart::ToolCall carries a string JSON value. acc.process(&StreamEvent::IndexedDelta(IndexedDelta { index: 0, delta: DeltaPart::ToolCall { @@ -1416,7 +1557,6 @@ mod tests { })) .unwrap(); - // Non-string JSON value — should be silently ignored. acc.process(&StreamEvent::IndexedDelta(IndexedDelta { index: 0, delta: DeltaPart::ToolCall { @@ -1468,7 +1608,6 @@ mod tests { 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("")), @@ -1483,7 +1622,6 @@ mod tests { .unwrap(); acc.process(&StreamEvent::PartStop).unwrap(); - // Second text part at index 1. acc.process(&StreamEvent::PartStart(PartStart { index: 1, part: Some(MessagePart::text("")), From 738b7a4cc667436910cfeeabd43df52b90d16ae2 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 27 Jul 2026 15:07:44 +1200 Subject: [PATCH 06/21] fix: openai multi-chunk streaming --- CHANGELOG.md | 26 ++++++++----- src/provider/openai.rs | 84 +++++++++++++++++++++++++++++------------- 2 files changed, 75 insertions(+), 35 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a7f8f7c..183ce0c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -383,15 +383,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. ### Fixed -- OpenAI streaming corrupted tool-call arguments that arrived across more - than one chunk. The `StreamEmitter` opened a `PartStart` for the tool - call on every chunk carrying a `function` field (which is all of them, - including argument-fragment chunks), so each fragment after the first - re-opened the part and wiped the previously buffered JSON in the - downstream `StreamAccumulator`. The emitter now tracks each tool call's - `index` and emits `PartStart` exactly once per call; argument fragments - append to the buffer as intended and the complete JSON parses at - `PartStop`. +- OpenAI streaming silently dropped every multi-chunk tool-call argument + fragment after the first, truncating the tool input JSON. The + deserialization structs declared `id` (on the tool-call delta) and + `name` (on the function object) as required `String` fields, but the + real OpenAI streaming protocol omits both on continuation chunks — + those carry only `index` and an `arguments` fragment. Serde rejected + those chunks with `missing field`, `OpenAiChunk::parse` returned + `None`, and the stream loop silently skipped them, leaving the + accumulated JSON incomplete. Both fields are now `Option` with + `#[serde(default)]`; the emitter latches `id` and `name` on the first + chunk for each `index` and ignores them on continuations. +- OpenAI streaming re-opened a tool-call part on every chunk carrying a + `function` field. Because every chunk (including argument-fragment + continuations) carries `function`, the `StreamEmitter` emitted + `PartStart` for each one, which wiped the downstream accumulator's + buffered JSON. The emitter now tracks each tool call's `index` and + emits `PartStart` exactly once per call. - `StreamAccumulator` dropped parallel tool calls whose argument fragments arrived interleaved (the shape OpenAI streams for `parallel_tool_calls`). The accumulator tracked a single in-progress diff --git a/src/provider/openai.rs b/src/provider/openai.rs index d126bd7..8f38bb4 100644 --- a/src/provider/openai.rs +++ b/src/provider/openai.rs @@ -947,9 +947,12 @@ struct OpenAiDelta { struct OpenAiToolCallDelta { /// Server-assigned identifier for the tool call. /// - /// Present on the first chunk for this `index`; the emitter forwards - /// it as the call id so the host can match the result later. - id: String, + /// Present only on the first chunk for this `index`; continuation + /// chunks omit it. The emitter latches it on the first chunk and + /// forwards it as the call id so the host can match the result + /// later. Defaults to `None` when the server omits the field. + #[serde(default)] + id: Option, /// Position of this tool call in the request's tool list. /// @@ -965,6 +968,7 @@ struct OpenAiToolCallDelta { /// the function `name` (first chunk for this `index`) or a fragment /// of the JSON `arguments` (subsequent chunks) — the /// [`StreamEmitter`] reassembles both per index. + #[serde(default)] function: Option, } @@ -973,23 +977,25 @@ struct OpenAiToolCallDelta { /// The `name` arrives on the first chunk for a tool call; `arguments` /// is a JSON string that may itself arrive in fragments across several /// chunks and must be concatenated before parsing. -#[derive(Deserialize)] +#[derive(Deserialize, Default)] struct OpenAiToolCallFunction { /// Accumulated JSON arguments for the tool call. /// /// A partial JSON string that grows across chunks; the emitter /// buffers it per `index` and hands the complete string to the /// caller once the part stops. + #[serde(default)] arguments: String, /// Fully-qualified name of the tool to invoke. /// /// Matches the `name` the tool was registered under in the request's /// `tools` array. Arrives on the first chunk for a given `index` - /// only; subsequent chunks for the same call carry just argument - /// fragments, so the emitter latches this value and reuses it for - /// every later chunk with the same index. - name: String, + /// only; continuation chunks for the same call omit it, so the + /// emitter latches this value on the first chunk and ignores it on + /// later ones. Defaults to `None` when the server omits the field. + #[serde(default)] + name: Option, } /// Stateful translator that converts a sequence of [`OpenAiChunk`]s @@ -1180,11 +1186,11 @@ impl StreamEmitter { self.push(StreamEvent::PartStart(PartStart { index: tc.index, part: Some(MessagePart::ToolCall { - id: tc.id.clone(), + id: tc.id.clone().unwrap_or_default(), name: tc .function .as_ref() - .map(|f| f.name.clone()) + .and_then(|f| f.name.clone()) .unwrap_or_default(), input: Value::Null, }), @@ -1393,7 +1399,6 @@ mod tests { 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"}"# @@ -1418,9 +1423,6 @@ mod tests { #[test] fn convert_message_multiple_tool_results_expand() { - // A single loopctl message with two tool-result parts must expand to - // two separate OpenAI tool messages (not one array-valued element, - // which Ollama/OpenAI reject). let m = Message::new( Role::User, vec![ @@ -1540,7 +1542,6 @@ mod tests { em.process_chunk(&chunk); let events = em.drain(); - // MessageStart, PartStart(0), IndexedDelta(Text) assert_eq!(events.len(), 3); assert!(matches!( events[1], @@ -1574,7 +1575,6 @@ mod tests { 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}]}"#, ) @@ -1582,7 +1582,6 @@ mod tests { em.process_chunk(&chunk); let events = em.drain(); - // PartStart(tool) + IndexedDelta(InputJson) assert_eq!(events.len(), 2); assert!(matches!( events[0], @@ -1612,7 +1611,7 @@ mod tests { em.drain(); let fragment = OpenAiChunk::parse( - r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":1,"id":"call_1","function":{"name":"echo","arguments":"\"hi\"}"}}]},"finish_reason":null}]}"#, + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":1,"function":{"arguments":"\"hi\"}"}}]},"finish_reason":null}]}"#, ) .unwrap(); em.process_chunk(&fragment); @@ -1642,7 +1641,7 @@ mod tests { let chunks = [ r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":null},"finish_reason":null}]}"#, r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":1,"id":"call_1","function":{"name":"echo","arguments":"{\"msg\":"}}]},"finish_reason":null}]}"#, - r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":1,"id":"call_1","function":{"name":"echo","arguments":"\"hi\"}"}}]},"finish_reason":null}]}"#, + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":1,"function":{"arguments":"\"hi\"}"}}]},"finish_reason":null}]}"#, r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":null,"finish_reason":"tool_calls"}]}"#, ]; for raw in chunks { @@ -1674,16 +1673,11 @@ mod tests { let mut em = StreamEmitter::default(); let mut acc = StreamAccumulator::new(); let chunks = [ - // Message start. r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":null},"finish_reason":null}]}"#, - // Call A (index 0): header + first fragment. r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_a","function":{"name":"echo","arguments":"{\"msg\":"}}]},"finish_reason":null}]}"#, - // Call B (index 1): header + first fragment. r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":1,"id":"call_b","function":{"name":"search","arguments":"{\"q\":"}}]},"finish_reason":null}]}"#, - // Call A: remaining fragment. - r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_a","function":{"name":"echo","arguments":"\"a\"}"}}]},"finish_reason":null}]}"#, - // Call B: remaining fragment. - r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":1,"id":"call_b","function":{"name":"search","arguments":"\"b\"}"}}]},"finish_reason":null}]}"#, + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"a\"}"}}]},"finish_reason":null}]}"#, + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":1,"function":{"arguments":"\"b\"}"}}]},"finish_reason":null}]}"#, r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":null,"finish_reason":"tool_calls"}]}"#, ]; for raw in chunks { @@ -1715,6 +1709,44 @@ mod tests { } } + #[test] + fn emitter_real_continuation_chunks_omit_id_and_name() { + use crate::stream::StreamAccumulator; + + let mut em = StreamEmitter::default(); + let mut acc = StreamAccumulator::new(); + let chunks = [ + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":null},"finish_reason":null}]}"#, + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"echo","arguments":"{\"msg\":"}}]},"finish_reason":null}]}"#, + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"hi\"}"}}]},"finish_reason":null}]}"#, + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":null,"finish_reason":"tool_calls"}]}"#, + ]; + for raw in chunks { + let chunk = OpenAiChunk::parse(raw).expect("real chunk shape must deserialize"); + em.process_chunk(&chunk); + for ev in em.drain() { + acc.process(&ev).expect("accumulator accepts events"); + } + } + for ev in em.finish() { + acc.process(&ev).expect("accumulator accepts finish events"); + } + + let msg = acc.build(); + assert_eq!(msg.parts.len(), 1, "one tool call expected"); + match &msg.parts[0] { + MessagePart::ToolCall { name, input, .. } => { + assert_eq!(name, "echo"); + assert_eq!( + input, + &serde_json::json!({"msg": "hi"}), + "continuation fragment must accumulate, not be dropped" + ); + } + other => panic!("expected ToolCall, got {other:?}"), + } + } + #[test] fn emitter_finish_emits_part_stops_and_message_delta() { let mut em = StreamEmitter::default(); From 49feb284b4545700006764db7549a2eccc74e92f Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 27 Jul 2026 15:25:37 +1200 Subject: [PATCH 07/21] fix: cancel signal re-arm --- CHANGELOG.md | 15 ++++++ src/cancel.rs | 121 ++++++++++++++++++++++++++++++++++++++++----- src/engine/bare.rs | 42 ++++++++++++++++ src/stream.rs | 3 +- 4 files changed, 168 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 183ce0c..414bba0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -205,6 +205,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. `Result` — chain with `?`), `with_observer`, `with_text_streamer`, `with_contributor`, `with_request_options`. The original `set_*`/`register_*` setters are unchanged and remain available. +- `CancelSignal::reset()` — re-arm a fired signal by swapping in a fresh + underlying `CancellationToken`. Required because the token is one-shot + by design; once fired it cannot be revived. All clones of an + `Arc` observe the new token, so a handle returned by + `BareLoop::cancel_signal()` keeps working across resets. `BareLoop` + calls this in `finalize()` so each `run()` starts with a clean signal. ### Changed @@ -410,6 +416,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. to the matching slot, and flushes slots in `PartStart` arrival order (FIFO) on `PartStop`. Anthropic (strictly sequential) and Gemini (atomic per-chunk tool calls) are unaffected. +- `BareLoop` was permanently dead after a single cancellation. Once + `cancel()` fired, the `CancelSignal` (a one-shot + `CancellationToken`) stayed cancelled forever, so every subsequent + `run()` returned `LoopError::Cancelled` immediately. `run()` now + re-arms the signal in `finalize()` — the single chokepoint every run + exit path passes through — so the next `run()` starts clean. A cancel + that arrives *before* a run still cancels that run (the signal is + cleared only after the run observes it and returns), preserving the + pre-run-cancel contract. ## [0.1.0] - 2025-07-01 diff --git a/src/cancel.rs b/src/cancel.rs index c606620..dea373d 100644 --- a/src/cancel.rs +++ b/src/cancel.rs @@ -21,6 +21,8 @@ //! } //! ``` +use std::sync::Mutex; + use tokio_util::sync::CancellationToken; /// Shared cancellation signal backed by a @@ -28,9 +30,18 @@ use 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`]. +/// instant notification with [`CancelSignal::notified`]. Re-arm between +/// runs with [`CancelSignal::reset`](Self::reset). +/// +/// The underlying token is intentionally one-shot — that is what makes it +/// race-free — so [`reset`](Self::reset) swaps in a fresh token rather +/// than trying to revive a fired one. Because every clone of an +/// `Arc` derefs through this same struct, all existing +/// handles observe the new token after a reset; a caller holding +/// [`cancel_signal`](crate::engine::BareLoop::cancel_signal) across runs +/// keeps working. pub struct CancelSignal { - inner: CancellationToken, + inner: Mutex, } impl CancelSignal { @@ -42,7 +53,7 @@ impl CancelSignal { #[must_use] pub fn new() -> Self { Self { - inner: CancellationToken::new(), + inner: Mutex::new(CancellationToken::new()), } } @@ -52,17 +63,25 @@ impl CancelSignal { /// awaiting [`Self::notified`]. Idempotent — calling multiple times is /// safe. pub fn cancel(&self) { - self.inner.cancel(); + self.inner + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .cancel(); } - /// Check whether the signal has been cancelled. + /// Check whether the signal has been cancelled since the last + /// [`reset`](Self::reset). /// - /// Performs a non-blocking check of the internal [`CancellationToken`]. - /// Returns `true` if [`cancel`](Self::cancel) has been called since - /// construction. + /// Performs a non-blocking check of the current + /// [`CancellationToken`]. Returns `true` if + /// [`cancel`](Self::cancel) has been called and no `reset` has + /// followed it. #[must_use] pub fn is_cancelled(&self) -> bool { - self.inner.is_cancelled() + self.inner + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_cancelled() } /// Return a future that completes **instantly** when [`Self::cancel`] @@ -72,7 +91,10 @@ impl CancelSignal { /// immediately on the first `.await`. /// /// This delegates to [`CancellationToken::cancelled`], which is - /// race-free by construction — no flag-then-wait loop required. + /// race-free by construction — no flag-then-wait loop required. The + /// lock is held only long enough to clone the token out, so awaiting + /// the returned future never blocks another caller from cancelling + /// or resetting. /// /// Use inside `tokio::select!` alongside the actual work future: /// @@ -83,7 +105,33 @@ impl CancelSignal { /// } /// ``` pub async fn notified(&self) { - self.inner.cancelled().await; + let token = self + .inner + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + token.cancelled().await; + } + + /// Re-arm the signal for a fresh run. + /// + /// Replaces the internal token with a brand-new, non-cancelled one. + /// Required at the start of each run because the underlying + /// [`CancellationToken`] is one-shot by design — once fired it stays + /// fired, so without a reset a single cancellation would make every + /// subsequent run return immediately as cancelled. + /// + /// All clones of an `Arc` observe the new token, + /// because they deref through this same struct; a handle returned by + /// [`cancel_signal`](crate::engine::BareLoop::cancel_signal) keeps + /// working across runs. A task already awaiting + /// [`notified`](Self::notified) on the previous token is unaffected — + /// its awaited run was cancelled, and it resolves against that token. + pub fn reset(&self) { + *self + .inner + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = CancellationToken::new(); } } @@ -160,4 +208,55 @@ mod tests { assert!(handle.await.unwrap()); } } + + #[test] + fn test_reset_clears_cancelled_state() { + let signal = CancelSignal::new(); + signal.cancel(); + assert!(signal.is_cancelled()); + + signal.reset(); + assert!( + !signal.is_cancelled(), + "reset must re-arm the signal to non-cancelled" + ); + } + + #[test] + fn test_reset_is_visible_through_existing_arc_clone() { + let signal = Arc::new(CancelSignal::new()); + let handle = Arc::clone(&signal); + + signal.cancel(); + assert!(handle.is_cancelled()); + + signal.reset(); + assert!(!handle.is_cancelled()); + } + + #[tokio::test] + async fn test_notified_after_reset_waits_for_new_cancel() { + let signal = Arc::new(CancelSignal::new()); + signal.cancel(); + signal.reset(); + + let handle = { + let s = Arc::clone(&signal); + tokio::spawn(async move { + s.notified().await; + "woke" + }) + }; + + tokio::task::yield_now().await; + let mut pending = tokio::time::interval(std::time::Duration::from_millis(5)); + pending.tick().await; + assert!( + !handle.is_finished(), + "notified must not fire after reset until a new cancel arrives" + ); + + signal.cancel(); + assert_eq!(handle.await.unwrap(), "woke"); + } } diff --git a/src/engine/bare.rs b/src/engine/bare.rs index 8ec90cd..ea3516a 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -1754,6 +1754,16 @@ impl crate::engine::core::Loop for BareLoop { !self.machine.is_terminal() } + /// Finalize the current run and return its [`Run`] accumulator. + /// + /// Every `run()` exit path — clean completion, error, max-turns, + /// cancellation — funnels through here. Records the run's end + /// timestamp, fires the session-end observers, and re-arms the + /// cancel signal so the next `run()` starts clean. Re-arming here + /// (rather than at the top of `run()`) preserves a cancel that + /// arrived before the run: the run observes it and returns + /// [`LoopError::Cancelled`], and only then is the signal cleared, + /// so the agent is never left permanently dead after one cancel. fn finalize<'a>( &'a mut self, success: bool, @@ -1767,6 +1777,7 @@ impl crate::engine::core::Loop for BareLoop { let duration = run.duration(); self.notify_session_end(&run, success, duration); + self.cancelled.reset(); Ok(run) }) @@ -3540,6 +3551,37 @@ mod tests { assert!(agent.is_cancelled()); } + #[tokio::test] + async fn test_second_run_after_cancel_is_not_dead() { + let client = MockClient::new("test-model"); + client.add_text_response("second run should reach me"); + + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + + agent.cancel(); + + let first = agent.run("first", &RunConfig::default()).await; + assert!( + matches!(first, Err(LoopError::Cancelled)), + "first run must be cancelled, got {first:?}" + ); + + let client2 = MockClient::new("test-model"); + client2.add_text_response("second run ok"); + agent.client = Arc::new(client2); + + let second = agent.run("second", &RunConfig::default()).await; + match &second { + Ok(run) => assert_eq!( + run.output.as_deref(), + Some("second run ok"), + "second run must complete after cancel, got run with output {:?}", + run.output + ), + Err(e) => panic!("second run after cancel must not fail, got {e:?}"), + } + } + #[tokio::test] async fn test_run_result_fields() { let client = MockClient::new("test-model"); diff --git a/src/stream.rs b/src/stream.rs index 6ad3742..9370e20 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -1004,8 +1004,7 @@ impl StreamAccumulator { /// fragment to the open slot whose `index` matches and appends it /// to that slot's buffer. /// - [`PartStop`](StreamEvent::PartStop) — Flushes the oldest - /// still-open slot into [`completed`](Self::completed) as a - /// [`MessagePart`] and drops it. + /// still-open slot into a finished [`MessagePart`] and drops it. /// - [`MessageDelta`](StreamEvent::MessageDelta) — Captures usage statistics. /// - [`MessageStop`](StreamEvent::MessageStop) / [`Ping`](StreamEvent::Ping) — Ignored. /// From 9344b7e910bb74fb53f42835209e47b1fd2a3ee8 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 27 Jul 2026 15:46:03 +1200 Subject: [PATCH 08/21] fix: remove git add -A for auto_commit, use list of files --- CHANGELOG.md | 14 +++ src/hooks/builtin/auto_commit.rs | 144 ++++++++++++++++++++++++++++--- 2 files changed, 144 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 414bba0..9210185 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -426,6 +426,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. cleared only after the run observes it and returns), preserving the pre-run-cancel contract. +### Security + +- The auto-commit hook's `GitExecutor::stage_files` ran `git add -A` + when called with an empty file list — the default configuration + (`AutoCommitConfig::files` defaults to `vec![]`, and a session with no + recorded modifications passes `None`, which falls back to that empty + list). This staged and committed the entire working tree on every + session end: unrelated user edits, scratch files, and secrets such as + `.env` and credentials. The empty-list branch now refuses with a + `GitExecutorError::GitError` instead of broadening scope; callers must + populate `AutoCommitConfig::files` or rely on the hook's per-session + modification tracking. Misconfiguration now fails loudly rather than + silently committing everything. + ## [0.1.0] - 2025-07-01 Initial crates.io release. diff --git a/src/hooks/builtin/auto_commit.rs b/src/hooks/builtin/auto_commit.rs index 36ef6ab..2f12c36 100644 --- a/src/hooks/builtin/auto_commit.rs +++ b/src/hooks/builtin/auto_commit.rs @@ -215,22 +215,32 @@ impl GitExecutor { /// Stage files for commit. /// - /// Each path in `files` is passed to `git add` individually. Passing - /// an empty slice stages every change in the working tree - /// (`git add -A`), which is convenient for "commit everything" - /// semantics but will include unrelated modifications. + /// Each path in `files` is passed to `git add -- ` individually, + /// scoping staging to exactly the listed paths. An empty slice is + /// refused with an error rather than falling through to `git add -A`: + /// the auto-commit hook is meant to commit only the agent's own + /// recorded modifications, so staging the entire working tree would + /// silently sweep up unrelated edits and secret files (`.env`, + /// credentials, scratch buffers). Misconfiguration must fail loudly, + /// not broaden scope. /// /// # Errors /// - /// Returns [`GitExecutorError`] if any `git add` invocation fails; - /// earlier files in the list may already be staged. + /// Returns [`GitExecutorError::GitError`] if `files` is empty — + /// populate `AutoCommitConfig::files` or record modifications through + /// the hook. Returns [`GitExecutorError`] if any individual `git add` + /// invocation fails; earlier files in the list may already be staged. pub fn stage_files(files: &[String]) -> Result<(), GitExecutorError> { if files.is_empty() { - Self::run_git(&["add", "-A"])?; - } else { - for file in files { - Self::run_git(&["add", "--", file])?; - } + return Err(GitExecutorError::GitError( + "refusing to stage with empty file list (would run `git add -A` \ + and commit unrelated changes); populate `AutoCommitConfig::files` \ + or record modifications via the hook" + .to_string(), + )); + } + for file in files { + Self::run_git(&["add", "--", file])?; } Ok(()) } @@ -483,10 +493,16 @@ pub struct AutoCommitConfig { /// session end. pub commit_on_tools: Vec, - /// Files to add before commit (empty = all files). + /// Files to add before commit. /// - /// Explicit allow-list of paths to `git add`. An empty vector stages every - /// change in the working tree (`git add -A`). + /// Explicit allow-list of paths passed to `git add`. When the hook + /// records no per-session modifications, this list is what gets + /// staged. An empty list is **not** a "commit everything" wildcard: + /// [`stage_files`](GitExecutor::stage_files) refuses it with an error + /// rather than running `git add -A`, so the default configuration + /// commits nothing instead of sweeping up unrelated or secret files. + /// Populate this list, or rely on the hook's per-session modification + /// tracking, to stage exactly the agent's own changes. pub files: Vec, } @@ -875,4 +891,104 @@ mod tests { hook.clear_modifications(); assert!(hook.modified_files.lock().unwrap().is_empty()); } + + /// An empty file list must be refused, not staged as `git add -A`. + /// + /// Reproduces the footgun: a real git repo with a tracked file that + /// has an uncommitted modification, then `stage_files(&[])`. The + /// fixed code returns `Err` and leaves the file unstaged; the buggy + /// code ran `git add -A` and staged the unrelated change. + struct RepoGuard(std::path::PathBuf); + impl Drop for RepoGuard { + fn drop(&mut self) { + std::fs::remove_dir_all(&self.0).ok(); + } + } + + #[test] + fn stage_files_empty_list_does_not_stage_unrelated_changes() { + use std::fs; + use std::process::Command; + + if Command::new("git") + .arg("--version") + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .is_err() + { + eprintln!("skipping: git not on PATH"); + return; + } + + let mut repo_dir = std::env::temp_dir(); + repo_dir.push(format!( + "loopctl-stage-files-test-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |d| d.as_nanos()) + )); + fs::create_dir_all(&repo_dir).expect("created temp repo dir"); + let guard = RepoGuard(repo_dir.clone()); + let prev_dir = std::env::current_dir().expect("cwd readable"); + + for (label, args) in [ + ("init", vec!["init"]), + ("config user.name", vec!["config", "user.name", "test"]), + ("config user.email", vec!["config", "user.email", "t@t"]), + ] { + let status = Command::new("git") + .args(&args) + .current_dir(&repo_dir) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .unwrap_or_else(|e| panic!("git {label} spawn failed: {e}")); + assert!(status.success(), "git {label} failed"); + } + + let sentinel = repo_dir.join("unrelated.txt"); + fs::write(&sentinel, "v1\n").expect("wrote sentinel v1"); + for (label, args) in [ + ("add", vec!["add", "unrelated.txt"]), + ("commit", vec!["commit", "-m", "init"]), + ] { + let status = Command::new("git") + .args(&args) + .current_dir(&repo_dir) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .unwrap_or_else(|e| panic!("git {label} spawn failed: {e}")); + assert!(status.success(), "git {label} failed"); + } + + fs::write(&sentinel, "v2-uncommitted\n").expect("wrote sentinel v2"); + + std::env::set_current_dir(&repo_dir).expect("cd into repo"); + let result = GitExecutor::stage_files(&[]); + std::env::set_current_dir(&prev_dir).expect("restored cwd"); + + assert!( + result.is_err(), + "stage_files(&[]) must refuse to run, got {result:?}" + ); + + let output = Command::new("git") + .args(["status", "--porcelain"]) + .current_dir(&repo_dir) + .output() + .unwrap_or_else(|e| panic!("git status spawn failed: {e}")); + let status_text = String::from_utf8_lossy(&output.stdout); + assert!( + status_text.contains(" M unrelated.txt"), + "unrelated.txt must remain unstaged after stage_files(&[]), \ + but git status was:\n{status_text}\n\ + (this means stage_files ran `git add -A` — the empty-list \ + footgun is still present)" + ); + + drop(guard); + } } From 11f0d3468dae7391758b8feff471543682578800 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 27 Jul 2026 16:24:43 +1200 Subject: [PATCH 09/21] fix: merge tool results into single MessagePart --- CHANGELOG.md | 12 +++++ src/engine/bare.rs | 108 +++++++++++++++++++++++++++++++------ src/engine/bare/message.rs | 29 ++++------ 3 files changed, 115 insertions(+), 34 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9210185..8017d49 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -425,6 +425,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. that arrives *before* a run still cancels that run (the signal is cleared only after the run observes it and returns), preserving the pre-run-cancel contract. +- A turn's tool results were split across multiple consecutive user + messages instead of merged into one. Unknown-tool results (preresolved + by the machine) each arrived as their own user `Message`, and the + dispatched known-tool results arrived as another, so the loop pushed + them into history as separate entries. `BareLoop::handle_call_tools` + now collects all tool-result parts for a turn — preresolved plus + dispatched — into a single user `Message` before feeding it to the + machine, so a turn with any mix of unknown and known tools produces + exactly one user message regardless of how the results were produced. + `build_tool_result_message` is renamed to `build_tool_result_parts` + and now returns `Vec` (it no longer wraps the parts in a + throwaway `Message`). ### Security diff --git a/src/engine/bare.rs b/src/engine/bare.rs index ea3516a..d82357c 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -1063,12 +1063,12 @@ impl BareLoop { turn_start: Instant, turn_input_tokens: u64, turn_output_tokens: u64, - ) -> Result { + ) -> Result, LoopError> { let result = self.dispatch_tools(tool_calls, turn_index).await; let turn_duration = turn_start.elapsed(); match result { Ok(results) => { - let tool_result_msg = Self::build_tool_result_message(results); + let parts = Self::build_tool_result_parts(results); self.notify_turn_end( turn_index, true, @@ -1077,7 +1077,7 @@ impl BareLoop { turn_input_tokens, turn_output_tokens, ); - Ok(tool_result_msg) + Ok(parts) } Err(e) => { let err_str = e.to_string(); @@ -1585,9 +1585,13 @@ impl BareLoop { /// Handle a tool-dispatch request from the machine. /// /// Fires `on_tool_call_received`, dispatches the calls that are not - /// preresolved, feeds every tool-result message (dispatched plus - /// preresolved) back to the machine, and keeps the run budget in sync. - /// Cancellation races the dispatch via a biased `select!`. + /// preresolved, then assembles every tool result for the turn — + /// preresolved unknown-tool results plus dispatched known-tool + /// results — into a single user [`Message`] and feeds it back to the + /// machine. One turn yields one user message regardless of how the + /// results were produced, which is the shape providers expect. Keeps + /// the run budget in sync. Cancellation races the dispatch via a + /// biased `select!`. /// /// # Errors /// @@ -1606,11 +1610,11 @@ impl BareLoop { .map_or((0, 0), |t| (t.input_tokens, t.output_tokens)); let mut tool_calls: Vec = Vec::with_capacity(calls.len()); let mut dispatch_calls: Vec = Vec::new(); - let mut preresolved: Vec = Vec::new(); + let mut preresolved_parts: Vec = Vec::new(); for pending in calls { tool_calls.push(pending.call.clone()); match &pending.preresolved_result { - Some(msg) => preresolved.push(msg.clone()), + Some(msg) => preresolved_parts.extend(msg.parts.iter().cloned()), None => dispatch_calls.push(pending.call.clone()), } } @@ -1623,7 +1627,7 @@ impl BareLoop { .await }; - let dispatched: Message = tokio::select! { + let mut parts: Vec = tokio::select! { biased; () = cancel.notified() => { self.machine.cancel(); @@ -1638,7 +1642,7 @@ impl BareLoop { return Err(LoopError::Cancelled); } result = dispatch => match result { - Ok(msg) => msg, + Ok(parts) => parts, Err(e) => { self.set_error_state(&e); return Err(e); @@ -1646,8 +1650,9 @@ impl BareLoop { }, }; - preresolved.push(dispatched); - self.machine.tool_results(preresolved); + preresolved_parts.append(&mut parts); + self.machine + .tool_results(vec![Message::new(Role::User, preresolved_parts)]); Ok(()) } @@ -2825,11 +2830,10 @@ mod tests { display_hint: None, }]; - let msg = BareLoop::::build_tool_result_message(results); - assert_eq!(msg.role, Role::User); - assert_eq!(msg.parts.len(), 1); + let parts = BareLoop::::build_tool_result_parts(results); + assert_eq!(parts.len(), 1); - match &msg.parts[0] { + match &parts[0] { MessagePart::ToolResult { call_id, output, @@ -2903,6 +2907,78 @@ mod tests { assert_eq!(result.tool_call_count(), 2); } + #[tokio::test] + async fn test_mixed_known_unknown_tools_merge_into_one_user_message() { + let client = MockClient::new("test-model"); + + // One known tool call (echo) and one unknown (nonexistent) in the + // same turn. The unknown result is preresolved; the known one is + // dispatched. Both must land in a single user Message in history. + let tool_events = vec![ + StreamEvent::MessageStart(MessageStart { + message: MessageMetadata { + id: "msg_mixed".into(), + role: "assistant".into(), + model: "test-model".into(), + }, + }), + StreamEvent::PartStart(PartStart { + index: 0, + part: Some(MessagePart::tool_call( + "t1", + "echo", + json!({"message": "hi"}), + )), + }), + StreamEvent::PartStop, + StreamEvent::PartStart(PartStart { + index: 1, + part: Some(MessagePart::tool_call("t2", "nonexistent", json!({}))), + }), + StreamEvent::PartStop, + StreamEvent::MessageDelta(MessageDelta { + delta: MessageDeltaPayload { + stop_reason: Some("tool_call".to_string()), + }, + usage: Some(Usage::new(50, 20)), + }), + StreamEvent::MessageStop, + ]; + crate::error::recover_guard(client.responses.lock()).push(tool_events); + client.add_text_response("done"); + + let mut registry = ToolRegistry::new(); + registry.register(EchoTool); + + let mut agent = BareLoop::new(Arc::new(client), registry, make_config()); + agent + .run("mixed tools", &RunConfig::default()) + .await + .unwrap(); + + let user_messages: Vec<&Message> = agent + .conversation() + .iter() + .filter(|m| m.role == Role::User) + .collect(); + assert_eq!( + user_messages.len(), + 2, + "expected [prompt, one merged tool-result message], got {} user messages", + user_messages.len() + ); + let tool_results: Vec<&MessagePart> = user_messages[1] + .parts + .iter() + .filter(|p| p.is_tool_result()) + .collect(); + assert_eq!( + tool_results.len(), + 2, + "merged user message must hold both tool-result parts" + ); + } + #[tokio::test] async fn test_text_streamer_fires_on_text_delta() { let client = MockClient::new("test-model"); diff --git a/src/engine/bare/message.rs b/src/engine/bare/message.rs index 7e09690..3d56d91 100644 --- a/src/engine/bare/message.rs +++ b/src/engine/bare/message.rs @@ -4,34 +4,27 @@ //! assembling tool-result messages, extracting text or tool calls from an //! assistant response, and computing token counts. -use super::{ - ApiClient, BareLoop, Message, MessagePart, Role, ToolContext, ToolDispatchResult, ToolSchema, -}; +use super::{ApiClient, BareLoop, MessagePart, ToolContext, ToolDispatchResult, ToolSchema}; impl BareLoop { - /// Build the tool result message from executed tool results. + /// Build the tool-result parts from executed tool results. /// - /// Per API convention, tool results are sent as a **user** message - /// containing `tool_result` content parts. Each part pairs the - /// `tool_call_id` with the tool's output (wrapped in a - /// [`ToolContent`](ToolContent)) and an `is_error` - /// flag so the model can distinguish successes from failures. + /// Each dispatch result becomes one `tool_result` [`MessagePart`]: + /// the `tool_call_id` paired with the tool's output (wrapped in a + /// [`ToolContent`](ToolContent)) and an `is_error` flag so the model + /// can distinguish successes from failures. The caller is responsible + /// for assembling these parts — alongside any preresolved results — + /// into the single user [`Message`] that represents the turn. /// /// # Parameters /// /// - `results` — The [`ToolDispatchResult`]s produced by /// [`dispatch_tools()`](BareLoop::dispatch_tools). - /// - /// # Returns - /// - /// A [`Message`] with [`Role::User`] and one `tool_result` - /// [`MessagePart`] per result. - pub(super) fn build_tool_result_message(results: Vec) -> Message { - let parts: Vec = results + pub(super) fn build_tool_result_parts(results: Vec) -> Vec { + results .into_iter() .map(|r| MessagePart::tool_result(r.tool_call_id, r.output, r.is_error)) - .collect(); - Message::new(Role::User, parts) + .collect() } /// Build tool schemas for the API request. From a0263d0e507961c996c8f4aff2f5a4a3b2db0ceb Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 27 Jul 2026 18:12:16 +1200 Subject: [PATCH 10/21] refactor: fallback manager, fix TOCTOU races --- CHANGELOG.md | 34 + Makefile | 14 +- src/capabilities.rs | 2 +- src/engine/bare.rs | 10 +- src/fallback.rs | 1460 +++++++++++++++++-------------------------- src/managers.rs | 4 +- 6 files changed, 628 insertions(+), 896 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8017d49..bf64f88 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -373,9 +373,43 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. - Internally-built `reqwest::Client`s now set `tcp_nodelay(true)` by default. SSE streaming emits many small packets; disabling Nagle's algorithm reduces per-delta latency. No correctness impact. +- `FallbackManager`'s circuit-breaker state is now unified behind a single + `Mutex`. Previously it was split across four independent + `Relaxed` atomics (`fallback_state`, `consecutive_failures`, + `primary_success_count`, `fallback_activated`) plus a `Mutex`, + which left a TOCTOU window between the state read and the + transition/counter-reset. The read-decide-transition in `record_failure` / + `record_success` now runs while holding the one lock, so the race is closed + by construction. Internal refactor — no public-API change beyond the method + merges noted under `### Removed`. ### Removed +- `FallbackManager::record_api_failure` and + `FallbackManager::record_model_failure` — merged into a single + [`FallbackManager::record_failure(FailureKind)`](crate::fallback::FallbackManager::record_failure). + The two methods differed only in how a failure during + [`Recovering`](crate::fallback::FallbackState::Recovering) was treated: + a sustained rate-limit re-trips the breaker, a transient error leaves + the half-open probe in place. That distinction is now an explicit + [`FailureKind`](crate::fallback::FailureKind) argument (`RateLimit` vs + `Transient`) instead of two near-identical methods. Migration: pass + `FailureKind::RateLimit` where you called `record_model_failure`, + `FailureKind::Transient` where you called `record_api_failure`. +- `FallbackManager::record_failure` / `record_success` zero-body aliases + for `record_api_failure` / `record_model_success` — removed alongside + the merge. `record_success` is the new canonical name for the success + path (was `record_model_success`). +- `FallbackEntry` and `AttemptRecord` are now `pub(crate)` — they were + `pub` with no external users and ~400 lines of speculative accessors. + The `fallback_entry(name)` lookup method (which leaked the internal + type) is removed. These types were never part of the documented public + API surface; only `FallbackManager`, `FallbackState`, `FallbackConfig`, + and `FailureKind` are public in the `fallback` module now. +- `FallbackState::From` impl and the `= 0`/`= 1`/`= 2` discriminants + — dead weight from when state round-tripped through an `AtomicU8`. + State is now a plain field on an internal struct; no `u8` casting + remains. - `Loop::process_turn` trait method and `BareLoop::run_turn_body` — the machine-driven `run()` replaces the old per-turn execution path. The `LoopMachine` is the new turn unit; drive it via `BareLoop::run()` (or diff --git a/Makefile b/Makefile index 41be7e1..ce184f7 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: check test clippy fmt docs ci lint examples +.PHONY: check test clippy fmt docs ci lint examples e2e e2e-providers e2e-ollama ci: fmt check clippy test docs examples @@ -23,3 +23,15 @@ docs: examples: cargo build --examples --all-features + +e2e: e2e-providers e2e-ollama + +e2e-providers: + LOOPCTL_E2E=1 cargo test --features ollama,openai,anthropic,gemini,grok,deepseek,zai --test provider_e2e -- --nocapture --test-threads=1 + +e2e-ollama: + @test -n "$(OLLAMA_MODEL)" || { echo "ERROR: set OLLAMA_MODEL (e.g. make e2e-ollama OLLAMA_MODEL=qwen2.5:7b)"; exit 1; } + LOOPCTL_E2E=1 cargo test --features ollama,grammar --test constrained_decode -- --nocapture + LOOPCTL_E2E=1 cargo test --features ollama --test examples_e2e -- --nocapture + LOOPCTL_E2E=1 cargo test --features ollama --test provider_survival -- --nocapture + LOOPCTL_E2E=1 cargo test --features ollama --test structured_output -- --nocapture diff --git a/src/capabilities.rs b/src/capabilities.rs index a6e2102..219a598 100644 --- a/src/capabilities.rs +++ b/src/capabilities.rs @@ -122,7 +122,7 @@ pub trait Detectable { /// /// ```rust,ignore /// fn handle_stream_error(runtime: &impl FallbackCapable) { -/// let tripped = runtime.fallback().record_api_failure(); +/// let tripped = runtime.fallback().record_failure(loopctl::fallback::FailureKind::Transient); /// if tripped { /// if let Some(model) = runtime.fallback().fallback_model() { /// // switch to fallback model diff --git a/src/engine/bare.rs b/src/engine/bare.rs index d82357c..029014d 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -1135,7 +1135,7 @@ impl BareLoop { /// `(Message, Option, StreamStopReason)` and just needs the /// side-effects. fn record_stream_success(&mut self, usage: Option<&Usage>) { - self.managers.fallback().record_model_success(); + self.managers.fallback().record_success(); let (in_tok, out_tok) = Self::usage_tokens(usage); self.managers.observers().on_stream_success(&StreamContext { @@ -1173,9 +1173,13 @@ impl BareLoop { /// handles cancellation before reaching [`do_stream`](Self::do_stream). fn record_stream_failure(&mut self, e: LoopError) -> LoopError { let tripped = if matches!(e, LoopError::RateLimitEscalation { .. }) { - self.managers.fallback().record_model_failure() + self.managers + .fallback() + .record_failure(crate::fallback::FailureKind::RateLimit) } else { - self.managers.fallback().record_api_failure() + self.managers + .fallback() + .record_failure(crate::fallback::FailureKind::Transient) }; if tripped { diff --git a/src/fallback.rs b/src/fallback.rs index 9d1b43a..29c1f35 100644 --- a/src/fallback.rs +++ b/src/fallback.rs @@ -19,31 +19,30 @@ //! //! - **[`FallbackState`]** — The three circuit-breaker states (`Primary`, `Fallback`, `Recovering`). //! - **[`FallbackConfig`]** — Configuration struct with thresholds and timeouts. -//! - **[`FallbackManager`]** — The state machine itself; thread-safe via atomics and `Mutex`. +//! - **[`FallbackManager`]** — The state machine itself; thread-safe via a single `Mutex`. //! //! # Quick Start //! //! ```rust -//! use loopctl::fallback::{FallbackManager, FallbackConfig}; +//! use loopctl::fallback::{FallbackManager, FallbackConfig, FailureKind}; //! //! // Create a manager with a trip threshold of 3 failures //! let mgr = FallbackManager::new(3, 2); //! //! // Simulate failures until the circuit trips -//! mgr.record_model_failure(); -//! mgr.record_model_failure(); -//! assert!(mgr.record_model_failure()); // 3rd failure trips the circuit +//! mgr.record_failure(FailureKind::Transient); +//! mgr.record_failure(FailureKind::Transient); +//! assert!(mgr.record_failure(FailureKind::Transient)); // 3rd failure trips the circuit //! assert!(mgr.is_using_fallback()); //! //! // Manually transition to recovering, then record successes to resume primary //! mgr.transition_to_recovering(); -//! mgr.record_model_success(); // 1st success -//! mgr.record_model_success(); // 2nd → back to Primary +//! mgr.record_success(); // 1st success +//! mgr.record_success(); // 2nd → back to Primary //! assert!(!mgr.is_using_fallback()); //! ``` use std::sync::Mutex; -use std::sync::atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering}; use std::time::{Duration, Instant}; use tracing::{debug, info, warn}; @@ -66,452 +65,171 @@ use tracing::{debug, info, warn}; /// use loopctl::fallback::FallbackState; /// /// let state = FallbackState::Primary; -/// assert_eq!(state as u8, 0); -/// assert_eq!(FallbackState::from(1), FallbackState::Fallback); +/// assert_eq!(state, FallbackState::Primary); /// ``` #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum FallbackState { - /// Operating on the primary model — no failures have tripped the - /// circuit breaker yet. - Primary = 0, - - /// A fallback model is active — the primary model failed and the - /// breaker tripped. Subsequent failures on the fallback model are - /// tracked separately. - Fallback = 1, - - /// 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, + /// The circuit is closed and the primary model is in use. + /// + /// This is the initial state and the steady-state target: requests + /// route to the primary model, the failure counter increments on + /// each failure, and a single success resets it. When + /// [`consecutive_failures`](FallbackManager::consecutive_failures) + /// reaches [`trip_threshold`](FallbackConfig::trip_threshold), the + /// manager transitions to [`Fallback`](Self::Fallback). + Primary, + + /// The circuit is open and a fallback model is in use. + /// + /// Entered when the primary model's consecutive failures exceeded + /// the trip threshold. Requests route to the active fallback model + /// (the first non-failed entry in the chain). The manager stays here + /// for at least + /// [`recovery_timeout`](FallbackConfig::recovery_timeout) before it + /// is willing to probe the primary again, at which point it + /// transitions to [`Recovering`](Self::Recovering). + Fallback, + + /// Half-open: the manager is probing whether the primary has recovered. + /// + /// Entered from [`Fallback`](Self::Fallback) after the cooldown + /// elapses (see + /// [`should_try_resume_primary`](FallbackManager::should_try_resume_primary)). + /// Requests route back to the primary while a recovery success + /// counter accrues; reaching + /// [`recovery_successes_needed`](FallbackConfig::recovery_successes_needed) + /// successes closes the circuit back to [`Primary`](Self::Primary). + /// A single sustained failure ([`FailureKind::RateLimit`]) + /// during this probe re-opens the circuit to [`Fallback`](Self::Fallback). + Recovering, } -/// Converts a raw `u8` back into a [`FallbackState`]. +/// What kind of failure triggered a [`FallbackManager::record_failure`] call. /// -/// 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 value cannot -/// cause a panic. -/// -/// # Example -/// -/// ```rust -/// use loopctl::fallback::FallbackState; -/// -/// assert_eq!(FallbackState::from(0u8), FallbackState::Primary); -/// assert_eq!(FallbackState::from(1u8), FallbackState::Fallback); -/// assert_eq!(FallbackState::from(2u8), FallbackState::Recovering); -/// assert_eq!(FallbackState::from(255u8), FallbackState::Primary); // unknown → safe default -/// ``` -impl From for FallbackState { - fn from(value: u8) -> Self { - match value { - 1 => FallbackState::Fallback, - 2 => FallbackState::Recovering, - _ => FallbackState::Primary, - } - } +/// The engine routes different failure causes through one recorder, and the +/// cause changes how a failure during [`FallbackState::Recovering`] is +/// treated: a sustained rate-limit means the primary is still degraded, so +/// the circuit re-trips to [`FallbackState::Fallback`]; a transient error +/// leaves the half-open probe in place. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FailureKind { + /// A transient API error (timeout, 5xx, network). + /// + /// During [`Recovering`](FallbackState::Recovering) the circuit stays + /// half-open — one transient blip does not by itself prove the primary + /// is still bad. + Transient, + + /// A sustained rate-limit escalation. + /// + /// During [`Recovering`](FallbackState::Recovering) the circuit + /// re-trips to [`Fallback`](FallbackState::Fallback): the primary is + /// still actively refusing load, so probing it further is pointless. + RateLimit, } -/// A single failed attempt on a fallback model. -/// -/// Records *when* a failure occurred and an optional reason string -/// (e.g. `"rate_limit"`, `"timeout"`, `"500 Internal Server Error"`). -/// Entries accumulate in [`FallbackEntry::attempts`]; once the count -/// reaches [`FallbackEntry::max_fail_count`], the model is considered -/// failed and skipped by [`FallbackManager::fallback_model`]. +/// A single model in the fallback chain, with a failure counter. /// -/// # Example +/// Each entry pairs a model name with a consecutive-failure count and a +/// `max_fail_count` threshold. When [`attempt_count`](Self::attempt_count) +/// reaches `max_fail_count`, the entry is considered [`failed`](Self::failed) +/// and is skipped by [`FallbackManager::fallback_model`]. /// -/// ```rust -/// use loopctl::fallback::AttemptRecord; -/// -/// let record = AttemptRecord::new("rate_limit"); -/// assert_eq!(record.reason(), Some("rate_limit")); -/// ``` +/// Internal to the fallback manager — not part of the public API. #[derive(Debug, Clone)] -pub struct AttemptRecord { - failed_at: Instant, - reason: Option, -} - -impl AttemptRecord { - /// Create a new attempt record with the given reason. - /// - /// # Example - /// - /// ```rust - /// use loopctl::fallback::AttemptRecord; - /// let record = AttemptRecord::new("rate_limit"); - /// assert_eq!(record.reason(), Some("rate_limit")); - /// ``` - #[must_use] - pub fn new(reason: impl Into) -> Self { - Self { - failed_at: Instant::now(), - reason: Some(reason.into()), - } - } - - /// Create a new attempt record without a reason. - /// - /// # Example - /// - /// ```rust - /// use loopctl::fallback::AttemptRecord; - /// let record = AttemptRecord::anonymous(); - /// assert!(record.reason().is_none()); - /// ``` - #[must_use] - pub fn anonymous() -> Self { - Self { - failed_at: Instant::now(), - reason: None, - } - } +pub(crate) struct FallbackEntry { + /// The model identifier this entry represents. + /// + /// Must match the routing identifier the API client expects (e.g. + /// `"llm-70b"`, `"gpt-4o"`). Used both to look the entry up in the + /// chain and as the value returned by + /// [`fallback_model`](FallbackManager::fallback_model) when this + /// entry is the active fallback. Set at construction and never + /// mutated. + name: String, - /// Set the reason for this failure record. - /// - /// Pass `None` for an anonymous record, or `Some("reason")` for - /// a labelled one. + /// Whether this model is eligible to serve requests. /// - /// # Example - /// - /// ```rust - /// 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(mut self, reason: Option) -> Self { - self.reason = reason; - self - } + /// An out-of-band kill switch separate from failure tracking: + /// `false` forces [`failed`](Self::failed) to return `true` + /// regardless of the attempt count, so the manager skips this + /// entry. Use [`set_available`](Self::set_available) to take a + /// model out of rotation (e.g. an API key was revoked, the model + /// was decommissioned) without recording spurious failures. + available: bool, - /// When this failure was recorded. + /// How many consecutive failures this model has recorded. /// - /// # Example - /// - /// ```rust - /// use loopctl::fallback::AttemptRecord; - /// let record = AttemptRecord::new("timeout"); - /// // failed_at is close to now - /// assert!(record.failed_at().elapsed().as_secs() < 1); - /// ``` - #[must_use] - pub fn failed_at(&self) -> Instant { - self.failed_at - } + /// Incremented by [`record_attempt`](Self::record_attempt) and + /// compared against [`max_fail_count`](Self::max_fail_count) to + /// decide [`failed`](Self::failed). Reset to zero by + /// [`clear_attempts`](Self::clear_attempts) when the circuit + /// recovers, giving a degraded model a fresh start. + attempt_count: usize, - /// The optional reason for this failure. - /// - /// # Example + /// The failure count at which this model is taken out of rotation. /// - /// ```rust - /// use loopctl::fallback::AttemptRecord; - /// let record = AttemptRecord::new("rate_limit"); - /// assert_eq!(record.reason(), Some("rate_limit")); - /// ``` - #[must_use] - pub fn reason(&self) -> Option<&str> { - self.reason.as_deref() - } -} - -/// A single model in the fallback chain, with attempt history. -/// -/// Each entry has a model name, a list of recorded failure attempts, -/// and a `max_fail_count` threshold. When [`attempts`](Self::attempts) -/// grows to `max_fail_count` entries, the entry is considered -/// [`failed`](Self::failed) and is automatically skipped by -/// [`FallbackManager::fallback_model`]. -/// -/// # Example -/// -/// ```rust -/// use loopctl::fallback::FallbackEntry; -/// -/// let mut entry = FallbackEntry::new("llm-70b"); -/// assert_eq!(entry.name(), "llm-70b"); -/// assert!(!entry.failed()); -/// -/// entry.record_attempt("timeout"); -/// entry.record_attempt("timeout"); -/// assert_eq!(entry.attempt_count(), 2); -/// assert!(entry.failed()); // max_fail_count defaults to 2 -/// ``` -#[derive(Debug, Clone)] -pub struct FallbackEntry { - /// Model identifier (e.g. `"llm-70b"`). Must match the API client's routing identifier. - name: String, - /// Set to `false` to take a model out of rotation independently of failure tracking. - available: bool, - /// Recorded failure attempts for this model. - attempts: Vec, - /// When `attempts.len()` reaches this threshold the model is taken out of rotation. + /// Once [`attempt_count`](Self::attempt_count) reaches this value + /// the entry is [`failed`](Self::failed) and the manager advances + /// to the next candidate in the chain. Defaults to the manager's + /// [`max_fail_count`](FallbackConfig::max_fail_count) at + /// construction; override per-entry with + /// [`with_max_fail_count`](Self::with_max_fail_count). max_fail_count: usize, } impl FallbackEntry { /// Create a new entry for the given model, not yet failed. /// - /// Uses a default `max_fail_count` of `2` - /// failure marks the model as failed. Use + /// Uses a default `max_fail_count` of `2`. Use /// [`with_max_fail_count`](Self::with_max_fail_count) to customise. - /// - /// # Example - /// - /// ```rust - /// use loopctl::fallback::FallbackEntry; - /// let entry = FallbackEntry::new("llm-70b"); - /// assert_eq!(entry.name(), "llm-70b"); - /// assert!(!entry.failed()); - /// assert_eq!(entry.max_fail_count(), 2); - /// ``` - #[must_use] - pub fn new(name: impl Into) -> Self { + pub(crate) fn new(name: impl Into) -> Self { Self { name: name.into(), available: true, - attempts: Vec::new(), + attempt_count: 0, max_fail_count: 2, } } /// 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). - /// - /// # Example - /// - /// ```rust - /// 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"); - /// entry.record_attempt("rate_limit"); - /// assert!(!entry.failed()); // only 2 of 3 - /// - /// entry.record_attempt("server_error"); - /// assert!(entry.failed()); // 3 of 3 - /// ``` + /// If raising the threshold on an already-failed entry, the attempt + /// count is padded up so it stays failed under the new threshold. #[must_use] - pub fn with_max_fail_count(mut self, max_fail_count: usize) -> Self { + pub(crate) fn with_max_fail_count(mut self, max_fail_count: usize) -> Self { 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()); + if self.failed() && self.attempt_count < new_max { + self.attempt_count = new_max; } self.max_fail_count = new_max; self } - /// Create a new entry already marked as failed. - /// - /// Useful when initializing from a known-degraded model. - /// [`failed()`](Self::failed) returns `true` immediately. - /// - /// # Example - /// - /// ```rust - /// use loopctl::fallback::FallbackEntry; - /// let entry = FallbackEntry::new_failed("llm-70b"); - /// assert!(entry.failed()); - /// ``` - #[must_use] - pub fn new_failed(name: impl Into) -> Self { - Self { - name: name.into(), - available: true, - attempts: vec![AttemptRecord::anonymous(), AttemptRecord::anonymous()], - max_fail_count: 2, - } - } - - /// The model name. - /// - /// # Example - /// - /// ```rust - /// use loopctl::fallback::FallbackEntry; - /// let entry = FallbackEntry::new("llm-70b"); - /// assert_eq!(entry.name(), "llm-70b"); - /// ``` - #[must_use] - pub fn name(&self) -> &str { - &self.name - } - /// Whether this model should be skipped. /// - /// Returns `true` when the model is not [`available`](Self::available) - /// or when [`attempt_count`](Self::attempt_count) has reached - /// [`max_fail_count`](Self::max_fail_count). - /// - /// # Example - /// - /// ```rust - /// use loopctl::fallback::FallbackEntry; - /// let mut entry = FallbackEntry::new("llm-70b"); - /// assert!(!entry.failed()); - /// entry.record_attempt("timeout"); - /// entry.record_attempt("timeout"); // max_fail_count defaults to 2 - /// assert!(entry.failed()); - /// ``` - #[must_use] - pub fn failed(&self) -> bool { - !self.available || self.attempts.len() >= self.max_fail_count - } - - /// The configured maximum failure count before this model is skipped. - /// - /// # Example - /// - /// ```rust - /// use loopctl::fallback::FallbackEntry; - /// let entry = FallbackEntry::new("llm-70b").with_max_fail_count(3); - /// assert_eq!(entry.max_fail_count(), 3); - /// ``` - #[must_use] - pub fn max_fail_count(&self) -> usize { - self.max_fail_count - } - - /// Whether this model is available for use. - /// - /// A model that is not available is always skipped by - /// [`FallbackManager::fallback_model`], regardless of failure count. - /// Set to `false` via [`set_available`](Self::set_available) to take - /// a model out of rotation (e.g. API key revoked, model decommissioned). - /// - /// # Example - /// - /// ```rust - /// use loopctl::fallback::FallbackEntry; - /// let mut entry = FallbackEntry::new("llm-70b"); - /// assert!(entry.available()); - /// entry.set_available(false); - /// assert!(!entry.available()); - /// assert!(entry.failed()); // unavailable ⇒ failed - /// ``` - #[must_use] - pub fn available(&self) -> bool { - self.available + /// `true` when the model is not `available` or when its attempt count + /// has reached `max_fail_count`. + pub(crate) fn failed(&self) -> bool { + !self.available || self.attempt_count >= self.max_fail_count } /// Set whether this model is available for use. /// - /// When set to `false`, [`failed()`](Self::failed) returns `true` + /// When set to `false`, [`failed`](Self::failed) returns `true` /// regardless of the attempt count, and the manager skips this model. - /// - /// # Example - /// - /// ```rust - /// use loopctl::fallback::FallbackEntry; - /// let mut entry = FallbackEntry::new("llm-70b"); - /// entry.set_available(false); - /// assert!(entry.failed()); - /// entry.set_available(true); - /// assert!(!entry.failed()); // no attempts recorded - /// ``` - pub fn set_available(&mut self, available: bool) { + pub(crate) fn set_available(&mut self, available: bool) { self.available = available; } - /// How many failure attempts have been recorded. - /// - /// # Example - /// - /// ```rust - /// use loopctl::fallback::FallbackEntry; - /// let mut entry = FallbackEntry::new("llm-70b"); - /// assert_eq!(entry.attempt_count(), 0); - /// entry.record_attempt("timeout"); - /// assert_eq!(entry.attempt_count(), 1); - /// ``` - #[must_use] - pub fn attempt_count(&self) -> usize { - self.attempts.len() - } - - /// Access the recorded failure attempts. - /// - /// Returns a slice of [`AttemptRecord`] in chronological order - /// (oldest first). - /// - /// # Example - /// - /// ```rust - /// use loopctl::fallback::FallbackEntry; - /// let mut entry = FallbackEntry::new("llm-70b"); - /// entry.record_attempt("timeout"); - /// entry.record_attempt("rate_limit"); - /// let attempts = entry.attempts(); - /// assert_eq!(attempts.len(), 2); - /// assert_eq!(attempts[0].reason(), Some("timeout")); - /// ``` - #[must_use] - pub fn attempts(&self) -> &[AttemptRecord] { - &self.attempts - } - - /// Record a new failure attempt with an optional reason. - /// - /// 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::fallback::FallbackEntry; - /// let mut entry = FallbackEntry::new("llm-70b"); - /// entry.record_attempt("timeout"); - /// entry.record_attempt("timeout"); // exceeds max_fail_count (default = 2) - /// assert!(entry.failed()); - /// ``` - pub fn record_attempt(&mut self, reason: impl Into) { - self.attempts.push(AttemptRecord::new(reason)); + /// Record a new failure attempt. + pub(crate) fn record_attempt(&mut self) { + self.attempt_count = self.attempt_count.saturating_add(1); } - /// Record a new failure attempt without a reason. - /// - /// # Example - /// - /// ```rust - /// 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) - /// assert!(entry.failed()); - /// ``` - pub fn record_attempt_anonymous(&mut self) { - self.attempts.push(AttemptRecord::anonymous()); - } - - /// Clear all recorded attempts, resetting [`failed()`](Self::failed) - /// to `false`. - /// - /// # Example - /// - /// ```rust - /// 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) - /// assert!(entry.failed()); - /// entry.clear_attempts(); - /// assert!(!entry.failed()); - /// ``` - pub fn clear_attempts(&mut self) { - self.attempts.clear(); + /// Clear all recorded attempts, resetting [`failed`](Self::failed) to `false`. + pub(crate) fn clear_attempts(&mut self) { + self.attempt_count = 0; } } @@ -570,7 +288,7 @@ pub struct FallbackConfig { /// Per-model failure threshold before a fallback model is skipped. Defaults to `2`. /// - /// Applied to each [`FallbackEntry`] in the fallback chain: once a single + /// Applied to each fallback model in the chain: once a single /// model accumulates this many recorded attempts, it is taken out of /// rotation and [`FallbackManager::active_model`] advances to the next /// candidate. @@ -611,33 +329,68 @@ impl Default for FallbackConfig { /// # Example /// /// ```rust -/// use loopctl::fallback::{FallbackManager, FallbackConfig}; +/// use loopctl::fallback::{FallbackManager, FallbackConfig, FailureKind}; /// use std::sync::Arc; /// use std::time::Duration; /// -/// let mgr = Arc::new(FallbackManager::new(5, 3).with_config(&FallbackConfig::default())); +/// let mgr = Arc::new(FallbackManager::new(5, 3).with_config(FallbackConfig::default())); /// /// // Simulate failures -/// assert!(!mgr.record_model_failure()); // 1 -/// assert!(!mgr.record_model_failure()); // 2 -/// assert!(mgr.record_model_failure()); // 3 → circuit trips +/// assert!(!mgr.record_failure(FailureKind::Transient)); // 1 +/// assert!(!mgr.record_failure(FailureKind::Transient)); // 2 +/// assert!(mgr.record_failure(FailureKind::Transient)); // 3 → circuit trips /// /// assert!(mgr.is_using_fallback()); /// /// // ... later, after cooldown ... /// if mgr.should_try_resume_primary(Duration::from_secs(60)) { /// mgr.transition_to_recovering(); -/// mgr.record_model_success(); -/// mgr.record_model_success(); // → back to Primary +/// mgr.record_success(); +/// mgr.record_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 { +/// The complete mutable circuit-breaker state. +/// +/// Every field that can change over the lifetime of a +/// [`FallbackManager`] lives here, behind a single +/// [`Mutex`](std::sync::Mutex). Holding the whole state under one lock +/// makes each public method's read-modify-write atomic with respect to +/// every other method: there is no window in which one caller can +/// observe a half-applied transition or race a counter reset. Plain +/// fields, no atomics — the lock is the only synchronization. +struct BreakerState { + /// Current circuit-breaker phase. + /// + /// `Primary` at construction; transitions through `Fallback` and + /// `Recovering` as failures and recoveries are recorded. + state: FallbackState, + + /// Consecutive failures on the current model. + /// + /// Incremented by the `record_*_failure` methods while in `Primary`; + /// reaching [`FallbackManager::fallback_threshold`] trips the + /// circuit. Reset to `0` on trip and on any success in `Primary`. + consecutive_failures: usize, + + /// Consecutive successes on the primary accumulated during + /// `Recovering`. + /// + /// Incremented by [`record_success`](FallbackManager::record_success) + /// while in `Recovering`; reaching + /// [`primary_resume_threshold`](FallbackManager::primary_resume_threshold) + /// closes the circuit back to `Primary`. Reset on entry to + /// `Recovering` and on trip. + primary_success_count: usize, + + /// Sticky flag recording whether fallback has ever been activated. + /// + /// `true` once the circuit has tripped at least once, and never + /// reset for the life of the manager. Lets observers tell "still on + /// primary, never failed" from "recovered back to primary". + fallback_activated: bool, + /// Original model name, before any fallback was activated. /// /// Captured when the manager first trips so the recovering state @@ -658,7 +411,7 @@ struct FallbackInner { /// Computed once and reused while the manager is in the fallback /// state, so callers asking "which model am I on now?" get an /// `O(1)` answer without re-scanning `fallback_models`. Recomputed - /// whenever the set of failed entries changes. + /// in place whenever the set of failed entries changes. active_fallback: Option, /// Instant at which fallback was activated. @@ -671,86 +424,44 @@ struct FallbackInner { fallback_switched_at: Option, } +impl Default for BreakerState { + fn default() -> Self { + Self { + state: FallbackState::Primary, + consecutive_failures: 0, + primary_success_count: 0, + fallback_activated: false, + original_model: None, + fallback_models: Vec::new(), + active_fallback: None, + fallback_switched_at: None, + } + } +} + /// 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 { - /// Number of consecutive failures on the primary that trips the - /// circuit. - /// - /// Once [`consecutive_failures`](Self::consecutive_failures) reaches - /// this value the manager transitions from `Primary` to `Fallback`. - /// Set at construction and immutable thereafter. - fallback_threshold: usize, - - /// Number of consecutive successes required on the primary during - /// recovery before transitioning back. - /// - /// Counts towards this via - /// [`record_model_success`](FallbackManager::record_model_success) - /// while in the `Recovering` state; reaching it transitions to - /// `Primary`. Set at construction and immutable thereafter. - primary_resume_threshold: usize, - - /// Per-model max failure count seeded into new - /// [`FallbackEntry`] instances. - /// - /// When a fallback model is added without an explicit cap, it - /// inherits this value as its `max_fail_count`. Set at construction - /// and immutable thereafter. - default_max_fail_count: usize, - - /// Consecutive API failure counter on the current model. - /// - /// Incremented by - /// [`record_model_failure`](FallbackManager::record_model_failure), - /// reset to zero on any success. Hitting `fallback_threshold` trips - /// the circuit. Atomic so it updates lock-free on the hot path. - consecutive_failures: AtomicUsize, - - /// Sticky flag recording whether fallback has ever been activated. - /// - /// `true` once the circuit has tripped at least once, and never - /// reset for the life of the manager. Lets observers tell "still - /// on primary, never failed" from "recovered back to primary". - fallback_activated: AtomicBool, - - /// Current circuit-breaker state, encoded as an atomic. - /// - /// `0` = `Primary`, `1` = `Fallback`, `2` = `Recovering` (see - /// [`FallbackState`]). Stored as `AtomicU8` for lock-free reads on - /// every request; mutated only on state transitions. - fallback_state: AtomicU8, - - /// Consecutive successes on the primary accumulated during the - /// `Recovering` state. - /// - /// Incremented by - /// [`record_model_success`](FallbackManager::record_model_success); - /// reaching `primary_resume_threshold` transitions back to - /// `Primary` and resets this to zero. Atomic so it updates - /// lock-free. - primary_success_count: AtomicUsize, - - /// 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. - /// - /// Compared against the elapsed time since - /// [`fallback_switched_at`](FallbackInner::fallback_switched_at) - /// by - /// [`should_try_resume_primary`](FallbackManager::should_try_resume_primary) - /// to decide when a cooldown has elapsed. Set at construction and - /// immutable thereafter. - recovery_timeout: Duration, + /// Immutable configuration: thresholds and timeouts. + /// + /// Single source of truth for the trip threshold, recovery + /// threshold, per-model max-fail count, and recovery timeout — + /// mirrors the [`DetectionManager::config`](crate::detection::DetectionManager::config) + /// pattern. Held by value, outside the state lock. + config: FallbackConfig, + + /// All mutable circuit-breaker state, behind a single lock. + /// + /// Holding every changeable field ([`BreakerState`]) under one + /// [`Mutex`](std::sync::Mutex) makes each public method's + /// read-modify-write atomic relative to every other method — no + /// window for a partial-state read or a counter/transition race. + /// Immutable config ([`config`](Self::config)) stays outside the + /// lock. + state: Mutex, } impl FallbackManager { @@ -772,7 +483,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new(5, 3); /// // Trip after 5 failures, resume after 3 consecutive successes @@ -780,27 +491,25 @@ impl FallbackManager { #[must_use] pub fn new(fallback_threshold: usize, primary_resume_threshold: usize) -> Self { Self { - fallback_threshold, - primary_resume_threshold, - default_max_fail_count: 2, - consecutive_failures: AtomicUsize::new(0), - fallback_activated: AtomicBool::new(false), - fallback_state: AtomicU8::new(FallbackState::Primary as u8), - primary_success_count: AtomicUsize::new(0), - inner: Mutex::new(FallbackInner::default()), - recovery_timeout: Duration::from_mins(1), + config: FallbackConfig { + trip_threshold: fallback_threshold, + recovery_successes_needed: primary_resume_threshold, + ..FallbackConfig::default() + }, + state: Mutex::new(BreakerState::default()), } } /// Apply configuration from a [`FallbackConfig`] struct. /// - /// Sets the failure threshold, recovery parameters, and per-model - /// max fail count from the config. + /// Replaces the manager's entire config. Prefer + /// [`new_with_config`](Self::new_with_config) for config-driven + /// construction; use this builder when chaining off [`new`](Self::new). /// /// # Example /// /// ```rust - /// use loopctl::fallback::{FallbackManager, FallbackConfig}; + /// use loopctl::fallback::{FallbackManager, FallbackConfig, FailureKind}; /// use std::time::Duration; /// /// let config = FallbackConfig { @@ -809,45 +518,55 @@ impl FallbackManager { /// recovery_successes_needed: 3, /// max_fail_count: 2, /// }; - /// let mgr = FallbackManager::new(5, 3).with_config(&config); + /// let mgr = FallbackManager::new(5, 3).with_config(config); /// ``` #[must_use] - 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.recovery_timeout = config.recovery_timeout; + pub fn with_config(mut self, config: FallbackConfig) -> Self { + self.config = config; self } - /// Set the recovery timeout (builder style). + /// The immutable [`FallbackConfig`] this manager was constructed with. /// - /// 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). + /// Read thresholds and timeouts (trip threshold, recovery timeout, + /// per-model max-fail count, recovery successes needed) from here + /// rather than from duplicated accessors. /// - /// Mirrors [`FallbackConfig::recovery_timeout`] for cases where a full - /// [`with_config`](Self::with_config) is not desired. + /// # Example + /// + /// ```rust + /// use loopctl::fallback::{FallbackManager, FallbackConfig}; + /// use std::time::Duration; + /// + /// let cfg = FallbackConfig { trip_threshold: 5, ..FallbackConfig::default() }; + /// let mgr = FallbackManager::new_with_config(cfg); + /// assert_eq!(mgr.config().trip_threshold, 5); + /// ``` #[must_use] - pub fn with_recovery_timeout(mut self, recovery_timeout: Duration) -> Self { - self.recovery_timeout = recovery_timeout; - self + pub fn config(&self) -> &FallbackConfig { + &self.config } - /// Configured recovery timeout. + /// Create a new manager directly from a [`FallbackConfig`]. /// - /// 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). + /// The config-driven counterpart of [`new`](Self::new): build the + /// whole config struct (typically from a file or env) and construct + /// in one step instead of overriding fields after the fact. + /// + /// # Example + /// + /// ```rust + /// use loopctl::fallback::{FallbackManager, FallbackConfig}; /// - /// Pass this to - /// [`should_try_resume_primary`](Self::should_try_resume_primary) to - /// honour the configured timeout without hard-coding a value. + /// let cfg = FallbackConfig::default(); + /// let mgr = FallbackManager::new_with_config(cfg); + /// ``` #[must_use] - pub fn recovery_timeout(&self) -> Duration { - self.recovery_timeout + pub fn new_with_config(config: FallbackConfig) -> Self { + Self { + config, + state: Mutex::new(BreakerState::default()), + } } /// Create with fallback already activated. @@ -872,7 +591,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new_with_fallback("llm-70b".into(), 3); /// assert!(mgr.is_using_fallback()); @@ -881,14 +600,10 @@ 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 inner) = mgr.inner.lock() { - inner.original_model = Some(original_model); - inner.fallback_switched_at = Some(Instant::now()); + if let Ok(mut state) = mgr.state.lock() { + state.original_model = Some(original_model); + Self::transition_to_fallback_impl(&mut state); } - mgr.fallback_activated.store(true, Ordering::Relaxed); - mgr.consecutive_failures.store(0, Ordering::Relaxed); - mgr.fallback_state - .store(FallbackState::Fallback as u8, Ordering::Relaxed); mgr } @@ -902,7 +617,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::for_model("llm-70b"); /// assert_eq!(mgr.original_model(), Some("llm-70b".to_string())); @@ -910,8 +625,8 @@ impl FallbackManager { /// ``` pub fn for_model(primary_model: impl Into) -> Self { let mgr = Self::new(3, 2); - if let Ok(mut inner) = mgr.inner.lock() { - inner.original_model = Some(primary_model.into()); + if let Ok(mut state) = mgr.state.lock() { + state.original_model = Some(primary_model.into()); } mgr } @@ -921,13 +636,15 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// assert_eq!(mgr.state(), FallbackState::Primary); /// ``` pub fn state(&self) -> FallbackState { - FallbackState::from(self.fallback_state.load(Ordering::Relaxed)) + self.state + .lock() + .map_or(FallbackState::Primary, |s| s.state) } /// Check if we're currently using a fallback model. @@ -941,7 +658,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState, FailureKind}; /// let mgr = FallbackManager::new(3, 2); /// assert!(!mgr.is_using_fallback()); /// ``` @@ -960,12 +677,12 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// let mgr = FallbackManager::new(3, 2); /// assert!(!mgr.is_fallback_active()); /// ``` pub fn is_fallback_active(&self) -> bool { - self.fallback_activated.load(Ordering::Relaxed) + self.state.lock().is_ok_and(|s| s.fallback_activated) } /// Get the number of consecutive failures. @@ -976,14 +693,14 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// let mgr = FallbackManager::new(3, 2); /// assert_eq!(mgr.consecutive_failures(), 0); - /// mgr.record_model_failure(); + /// mgr.record_failure(FailureKind::Transient); /// assert_eq!(mgr.consecutive_failures(), 1); /// ``` pub fn consecutive_failures(&self) -> usize { - self.consecutive_failures.load(Ordering::Relaxed) + self.state.lock().map_or(0, |s| s.consecutive_failures) } /// Get the original model name (before fallback). @@ -997,15 +714,15 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// let mgr = FallbackManager::for_model("llm-70b"); /// assert_eq!(mgr.original_model(), Some("llm-70b".to_string())); /// ``` pub fn original_model(&self) -> Option { - self.inner + self.state .lock() .ok() - .and_then(|i| i.original_model.clone()) + .and_then(|s| s.original_model.clone()) } /// Set the original model name. @@ -1016,15 +733,15 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// let mgr = FallbackManager::new(3, 2); /// assert_eq!(mgr.original_model(), None); /// mgr.set_original_model("llm-70b".into()); /// assert_eq!(mgr.original_model(), Some("llm-70b".to_string())); /// ``` pub fn set_original_model(&self, model: String) { - if let Ok(mut inner) = self.inner.lock() { - inner.original_model = Some(model); + if let Ok(mut state) = self.state.lock() { + state.original_model = Some(model); } } @@ -1040,14 +757,14 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// use std::time::Duration; /// /// let mgr = FallbackManager::new(3, 2); /// assert!(mgr.fallback_switched_at().is_none()); /// ``` pub fn fallback_switched_at(&self) -> Option { - self.inner.lock().ok().and_then(|i| i.fallback_switched_at) + self.state.lock().ok().and_then(|s| s.fallback_switched_at) } /// Get the model that should be used for the next request. @@ -1071,14 +788,20 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// let mgr = FallbackManager::for_model("llm-70b"); /// assert_eq!(mgr.active_model(), Some("llm-70b".to_string())); /// ``` pub fn active_model(&self) -> Option { - match self.state() { - FallbackState::Primary | FallbackState::Recovering => self.original_model(), - FallbackState::Fallback => self.fallback_model().or_else(|| self.original_model()), + let Ok(state) = self.state.lock() else { + return None; + }; + match state.state { + FallbackState::Primary | FallbackState::Recovering => state.original_model.clone(), + FallbackState::Fallback => state + .active_fallback + .clone() + .or_else(|| state.original_model.clone()), } } @@ -1097,7 +820,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState, FailureKind}; /// /// let mgr = FallbackManager::for_model("llm-1"); /// mgr.add_fallback_model("llm-2"); @@ -1108,21 +831,21 @@ impl FallbackManager { /// assert_eq!(mgr.fallback_models(), vec!["llm-4"]); /// /// // Simulate failures until circuit trips - /// mgr.record_model_failure(); - /// mgr.record_model_failure(); - /// mgr.record_model_failure(); // trips to Fallback + /// mgr.record_failure(FailureKind::Transient); + /// mgr.record_failure(FailureKind::Transient); + /// mgr.record_failure(FailureKind::Transient); // trips to Fallback /// /// assert_eq!(mgr.state(), FallbackState::Fallback); /// assert_eq!(mgr.active_model(), Some("llm-4".to_string())); /// ``` pub fn set_fallback_model(&self, model: impl Into) { - if let Ok(mut inner) = self.inner.lock() { - inner.fallback_models.clear(); - inner + if let Ok(mut state) = self.state.lock() { + state.fallback_models.clear(); + state .fallback_models - .push(FallbackEntry::new(model).with_max_fail_count(self.default_max_fail_count)); + .push(FallbackEntry::new(model).with_max_fail_count(self.config.max_fail_count)); + Self::recompute_active_fallback_impl(&mut state); } - self.recompute_active_fallback(); } /// Get the first fallback model, if any are configured. @@ -1134,7 +857,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// assert!(mgr.fallback_model().is_none()); @@ -1144,10 +867,10 @@ impl FallbackManager { /// assert_eq!(mgr.fallback_model(), Some("llm-70b".to_string())); // first in chain /// ``` pub fn fallback_model(&self) -> Option { - self.inner + self.state .lock() .ok() - .and_then(|i| i.active_fallback.clone()) + .and_then(|s| s.active_fallback.clone()) } /// Get the full fallback model chain. @@ -1158,7 +881,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-70b"); @@ -1169,10 +892,10 @@ impl FallbackManager { /// assert_eq!(chain, vec!["llm-70b", "llm-120b", "llm-32b"]); /// ``` pub fn fallback_models(&self) -> Vec { - self.inner + self.state .lock() .ok() - .map(|i| i.fallback_models.iter().map(|e| e.name.clone()).collect()) + .map(|s| s.fallback_models.iter().map(|e| e.name.clone()).collect()) .unwrap_or_default() } @@ -1189,7 +912,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-70b"); @@ -1199,12 +922,12 @@ impl FallbackManager { /// assert_eq!(chain, vec!["llm-70b", "llm-120b"]); /// ``` pub fn add_fallback_model(&self, model: impl Into) { - if let Ok(mut inner) = self.inner.lock() { - inner + if let Ok(mut state) = self.state.lock() { + state .fallback_models - .push(FallbackEntry::new(model).with_max_fail_count(self.default_max_fail_count)); + .push(FallbackEntry::new(model).with_max_fail_count(self.config.max_fail_count)); + Self::recompute_active_fallback_impl(&mut state); } - self.recompute_active_fallback(); } /// Insert a fallback model at a specific position in the chain. @@ -1220,7 +943,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-70b"); @@ -1231,15 +954,15 @@ 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 inner) = self.inner.lock() { - let entry = FallbackEntry::new(model).with_max_fail_count(self.default_max_fail_count); - if index >= inner.fallback_models.len() { - inner.fallback_models.push(entry); + if let Ok(mut state) = self.state.lock() { + let entry = FallbackEntry::new(model).with_max_fail_count(self.config.max_fail_count); + if index >= state.fallback_models.len() { + state.fallback_models.push(entry); } else { - inner.fallback_models.insert(index, entry); + state.fallback_models.insert(index, entry); } + Self::recompute_active_fallback_impl(&mut state); } - self.recompute_active_fallback(); } /// Remove a fallback model by name. @@ -1251,7 +974,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-70b"); @@ -1264,18 +987,15 @@ impl FallbackManager { /// assert_eq!(chain, vec!["llm-120b"]); /// ``` pub fn remove_fallback_model(&self, model: &str) -> bool { - 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 + let mut removed = false; + if let Ok(mut state) = self.state.lock() { + if let Some(pos) = state.fallback_models.iter().position(|x| x.name == model) { + state.fallback_models.remove(pos); + removed = true; + } + if removed { + Self::recompute_active_fallback_impl(&mut state); } - } else { - false - }; - if removed { - self.recompute_active_fallback(); } removed } @@ -1292,7 +1012,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.set_fallback_models(vec!["llm-70b".into(), "llm-120b".into(), "llm-32b".into()]); @@ -1301,14 +1021,14 @@ impl FallbackManager { /// assert_eq!(chain, vec!["llm-70b", "llm-120b", "llm-32b"]); /// ``` pub fn set_fallback_models(&self, models: Vec) { - let max_fc = self.default_max_fail_count; - if let Ok(mut inner) = self.inner.lock() { - inner.fallback_models = models + let max_fc = self.config.max_fail_count; + if let Ok(mut state) = self.state.lock() { + state.fallback_models = models .into_iter() .map(|name| FallbackEntry::new(name).with_max_fail_count(max_fc)) .collect(); + Self::recompute_active_fallback_impl(&mut state); } - self.recompute_active_fallback(); } /// Mark a fallback model as failed by name. @@ -1322,7 +1042,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-2"); @@ -1338,18 +1058,15 @@ 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 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 { - false + let mut found = false; + if let Ok(mut state) = self.state.lock() { + if let Some(entry) = state.fallback_models.iter_mut().find(|e| e.name == model) { + entry.record_attempt(); + found = true; + } + if found { + Self::recompute_active_fallback_impl(&mut state); } - } else { - false - }; - if found { - self.recompute_active_fallback(); } found } @@ -1363,7 +1080,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-2"); @@ -1375,18 +1092,15 @@ impl FallbackManager { /// assert!(mgr.failed_fallbacks().is_empty()); /// ``` pub fn clear_fallback_failed(&self, model: &str) -> bool { - let found = if let Ok(mut inner) = self.inner.lock() { - if let Some(entry) = inner.fallback_models.iter_mut().find(|e| e.name == model) { + let mut found = false; + if let Ok(mut state) = self.state.lock() { + if let Some(entry) = state.fallback_models.iter_mut().find(|e| e.name == model) { entry.clear_attempts(); - true - } else { - false + found = true; + } + if found { + Self::recompute_active_fallback_impl(&mut state); } - } else { - false - }; - if found { - self.recompute_active_fallback(); } found } @@ -1399,7 +1113,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-2"); @@ -1414,12 +1128,12 @@ impl FallbackManager { /// assert!(mgr.failed_fallbacks().is_empty()); /// ``` pub fn clear_all_fallback_failed(&self) { - if let Ok(mut inner) = self.inner.lock() { - for entry in &mut inner.fallback_models { + if let Ok(mut state) = self.state.lock() { + for entry in &mut state.fallback_models { entry.clear_attempts(); } + Self::recompute_active_fallback_impl(&mut state); } - self.recompute_active_fallback(); } /// Get the names of all fallback models marked as failed. @@ -1430,7 +1144,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-2"); @@ -1442,11 +1156,11 @@ impl FallbackManager { /// assert_eq!(failed, vec!["llm-3"]); /// ``` pub fn failed_fallbacks(&self) -> Vec { - self.inner + self.state .lock() .ok() - .map(|i| { - i.fallback_models + .map(|s| { + s.fallback_models .iter() .filter(|e| e.failed()) .map(|e| e.name.clone()) @@ -1463,7 +1177,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-2"); @@ -1475,11 +1189,11 @@ impl FallbackManager { /// assert_eq!(available, vec!["llm-2"]); /// ``` pub fn available_fallbacks(&self) -> Vec { - self.inner + self.state .lock() .ok() - .map(|i| { - i.fallback_models + .map(|s| { + s.fallback_models .iter() .filter(|e| !e.failed()) .map(|e| e.name.clone()) @@ -1488,43 +1202,13 @@ impl FallbackManager { .unwrap_or_default() } - /// Look up a fallback entry by name. - /// - /// Returns a cloned [`FallbackEntry`] for the first model in the chain - /// whose name matches, or `None` if the name is not present. Use this - /// to inspect a specific model's failure status without iterating the - /// full chain. - /// - /// # Example - /// - /// ```rust - /// use loopctl::fallback::FallbackManager; - /// - /// let mgr = FallbackManager::new(3, 2); - /// mgr.add_fallback_model("llm-70b"); - /// mgr.add_fallback_model("llm-120b"); + /// Set the `available` flag on a fallback model. /// - /// let entry = mgr.fallback_entry("llm-70b"); - /// assert!(entry.is_some()); - /// let e = entry.unwrap(); - /// assert_eq!(e.name(), "llm-70b"); - /// assert!(!e.failed()); // not failed yet - /// - /// assert!(mgr.fallback_entry("nonexistent").is_none()); - /// ``` - pub fn fallback_entry(&self, name: &str) -> Option { - self.inner - .lock() - .ok() - .and_then(|i| i.fallback_models.iter().find(|e| e.name == name).cloned()) - } - - /// Set the [`available`](FallbackEntry::available) flag on a fallback model. - /// - /// When set to `false`, the entry is considered [`failed`](FallbackEntry::failed) - /// regardless of its attempt count, and [`active_model`](Self::active_model) - /// will skip it. When set back to `true`, the entry becomes eligible again - /// (unless its attempt count has also reached [`FallbackEntry::max_fail_count`]). + /// When set to `false`, the model is considered failed regardless of + /// its attempt count, and [`active_model`](Self::active_model) will + /// skip it. When set back to `true`, the model becomes eligible again + /// (unless its attempt count has also reached the per-model + /// `max_fail_count`). /// /// Returns `true` if the model was found in the chain and updated, /// `false` if no model with that name exists. @@ -1532,7 +1216,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-2"); @@ -1549,18 +1233,15 @@ 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 inner) = self.inner.lock() { - if let Some(entry) = inner.fallback_models.iter_mut().find(|e| e.name == model) { + let mut found = false; + if let Ok(mut state) = self.state.lock() { + if let Some(entry) = state.fallback_models.iter_mut().find(|e| e.name == model) { entry.set_available(available); - true - } else { - false + found = true; + } + if found { + Self::recompute_active_fallback_impl(&mut state); } - } else { - false - }; - if found { - self.recompute_active_fallback(); } found } @@ -1583,50 +1264,98 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// 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, now activated - /// assert!(!mgr.record_api_failure()); // 4 — already activated, no re-trip + /// assert!(!mgr.record_failure(FailureKind::Transient)); // 1 + /// assert!(!mgr.record_failure(FailureKind::Transient)); // 2 + /// assert!(mgr.record_failure(FailureKind::Transient)); // 3 — threshold reached, now activated + /// assert!(!mgr.record_failure(FailureKind::Transient)); // 4 — already activated, no re-trip /// ``` - pub fn record_api_failure(&self) -> bool { - match self.state() { + /// Record a failure on the current model and trip the circuit if the + /// threshold is reached. + /// + /// Called by the agent loop each time an LLM request fails. The + /// effect depends on the current circuit state and the failure + /// [`kind`](FailureKind): + /// + /// - **[`Primary`](FallbackState::Primary)**: Increments the failure + /// counter. If it reaches + /// [`fallback_threshold`](Self::new), transitions to + /// [`Fallback`](FallbackState::Fallback) and returns `true`. + /// - **[`Fallback`](FallbackState::Fallback)**: Logs a warning — the + /// fallback model itself is failing. Consider calling + /// [`mark_fallback_failed`](Self::mark_fallback_failed) to skip it. + /// Returns `false`; the counter is unchanged. + /// - **[`Recovering`](FallbackState::Recovering)**: A + /// [`FailureKind::RateLimit`] re-trips the circuit to + /// [`Fallback`](FallbackState::Fallback) — the primary is still + /// rate-limited, so probing it further is pointless. A + /// [`FailureKind::Transient`] error leaves the half-open probe in + /// place. Returns `false` either way. + /// + /// # Returns + /// + /// `true` only when this call tripped the circuit from `Primary` to + /// `Fallback`; `false` otherwise. + /// + /// # Example + /// + /// ```rust + /// use loopctl::fallback::{FallbackManager, FallbackState, FailureKind}; + /// let mgr = FallbackManager::new(3, 2); + /// assert!(!mgr.record_failure(FailureKind::Transient)); // 1 + /// assert!(!mgr.record_failure(FailureKind::Transient)); // 2 + /// assert!(mgr.record_failure(FailureKind::Transient)); // 3 → trips to Fallback + /// assert_eq!(mgr.state(), FallbackState::Fallback); + /// ``` + pub fn record_failure(&self, kind: FailureKind) -> bool { + let Ok(mut state) = self.state.lock() else { + return false; + }; + match state.state { FallbackState::Primary => { - let failures = self - .consecutive_failures - .fetch_add(1, Ordering::Relaxed) - .saturating_add(1); - if failures >= self.fallback_threshold { + state.consecutive_failures = state.consecutive_failures.saturating_add(1); + if state.consecutive_failures >= self.config.trip_threshold { warn!( - consecutive_failures = failures, - threshold = self.fallback_threshold, + consecutive_failures = state.consecutive_failures, + threshold = self.config.trip_threshold, "Fallback threshold reached" ); - self.transition_to_fallback(); - return true; + Self::transition_to_fallback_impl(&mut state); + true + } else { + false } - false } - FallbackState::Fallback | FallbackState::Recovering => { - let failures = self.consecutive_failures.load(Ordering::Relaxed); + FallbackState::Fallback => { + let fb_name = state + .active_fallback + .clone() + .unwrap_or_else(|| "unknown".to_string()); warn!( - consecutive_failures = failures, - "API failure recorded while not in Primary state; counter unchanged" + "Fallback model \"{fb_name}\" also experiencing failures; consider calling mark_fallback_failed(\"{fb_name}\") to skip it" ); false } + FallbackState::Recovering => match kind { + FailureKind::RateLimit => { + warn!( + "Primary model rate-limited during recovery test, re-tripping to fallback" + ); + Self::transition_to_fallback_impl(&mut state); + false + } + FailureKind::Transient => { + warn!( + consecutive_failures = state.consecutive_failures, + "Transient failure recorded during recovery; staying half-open" + ); + false + } + }, } } - /// Alias for [`record_api_failure`](Self::record_api_failure) — framework-style name. - /// - /// Provided for callers that prefer the shorter `record_failure` name. - /// Delegates directly to [`record_api_failure`](Self::record_api_failure). - pub fn record_failure(&self) -> bool { - self.record_api_failure() - } - /// Reset the consecutive failure counter. /// /// Sets [`consecutive_failures`](Self::consecutive_failures) back to @@ -1637,16 +1366,18 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// let mgr = FallbackManager::new(3, 2); - /// mgr.record_api_failure(); - /// mgr.record_api_failure(); + /// mgr.record_failure(FailureKind::Transient); + /// mgr.record_failure(FailureKind::Transient); /// assert_eq!(mgr.consecutive_failures(), 2); /// mgr.reset_failure_counter(); /// assert_eq!(mgr.consecutive_failures(), 0); /// ``` pub fn reset_failure_counter(&self) { - self.consecutive_failures.store(0, Ordering::Relaxed); + if let Ok(mut state) = self.state.lock() { + state.consecutive_failures = 0; + } } /// Record a success on the current model. @@ -1666,104 +1397,32 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// let mgr = FallbackManager::new(3, 2); - /// mgr.record_api_failure(); + /// mgr.record_failure(FailureKind::Transient); /// assert_eq!(mgr.consecutive_failures(), 1); - /// mgr.record_model_success(); // resets failures to 0 + /// mgr.record_success(); // resets failures to 0 /// assert_eq!(mgr.consecutive_failures(), 0); /// ``` - pub fn record_model_success(&self) { - match self.state() { + pub fn record_success(&self) { + let Ok(mut state) = self.state.lock() else { + return; + }; + match state.state { FallbackState::Primary => { - self.consecutive_failures.store(0, Ordering::Relaxed); - } - FallbackState::Fallback => { - // Success on fallback, stay in fallback + state.consecutive_failures = 0; } + FallbackState::Fallback => {} FallbackState::Recovering => { - let successes = self - .primary_success_count - .fetch_add(1, Ordering::Relaxed) - .saturating_add(1); + state.primary_success_count = state.primary_success_count.saturating_add(1); debug!( - successes, - threshold = self.primary_resume_threshold, + successes = state.primary_success_count, + threshold = self.config.recovery_successes_needed, "Primary model success during recovery test" ); - if successes >= self.primary_resume_threshold { - self.transition_to_primary(); - } - } - } - } - - /// Alias for [`record_model_success`](Self::record_model_success) — framework-style name. - /// - /// Provided for callers that prefer the shorter `record_success` name. - /// Delegates directly to [`record_model_success`](Self::record_model_success). - pub fn record_success(&self) { - self.record_model_success(); - } - - /// Record a failure on the current model. - /// - /// Called by the agent loop when an LLM request fails. The effect - /// depends on the current circuit state: - /// - /// - **[`Primary`](FallbackState::Primary)**: Increments the failure - /// counter. If the count reaches the threshold, transitions to - /// [`Fallback`](FallbackState::Fallback) via - /// [`transition_to_fallback`](Self::transition_to_fallback). - /// - **[`Fallback`](FallbackState::Fallback)**: Logs a warning — - /// the fallback model itself is experiencing failures. The circuit - /// stays open. - /// - **[`Recovering`](FallbackState::Recovering)**: Immediately - /// reopens the circuit back to [`Fallback`](FallbackState::Fallback) - /// — the primary is not yet healthy. - /// - /// # Returns - /// - /// `true` if this specific failure caused the circuit to trip from - /// `Primary` to `Fallback`; `false` otherwise. - /// - /// # Example - /// - /// ```rust - /// use loopctl::fallback::{FallbackManager, FallbackState}; - /// - /// let mgr = FallbackManager::new(3, 2); - /// assert!(!mgr.record_model_failure()); // 1 - /// assert!(!mgr.record_model_failure()); // 2 - /// assert!(mgr.record_model_failure()); // 3 → trips to Fallback - /// assert_eq!(mgr.state(), FallbackState::Fallback); - /// ``` - pub fn record_model_failure(&self) -> bool { - match self.state() { - FallbackState::Primary => { - let failures = self - .consecutive_failures - .fetch_add(1, Ordering::Relaxed) - .saturating_add(1); - if failures >= self.fallback_threshold { - self.transition_to_fallback(); - return true; + if state.primary_success_count >= self.config.recovery_successes_needed { + Self::transition_to_primary_impl(&mut state); } - false - } - FallbackState::Fallback => { - let fb_name = self - .fallback_model() - .unwrap_or_else(|| "unknown".to_string()); - warn!( - "Fallback model \"{fb_name}\" also experiencing failures; consider calling mark_fallback_failed(\"{fb_name}\") to skip it" - ); - false - } - FallbackState::Recovering => { - warn!("Primary model failed during recovery test, staying on fallback"); - self.transition_to_fallback(); - false } } } @@ -1791,7 +1450,7 @@ impl FallbackManager { /// /// ```rust /// use std::time::Duration; - /// use loopctl::fallback::FallbackManager; + /// use loopctl::fallback::{FallbackManager, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// // Not in fallback state → false @@ -1813,7 +1472,7 @@ impl FallbackManager { /// Moves the circuit breaker to [`FallbackState::Fallback`], records /// the current time as [`fallback_switched_at`](Self::fallback_switched_at), /// and resets the recovery success counter. Called automatically by - /// [`record_model_failure`](Self::record_model_failure) when the + /// [`record_failure`](Self::record_failure) when the /// failure threshold is reached, or manually when the circuit needs /// to trip immediately. /// @@ -1823,20 +1482,29 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.transition_to_fallback(); /// assert_eq!(mgr.state(), FallbackState::Fallback); /// ``` pub fn transition_to_fallback(&self) { - self.fallback_state - .store(FallbackState::Fallback as u8, Ordering::Relaxed); - self.fallback_activated.store(true, Ordering::Relaxed); - if let Ok(mut inner) = self.inner.lock() { - inner.fallback_switched_at = Some(Instant::now()); + if let Ok(mut state) = self.state.lock() { + Self::transition_to_fallback_impl(&mut state); } - self.primary_success_count.store(0, Ordering::Relaxed); + } + + /// [`transition_to_fallback`](Self::transition_to_fallback) assuming + /// the caller already holds the state lock. + /// + /// Called from the `record_*` methods so the trip runs while their + /// guard is held, making the read-decide-transition atomic with + /// respect to other callers. + fn transition_to_fallback_impl(state: &mut BreakerState) { + state.state = FallbackState::Fallback; + state.fallback_activated = true; + state.fallback_switched_at = Some(Instant::now()); + state.primary_success_count = 0; info!("Circuit breaker: transitioned to Fallback state"); } @@ -1847,13 +1515,13 @@ impl FallbackManager { /// counter. Called by the agent loop after /// [`should_try_resume_primary`](Self::should_try_resume_primary) /// returns `true`. Subsequent calls to - /// [`record_model_success`](Self::record_model_success) will count + /// [`record_success`](Self::record_success) will count /// toward the recovery threshold. /// /// # Example /// /// ```rust - /// use loopctl::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.transition_to_fallback(); @@ -1861,9 +1529,16 @@ impl FallbackManager { /// assert_eq!(mgr.state(), FallbackState::Recovering); /// ``` pub fn transition_to_recovering(&self) { - self.fallback_state - .store(FallbackState::Recovering as u8, Ordering::Relaxed); - self.primary_success_count.store(0, Ordering::Relaxed); + if let Ok(mut state) = self.state.lock() { + Self::transition_to_recovering_impl(&mut state); + } + } + + /// [`transition_to_recovering`](Self::transition_to_recovering) + /// assuming the caller already holds the state lock. + fn transition_to_recovering_impl(state: &mut BreakerState) { + state.state = FallbackState::Recovering; + state.primary_success_count = 0; info!("Circuit breaker: transitioned to Recovering state (testing primary)"); } @@ -1873,7 +1548,7 @@ impl FallbackManager { /// the fallback timestamp, and resets all counters (failures, /// successes, and the `fallback_activated` flag). Called /// automatically when the recovery success threshold is reached - /// inside [`record_model_success`](Self::record_model_success), or + /// inside [`record_success`](Self::record_success), or /// manually to force an immediate return to primary. /// /// After this call, [`is_using_fallback`](Self::is_using_fallback) @@ -1882,7 +1557,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.transition_to_fallback(); @@ -1891,15 +1566,28 @@ impl FallbackManager { /// assert!(!mgr.is_using_fallback()); /// ``` pub fn transition_to_primary(&self) { - self.fallback_state - .store(FallbackState::Primary as u8, Ordering::Relaxed); - if let Ok(mut inner) = self.inner.lock() { - inner.fallback_switched_at = None; + if let Ok(mut state) = self.state.lock() { + Self::transition_to_primary_impl(&mut state); + } + } + + /// [`transition_to_primary`](Self::transition_to_primary) assuming + /// the caller already holds the state lock. + /// + /// Inlines the [`clear_all_fallback_failed`](Self::clear_all_fallback_failed) + /// and active-fallback recompute against the held guard, because the + /// public versions acquire the same lock themselves and + /// `std::sync::Mutex` is not reentrant. + fn transition_to_primary_impl(state: &mut BreakerState) { + state.state = FallbackState::Primary; + state.fallback_switched_at = None; + state.primary_success_count = 0; + state.consecutive_failures = 0; + state.fallback_activated = false; + for entry in &mut state.fallback_models { + entry.clear_attempts(); } - self.primary_success_count.store(0, Ordering::Relaxed); - self.consecutive_failures.store(0, Ordering::Relaxed); - self.fallback_activated.store(false, Ordering::Relaxed); - self.clear_all_fallback_failed(); + Self::recompute_active_fallback_impl(state); info!("Circuit breaker: transitioned to Primary state (primary model recovered)"); } @@ -1917,10 +1605,10 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState, FailureKind}; /// /// let mgr = FallbackManager::new(3, 2); - /// for _ in 0..3 { mgr.record_model_failure(); } + /// for _ in 0..3 { mgr.record_failure(FailureKind::Transient); } /// assert_eq!(mgr.state(), FallbackState::Fallback); /// /// mgr.reset(); @@ -1929,33 +1617,26 @@ impl FallbackManager { /// assert_eq!(mgr.consecutive_failures(), 0); /// ``` pub fn reset(&self) { - self.fallback_state - .store(FallbackState::Primary as u8, Ordering::Relaxed); - 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 inner) = self.inner.lock() { - inner.fallback_switched_at = None; - } - self.clear_all_fallback_failed(); - } - - /// 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.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; + if let Ok(mut state) = self.state.lock() { + Self::transition_to_primary_impl(&mut state); } } + + /// Recompute the cached [`active_fallback`](BreakerState::active_fallback) + /// from the current fallback chain, in place against a held guard. + /// + /// Walks [`fallback_models`](BreakerState::fallback_models) for the + /// first non-failed entry and stores its name, so + /// [`fallback_model`](Self::fallback_model) stays `O(1)`. Called by + /// every method that mutates the chain or a model's failed flag, and + /// by [`transition_to_primary_impl`](Self::transition_to_primary_impl). + fn recompute_active_fallback_impl(state: &mut BreakerState) { + state.active_fallback = state + .fallback_models + .iter() + .find(|e| !e.failed()) + .map(|e| e.name.clone()); + } } impl Default for FallbackManager { @@ -1980,17 +1661,17 @@ mod tests { #[test] fn test_failure_threshold() { 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_failure(FailureKind::Transient)); // 1 + assert!(!mgr.record_failure(FailureKind::Transient)); // 2 + assert!(mgr.record_failure(FailureKind::Transient)); // 3 — threshold reached } #[test] fn test_model_failure_triggers_fallback() { let mgr = FallbackManager::new(3, 2); - assert!(!mgr.record_model_failure()); // 1 - assert!(!mgr.record_model_failure()); // 2 - assert!(mgr.record_model_failure()); // 3 — triggers fallback + assert!(!mgr.record_failure(FailureKind::Transient)); // 1 + assert!(!mgr.record_failure(FailureKind::Transient)); // 2 + assert!(mgr.record_failure(FailureKind::Transient)); // 3 — triggers fallback assert_eq!(mgr.state(), FallbackState::Fallback); } @@ -1999,7 +1680,7 @@ mod tests { let mgr = FallbackManager::new(3, 2); // Trigger fallback for _ in 0..3 { - mgr.record_model_failure(); + mgr.record_failure(FailureKind::Transient); } assert_eq!(mgr.state(), FallbackState::Fallback); @@ -2008,8 +1689,8 @@ mod tests { assert_eq!(mgr.state(), FallbackState::Recovering); // Recover after enough successes - mgr.record_model_success(); // 1 - mgr.record_model_success(); // 2 — threshold reached + mgr.record_success(); // 1 + mgr.record_success(); // 2 — threshold reached assert_eq!(mgr.state(), FallbackState::Primary); } @@ -2017,10 +1698,10 @@ mod tests { fn test_recovery_failure_goes_back_to_fallback() { let mgr = FallbackManager::new(3, 2); for _ in 0..3 { - mgr.record_model_failure(); + mgr.record_failure(FailureKind::Transient); } mgr.transition_to_recovering(); - mgr.record_model_failure(); // failure during recovery + mgr.record_failure(FailureKind::RateLimit); // sustained failure during recovery re-trips assert_eq!(mgr.state(), FallbackState::Fallback); } @@ -2031,7 +1712,7 @@ mod tests { // Trigger fallback for _ in 0..3 { - mgr.record_model_failure(); + mgr.record_failure(FailureKind::Transient); } // Not enough time assert!(!mgr.should_try_resume_primary(Duration::from_hours(1))); @@ -2051,7 +1732,7 @@ mod tests { fn test_reset() { let mgr = FallbackManager::new(3, 2); for _ in 0..3 { - mgr.record_model_failure(); + mgr.record_failure(FailureKind::Transient); } assert_eq!(mgr.state(), FallbackState::Fallback); @@ -2066,24 +1747,23 @@ mod tests { let mgr = FallbackManager::new(3, 2); // Trip the circuit for _ in 0..3 { - mgr.record_api_failure(); + mgr.record_failure(FailureKind::Transient); } // Activate fallback mgr.transition_to_fallback(); - mgr.fallback_activated.store(true, Ordering::Relaxed); // Further failures should not return true (already activated) - assert!(!mgr.record_api_failure()); + assert!(!mgr.record_failure(FailureKind::Transient)); } #[test] fn test_record_success_resets_on_primary() { let mgr = FallbackManager::new(3, 2); - mgr.record_api_failure(); - mgr.record_api_failure(); + mgr.record_failure(FailureKind::Transient); + mgr.record_failure(FailureKind::Transient); assert_eq!(mgr.consecutive_failures(), 2); - mgr.record_model_success(); + mgr.record_success(); assert_eq!(mgr.consecutive_failures(), 0); } @@ -2105,8 +1785,8 @@ mod tests { for _ in 0..10 { let mgr = Arc::clone(&mgr); handles.push(thread::spawn(move || { - mgr.record_api_failure(); - mgr.record_model_success(); + mgr.record_failure(FailureKind::Transient); + mgr.record_success(); mgr.state(); mgr.consecutive_failures(); })); @@ -2163,7 +1843,7 @@ mod tests { let mgr = FallbackManager::for_model("primary-model"); mgr.add_fallback_model("fallback-model"); // Record a failure while in Primary to dirty the counter, then trip. - mgr.record_failure(); + mgr.record_failure(FailureKind::Transient); mgr.transition_to_fallback(); // State is dirty. @@ -2182,46 +1862,46 @@ mod tests { 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.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"); + entry.record_attempt(); + entry.record_attempt(); assert!(entry.failed()); - assert_eq!(entry.attempt_count(), 2); + 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_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"); + entry.record_attempt(); + entry.record_attempt(); assert!(entry.failed()); let entry = entry.with_max_fail_count(3); - assert_eq!(entry.attempt_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"); + entry.record_attempt(); + entry.record_attempt(); 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_eq!(entry.attempt_count, 2); assert!(entry.failed()); } diff --git a/src/managers.rs b/src/managers.rs index 7504395..d195c12 100644 --- a/src/managers.rs +++ b/src/managers.rs @@ -710,7 +710,9 @@ mod tests { #[test] fn test_reset_all_clears_fallback() { let managers = LoopManagers::new(); - let _ = managers.fallback().record_api_failure(); + let _ = managers + .fallback() + .record_failure(crate::fallback::FailureKind::Transient); managers.reset_all(); assert!(managers.fallback().active_model().is_none()); } From e9e21a62b164010764c8663def3c7942b3dbbf11 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 27 Jul 2026 19:57:12 +1200 Subject: [PATCH 11/21] refactor: bare loop dispatch, fix parallel dispatch recovery drop --- CHANGELOG.md | 17 ++ src/engine/bare/dispatch.rs | 412 +++++++++++++----------------------- 2 files changed, 159 insertions(+), 270 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bf64f88..de1aa9f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -423,6 +423,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. ### Fixed +- Tool dispatch had three separate code paths (sequential, small-batch + parallel, and wave-parallel) each with its own copy of the recovery + loop, PRE/POST side-effect logic, and cancel handling. Two bugs came + from this split: the small-batch path called `dispatch_tool` directly + with no recovery (a flaky tool failed permanently where it would + recover elsewhere), and the parallel paths fired observer/detection/ + hook side-effects once on the final result while sequential fired + them per retry attempt — diverging loop-detection sensitivity, health + inputs, and observer event counts by dispatch mode. The entire + dispatch machinery is now one function (`execute_tool_call`) that owns + the full lifecycle (PRE → dispatch → POST → recovery loop) with + mid-flight cancel via `tokio::select!`. Sequential and parallel both + call it, so there is exactly one definition of "execute a tool call" + — no divergence is possible. Six functions (`parallel_pre_phase`, + `parallel_exec_phase`, `parallel_post_phase`, `parallel_run_remaining`, + `run_parallel_task`, `dispatch_tool_with_recovery`) collapsed into one + `execute_tool_call` + two thin dispatch loops. - OpenAI streaming silently dropped every multi-chunk tool-call argument fragment after the first, truncating the tool input JSON. The deserialization structs declared `id` (on the tool-call delta) and diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index b23bf6f..e01176b 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -229,17 +229,15 @@ impl BareLoop { /// Dispatch tool calls one at a time. /// /// Each [`ToolCall`] runs to completion via - /// [`dispatch_tool_with_recovery`](Self::dispatch_tool_with_recovery) - /// before the next begins. Observers see strictly paired - /// `on_tool_pre` / `on_tool_post` events per call, in input order. The - /// cancel signal is checked between calls so a Ctrl-C mid-batch aborts the - /// remaining calls rather than running them all. + /// [`execute_tool_call`](Self::execute_tool_call) before the next + /// begins. The cancel signal is checked between calls so a Ctrl-C + /// mid-batch aborts the remaining calls rather than running them all. /// /// # Errors /// /// Returns [`LoopError::Cancelled`] if the cancel signal fires between /// calls, or any hard error from - /// [`dispatch_tool_with_recovery`](Self::dispatch_tool_with_recovery). + /// [`execute_tool_call`](Self::execute_tool_call). async fn dispatch_tools_sequential( &self, tool_calls: &[ToolCall], @@ -250,42 +248,31 @@ impl BareLoop { if self.is_cancelled() { return Err(LoopError::Cancelled); } - let result = self.dispatch_tool_with_recovery(call, turn_idx).await?; + let result = self.execute_tool_call(call.clone(), turn_idx).await?; results.push(result); } Ok(results) } - /// Dispatch independent tool calls concurrently, preserving the ordering - /// invariant. + /// Dispatch independent tool calls concurrently. /// - /// The batch flows through five phases: - /// - /// 1. **PRE** (sequential, input order): for each call — cancel-check, - /// `on_tool_pre`, pre-hooks, pre-detection. Blocked calls get a - /// pre-computed soft-error result and are excluded from parallel - /// execution. A detection hard-stop returns `Err` immediately, before - /// any task spawns. - /// 2. **PLAN**: [`ToolDependencyGraph`] partitions eligible calls into - /// waves based on `is_safe_for_concurrent_execution` and - /// `resource_key` conflicts. - /// 3. **EXEC** (concurrent): each wave's calls run under a - /// [`Semaphore`](tokio::sync::Semaphore) capped at `max_concurrency`, - /// each racing the shared cancel signal via `tokio::select!`. - /// 4. **JOIN**: results are collected indexed by original position so they - /// come home in input order regardless of completion order. - /// 5. **POST** (sequential, input order): `post_detection`, - /// `on_tool_post`, post-hooks, `record_tool_health` for each call. + /// Builds a [`DispatchPlan`] from the registry's + /// [`is_safe_for_concurrent_execution`](crate::tool::Tool::is_safe_for_concurrent_execution) + /// and [`resource_key`](crate::tool::Tool::resource_key) metadata, then + /// runs each wave of calls concurrently under a semaphore capped at + /// [`max_concurrency`](crate::config::ParallelDispatchConfig::max_concurrency). + /// Each call runs through [`execute_tool_call`](Self::execute_tool_call) + /// — the same end-to-end path as sequential — so observers, + /// detection, hooks, and health all fire identically regardless of + /// dispatch mode. /// - /// Falls back to the sequential path when parallelism yields no benefit - /// (fewer than 2 eligible calls). Soft errors are collected (not fatal); - /// the model sees all N results in input order. + /// Falls back to the sequential path when there are fewer than 2 calls. /// /// # Errors /// - /// [`LoopError::Cancelled`] if the cancel signal fires during dispatch - /// (either in PRE or inside an EXEC task). [`LoopError::LoopDetected`] - /// (propagated from PRE detection) on a hard stop. + /// [`LoopError::Cancelled`] if the cancel signal fires during dispatch. + /// [`LoopError::LoopDetected`] on a hard stop from detection. Any hard + /// error from an individual [`execute_tool_call`](Self::execute_tool_call). async fn dispatch_tools_parallel( &self, tool_calls: &[ToolCall], @@ -295,124 +282,38 @@ impl BareLoop { return self.dispatch_tools_sequential(tool_calls, turn_idx).await; } - let tool_context = self.build_tool_context(); - let (pre_results, eligible) = self.parallel_pre_phase(tool_calls, turn_idx)?; - - if eligible.len() < 2 { - return self - .parallel_run_remaining(tool_calls, &pre_results, &tool_context, turn_idx) - .await; - } - - let exec_results = self - .parallel_exec_phase(tool_calls, &eligible, &pre_results, &tool_context, turn_idx) - .await?; - Ok(self.parallel_post_phase(tool_calls, &pre_results, &exec_results, turn_idx)) - } - - /// PRE phase: for each call (input order), fire `on_tool_pre`, check - /// pre-hooks, run pre-detection. Blocked calls get a pre-computed - /// soft-error result stashed in `pre_results`; eligible calls are collected - /// into the returned index list. - /// - /// # Errors - /// - /// [`LoopError::Cancelled`] if the cancel signal fires, or - /// [`LoopError::LoopDetected`] from pre-detection. - fn parallel_pre_phase( - &self, - tool_calls: &[ToolCall], - turn_idx: usize, - ) -> Result<(Vec>, Vec), LoopError> { - let mut pre_results: Vec> = - (0..tool_calls.len()).map(|_| None).collect(); - let mut eligible: Vec = Vec::new(); - - for (i, tc) in tool_calls.iter().enumerate() { - if self.is_cancelled() { - return Err(LoopError::Cancelled); - } - self.fire_tool_pre(turn_idx, tc); - - if let Some(blocked) = self.check_pre_tool_use_hooks(tc, turn_idx) { - self.fire_tool_post(turn_idx, tc, &blocked); - if let Some(slot) = pre_results.get_mut(i) { - *slot = Some(blocked); - } - continue; - } - - if let Some(blocked) = self.pre_detection(tc, turn_idx)? { - self.fire_tool_post(turn_idx, tc, &blocked); - if let Some(slot) = pre_results.get_mut(i) { - *slot = Some(blocked); - } - continue; - } - eligible.push(i); - } - Ok((pre_results, eligible)) - } - - /// EXEC phase: run eligible calls concurrently in waves under a semaphore. - /// - /// Each wave's tasks race the cancel signal via `tokio::select!`. Results - /// are placed in a position-indexed vec so they land in input order. - /// - /// # Errors - /// - /// [`LoopError::Cancelled`] if any task is cancelled, or a hard - /// [`LoopError`] propagated from `dispatch_tool`. - async fn parallel_exec_phase( - &self, - tool_calls: &[ToolCall], - eligible: &[usize], - pre_results: &[Option], - tool_context: &ToolContext, - turn_idx: usize, - ) -> Result>, LoopError> { - let graph = ToolDependencyGraph::from_calls(tool_calls, &self.tools); - let plan = graph.plan(); + let plan = ToolDependencyGraph::from_calls(tool_calls, &self.tools).plan(); let max_concurrency = self .run_config() .parallel_tool_dispatch .max_concurrency - .clamp(1, eligible.len()); + .clamp(1, tool_calls.len()); let semaphore = Arc::new(tokio::sync::Semaphore::new(max_concurrency)); - let mut exec_results: Vec> = + let mut results: Vec> = (0..tool_calls.len()).map(|_| None).collect(); + for wave in &plan.waves { if self.is_cancelled() { return Err(LoopError::Cancelled); } - let wave_eligible: Vec = wave - .iter() - .copied() - .filter(|idx| pre_results.get(*idx).is_some_and(Option::is_none)) - .collect(); - if wave_eligible.is_empty() { - continue; - } - let mut tasks = Vec::with_capacity(wave_eligible.len()); - for &idx in &wave_eligible { + let mut tasks = Vec::with_capacity(wave.len()); + for &idx in wave { let tc = tool_calls.get(idx).cloned(); - let ctx = tool_context.clone(); let sem = Arc::clone(&semaphore); - let cancel = Arc::clone(&self.cancelled); let turn = turn_idx; tasks.push(async move { let _permit = sem.acquire_owned().await.ok()?; - self.run_parallel_task(tc?, &ctx, &cancel, turn).await + Some(self.execute_tool_call(tc?, turn).await) }); } let outcomes = futures::future::join_all(tasks).await; - for (outcome, &idx) in outcomes.into_iter().zip(&wave_eligible) { + for (outcome, &idx) in outcomes.into_iter().zip(wave) { match outcome { None => return Err(LoopError::Cancelled), Some(Ok(result)) => { - if let Some(slot) = exec_results.get_mut(idx) { + if let Some(slot) = results.get_mut(idx) { *slot = Some(result); } } @@ -420,105 +321,53 @@ impl BareLoop { } } } - Ok(exec_results) - } - /// POST phase: for each call (input order), run `post_detection`, - /// `on_tool_post`, post-hooks, and `record_tool_health`. Blocked calls - /// (from PRE) emit their stashed soft-error result; executed calls emit - /// their exec result. - fn parallel_post_phase( - &self, - tool_calls: &[ToolCall], - pre_results: &[Option], - exec_results: &[Option], - turn_idx: usize, - ) -> Vec { - let mut results = Vec::with_capacity(tool_calls.len()); - for (i, tc) in tool_calls.iter().enumerate() { - if let Some(Some(pr)) = pre_results.get(i) { - results.push(pr.clone()); - } else if let Some(er) = exec_results.get(i).and_then(|r| r.as_ref()) { - self.post_detection(tc, er); - self.fire_tool_post(turn_idx, tc, er); - self.notify_post_tool_use_hooks(tc, er, turn_idx); - // Health was already recorded inside run_parallel_task on - // every attempt — don't double-count here. - results.push(er.clone()); - } else { - let soft = ToolDispatchResult { - tool_call_id: tc.id.clone(), + Ok(results + .into_iter() + .map(|r| { + r.unwrap_or_else(|| ToolDispatchResult { + tool_call_id: String::new(), output: ToolContent::Text("dispatch produced no result".to_string()), is_error: true, duration: Duration::ZERO, - resolved_tool_name: tc.tool.clone(), + resolved_tool_name: String::new(), display_hint: None, - }; - results.push(soft); - } - } - results - } - - /// Run remaining (non-blocked, non-eligible-for-parallel) calls - /// sequentially when there aren't enough to justify concurrency. - /// - /// # Errors - /// - /// [`LoopError::Cancelled`] or any hard error from - /// [`dispatch_tool`](Self::dispatch_tool). - async fn parallel_run_remaining( - &self, - tool_calls: &[ToolCall], - pre_results: &[Option], - tool_context: &ToolContext, - turn_idx: usize, - ) -> Result, LoopError> { - let mut results = Vec::with_capacity(tool_calls.len()); - for (i, tc) in tool_calls.iter().enumerate() { - if let Some(Some(pr)) = pre_results.get(i) { - results.push(pr.clone()); - } else { - if self.is_cancelled() { - return Err(LoopError::Cancelled); - } - let start = Instant::now(); - let r = self - .dispatch_tool(tc, tool_context, start, turn_idx) - .await?; - self.post_detection(tc, &r); - self.fire_tool_post(turn_idx, tc, &r); - self.notify_post_tool_use_hooks(tc, &r, turn_idx); - self.record_tool_health(tc.tool.as_str(), &r); - results.push(r); - } - } - Ok(results) + }) + }) + .collect()) } - /// Dispatch a single tool call with reflection and recovery on errors. + /// Execute a single tool call end-to-end. /// - /// If the tool succeeds, returns immediately. On failure, consults the - /// [`Reflector`](crate::reflection::Reflector) and - /// [`RecoveryStrategy`](crate::reflection::RecoveryStrategy) to decide - /// whether to retry, skip, ask the user, or fail. + /// The single function that owns the full lifecycle of one tool call: + /// PRE (observer pre-notification, pre-hooks, pre-detection) → dispatch + /// → POST (post-detection, observer post-notification, post-hooks, + /// health recording) → recovery decision. On failure, consults + /// [`recovery_wait_or_return`](Self::recovery_wait_or_return) and loops + /// — re-firing PRE and POST on every retry attempt — until the tool + /// succeeds, the strategy gives up (returns a soft error), or the user + /// cancels. /// - /// Each attempt fires [`on_tool_pre`](crate::observer::LoopObserver::on_tool_pre) - /// before dispatch and [`on_tool_post`](crate::observer::LoopObserver::on_tool_post) - /// after, so observers always see a complete lifecycle pair. + /// Safe to run concurrently: every side-effect target (`LoopObserver`, + /// `DetectionManager`, `HookExecutor`, `HealthRegistry`) is `Send + + /// Sync`. Sequential and parallel dispatch both call this function, so + /// there is exactly one definition of what "execute a tool call" + /// means — no divergence in side-effect granularity, observer event + /// counts, or recovery behaviour between modes. /// /// # Errors /// - /// Returns [`LoopError`] if the detection manager signals a hard stop - /// (e.g. [`LoopError::LoopDetected`]). - async fn dispatch_tool_with_recovery( + /// Returns [`LoopError::Cancelled`] when the recovery strategy aborts + /// on cancellation, [`LoopError::LoopDetected`] on a hard detection + /// stop, or any hard error from + /// [`dispatch_tool`](Self::dispatch_tool). + async fn execute_tool_call( &self, - tc: &ToolCall, + mut tc: ToolCall, turn_idx: usize, ) -> Result { let tool_context = self.build_tool_context(); let mut attempt: u32 = 0; - let mut tc = tc.clone(); loop { self.fire_tool_pre(turn_idx, &tc); @@ -534,10 +383,11 @@ impl BareLoop { } let start = Instant::now(); - let tool_result = self - .dispatch_tool(&tc, &tool_context, start, turn_idx) - .await?; - + let tool_result = tokio::select! { + biased; + () = self.cancelled.notified() => return Err(LoopError::Cancelled), + r = self.dispatch_tool(&tc, &tool_context, start, turn_idx) => r?, + }; self.post_detection(&tc, &tool_result); self.fire_tool_post(turn_idx, &tc, &tool_result); self.notify_post_tool_use_hooks(&tc, &tool_result, turn_idx); @@ -561,61 +411,6 @@ impl BareLoop { } } - /// Execute one parallel-wave task: the retry/recovery loop with cancellation, - /// no observer/hook/detection side-effects. - /// - /// Runs [`dispatch_tool`](Self::dispatch_tool) once, racing the shared - /// cancel signal. On error, consults - /// [`recovery_wait_or_return`](Self::recovery_wait_or_return) — retrying - /// with any correction, returning the soft error, or aborting on cancel. - /// The parallel PRE/POST phases own the observer/hook/detection - /// side-effects. - /// - /// Returns `None` if cancelled, `Some(Ok(result))` for success or soft - /// error, `Some(Err(e))` for a hard error. - async fn run_parallel_task( - &self, - mut tc: ToolCall, - ctx: &ToolContext, - cancel: &Arc, - turn: usize, - ) -> Option> { - let mut attempt: u32 = 0; - loop { - let start = Instant::now(); - let result = tokio::select! { - biased; - () = cancel.notified() => return None, - r = self.dispatch_tool(&tc, ctx, start, turn) => r, - }; - let tool_result = match result { - Ok(r) => r, - Err(e) => return Some(Err(e)), - }; - // Record health on every attempt — the registry is thread-safe - // (atomic counters), so this is safe to call from concurrent - // tasks. Detection/observer side-effects stay in the POST phase - // (they are not thread-safe). - self.record_tool_health(tc.tool.as_str(), &tool_result); - if !tool_result.is_error { - return Some(Ok(tool_result)); - } - match self - .recovery_wait_or_return(&tc, &tool_result, attempt) - .await - { - Ok((next_attempt, correction)) => { - attempt = next_attempt; - Self::apply_correction_if_present(&mut tc, correction); - } - Err(RecoveryOutcome::SoftError(returned_result)) => { - return Some(Ok(returned_result)); - } - Err(RecoveryOutcome::Cancelled) => return None, - } - } - } - /// Notify observers that a tool call is about to be dispatched. /// /// Fires [`on_tool_pre`](crate::observer::LoopObserver::on_tool_pre) with @@ -718,7 +513,7 @@ impl BareLoop { /// error results. A tool not in the registry produces a soft error. /// /// Observer notifications are handled by the caller - /// ([`dispatch_tool_with_recovery`](Self::dispatch_tool_with_recovery)). + /// ([`execute_tool_call`](Self::execute_tool_call)). /// /// # Errors /// @@ -993,7 +788,7 @@ impl BareLoop { /// middleware chain (timeout, permissions, output limits, etc.). /// /// Observer notifications are handled by the caller - /// ([`dispatch_tool_with_recovery`](Self::dispatch_tool_with_recovery)). + /// ([`execute_tool_call`](Self::execute_tool_call)). /// /// # Errors /// @@ -1653,4 +1448,81 @@ mod tests { "cancel should interrupt the 10s backoff promptly; elapsed {elapsed:?}" ); } + + #[tokio::test] + async fn execute_tool_call_runs_recovery_on_failure() { + use crate::reflection::{ + FailureAnalysis, FailureSeverity, RecoveryAction, RecoveryStrategy, + }; + use std::sync::atomic::{AtomicU32, Ordering}; + + struct AlwaysRecoverable; + impl crate::reflection::Reflector for AlwaysRecoverable { + fn analyze( + &self, + error: &str, + tool_name: &str, + _tool_input: &Value, + _tool_schema: Option<&crate::tool::ToolSchema>, + _context: &crate::reflection::ReflectionContext, + ) -> Pin< + Box< + dyn Future> + + Send + + '_, + >, + > { + let error = error.to_string(); + let tool_name = tool_name.to_string(); + Box::pin(async move { + Ok(FailureAnalysis { + is_recoverable: true, + root_cause: error, + severity: FailureSeverity::Medium, + correction: None, + context: format!("tool: {tool_name}"), + }) + }) + } + } + + struct CountingRetry { + calls: Arc, + } + impl RecoveryStrategy for CountingRetry { + fn decide( + &self, + _analysis: &FailureAnalysis, + _attempt: u32, + _max_attempts: u32, + ) -> Pin + Send + '_>> { + self.calls.fetch_add(1, Ordering::Relaxed); + Box::pin(async { RecoveryAction::Skip("counted".into()) }) + } + } + + let decide_calls = Arc::new(AtomicU32::new(0)); + let error_tool = crate::tool::FnTool::new( + "error_tool".into(), + "Always errors".into(), + Value::Object(serde_json::Map::new()), + |_, _| Box::pin(async { Err(ToolError::Execution("boom".to_string())) }), + ); + + let mut registry = ToolRegistry::new(); + registry.register(error_tool); + let mut bare = make_loop(registry); + bare.set_reflector(Arc::new(AlwaysRecoverable)); + bare.set_recovery_strategy(Arc::new(CountingRetry { + calls: Arc::clone(&decide_calls), + })); + + let tc = make_call("1", "error_tool", Value::Null); + let _ = bare.execute_tool_call(tc, 0).await.ok(); + + assert!( + decide_calls.load(Ordering::Relaxed) > 0, + "execute_tool_call must run recovery on failure" + ); + } } From 42aa9476ded7aec9f3ebcd76fe5269958a4428b5 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 27 Jul 2026 19:58:21 +1200 Subject: [PATCH 12/21] refactor: notify prefix consistency --- src/engine/bare.rs | 16 ++++++++-------- src/engine/bare/dispatch.rs | 12 ++++++------ 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/src/engine/bare.rs b/src/engine/bare.rs index 029014d..eb537fe 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -1234,7 +1234,7 @@ impl BareLoop { /// per-turn counter increment). `query` is the user's message on the first /// turn and `""` on continuation turns (the previous turn's tool results /// are already in the conversation history). - fn fire_turn_start(&self, turn: usize, query: &str) { + fn notify_turn_start(&self, turn: usize, query: &str) { self.managers.observers().on_turn_start(&TurnStartContext { turn, query: query.to_string(), @@ -1245,9 +1245,9 @@ impl BareLoop { /// the assembled assistant text for the turn that just streamed. /// /// `turn` is the same 0-indexed current turn passed to - /// [`fire_turn_start`](Self::fire_turn_start). `usage` is `None` when the + /// [`notify_turn_start`](Self::notify_turn_start). `usage` is `None` when the /// provider did not report token counts for the turn. - fn fire_response(&self, turn: usize, text: &str, usage: Option) { + fn notify_response(&self, turn: usize, text: &str, usage: Option) { self.managers.observers().on_response(&ResponseContext { turn, text: text.to_string(), @@ -1261,9 +1261,9 @@ impl BareLoop { /// Fires once per call regardless of how many recovery retries the call /// later undergoes (the retry loop lives in dispatch and re-fires only /// `on_tool_pre`/`on_tool_post`). `turn` is the same 0-indexed current - /// turn passed to [`fire_turn_start`](Self::fire_turn_start); it reaches + /// turn passed to [`notify_turn_start`](Self::notify_turn_start); it reaches /// the dispatch path as `turn_idx`, so the two events correlate. - fn fire_tool_calls_received(&self, turn: usize, tool_calls: &[ToolCall]) { + fn notify_tool_calls_received(&self, turn: usize, tool_calls: &[ToolCall]) { for tc in tool_calls { self.managers .observers() @@ -1491,7 +1491,7 @@ impl BareLoop { .unwrap_or_default() }; - self.fire_turn_start(current_turn, &turn_input); + self.notify_turn_start(current_turn, &turn_input); self.inject_contributors(current_turn); let cancel = Arc::clone(&self.cancelled); @@ -1529,7 +1529,7 @@ impl BareLoop { let text = msg.text_content(); let (turn_in, turn_out) = Self::usage_tokens(usage.as_ref()); let pattern = self.managers.detection().record_response(&text); - self.fire_response(current_turn, &text, usage); + self.notify_response(current_turn, &text, usage); if let Some(e) = self.apply_loop_detection(current_turn, &pattern) { return Err(e); @@ -1623,7 +1623,7 @@ impl BareLoop { } } - self.fire_tool_calls_received(current_turn, &tool_calls); + self.notify_tool_calls_received(current_turn, &tool_calls); let cancel = Arc::clone(&self.cancelled); let dispatch = async { diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index e01176b..d51079b 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -370,15 +370,15 @@ impl BareLoop { let mut attempt: u32 = 0; loop { - self.fire_tool_pre(turn_idx, &tc); + self.notify_tool_pre(turn_idx, &tc); if let Some(blocked) = self.check_pre_tool_use_hooks(&tc, turn_idx) { - self.fire_tool_post(turn_idx, &tc, &blocked); + self.notify_tool_post(turn_idx, &tc, &blocked); return Ok(blocked); } if let Some(blocked) = self.pre_detection(&tc, turn_idx)? { - self.fire_tool_post(turn_idx, &tc, &blocked); + self.notify_tool_post(turn_idx, &tc, &blocked); return Ok(blocked); } @@ -389,7 +389,7 @@ impl BareLoop { r = self.dispatch_tool(&tc, &tool_context, start, turn_idx) => r?, }; self.post_detection(&tc, &tool_result); - self.fire_tool_post(turn_idx, &tc, &tool_result); + self.notify_tool_post(turn_idx, &tc, &tool_result); self.notify_post_tool_use_hooks(&tc, &tool_result, turn_idx); self.record_tool_health(tc.tool.as_str(), &tool_result); @@ -417,7 +417,7 @@ impl BareLoop { /// the turn index, tool name, and tool-call ID. Called once per call before /// any hook checks, detection, or execution — observers always see this /// first, regardless of whether the call is later blocked or retried. - fn fire_tool_pre(&self, turn_idx: usize, tc: &ToolCall) { + fn notify_tool_pre(&self, turn_idx: usize, tc: &ToolCall) { self.managers.observers().on_tool_pre(&ToolPreContext { turn: turn_idx, tool: tc.tool.clone(), @@ -432,7 +432,7 @@ impl BareLoop { /// successful execution, hook block, detection block, or soft error — so /// that every `on_tool_pre` has a matching `on_tool_post`, regardless of /// the path taken. Observers can pair the two by `tool_call_id`. - fn fire_tool_post(&self, turn_idx: usize, tc: &ToolCall, result: &ToolDispatchResult) { + fn notify_tool_post(&self, turn_idx: usize, tc: &ToolCall, result: &ToolDispatchResult) { self.managers.observers().on_tool_post(&ToolPostContext { turn: turn_idx, tool: tc.tool.clone(), From 087b2dc8e8f953ac5523f39be767581e8f0ecc6b Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 27 Jul 2026 20:28:56 +1200 Subject: [PATCH 13/21] fix: rate limit hard stop check fix --- src/stream/handler.rs | 85 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 69 insertions(+), 16 deletions(-) diff --git a/src/stream/handler.rs b/src/stream/handler.rs index a58e1b2..918d41b 100644 --- a/src/stream/handler.rs +++ b/src/stream/handler.rs @@ -1456,19 +1456,27 @@ impl StreamHandler { /// Set a custom [`RateLimitConfig`]. Consuming builder. /// + /// Validates the config: if it violates any constraint (zero delay, + /// inverted ceilings, etc.), the invalid value is logged and the + /// default config is used instead. This prevents silently storing a + /// config that would invert retry/escalation behavior. + /// /// # Example /// /// ``` /// use loopctl::stream::handler::{StreamHandler, RateLimitConfig}; - /// use std::time::Duration; /// /// let handler = StreamHandler::new().with_rate_limit_config( - /// RateLimitConfig { max_retries: 2, ..Default::default() }, + /// RateLimitConfig { max_retries: 5, fallback_after_retries: 2, ..Default::default() }, /// ); - /// assert_eq!(handler.rate_limit_config().max_retries, 2); + /// assert_eq!(handler.rate_limit_config().max_retries, 5); /// ``` #[must_use] pub fn with_rate_limit_config(mut self, rl: RateLimitConfig) -> Self { + if let Err(e) = rl.validate() { + tracing::warn!(error = %e, "invalid RateLimitConfig, falling back to default"); + return self; + } self.rate_limit_config = rl; self } @@ -1693,13 +1701,16 @@ impl StreamHandler { /// Decide how to handle a rate-limit failure on the current model. /// /// Bumps `count` and returns one of: - /// - [`RateLimitRetry::Escalate`] once `count` exceeds - /// [`fallback_after_retries`](RateLimitConfig::fallback_after_retries) — the - /// caller escalates to the model circuit breaker; /// - [`RateLimitRetry::HardStop`] once `count` exceeds - /// [`max_retries`](RateLimitConfig::max_retries) — for when escalation is - /// unavailable (e.g. no fallback model configured); + /// [`max_retries`](RateLimitConfig::max_retries) — the absolute ceiling; + /// - [`RateLimitRetry::Escalate`] once `count` exceeds + /// [`fallback_after_retries`](RateLimitConfig::fallback_after_retries) + /// but is still within `max_retries` — the caller escalates to the + /// model circuit breaker; /// - [`RateLimitRetry::Retry`] with the deadline-clamped backoff otherwise. + /// + /// `max_retries` is checked first so it is always enforced as the hard + /// ceiling, regardless of `fallback_after_retries`. fn rate_limit_retry( &self, detail: &DetectedRateLimit, @@ -1707,15 +1718,15 @@ impl StreamHandler { deadline: Option, ) -> RateLimitRetry { *count = count.saturating_add(1); + if *count > self.rate_limit_config.max_retries { + return RateLimitRetry::HardStop; + } if *count > self.rate_limit_config.fallback_after_retries { return RateLimitRetry::Escalate { attempts: *count, retry_after: detail.retry_after, }; } - if *count > self.rate_limit_config.max_retries { - return RateLimitRetry::HardStop; - } let delay = clamp_delay_to_deadline(self.rate_limit_config.backoff(detail.retry_after), deadline); RateLimitRetry::Retry(delay) @@ -3274,10 +3285,8 @@ mod tests { #[test] fn rate_limit_retry_hard_stops_after_max_retries() { - // fallback_after_retries high so escalation never fires; the hard-stop - // backstop kicks in once max_retries is exceeded. let handler = StreamHandler::new().with_rate_limit_config(RateLimitConfig { - fallback_after_retries: 100, + fallback_after_retries: 1, max_retries: 2, ..Default::default() }); @@ -3294,6 +3303,49 @@ mod tests { assert_eq!(count, 3); } + #[test] + fn rate_limit_retry_max_retries_should_be_enforced_under_valid_config() { + let handler = StreamHandler::new(); + let detail = detected_limit(None); + let mut count = 0u32; + + for _ in 0..(handler.rate_limit_config().max_retries + 2) { + let _ = handler.rate_limit_retry(&detail, &mut count, None); + } + let max = handler.rate_limit_config().max_retries; + assert!( + count > max, + "count {count} must exceed max_retries {max} after enough calls" + ); + let decision = handler.rate_limit_retry(&detail, &mut count, None); + assert!( + matches!(decision, RateLimitRetry::HardStop), + "max_retries={max} should be enforced as a hard ceiling, \ + but Escalate shadows it — HardStop is dead code under valid config" + ); + } + + #[test] + fn with_rate_limit_config_should_reject_invalid() { + let invalid = RateLimitConfig { + fallback_after_retries: 10, + max_retries: 3, + ..Default::default() + }; + let result = StreamHandler::new().with_rate_limit_config(invalid); + let detail = detected_limit(None); + let mut count = 0u32; + for _ in 0..4 { + let _ = result.rate_limit_retry(&detail, &mut count, None); + } + let decision = result.rate_limit_retry(&detail, &mut count, None); + assert!( + !matches!(decision, RateLimitRetry::HardStop), + "invalid config (fallback_after=10 > max_retries=3) must not \ + silently invert behavior — HardStop should never fire before Escalation" + ); + } + #[test] fn detected_rate_limit_from_structured_variant() { let err = ApiError::RateLimit { @@ -3381,6 +3433,7 @@ mod tests { let custom = RateLimitConfig { max_retries: 2, + fallback_after_retries: 1, default_delay: Duration::from_secs(1), ..Default::default() }; @@ -3896,7 +3949,7 @@ mod tests { } let handler = StreamHandler::new().with_rate_limit_config(RateLimitConfig { - fallback_after_retries: 100, + fallback_after_retries: 2, max_retries: 2, default_delay: Duration::from_millis(1), max_delay: Duration::from_millis(1), @@ -3913,7 +3966,7 @@ mod tests { StreamHandlerError::StreamFailed(StreamOutcome::RateLimited { .. }) | StreamHandlerError::InitFailed(StreamOutcome::RateLimited { .. }) => {} StreamHandlerError::RateLimitEscalation { .. } => { - panic!("escalation must not fire when max_retries < fallback_after_retries") + panic!("escalation must not fire when max_retries == fallback_after_retries") } other => panic!("expected rate-limit outcome, got {other:?}"), } From bd73813833e3923358b68c34b358d27b28a04831 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 27 Jul 2026 20:42:55 +1200 Subject: [PATCH 14/21] fix: circuit breaker mutex for tool health --- src/tool/health.rs | 185 ++++++++++++++++++++++++--------------------- 1 file changed, 99 insertions(+), 86 deletions(-) diff --git a/src/tool/health.rs b/src/tool/health.rs index a0c7980..dd5cf94 100644 --- a/src/tool/health.rs +++ b/src/tool/health.rs @@ -59,7 +59,7 @@ use std::collections::HashMap; use std::fmt; -use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -444,17 +444,17 @@ impl Default for CircuitBreakerConfig { } } -/// Per-tool circuit breaker using atomic state. +/// Per-tool circuit breaker. /// /// When a tool fails `failure_threshold` times consecutively the breaker /// opens. After `recovery_duration` it transitions to `HalfOpen` and allows /// one probe call. If the probe succeeds the breaker closes; if it fails /// the breaker reopens. /// -/// All state is atomic — no locks needed on the hot path. The only -/// `Mutex` is for `last_failure_time`, which is updated only when a -/// failure occurs (cold path relative to the read-only `allow_request` -/// check). +/// All mutable state is behind a single `Mutex`, making every method's +/// read-modify-write atomic with respect to every other method — no +/// window for a counter/state race between concurrent `record_success` +/// and `record_failure` calls. /// /// # Example /// @@ -476,38 +476,64 @@ impl Default for CircuitBreakerConfig { /// assert!(!breaker.allow_request()); /// ``` pub struct ToolCircuitBreaker { - /// Current circuit-breaker state, encoded as a [`CircuitState`] in - /// an `AtomicU32`. + /// Number of consecutive failures that trips the breaker. /// - /// Read on every `allow_request` check and mutated only on state - /// transitions; atomic so the hot path is lock-free. - state: AtomicU32, + /// Compared against [`consecutive_failures`](BreakerState::consecutive_failures) + /// on each `record_failure` call; reaching this value while in + /// `Closed` transitions to `Open`. Set at construction and immutable + /// thereafter. + failure_threshold: u64, - /// Consecutive failures recorded since the last success. + /// How long the breaker stays `Open` before allowing a `HalfOpen` probe. /// - /// Reset to zero by [`record_success`](Self::record_success); - /// reaching `failure_threshold` while in `Closed` transitions the - /// breaker to `Open`. - consecutive_failures: AtomicU64, + /// Compared against the elapsed time since the last failure in + /// [`allow_request`](Self::allow_request). Set at construction and + /// immutable thereafter. + recovery_duration: Duration, - /// Number of consecutive failures that trips the breaker. + /// All mutable state, behind a single lock. /// - /// Set at construction and immutable thereafter; compared against - /// `consecutive_failures` on each `record_failure`. - failure_threshold: u64, + /// Holding the counter, state, and last-failure timestamp together + /// prevents the TOCTOU race where a concurrent `record_success` + /// could reset the counter while a `record_failure` is mid-transition, + /// or vice versa. + state: Mutex, +} - /// How long to remain `Open` before allowing a `HalfOpen` probe. +/// The complete mutable state of a [`ToolCircuitBreaker`]. +/// +/// Lives behind a single [`Mutex`] on the breaker so that every method's +/// read-modify-write is atomic. All three fields are always observed +/// together — no partial-state reads. +struct BreakerState { + /// Current circuit-breaker phase. /// - /// Set at construction and immutable thereafter; compared against - /// the elapsed time since `last_failure_time` in `allow_request`. - recovery_duration: Duration, + /// `Closed` at construction; transitions through `Open` (failures + /// exceeded threshold) and `HalfOpen` (recovery probe in flight) as + /// failures and successes are recorded. + circuit: CircuitState, - /// Instant of the most recent failure, or `None` if none yet. + /// Consecutive failures since the last success. /// - /// Guarded by a `Mutex` because `Instant` is not atomic and this is - /// only written on the failure cold path; the read-only - /// `allow_request` check acquires the lock briefly. - last_failure_time: Mutex>, + /// Reset to zero by success; reaching the threshold while in + /// `Closed` transitions to `Open`. + consecutive_failures: u64, + + /// When the most recent failure occurred, or `None` if none yet. + /// + /// Compared against `recovery_duration` in `allow_request` to decide + /// whether enough time has passed to allow a `HalfOpen` probe. + last_failure_time: Option, +} + +impl Default for BreakerState { + fn default() -> Self { + Self { + circuit: CircuitState::Closed, + consecutive_failures: 0, + last_failure_time: None, + } + } } impl ToolCircuitBreaker { @@ -518,11 +544,9 @@ impl ToolCircuitBreaker { #[must_use] pub fn new(recovery_duration: Duration, failure_threshold: u64) -> Self { Self { - state: AtomicU32::new(CircuitState::Closed as u32), - consecutive_failures: AtomicU64::new(0), failure_threshold, recovery_duration, - last_failure_time: Mutex::new(None), + state: Mutex::new(BreakerState::default()), } } @@ -540,28 +564,25 @@ impl ToolCircuitBreaker { /// /// - **Closed**: always allowed. /// - **Open**: allowed only if `recovery_duration` has elapsed since - /// the last failure, in which case the calling thread atomically - /// transitions to `HalfOpen` and becomes the sole probe. + /// the last failure, in which case the breaker transitions to + /// `HalfOpen` and the caller becomes the sole probe. /// - **`HalfOpen`**: already probing — no additional probes allowed /// (returns `false` to prevent thundering-herd). #[must_use] pub fn allow_request(&self) -> bool { - match self.state.load(Ordering::Acquire).into() { + let Ok(mut state) = self.state.lock() else { + return false; + }; + match state.circuit { CircuitState::Closed => true, CircuitState::HalfOpen => false, CircuitState::Open => { - let recovered = crate::error::recover_guard(self.last_failure_time.lock()) - .map(|t| t.elapsed() >= self.recovery_duration) - .unwrap_or(false); + let recovered = state + .last_failure_time + .is_some_and(|t| t.elapsed() >= self.recovery_duration); if recovered { - self.state - .compare_exchange( - CircuitState::Open as u32, - CircuitState::HalfOpen as u32, - Ordering::AcqRel, - Ordering::Acquire, - ) - .is_ok() + state.circuit = CircuitState::HalfOpen; + true } else { false } @@ -574,9 +595,10 @@ impl ToolCircuitBreaker { /// Resets consecutive failures to zero and transitions the breaker /// to Closed. pub fn record_success(&self) { - self.consecutive_failures.store(0, Ordering::Release); - self.state - .store(CircuitState::Closed as u32, Ordering::Release); + if let Ok(mut state) = self.state.lock() { + state.consecutive_failures = 0; + state.circuit = CircuitState::Closed; + } } /// Record a failed call. @@ -585,29 +607,21 @@ impl ToolCircuitBreaker { /// `failure_threshold`, the breaker transitions to Open. In the /// `HalfOpen` state, a single failure reopens the breaker. pub fn record_failure(&self) { - let failures = self - .consecutive_failures - .fetch_add(1, Ordering::AcqRel) - .saturating_add(1); - if let Ok(mut guard) = self.last_failure_time.lock() { - *guard = Some(Instant::now()); - } - let current_state = self.state.load(Ordering::Acquire).into(); - match current_state { + let Ok(mut state) = self.state.lock() else { + return; + }; + state.consecutive_failures = state.consecutive_failures.saturating_add(1); + state.last_failure_time = Some(Instant::now()); + match state.circuit { CircuitState::Closed => { - if failures >= self.failure_threshold { - self.state - .store(CircuitState::Open as u32, Ordering::Release); + if state.consecutive_failures >= self.failure_threshold { + state.circuit = CircuitState::Open; } } CircuitState::HalfOpen => { - // Probe failed — reopen immediately - self.state - .store(CircuitState::Open as u32, Ordering::Release); - } - CircuitState::Open => { - // Already open; no state change needed + state.circuit = CircuitState::Open; } + CircuitState::Open => {} } } @@ -618,7 +632,11 @@ impl ToolCircuitBreaker { /// numeric encoding. #[must_use] pub fn state_label(&self) -> &'static str { - match self.state.load(Ordering::Acquire).into() { + match self + .state + .lock() + .map_or(CircuitState::Closed, |s| s.circuit) + { CircuitState::Closed => "closed", CircuitState::Open => "open", CircuitState::HalfOpen => "half-open", @@ -627,24 +645,21 @@ impl ToolCircuitBreaker { /// Number of consecutive failures recorded since the last success. /// - /// Lock-free relaxed load of the counter reset to zero on every - /// success; reaching `failure_threshold` trips the breaker. + /// Reset to zero on every success; reaching `failure_threshold` + /// trips the breaker. #[must_use] pub fn consecutive_failures(&self) -> u64 { - self.consecutive_failures.load(Ordering::Relaxed) + self.state.lock().map_or(0, |s| s.consecutive_failures) } /// Whether the breaker is currently in the Closed (healthy) state. /// - /// `true` when requests are allowed unconditionally. Note a breaker - /// that just closed may still have a low health score, so callers - /// that want the composite picture should consult the registry. + /// `true` when requests are allowed unconditionally. #[must_use] pub fn is_closed(&self) -> bool { - matches!( - self.state.load(Ordering::Acquire).into(), - CircuitState::Closed - ) + self.state + .lock() + .is_ok_and(|s| s.circuit == CircuitState::Closed) } /// Whether the breaker is currently in the Open (blocking) state. @@ -654,10 +669,9 @@ impl ToolCircuitBreaker { /// [`allow_request`](Self::allow_request)). #[must_use] pub fn is_open(&self) -> bool { - matches!( - self.state.load(Ordering::Acquire).into(), - CircuitState::Open - ) + self.state + .lock() + .is_ok_and(|s| s.circuit == CircuitState::Open) } /// Whether the breaker is currently in the `HalfOpen` (probing) @@ -667,10 +681,9 @@ impl ToolCircuitBreaker { /// probes are refused to avoid a thundering herd. #[must_use] pub fn is_half_open(&self) -> bool { - matches!( - self.state.load(Ordering::Acquire).into(), - CircuitState::HalfOpen - ) + self.state + .lock() + .is_ok_and(|s| s.circuit == CircuitState::HalfOpen) } } From 4ab881827f9254b84f5f4cb01439ae09cc8d5047 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 28 Jul 2026 00:18:41 +1200 Subject: [PATCH 15/21] fix: message accumulation, move history to machine --- src/api.rs | 18 +-- src/engine/bare.rs | 257 +++++++++++++++++++++++++++-------- src/engine/bare/compact.rs | 2 +- src/engine/bare/dispatch.rs | 4 +- src/engine/bare/stream.rs | 14 +- src/engine/core/lifecycle.rs | 6 +- src/engine/core/machine.rs | 151 ++++++++++++++++++-- src/provider/anthropic.rs | 26 ++-- src/provider/gemini.rs | 44 +++--- src/provider/grammar.rs | 2 +- src/provider/openai.rs | 44 +++--- src/reflection/llm.rs | 23 ++-- src/stream/handler.rs | 133 ++++++++---------- src/structured.rs | 22 +-- src/testing.rs | 22 +-- tests/constrained_decode.rs | 2 +- tests/provider_e2e.rs | 2 +- 17 files changed, 500 insertions(+), 272 deletions(-) diff --git a/src/api.rs b/src/api.rs index f5ff8ab..db674cf 100644 --- a/src/api.rs +++ b/src/api.rs @@ -6,10 +6,10 @@ //! See the sub-modules for detailed documentation. pub mod error; -use crate::api::error::ApiError; use crate::message::Message; use crate::stream::StreamEvent; use crate::tool::ToolSchema; +use error::ApiError; use futures::Stream; use std::future::Future; use std::pin::Pin; @@ -229,7 +229,7 @@ pub trait ApiClient: Send + Sync { /// A pinned, boxed stream of [`Result`]. fn stream_messages( &self, - request: StreamRequest, + request: &StreamRequest, ) -> Pin> + Send + 'static>>; /// Non-streaming message request (fallback). @@ -249,7 +249,7 @@ pub trait ApiClient: Send + Sync { /// from the provider, or an [`ApiError`] if the request fails. fn create_message( &self, - request: StreamRequest, + request: &StreamRequest, ) -> Pin> + Send + '_>>; /// Streaming variant that honors [`RequestOptions`](crate::structured::RequestOptions). @@ -263,7 +263,7 @@ pub trait ApiClient: Send + Sync { /// this to inject the schema. fn stream_messages_with_options( &self, - request: StreamRequest, + request: &StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { if options.response_format.is_none() { @@ -290,7 +290,7 @@ pub trait ApiClient: Send + Sync { /// on the extracted payload. fn create_message_with_options( &self, - request: StreamRequest, + request: &StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + '_>> { if options.response_format.is_none() { @@ -378,7 +378,7 @@ mod tests { fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &StreamRequest, ) -> Pin> + Send + 'static>> { // Return a simple stream with one event let events: Vec> = vec![ @@ -406,7 +406,7 @@ mod tests { fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { @@ -426,7 +426,7 @@ mod tests { #[tokio::test] async fn test_mock_client_stream() { let client = MockClient::new("test-model"); - let stream = client.stream_messages(crate::api::StreamRequest { + let stream = client.stream_messages(&StreamRequest { messages: vec![Message::user("Hi")], system: None, tools: None, @@ -456,7 +456,7 @@ mod tests { async fn test_mock_client_create_message() { let client = MockClient::new("test-model"); let result = client - .create_message(crate::api::StreamRequest { + .create_message(&StreamRequest { messages: vec![Message::user("Hi")], system: None, tools: None, diff --git a/src/engine/bare.rs b/src/engine/bare.rs index eb537fe..8b15b15 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -336,7 +336,7 @@ impl BareLoop { s.runs.push(Run::default()); s }, - machine: LoopMachine::from_history(vec![Message::user("")]), + machine: LoopMachine::from_history(Vec::new()), managers, reflector: Arc::new(NoopReflector), recovery: Arc::new(ExponentialBackoffRecovery::new(3)), @@ -356,11 +356,8 @@ impl BareLoop { /// driving state machine; it is empty until the first /// [`run()`](crate::engine::core::Loop::run) call mints a machine for /// the run. - pub fn conversation(&self) -> &[Message] { - if self.session.session_start.is_none() { - return &[]; - } - self.machine.history() + pub fn conversation(&self) -> Vec { + self.machine.full_history() } /// Get the session configuration. @@ -1112,8 +1109,11 @@ impl BareLoop { /// # Errors /// /// Propagates whatever [`stream_turn`](Self::stream_turn) returned. - async fn do_stream(&mut self) -> Result<(Message, Option, StreamStopReason), LoopError> { - match self.stream_turn().await { + async fn do_stream( + &mut self, + contributor_messages: Vec, + ) -> Result<(Message, Option, StreamStopReason), LoopError> { + match self.stream_turn(contributor_messages).await { Ok((msg, usage, stop)) => { self.record_stream_success(usage.as_ref()); Ok((msg, usage, stop)) @@ -1437,25 +1437,24 @@ impl ModelSwitch<'_, C> { } impl BareLoop { - /// Inject any contributor messages ahead of the next model call. + /// Collect transient contributor messages for the current turn. /// /// Each registered [`ContextContributor`] is consulted against the - /// machine-owned history snapshot; any message it returns is fed into the - /// machine via [`LoopMachine::inject`] in registration order. No-op when no - /// contributors are registered. - fn inject_contributors(&mut self, current_turn: usize) { + /// machine-owned history snapshot; the returned messages are prepended + /// to the outbound [`StreamRequest`](crate::api::StreamRequest) so the + /// model sees them, but they are **not** persisted into history — they + /// appear fresh each turn and never accumulate. Returns an empty vec + /// when no contributors are registered. + fn collect_contributor_messages(&self, current_turn: usize) -> Vec { if self.contributors.is_empty() { - return; + return Vec::new(); } - let ctx = ContributorContext::new(current_turn, self.machine.history()); - let injected: Vec = self - .contributors + let full = self.machine.full_history(); + let ctx = ContributorContext::new(current_turn, &full); + self.contributors .iter() .filter_map(|contributor| contributor.contribute(&ctx)) - .collect(); - for message in injected { - self.machine.inject(message); - } + .collect() } /// Handle a model-call request from the machine. @@ -1492,7 +1491,8 @@ impl BareLoop { }; self.notify_turn_start(current_turn, &turn_input); - self.inject_contributors(current_turn); + + let contributor_messages = self.collect_contributor_messages(current_turn); let cancel = Arc::clone(&self.cancelled); tokio::select! { @@ -1509,7 +1509,7 @@ impl BareLoop { ); Err(LoopError::Cancelled) } - stream_outcome = self.do_stream() => { + stream_outcome = self.do_stream(contributor_messages) => { let (msg, usage, _stream_stop) = match stream_outcome { Ok(triple) => triple, Err(LoopError::Cancelled) => { @@ -1696,10 +1696,9 @@ impl crate::engine::core::Loop for BareLoop { self.notify_session_start(); } - self.session.history.push(Message::user(input)); self.session.runs.push(Run::new(input, run_config)); self.managers.reset_all(); - self.machine = LoopMachine::from_history(self.session.history.clone()); + self.machine.accept_input(input); loop { let policy = self.machine_policy(); @@ -1729,11 +1728,9 @@ impl crate::engine::core::Loop for BareLoop { MachineStep::Done(outcome) => match outcome { MachineOutcome::Completed { final_text } => { if let Some(run) = self.current_run_mut() { - run.output = Some(final_text.clone()); + run.output = Some(final_text); } - self.session.history.push(Message::assistant(final_text)); - break; } MachineOutcome::MaxTurnsExceeded => { @@ -1782,6 +1779,12 @@ impl crate::engine::core::Loop for BareLoop { run.end = Some(Instant::now()); } + if success { + self.machine.commit_pending(); + } else { + self.machine.discard_pending(); + } + let run = self.current_run().cloned().unwrap_or_default(); let duration = run.duration(); @@ -2050,7 +2053,7 @@ mod tests { fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { let mut guard = crate::error::recover_guard(self.responses.lock()); @@ -2067,7 +2070,7 @@ mod tests { fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { Ok(json!({"content": []})) }) } @@ -2470,6 +2473,95 @@ mod tests { ); } + #[tokio::test] + async fn compaction_sees_pending_messages() { + let client = MockClient::new("test-model"); + client.add_text_response(&"x".repeat(200)); + client.add_text_response("done"); + + let config = make_config() + .with_context_window(100) + .with_compact_threshold(10); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + agent.set_context_manager(Arc::new( + crate::compact::ContextManager::new(Arc::new( + crate::compact::TruncatingCompactor::new(), + )) + .with_context_window(100) + .with_threshold(10), + )); + + agent + .run("fill it up", &RunConfig::default()) + .await + .unwrap(); + + let conv_before = agent.conversation(); + let size_before = conv_before.len(); + + agent + .run("second run", &RunConfig::default()) + .await + .unwrap(); + + let conv_after = agent.conversation(); + let size_after = conv_after.len(); + + assert!( + size_after < size_before + 4, + "compaction must have reduced history during second run; before={size_before} after={size_after}" + ); + assert!( + conv_after.iter().any(|m| m.role == Role::User + && m.parts.iter().any(|p| matches!( + p, + MessagePart::Text { text } if text == "second run" + ))), + "second run's user input must be in committed history after success" + ); + } + + #[tokio::test] + async fn compaction_then_failure_leaves_history_compacted() { + let client = MockClient::new("test-model"); + client.add_text_response(&"x".repeat(200)); + client.add_text_response("done"); + client.add_text_response("second done"); + + let config = make_config() + .with_context_window(100) + .with_compact_threshold(10); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + agent.set_context_manager(Arc::new( + crate::compact::ContextManager::new(Arc::new( + crate::compact::TruncatingCompactor::new(), + )) + .with_context_window(100) + .with_threshold(10), + )); + + agent.run("first run", &RunConfig::default()).await.unwrap(); + + agent.cancel(); + let _ = agent.run("will fail", &RunConfig::default()).await.ok(); + agent.cancelled.reset(); + + let history = agent.conversation(); + assert!( + !history.is_empty(), + "history must contain messages from the first successful run" + ); + assert!( + !history.iter().any(|m| m.role == Role::User + && m.parts + .iter() + .any(|p| matches!(p, MessagePart::Text { text } if text == "will fail"))), + "failed run's user input must not persist in history" + ); + + agent.run("third run", &RunConfig::default()).await.unwrap(); + } + #[tokio::test] async fn observer_sequence_compaction_turn() { let client = MockClient::new("test-model"); @@ -2960,8 +3052,8 @@ mod tests { .await .unwrap(); - let user_messages: Vec<&Message> = agent - .conversation() + let conversation = agent.conversation(); + let user_messages: Vec<&Message> = conversation .iter() .filter(|m| m.role == Role::User) .collect(); @@ -3986,7 +4078,7 @@ mod tests { fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { let rx = crate::error::recover_guard(self.rx.lock()) @@ -3997,7 +4089,7 @@ mod tests { fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { Err(ApiError::api("not implemented")) }) } @@ -4196,7 +4288,7 @@ mod tests { fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin< Box< dyn futures::stream::Stream> @@ -4209,7 +4301,7 @@ mod tests { fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin< Box< dyn std::future::Future> @@ -4730,7 +4822,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { Box::pin(futures::stream::once(async { @@ -4742,7 +4834,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { Ok(json!({})) }) } @@ -4932,10 +5024,10 @@ mod tests { fn stream_messages( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { - let messages = request.messages; + let messages = request.messages.clone(); crate::error::recover_guard(self.seen.lock()).push(messages); let mut guard = crate::error::recover_guard(self.responses.lock()); if let Some(events) = guard.pop_front() { @@ -4950,11 +5042,11 @@ mod tests { fn stream_messages_with_options( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { - let messages = request.messages; + let messages = request.messages.clone(); crate::error::recover_guard(self.seen.lock()).push(messages); crate::error::recover_guard(self.seen_options.lock()).push(options); let mut guard = crate::error::recover_guard(self.responses.lock()); @@ -4970,7 +5062,7 @@ mod tests { fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { Ok(json!({"content": []})) }) } @@ -5026,8 +5118,6 @@ mod tests { agent.add_contributor(Box::new(StaticReminder("stay on task".into()))); let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); - // The injected message reached the model: it's in the captured inbound - // conversation after the user message. let seen = client.first_seen(); let texts: Vec<&str> = seen .iter() @@ -5039,14 +5129,19 @@ mod tests { }) }) .collect(); - assert!(texts.iter().any(|t| t.contains("stay on task"))); + assert!( + texts.iter().any(|t| t.contains("stay on task")), + "contributor message must reach the model in the outbound request" + ); - // And it persists in the loop's own conversation history. let persisted = agent.conversation(); - assert!(persisted.iter().any(|m| m.role == Role::System - && m.parts.iter().any( - |p| matches!(p, MessagePart::Text { text } if text.contains("stay on task")) - ))); + assert!( + !persisted.iter().any(|m| m.role == Role::System + && m.parts.iter().any( + |p| matches!(p, MessagePart::Text { text } if text.contains("stay on task")) + )), + "contributor message must NOT persist in history" + ); } #[tokio::test] @@ -5069,6 +5164,61 @@ mod tests { assert_eq!(user_count, 1, "baseline conversation has one user message"); } + #[tokio::test] + async fn failed_run_leaves_history_clean() { + let client = MockClient::new("test-model"); + client.add_text_response("done"); + + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + + agent.cancel(); + let result = agent.run("first", &RunConfig::default()).await; + assert!(result.is_err(), "run must fail"); + + let history_after_fail = agent.conversation(); + assert!( + history_after_fail.is_empty(), + "failed run must not leave messages in committed history; \ + got {} messages", + history_after_fail.len() + ); + + agent.cancelled.reset(); + agent.run("second", &RunConfig::default()).await.unwrap(); + } + + #[tokio::test] + async fn contributor_messages_must_not_accumulate_across_turns() { + let client = RecordingClient::new("test-model"); + client.add_text_response("turn 1 done"); + client.add_text_response("turn 2 done"); + let config = contributor_config(); + let mut agent = BareLoop::new(Arc::new(client.clone()), ToolRegistry::new(), config); + agent.add_contributor(Box::new(StaticReminder("stay on task".into()))); + + agent.run("first run", &RunConfig::default()).await.unwrap(); + agent + .run("second run", &RunConfig::default()) + .await + .unwrap(); + + let system_count = agent + .conversation() + .iter() + .filter(|m| m.role == Role::System) + .filter(|m| { + m.parts + .iter() + .any(|p| matches!(p, MessagePart::Text { text } if text == "stay on task")) + }) + .count(); + assert_eq!( + system_count, 0, + "contributor messages must NOT persist in history; \ + found {system_count} copies (accumulated across turns)" + ); + } + #[tokio::test] async fn test_contributor_returning_none_injects_nothing() { let client = RecordingClient::new("test-model"); @@ -5095,10 +5245,9 @@ mod tests { agent.add_contributor(Box::new(StaticReminder("second".into()))); let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); - // Find positions of the two reminders in the persisted history. - let persisted = agent.conversation(); + let seen = client.first_seen(); let pos = |needle: &str| -> Option { - persisted.iter().position(|m| { + seen.iter().position(|m| { m.role == Role::System && m.parts .iter() diff --git a/src/engine/bare/compact.rs b/src/engine/bare/compact.rs index 6dc88eb..85020c2 100644 --- a/src/engine/bare/compact.rs +++ b/src/engine/bare/compact.rs @@ -48,7 +48,7 @@ impl BareLoop { turn: usize, reason: crate::compact::types::CompactReason, ) -> Result<(Vec, u64), LoopError> { - let history = self.machine.history().to_vec(); + let history = self.machine.full_history(); let Some(ctx_manager) = self.managers.context_manager() else { return Ok((history, 0)); }; diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index d51079b..ecac203 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -933,7 +933,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin< Box< dyn futures::Stream> @@ -945,7 +945,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { Err(ApiError::http("not implemented")) }) diff --git a/src/engine/bare/stream.rs b/src/engine/bare/stream.rs index 6c4d557..8d3ddd7 100644 --- a/src/engine/bare/stream.rs +++ b/src/engine/bare/stream.rs @@ -49,13 +49,17 @@ impl BareLoop { /// fires mid-stream. pub(super) async fn stream_turn( &self, + contributor_messages: Vec, ) -> Result<(Message, Option, StreamStopReason), LoopError> { let handler = self.managers.stream_handler(); + let mut messages = contributor_messages; + messages.extend(self.machine.full_history()); + let request = crate::api::StreamRequest::new(messages) + .with_system(self.session.config.system_prompt.clone()) + .with_tools(self.build_tool_schemas()); let mut stream = handler.stream_turn( &*self.client, - crate::api::StreamRequest::new(self.machine.history().to_vec()) - .with_system(self.session.config.system_prompt.clone()) - .with_tools(self.build_tool_schemas()), + &request, self.request_options.clone(), &self.cancelled, ); @@ -183,7 +187,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -191,7 +195,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + '_>, > { diff --git a/src/engine/core/lifecycle.rs b/src/engine/core/lifecycle.rs index 90a0a49..961e853 100644 --- a/src/engine/core/lifecycle.rs +++ b/src/engine/core/lifecycle.rs @@ -45,7 +45,7 @@ use uuid::Uuid; use crate::config::ParallelDispatchConfig; use crate::engine::core::machine::MachineState; use crate::error::LoopError; -use crate::message::Message; + use crate::reflection::{Correction, CorrectionResult, CorrectionType}; /// Provide a fresh `Instant` for serde-deserialized fields. @@ -501,9 +501,6 @@ pub struct Session { /// (turns, tokens, duration) sum over this list. The last entry is the /// in-flight run (or the most recently completed one). pub runs: Vec, - - /// TODO: add docs - pub history: Vec, } impl Session { @@ -518,7 +515,6 @@ impl Session { config, session_start: None, runs: Vec::new(), - history: Vec::new(), } } diff --git a/src/engine/core/machine.rs b/src/engine/core/machine.rs index 7440f47..5f13572 100644 --- a/src/engine/core/machine.rs +++ b/src/engine/core/machine.rs @@ -337,13 +337,27 @@ pub struct MachinePolicy { /// ``` #[derive(Debug, Clone, Serialize, Deserialize)] pub struct LoopMachine { - /// The append-only conversation history. + /// The committed conversation history from previous runs. /// - /// The complete record of the run: the opening user message, every - /// assistant response, and every tool result. The driver derives the LLM - /// feed from it on each call; compaction replaces it wholesale. + /// Immutable during a run — the driver derives the LLM feed from + /// `history + pending`. Only modified by compaction (which replaces it + /// wholesale with a compacted version) and by [`commit_pending`] + /// (which appends the current run's messages on success). + /// + /// [`commit_pending`]: Self::commit_pending history: Vec, + /// Messages accumulated by the current run. + /// + /// The user input, assistant responses, and tool results from the + /// in-flight run. Cleared by [`accept_input`] at the start of each + /// run. On success, [`commit_pending`] moves these into `history`; on + /// failure they are discarded so `history` stays clean. + /// + /// [`accept_input`]: Self::accept_input + /// [`commit_pending`]: Self::commit_pending + pending: Vec, + /// Where the machine is in its request/respond cycle. /// /// One of [`MachineState`]: ready to call the model, awaiting a model @@ -396,6 +410,7 @@ impl LoopMachine { pub fn from_history(history: Vec) -> Self { Self { history, + pending: Vec::new(), state: MachineState::Start, turns_taken: 0, context_tokens: 0, @@ -404,6 +419,27 @@ impl LoopMachine { } } + /// Begin a new run on this machine, preserving the accumulated history. + /// + /// Pushes the user's input message onto the history, then resets the + /// per-run state (turn counter, context tokens, cancellation flag, + /// pending tools, state machine position) so the machine is ready for + /// a fresh `CallLLM` → tool → ... → `Done` cycle. The conversation + /// history from previous runs is untouched — the model sees the full + /// cross-run context. + /// + /// This replaces the old pattern of cloning `session.history` into a + /// fresh `LoopMachine::from_history` at every `run()` call. + pub fn accept_input(&mut self, input: &str) { + self.pending.clear(); + self.pending.push(Message::user(input)); + self.state = MachineState::Start; + self.turns_taken = 0; + self.context_tokens = 0; + self.cancelled = false; + self.pending_tools.clear(); + } + /// Return the next step the driver must perform. /// /// `policy` supplies the turn budget and compaction thresholds the machine @@ -530,13 +566,13 @@ impl LoopMachine { input: input.clone(), }) .collect(); - self.history.push(message); + self.pending.push(message); self.context_tokens = response.input_tokens; self.turns_taken = self.turns_taken.saturating_add(1); if tool_calls.is_empty() { let final_text = self - .history + .pending .last() .map(Message::text_content) .unwrap_or_default(); @@ -563,7 +599,7 @@ impl LoopMachine { if self.is_terminal() { return; } - self.history.extend(messages); + self.pending.extend(messages); self.pending_tools.clear(); self.state = MachineState::Start; } @@ -582,7 +618,7 @@ impl LoopMachine { if self.is_terminal() { return; } - self.history.push(message); + self.pending.push(message); } /// Feed compacted history back into the machine. @@ -598,6 +634,7 @@ impl LoopMachine { return; } self.history = compacted; + self.pending.clear(); self.context_tokens = tokens_after; self.state = MachineState::Start; } @@ -639,15 +676,47 @@ impl LoopMachine { self.cancelled } - /// The append-only conversation history. + /// The committed conversation history from previous runs. /// - /// Read-only access to the record the machine owns. The driver builds the - /// LLM feed from this on each [`MachineStep::CallLLM`]. + /// Does not include the current run's pending messages. Use + /// [`full_history`](Self::full_history) when you need the complete + /// conversation for an API call. #[must_use] pub fn history(&self) -> &[Message] { &self.history } + /// The full conversation: committed history plus the current run's pending messages. + /// + /// Allocates a merged `Vec` each call — use when sending the outbound + /// `StreamRequest` or when a contributor needs to see the full context. + #[must_use] + pub fn full_history(&self) -> Vec { + let mut merged = self.history.clone(); + merged.extend_from_slice(&self.pending); + merged + } + + /// Move the current run's pending messages into the committed history. + /// + /// Called by the driver on successful run completion. After this, + /// the pending buffer is empty and the committed history contains the + /// full conversation. On failure the driver simply clears pending + /// via [`accept_input`](Self::accept_input) on the next run, + /// leaving the committed history untouched. + pub fn commit_pending(&mut self) { + self.history.append(&mut self.pending); + } + + /// Discard the current run's pending messages without committing. + /// + /// Called by the driver on run failure. The committed history stays + /// clean — no orphaned tool calls, model responses, or tool results + /// from the abandoned run leak into the next run's context. + pub fn discard_pending(&mut self) { + self.pending.clear(); + } + /// The current state of the step protocol. /// /// A snapshot of where the machine is in its request/respond cycle — useful @@ -785,7 +854,7 @@ mod tests { // A text-only turn completes the run, so the machine is terminal. assert!(machine.is_terminal()); assert_eq!(machine.turns_taken(), 1); - assert_eq!(machine.history().len(), 2); + assert_eq!(machine.full_history().len(), 2); assert!(matches!( machine.state(), MachineState::Terminal(MachineOutcome::Completed { .. }) @@ -1049,7 +1118,7 @@ mod tests { let compacted = vec![Message::user("compacted-only")]; machine.compaction_result(compacted.clone(), 0); // Compare by serialized form: Message is not PartialEq. - let got = serde_json::to_string(machine.history()).expect("serialize history"); + let got = serde_json::to_string(&machine.full_history()).expect("serialize history"); let want = serde_json::to_string(&compacted).expect("serialize expected"); assert_eq!(got, want, "history must be replaced by the compacted slice"); // After compaction the next step is the deferred CallLLM. @@ -1059,6 +1128,60 @@ mod tests { )); } + #[test] + fn compaction_includes_pending_messages() { + let policy = MachinePolicy { + max_turns: 5, + context_window: 100, + compact_threshold: 50, + auto_compact: true, + }; + let mut machine = LoopMachine::from_history(vec![Message::user("from previous run")]); + machine.accept_input("current run input"); + machine.model_response(tool_response("echo", &["echo"], 60)); + let _ = machine.next_step(policy); + machine.tool_results(vec![Message::user("tool-out")]); + + let full = machine.full_history(); + assert!( + full.len() >= 4, + "full_history must include both committed history and pending messages; \ + got {} messages", + full.len() + ); + } + + #[test] + fn compaction_result_clears_pending() { + let policy = MachinePolicy { + max_turns: 5, + context_window: 100, + compact_threshold: 50, + auto_compact: true, + }; + let mut machine = LoopMachine::from_history(vec![Message::user("hello")]); + let _ = machine.next_step(policy); + machine.model_response(tool_response("echo", &["echo"], 60)); + let _ = machine.next_step(policy); + machine.tool_results(vec![Message::user("tool-out")]); + assert!(matches!( + machine.next_step(policy), + MachineStep::Compact { .. } + )); + machine.compaction_result(vec![Message::user("compacted")], 0); + + assert_eq!( + machine.history().len(), + 1, + "history must contain only the compacted message" + ); + assert_eq!( + machine.full_history().len(), + 1, + "pending must be cleared after compaction" + ); + } + #[test] fn history_accumulates_user_assistant_tool_round() { let mut machine = small_machine(5); @@ -1083,7 +1206,7 @@ mod tests { ); machine.tool_results(vec![result]); - let roles: Vec = machine.history().iter().map(|m| m.role).collect(); + let roles: Vec = machine.full_history().iter().map(|m| m.role).collect(); assert_eq!( roles, vec![Role::User, Role::Assistant, Role::User], diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs index a167e75..53d9a36 100644 --- a/src/provider/anthropic.rs +++ b/src/provider/anthropic.rs @@ -212,34 +212,31 @@ impl ApiClient for AnthropicClient { fn stream_messages( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { self.stream_messages_with_options(request, crate::structured::RequestOptions::default()) } fn create_message( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, ) -> Pin> + Send + '_>> { self.create_message_with_options(request, crate::structured::RequestOptions::default()) } fn stream_messages_with_options( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { - let crate::api::StreamRequest { - messages, - system, - tools, - } = request; + let system = request.system.clone(); + let tools = request.tools.clone(); let model = crate::error::recover_guard(self.model.lock()).clone(); let rf = options.response_format.as_ref(); let body = build_request_body( &RequestBodySpec { model: &model, - messages: &messages, + messages: &request.messages, system: system.as_deref(), tools: tools.as_deref(), response_format: rf, @@ -272,20 +269,17 @@ impl ApiClient for AnthropicClient { fn create_message_with_options( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + '_>> { - let crate::api::StreamRequest { - messages, - system, - tools, - } = request; + let system = request.system.clone(); + let tools = request.tools.clone(); let model = crate::error::recover_guard(self.model.lock()).clone(); let rf = options.response_format.as_ref(); let body = build_request_body( &RequestBodySpec { model: &model, - messages: &messages, + messages: &request.messages, system: system.as_deref(), tools: tools.as_deref(), response_format: rf, diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs index 2552b0a..cbc05fa 100644 --- a/src/provider/gemini.rs +++ b/src/provider/gemini.rs @@ -222,15 +222,12 @@ impl ApiClient for GeminiClient { fn stream_messages( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { - let crate::api::StreamRequest { - messages, - system, - tools, - } = request; + let system = request.system.clone(); + let tools = request.tools.clone(); let body = build_request_body( - &messages, + &request.messages, system.as_deref(), tools.as_deref(), None, @@ -261,15 +258,12 @@ impl ApiClient for GeminiClient { fn create_message( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, ) -> Pin> + Send + '_>> { - let crate::api::StreamRequest { - messages, - system, - tools, - } = request; + let system = request.system.clone(); + let tools = request.tools.clone(); let body = build_request_body( - &messages, + &request.messages, system.as_deref(), tools.as_deref(), None, @@ -291,17 +285,14 @@ impl ApiClient for GeminiClient { fn stream_messages_with_options( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { - let crate::api::StreamRequest { - messages, - system, - tools, - } = request; + let system = request.system.clone(); + let tools = request.tools.clone(); let rf = options.response_format.as_ref(); let body = build_request_body( - &messages, + &request.messages, system.as_deref(), tools.as_deref(), rf, @@ -332,17 +323,14 @@ impl ApiClient for GeminiClient { fn create_message_with_options( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + '_>> { - let crate::api::StreamRequest { - messages, - system, - tools, - } = request; + let system = request.system.clone(); + let tools = request.tools.clone(); let response_format = options.response_format.as_ref(); let body = build_request_body( - &messages, + &request.messages, system.as_deref(), tools.as_deref(), response_format, diff --git a/src/provider/grammar.rs b/src/provider/grammar.rs index d6498ae..4922c3e 100644 --- a/src/provider/grammar.rs +++ b/src/provider/grammar.rs @@ -178,7 +178,7 @@ mod tests { let mut total = 0usize; for prompt in corpus { let stream = client.stream_messages_with_options( - crate::api::StreamRequest::new(vec![crate::message::Message::user(prompt)]) + &crate::api::StreamRequest::new(vec![crate::message::Message::user(prompt)]) .with_tools(Some(schemas.clone())), opts.clone(), ); diff --git a/src/provider/openai.rs b/src/provider/openai.rs index 8f38bb4..58bb239 100644 --- a/src/provider/openai.rs +++ b/src/provider/openai.rs @@ -210,17 +210,14 @@ impl ApiClient for OpenAiClient { fn stream_messages( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { - let crate::api::StreamRequest { - messages, - system, - tools, - } = request; + let system = request.system.clone(); + let tools = request.tools.clone(); let model = crate::error::recover_guard(self.model.lock()).clone(); let body = RequestBody::build( &model, - &messages, + &request.messages, system.as_deref(), tools.as_deref(), None, @@ -253,17 +250,14 @@ impl ApiClient for OpenAiClient { fn create_message( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, ) -> Pin> + Send + '_>> { - let crate::api::StreamRequest { - messages, - system, - tools, - } = request; + let system = request.system.clone(); + let tools = request.tools.clone(); let model = crate::error::recover_guard(self.model.lock()).clone(); let body = RequestBody::build( &model, - &messages, + &request.messages, system.as_deref(), tools.as_deref(), None, @@ -286,19 +280,16 @@ impl ApiClient for OpenAiClient { fn stream_messages_with_options( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { - let crate::api::StreamRequest { - messages, - system, - tools, - } = request; + let system = request.system.clone(); + let tools = request.tools.clone(); let model = crate::error::recover_guard(self.model.lock()).clone(); let rf = options.response_format.as_ref(); let body = RequestBody::build( &model, - &messages, + &request.messages, system.as_deref(), tools.as_deref(), rf, @@ -331,19 +322,16 @@ impl ApiClient for OpenAiClient { fn create_message_with_options( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + '_>> { - let crate::api::StreamRequest { - messages, - system, - tools, - } = request; + let system = request.system.clone(); + let tools = request.tools.clone(); let model = crate::error::recover_guard(self.model.lock()).clone(); let rf = options.response_format.as_ref(); let body = RequestBody::build( &model, - &messages, + &request.messages, system.as_deref(), tools.as_deref(), rf, diff --git a/src/reflection/llm.rs b/src/reflection/llm.rs index 1f3bc9e..47002a7 100644 --- a/src/reflection/llm.rs +++ b/src/reflection/llm.rs @@ -291,7 +291,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin< Box< dyn futures::Stream> @@ -303,14 +303,14 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { Ok(serde_json::json!({})) }) } fn create_message_with_options( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, _options: RequestOptions, ) -> Pin> + Send + '_>> { @@ -325,7 +325,10 @@ mod tests { _ => None, }) }); - *self.captured.lock().unwrap() = Some(Captured { system, user }); + *self.captured.lock().unwrap() = Some(Captured { + system: system.clone(), + user, + }); let response = self.response.clone(); Box::pin(async move { Ok(response) }) } @@ -342,7 +345,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin< Box< dyn futures::Stream> @@ -354,14 +357,14 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { Ok(serde_json::json!({})) }) } fn create_message_with_options( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, _options: RequestOptions, ) -> Pin> + Send + '_>> { @@ -378,7 +381,7 @@ mod tests { } fn stream_messages( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, ) -> Pin< Box< dyn futures::Stream> @@ -390,14 +393,14 @@ mod tests { } fn create_message( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, ) -> Pin> + Send + '_>> { self.0.create_message(request) } fn create_message_with_options( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, _options: RequestOptions, ) -> Pin> + Send + '_>> { diff --git a/src/stream/handler.rs b/src/stream/handler.rs index 918d41b..622ddf2 100644 --- a/src/stream/handler.rs +++ b/src/stream/handler.rs @@ -1554,19 +1554,10 @@ impl StreamHandler { pub fn stream_turn<'a, C: ApiClient>( &'a self, client: &'a C, - request: crate::api::StreamRequest, + request: &'a crate::api::StreamRequest, options: crate::structured::RequestOptions, cancel: &'a Arc, ) -> Pin> + Send + 'a>> { - let crate::api::StreamRequest { - messages: conversation, - system, - tools: tool_schemas, - } = request; - // `checked_add` returns `None` when the timeout (e.g. `Duration::MAX` - // in the passthrough handler) would overflow `Instant`. Treat that as - // "no total deadline" (`None`) rather than wrapping to `Instant::now()` - // (which would immediately fire a spurious TotalTimeout). let total_deadline = Instant::now().checked_add(self.timeout_config.total_stream_timeout); let stream_start = Instant::now(); let max_attempts = self.retry_config.max_retries.saturating_add(1); @@ -1598,11 +1589,6 @@ impl StreamHandler { first_attempt = false; self.gate_on_rate_limit(client, cancel, total_deadline).await?; - let request = crate::api::StreamRequest { - messages: conversation.clone(), - system: system.clone(), - tools: tool_schemas.clone(), - }; let mut stream = client.stream_messages_with_options(request, options.clone()); @@ -1667,11 +1653,6 @@ impl StreamHandler { if transport_attempts >= max_attempts.saturating_sub(1) { if self.timeout_config.fallback_to_non_streaming { - let request = crate::api::StreamRequest { - messages: conversation.clone(), - system: system.clone(), - tools: tool_schemas.clone(), - }; let (message, fallback_stop_reason) = self.fallback_non_streaming( client, request, @@ -1939,7 +1920,7 @@ impl StreamHandler { async fn fallback_non_streaming( &self, client: &C, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, cancel: &Arc, stream_outcome: Option, ) -> Result<(Message, StreamStopReason), StreamHandlerError> { @@ -2072,7 +2053,7 @@ mod tests { async fn drive_turn( &self, client: &C, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, cancel: &Arc, ) -> Result { let mut stream = self.stream_turn( @@ -2597,7 +2578,7 @@ mod tests { fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -2607,7 +2588,7 @@ mod tests { fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + '_>, > { @@ -2638,7 +2619,7 @@ mod tests { let (message, stop_reason) = handler .fallback_non_streaming( &client, - crate::api::StreamRequest::new(vec![]), + &crate::api::StreamRequest::new(vec![]), &cancel, Some(StreamOutcome::InitFailed { last_error: "stream failed".to_string(), @@ -2679,7 +2660,7 @@ mod tests { let err = handler .fallback_non_streaming( &client, - crate::api::StreamRequest::new(vec![]), + &crate::api::StreamRequest::new(vec![]), &cancel, None, ) @@ -2707,7 +2688,7 @@ mod tests { let err = handler .fallback_non_streaming( &client, - crate::api::StreamRequest::new(vec![]), + &crate::api::StreamRequest::new(vec![]), &cancel, Some(StreamOutcome::InitFailed { last_error: "stream timeout".to_string(), @@ -2744,9 +2725,10 @@ mod tests { let client = HandlerMock::new().with_text_response("hello"); let cancel = Arc::new(CancelSignal::new()); + let req = crate::api::StreamRequest::new(vec![]); let mut stream = handler.stream_turn( &client, - crate::api::StreamRequest::new(vec![]), + &req, crate::structured::RequestOptions::default(), &cancel, ); @@ -2784,13 +2766,13 @@ mod tests { } fn stream_messages( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { self.stream_messages_with_options(request, crate::structured::RequestOptions::default()) } fn create_message( &self, - request: crate::api::StreamRequest, + request: &crate::api::StreamRequest, ) -> Pin< Box> + Send + '_>, > { @@ -2798,7 +2780,7 @@ mod tests { } fn stream_messages_with_options( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, _options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { use std::sync::atomic::Ordering; @@ -2822,7 +2804,7 @@ mod tests { } fn create_message_with_options( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, _options: crate::structured::RequestOptions, ) -> Pin< Box> + Send + '_>, @@ -2845,9 +2827,10 @@ mod tests { let client = RetryingMock { attempts }; let cancel = Arc::new(CancelSignal::new()); + let req = crate::api::StreamRequest::new(vec![]); let mut stream = handler.stream_turn( &client, - crate::api::StreamRequest::new(vec![]), + &req, crate::structured::RequestOptions::default(), &cancel, ); @@ -2870,7 +2853,7 @@ mod tests { let cancel = Arc::new(CancelSignal::new()); let result = handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("stream_turn should succeed"); @@ -2886,7 +2869,7 @@ mod tests { cancel.cancel(); let err = handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect_err("should fail on cancellation"); @@ -2915,7 +2898,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -2925,7 +2908,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -2952,7 +2935,7 @@ mod tests { let cancel = Arc::new(CancelSignal::new()); let err = handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect_err("should fail when streaming errors and fallback is disabled"); @@ -2983,7 +2966,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { // Always fail streaming — forces the handler to retry, then // fall back to create_message. @@ -2993,7 +2976,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin< Box> + Send + '_>, > { @@ -3006,7 +2989,7 @@ mod tests { } fn stream_messages_with_options( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, _options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { Box::pin(futures::stream::once(async { @@ -3015,7 +2998,7 @@ mod tests { } fn create_message_with_options( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, _options: crate::structured::RequestOptions, ) -> Pin< Box> + Send + '_>, @@ -3053,7 +3036,7 @@ mod tests { let cancel = Arc::new(CancelSignal::new()); let result = handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("fallback should succeed"); @@ -3458,7 +3441,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -3466,7 +3449,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + '_>, > { @@ -3601,7 +3584,7 @@ mod tests { let start = Instant::now(); for _ in 0..3 { handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("turn should succeed"); } @@ -3628,7 +3611,7 @@ mod tests { // First turn consumes the only token. let cancel = Arc::new(CancelSignal::new()); handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("first turn should succeed"); @@ -3637,7 +3620,7 @@ mod tests { cancel2.cancel(); let start = Instant::now(); let err = handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel2) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel2) .await .expect_err("should be cancelled, not hang for 60s"); let elapsed = start.elapsed(); @@ -3666,14 +3649,14 @@ mod tests { // First turn consumes the only token. handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("first turn should succeed"); // Second turn: bucket empty, wait capped at 50ms, then proceeds. let start = Instant::now(); let result = handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await; let elapsed = start.elapsed(); @@ -3700,7 +3683,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -3718,7 +3701,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -3752,7 +3735,7 @@ mod tests { let start = Instant::now(); let result = handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("second attempt should succeed"); let elapsed = start.elapsed(); @@ -3775,7 +3758,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -3788,7 +3771,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -3810,7 +3793,7 @@ mod tests { let cancel = Arc::new(CancelSignal::new()); let err = handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect_err("should escalate, not succeed"); match err { @@ -3845,7 +3828,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -3862,7 +3845,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -3899,7 +3882,7 @@ mod tests { let cancel = Arc::new(CancelSignal::new()); let err = handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect_err("should escalate after the rate-limit budget, not fall through"); match err { @@ -3923,7 +3906,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -3936,7 +3919,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -3959,7 +3942,7 @@ mod tests { let cancel = Arc::new(CancelSignal::new()); let err = handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect_err("hard-stop should fail the turn"); match err { @@ -3987,7 +3970,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -4005,7 +3988,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -4029,7 +4012,7 @@ mod tests { }; let cancel = Arc::new(CancelSignal::new()); handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("first call should succeed after one rate-limit retry"); @@ -4039,7 +4022,7 @@ mod tests { attempts: AtomicUsize::new(0), }; handler - .drive_turn(&client2, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client2, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("second call should not see leaked rate-limit state"); } @@ -4055,7 +4038,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -4065,7 +4048,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -4092,7 +4075,7 @@ mod tests { let cancel = Arc::new(CancelSignal::new()); let err = handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect_err("transport errors should fail"); assert!( @@ -4112,7 +4095,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -4125,7 +4108,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -4163,7 +4146,7 @@ mod tests { let start = Instant::now(); let err = handler - .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .drive_turn(&client, &crate::api::StreamRequest::new(vec![]), &cancel) .await .expect_err("tight total timeout should fail the turn"); let elapsed = start.elapsed(); @@ -4189,7 +4172,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -4199,7 +4182,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -4230,7 +4213,7 @@ mod tests { let err = handler .drive_turn( &AlwaysFailingMock, - crate::api::StreamRequest::new(vec![]), + &crate::api::StreamRequest::new(vec![]), &cancel, ) .await diff --git a/src/structured.rs b/src/structured.rs index 17c9dbc..033447c 100644 --- a/src/structured.rs +++ b/src/structured.rs @@ -604,7 +604,7 @@ pub async fn request_structured Pin< Box< dyn futures::Stream< @@ -761,7 +761,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin< Box< dyn Future> @@ -779,7 +779,7 @@ mod tests { let opts = RequestOptions::new().with_response_format(ResponseFormat::from_type::()); let request = crate::api::StreamRequest::new(vec![]); - let result = client.create_message_with_options(request, opts).await; + let result = client.create_message_with_options(&request, opts).await; assert!( result.is_err(), "client without structured-output support should reject response_format" @@ -796,7 +796,7 @@ mod tests { let client = PlainMockClient; let opts = RequestOptions::new(); let request = crate::api::StreamRequest::new(vec![]); - let result = client.create_message_with_options(request, opts).await; + let result = client.create_message_with_options(&request, opts).await; assert!(result.is_ok(), "empty options should delegate normally"); } @@ -807,7 +807,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin< Box< dyn futures::Stream< @@ -820,7 +820,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin< Box< dyn Future> @@ -832,7 +832,7 @@ mod tests { } fn create_message_with_options( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, _options: RequestOptions, ) -> Pin< Box< @@ -866,7 +866,7 @@ mod tests { } fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin< Box< dyn futures::Stream< @@ -879,7 +879,7 @@ mod tests { } fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin< Box< dyn Future> @@ -891,7 +891,7 @@ mod tests { } fn create_message_with_options( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, _options: RequestOptions, ) -> Pin< Box< diff --git a/src/testing.rs b/src/testing.rs index b658b75..af287d9 100644 --- a/src/testing.rs +++ b/src/testing.rs @@ -632,14 +632,14 @@ impl ApiClient for MockApiClient { /// use loopctl::testing::MockApiClient; /// /// let client = MockApiClient::new("test-model").with_text_response("Hi!"); - /// let stream = client.stream_messages(StreamRequest::new(vec![])); + /// let stream = client.stream_messages(&StreamRequest::new(vec![])); /// let events: Vec<_> = futures::StreamExt::collect(stream).await; /// assert!(events.len() >= 4); /// # }); /// ``` fn stream_messages( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { if let Some(ref err) = self.error { let err = err.clone(); @@ -716,13 +716,13 @@ impl ApiClient for MockApiClient { /// use loopctl::testing::MockApiClient; /// /// let client = MockApiClient::new("test-model").with_text_response("Hi!"); - /// let result = client.create_message(StreamRequest::new(vec![])).await; + /// let result = client.create_message(&StreamRequest::new(vec![])).await; /// assert_eq!(result.unwrap()["content"][0]["text"], "Hi!"); /// # }); /// ``` fn create_message( &self, - _request: crate::api::StreamRequest, + _request: &crate::api::StreamRequest, ) -> Pin> + Send + '_>> { if let Some(ref err) = self.error { let err = err.clone(); @@ -1371,7 +1371,7 @@ mod tests { #[tokio::test] async fn test_mock_client_default_response() { let client = MockApiClient::new("test-model"); - let stream = client.stream_messages(crate::api::StreamRequest { + let stream = client.stream_messages(&crate::api::StreamRequest { messages: vec![Message::user("Hi")], system: None, tools: None, @@ -1384,7 +1384,7 @@ mod tests { #[tokio::test] async fn test_mock_client_custom_text() { let client = MockApiClient::new("test-model").with_text_response("Custom response"); - let stream = client.stream_messages(crate::api::StreamRequest { + let stream = client.stream_messages(&crate::api::StreamRequest { messages: vec![Message::user("Hi")], system: None, tools: None, @@ -1408,7 +1408,7 @@ mod tests { "echo", json!({"message": "hi"}), ); - let stream = client.stream_messages(crate::api::StreamRequest { + let stream = client.stream_messages(&crate::api::StreamRequest { messages: vec![Message::user("Hi")], system: None, tools: None, @@ -1438,7 +1438,7 @@ mod tests { #[tokio::test] async fn test_mock_client_error() { let client = MockApiClient::new("test-model").with_error("API error"); - let stream = client.stream_messages(crate::api::StreamRequest { + let stream = client.stream_messages(&crate::api::StreamRequest { messages: vec![Message::user("Hi")], system: None, tools: None, @@ -1463,7 +1463,7 @@ mod tests { }, ]); - let stream1 = client.stream_messages(crate::api::StreamRequest { + let stream1 = client.stream_messages(&crate::api::StreamRequest { messages: vec![Message::user("Hi")], system: None, tools: None, @@ -1479,7 +1479,7 @@ mod tests { }); assert!(has_first); - let stream2 = client.stream_messages(crate::api::StreamRequest { + let stream2 = client.stream_messages(&crate::api::StreamRequest { messages: vec![Message::user("Hi")], system: None, tools: None, @@ -1500,7 +1500,7 @@ mod tests { async fn test_mock_client_create_message() { let client = MockApiClient::new("test-model").with_text_response("Hello!"); let result = client - .create_message(crate::api::StreamRequest { + .create_message(&crate::api::StreamRequest { messages: vec![Message::user("Hi")], system: None, tools: None, diff --git a/tests/constrained_decode.rs b/tests/constrained_decode.rs index b96bfc6..38f73b6 100644 --- a/tests/constrained_decode.rs +++ b/tests/constrained_decode.rs @@ -64,7 +64,7 @@ async fn strict_mode_produces_valid_tool_calls() { let req = loopctl::api::StreamRequest::new(vec![loopctl::message::Message::user(prompt)]) .with_tools(Some(tools.clone())); let opts = RequestOptions::default().with_tool_constraint(ToolConstraint::Strict); - let stream = client.stream_messages_with_options(req, opts); + let stream = client.stream_messages_with_options(&req, opts); let events = match helpers::collect_events(Box::pin(stream)).await { Some(events) => events, diff --git a/tests/provider_e2e.rs b/tests/provider_e2e.rs index 6389ce7..4b0be98 100644 --- a/tests/provider_e2e.rs +++ b/tests/provider_e2e.rs @@ -43,7 +43,7 @@ async fn run_provider_test(client: &dyn ApiClient, name: &str) { let req = loopctl::api::StreamRequest::new(vec![loopctl::message::Message::user( "Say hello in exactly 3 words.", )]); - let stream = client.stream_messages(req); + let stream = client.stream_messages(&req); let mut stream = std::pin::pin!(stream); let mut events = Vec::new(); From 0fbe2af0ab83c851ba602655abbd5ce1af41826d Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 28 Jul 2026 00:31:29 +1200 Subject: [PATCH 16/21] fix: gemini provider stream, token usage --- src/provider/gemini.rs | 219 ++++++++++++++++++++++++++++++++++------- 1 file changed, 186 insertions(+), 33 deletions(-) diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs index cbc05fa..123e434 100644 --- a/src/provider/gemini.rs +++ b/src/provider/gemini.rs @@ -31,7 +31,7 @@ use crate::api::error::ApiError; use crate::message::{Message, MessagePart, Role}; use crate::stream::{ DeltaPart, IndexedDelta, MessageDelta, MessageDeltaPayload, MessageMetadata, MessageStart, - PartStart, StreamEvent, StreamStopReason, + PartStart, StreamEvent, StreamStopReason, Usage, }; use crate::structured::ToolConstraint; use crate::structured::tighten_json_schema; @@ -757,6 +757,21 @@ struct StreamEmitter { /// symmetric to the text lane. thinking_part_open: bool, + /// Whether a tool-call part is currently open. + /// + /// Each `functionCall` part emits a [`PartStart`]. Without tracking, + /// [`extract_finish_reason`](Self::extract_finish_reason) would never + /// close tool parts, leaving the downstream accumulator with unbalanced + /// `PartStart`/`PartStop` events. + tool_part_open: bool, + + /// Next tool-call index for multiple function calls in one response. + /// + /// Gemini can emit several `functionCall` parts in a single chunk. + /// Each gets its own `PartStart` at an incrementing index so the + /// accumulator can distinguish them. + next_tool_index: usize, + /// Whether the terminal stop signal has been processed. /// /// Set by [`finish`](Self::finish) when it appends the synthetic @@ -906,41 +921,54 @@ impl StreamEmitter { return; }; - let Some(func_call) = parts.iter().find_map(|p| p.get("functionCall")) else { + let func_calls: Vec<(&Value, usize)> = parts + .iter() + .enumerate() + .filter_map(|(i, p)| p.get("functionCall").map(|fc| (fc, i))) + .collect(); + if func_calls.is_empty() { return; - }; - 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(); - - if self.thinking_part_open { - self.thinking_part_open = false; - self.push(StreamEvent::PartStop); } - if self.text_part_open { - self.text_part_open = false; + for (func_call, _part_idx) in func_calls { + 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(); + + if self.thinking_part_open { + self.thinking_part_open = false; + self.push(StreamEvent::PartStop); + } + if self.text_part_open { + self.text_part_open = false; + self.push(StreamEvent::PartStop); + } + + let idx = self.next_tool_index; + self.next_tool_index = self.next_tool_index.saturating_add(1); + self.tool_part_open = true; + + self.push(StreamEvent::PartStart(PartStart { + index: idx, + part: Some(MessagePart::ToolCall { + id: String::new(), + name, + input: args, + }), + })); + self.push(StreamEvent::IndexedDelta(IndexedDelta { + index: idx, + delta: DeltaPart::InputJson { + partial_json: args_str, + }, + })); self.push(StreamEvent::PartStop); + self.tool_part_open = false; } - - 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 the finish reason from the chunk and emit stop events. @@ -972,16 +1000,40 @@ impl StreamEmitter { self.push(StreamEvent::PartStop); } + if self.tool_part_open { + self.tool_part_open = false; + self.push(StreamEvent::PartStop); + } + if self.text_part_open { self.text_part_open = false; self.push(StreamEvent::PartStop); } + let usage = json.pointer("/usageMetadata").and_then(|u| { + let input = u + .get("promptTokenCount") + .and_then(Value::as_u64) + .unwrap_or(0); + let output = u + .get("candidatesTokenCount") + .and_then(Value::as_u64) + .unwrap_or(0); + if input == 0 && output == 0 { + None + } else { + Some(Usage::new( + u32::try_from(input).unwrap_or(0), + u32::try_from(output).unwrap_or(0), + )) + } + }); + self.push(StreamEvent::MessageDelta(MessageDelta { delta: MessageDeltaPayload { stop_reason: Some(stop.to_api_str().into()), }, - usage: None, + usage, })); } @@ -1732,7 +1784,7 @@ mod tests { .iter() .filter(|e| matches!(e, StreamEvent::PartStop)) .count(); - assert_eq!(stops, 1, "text lane must be closed before tool PartStart"); + assert_eq!(stops, 2, "text lane closed + tool part closed"); // Ordering: TextDelta → PartStop → tool PartStart. let text_delta_idx = events @@ -2282,4 +2334,105 @@ mod tests { assert_eq!(contents[0]["parts"][0]["text"], "first"); assert_eq!(contents[2]["parts"][0]["text"], "third"); } + + #[test] + fn emitter_multiple_function_calls_per_chunk() { + let mut em = StreamEmitter::default(); + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": [ + {"functionCall": {"name": "search", "args": {"q": "rust"}}}, + {"functionCall": {"name": "write", "args": {"path": "/tmp"}}} + ] + } + }] + })); + let events = em.drain(); + + let tool_starts: Vec<_> = events + .iter() + .filter_map(|e| match e { + StreamEvent::PartStart(PartStart { + index, + part: Some(MessagePart::ToolCall { name, .. }), + }) => Some((*index, name.clone())), + _ => None, + }) + .collect(); + assert_eq!( + tool_starts.len(), + 2, + "both function calls must emit PartStart" + ); + assert_eq!(tool_starts[0].1, "search"); + assert_eq!(tool_starts[1].1, "write"); + assert_ne!( + tool_starts[0].0, tool_starts[1].0, + "each tool call must get a distinct index" + ); + + let stops = events + .iter() + .filter(|e| matches!(e, StreamEvent::PartStop)) + .count(); + assert_eq!(stops, 2, "each tool call must emit its own PartStop"); + } + + #[test] + fn emitter_finish_closes_open_tool_part() { + let mut em = StreamEmitter::default(); + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": [{"functionCall": {"name": "search", "args": {"q": "rust"}}}] + } + }] + })); + em.drain(); + em.process_chunk(&serde_json::json!({ + "candidates": [{"finishReason": "STOP"}] + })); + let events = em.drain(); + + let has_message_delta = events.iter().any(|e| { + matches!( + e, + StreamEvent::MessageDelta(MessageDelta { + delta: MessageDeltaPayload { + stop_reason: Some(_), + }, + .. + }) + ) + }); + assert!(has_message_delta, "finish must emit a MessageDelta"); + } + + #[test] + fn emitter_finish_extracts_usage_metadata() { + let mut em = StreamEmitter::default(); + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": {"parts": [{"text": "hi"}]} + }] + })); + em.drain(); + em.process_chunk(&serde_json::json!({ + "candidates": [{"finishReason": "STOP"}], + "usageMetadata": { + "promptTokenCount": 42, + "candidatesTokenCount": 7 + } + })); + let events = em.drain(); + + let usage = events.iter().find_map(|e| match e { + StreamEvent::MessageDelta(MessageDelta { usage: Some(u), .. }) => Some(*u), + _ => None, + }); + let usage = usage.expect("finish must include Usage from usageMetadata"); + assert_eq!(usage.input_tokens, 42); + assert_eq!(usage.output_tokens, 7); + } } From d0753bb5adc1245d9defbd8f3f72d65e6aacfb3d Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 28 Jul 2026 20:34:45 +1200 Subject: [PATCH 17/21] refactor: replace session start/end events with run start/end, add reset flag to run config --- CHANGELOG.md | 47 +++++++++++++ README.md | 19 +++-- examples/chat.rs | 5 +- src/capabilities.rs | 2 +- src/engine/bare.rs | 108 ++++++++++++++++++---------- src/engine/bare/emission.rs | 113 ++++++++++++++---------------- src/engine/core/lifecycle.rs | 41 ++++++++++- src/hooks.rs | 14 ++-- src/hooks/builtin.rs | 2 +- src/hooks/builtin/auto_commit.rs | 20 +++--- src/hooks/builtin/logging_hook.rs | 18 ++--- src/hooks/context.rs | 72 +++++++++---------- src/hooks/executor.rs | 70 +++++++++--------- src/managers.rs | 16 +++-- src/observer.rs | 53 +++++++------- src/observer/context.rs | 34 ++++----- 16 files changed, 373 insertions(+), 261 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index de1aa9f..7f8663a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -214,6 +214,53 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. ### Changed +- **Breaking (session→run lifecycle rename):** the observer and hook events + formerly named `on_session_start` / `on_session_end` are renamed to + `on_run_start` / `on_run_end`. These events fire once per `run()` call and + carry per-run data (turn count, per-run duration, per-run tokens), so the + new names match their actual semantics. The Session ⊃ Run ⊃ Turn hierarchy + is unchanged: a Session spans the agent's lifetime, a Run is one `run()` + call, a Turn is one loop iteration. Affected APIs (old → new): + `LoopObserver::on_session_start` → `on_run_start`, + `LoopObserver::on_session_end` → `on_run_end`, + `observer::SessionStartContext` → `observer::RunStartContext`, + `observer::SessionEndContext` → `observer::RunEndContext`, + `Hook::on_session_start` → `on_run_start`, + `Hook::on_session_end` → `on_run_end`, + `hooks::context::SessionStartContext` → `hooks::context::RunStartContext`, + `hooks::context::SessionEndContext` → `hooks::context::RunEndContext`, + `hooks::context::SessionEndReason` → `hooks::context::RunEndReason`, + `HookExecutor::notify_session_start` → `notify_run_start`, + `HookExecutor::notify_session_end` → `notify_run_end`, + `HookExecutor::notify_session_start_async` → `notify_run_start_async`, + `HookExecutor::notify_session_end_async` → `notify_run_end_async`, + `ObserverHost::on_session_start` → `on_run_start`, + `ObserverHost::on_session_end` → `on_run_end`. + Migration: rename the method/trait impl in every `impl LoopObserver` and + `impl Hook`; rename every `SessionStartContext` / `SessionEndContext` / + `SessionEndReason` reference. The `session_id` field on the renamed context + types is unchanged (a run belongs to a session). +- **Breaking (per-run manager reset):** `LoopManagers::reset_all` is no + longer called automatically from `run()`. Previously it fired on every + `run()` call, wiping session-scoped manager state (fallback circuit + breaker, loop detection, observers) between runs. On a fresh session + the managers are already in their default state, so the call was a + no-op on the first run and a correctness bug on subsequent runs (it + discarded accumulated circuit-breaker and detection state). To + reinitialise mid-session, call `managers.reset_all()` explicitly. +- `RunConfig` gains a `reset_managers: bool` field (default `false`). Set it + to `true` when a run is logically independent from the previous one and you + want fresh circuit-breaker / detection / observer state for that run. This + replaces the removed automatic `reset_all` with explicit, per-run control. + Because `RunConfig` is `#[non_exhaustive]`, existing code that constructs it + via `Default::default()` or `RunConfig { .. }` continues to work unchanged. +- `on_run_start` (the renamed `on_session_start`) now fires at the start of + **every** `run()` call, matching `on_run_end` which already fired on every + `run()` call. Previously it fired only once (on the first `run()`), creating + an asymmetry where a multi-run session saw 1 start event but N end events. + The session-scoped bookkeeping (`session_start` timestamp, `reset_all`) is + unaffected — it still runs only once, on the first `run()`. + - **Breaking (machine-driven engine):** `BareLoop::run()` is now a `match machine.next_step()` loop driving a `LoopMachine`. The machine owns the conversation history and every loop decision (turn count, max-turn, tool-call diff --git a/README.md b/README.md index 507e3fb..bca04b0 100644 --- a/README.md +++ b/README.md @@ -80,8 +80,9 @@ impl Tool for EchoTool { ```rust,no_run use loopctl::engine::BareLoop; use loopctl::engine::core::Loop; +use loopctl::engine::RunConfig; use loopctl::tool::ToolRegistry; -use loopctl::config::LoopConfig; +use loopctl::config::SessionConfig; use std::sync::Arc; // 1. Bring your own API client (implements ApiClient trait) @@ -89,8 +90,8 @@ use std::sync::Arc; # 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>> { +# fn stream_messages(&self, _req: &loopctl::api::StreamRequest) +# -> std::pin::Pin> + Send>> { # unimplemented!() # } # } @@ -101,16 +102,12 @@ let mut registry = ToolRegistry::new(); // registry.register(EchoTool); // 3. Configure -let config = LoopConfig { - max_turns: 50, - model: "llm-70b".into(), - ..Default::default() -}; +let config = SessionConfig::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 mut agent = BareLoop::new(client, registry, config); +// let result = agent.run("Use the echo tool to say hello", &RunConfig::default()).await?; +// println!("Completed in {} turns", result.turn_count()); ``` ### Use the Testing Module diff --git a/examples/chat.rs b/examples/chat.rs index 9eb3f84..80eccd5 100644 --- a/examples/chat.rs +++ b/examples/chat.rs @@ -329,10 +329,7 @@ async fn run_repl(client: Arc) { // Create the agent once — conversation history persists across inputs. let config = SessionConfig::default(); - let run_config = RunConfig { - max_turns: 10, - ..RunConfig::default() - }; + let run_config = RunConfig::default().with_max_turns(10); let mut agent = BareLoop::new(client, build_tools(), config); agent.register_observer(Arc::new(PrintingObserver)); diff --git a/src/capabilities.rs b/src/capabilities.rs index 219a598..452a0c6 100644 --- a/src/capabilities.rs +++ b/src/capabilities.rs @@ -204,7 +204,7 @@ pub trait StreamCapable { /// # When to use /// /// Use this trait bound when you need to run hooks before or after -/// tool dispatch, compaction, or session start/end. +/// tool dispatch, compaction, or run start/end. #[cfg(feature = "hooks")] pub trait Hookable { /// Returns the hook executor, if hooks are configured. diff --git a/src/engine/bare.rs b/src/engine/bare.rs index 8b15b15..69011a6 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -77,7 +77,7 @@ use crate::hooks::context::{ CompactTrigger, PostCompactContext, PostToolUseContext, PreCompactContext, PreToolUseContext, }; #[cfg(all(test, feature = "hooks"))] -use crate::hooks::context::{SessionEndContext as HookSessionEndContext, SessionEndReason}; +use crate::hooks::context::{RunEndContext as HookRunEndContext, RunEndReason}; #[cfg(feature = "hooks")] use crate::hooks::{HookAction, HookExecutor}; use crate::managers::LoopManagers; @@ -196,7 +196,8 @@ pub struct BareLoop { /// - Optional [`HookExecutor`] — bidirectional lifecycle hooks. /// - Optional [`ToolHealthRegistry`] — per-tool health tracking. /// - /// Reset at the start of every session via [`LoopManagers::reset_all`]. + /// Fresh on construction; call [`LoopManagers::reset_all`] to + /// reinitialise mid-session. /// /// [`FallbackManager`]: crate::fallback::FallbackManager /// [`DetectionManager`]: crate::detection::DetectionManager @@ -663,7 +664,7 @@ impl BareLoop { /// Set the [`HookExecutor`] for lifecycle interception. /// /// When set, the executor runs registered hooks before and after - /// tool dispatch, compaction, and session start/end. Hooks can + /// tool dispatch, compaction, and run 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). @@ -1693,11 +1694,14 @@ impl crate::engine::core::Loop for BareLoop { let session_is_new = self.session.session_start.is_none(); if session_is_new { self.session.session_start = Some(Instant::now()); - self.notify_session_start(); + } + + if run_config.reset_managers { + self.managers.reset_all(); } self.session.runs.push(Run::new(input, run_config)); - self.managers.reset_all(); + self.notify_run_start(); self.machine.accept_input(input); loop { @@ -1764,7 +1768,7 @@ impl crate::engine::core::Loop for BareLoop { /// /// Every `run()` exit path — clean completion, error, max-turns, /// cancellation — funnels through here. Records the run's end - /// timestamp, fires the session-end observers, and re-arms the + /// timestamp, fires the run-end observers, and re-arms the /// cancel signal so the next `run()` starts clean. Re-arming here /// (rather than at the top of `run()`) preserves a cancel that /// arrived before the run: the run observes it and returns @@ -1788,7 +1792,7 @@ impl crate::engine::core::Loop for BareLoop { let run = self.current_run().cloned().unwrap_or_default(); let duration = run.duration(); - self.notify_session_end(&run, success, duration); + self.notify_run_end(&run, success, duration); self.cancelled.reset(); Ok(run) @@ -2203,8 +2207,8 @@ mod tests { } struct CountingObserver { - session_starts: AtomicUsize, - session_ends: AtomicUsize, + run_starts: AtomicUsize, + run_ends: AtomicUsize, turn_starts: AtomicUsize, turn_ends: AtomicUsize, tool_calls_received: AtomicUsize, @@ -2215,8 +2219,8 @@ mod tests { impl CountingObserver { fn new() -> Self { Self { - session_starts: AtomicUsize::new(0), - session_ends: AtomicUsize::new(0), + run_starts: AtomicUsize::new(0), + run_ends: AtomicUsize::new(0), turn_starts: AtomicUsize::new(0), turn_ends: AtomicUsize::new(0), tool_calls_received: AtomicUsize::new(0), @@ -2231,12 +2235,12 @@ mod tests { "counting" } - fn on_session_start(&self, _ctx: &crate::observer::SessionStartContext) { - self.session_starts.fetch_add(1, Ordering::SeqCst); + fn on_run_start(&self, _ctx: &crate::observer::RunStartContext) { + self.run_starts.fetch_add(1, Ordering::SeqCst); } - fn on_session_end(&self, _ctx: &crate::observer::SessionEndContext) { - self.session_ends.fetch_add(1, Ordering::SeqCst); + fn on_run_end(&self, _ctx: &crate::observer::RunEndContext) { + self.run_ends.fetch_add(1, Ordering::SeqCst); } fn on_turn_start(&self, _ctx: &crate::observer::TurnStartContext) { @@ -2833,12 +2837,40 @@ mod tests { let _result = agent.run("Hi", &RunConfig::default()).await.unwrap(); - assert_eq!(plugin.session_starts.load(Ordering::SeqCst), 1); - assert_eq!(plugin.session_ends.load(Ordering::SeqCst), 1); + assert_eq!(plugin.run_starts.load(Ordering::SeqCst), 1); + assert_eq!(plugin.run_ends.load(Ordering::SeqCst), 1); assert_eq!(plugin.turn_starts.load(Ordering::SeqCst), 1); assert_eq!(plugin.turn_ends.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn test_observer_run_start_end_symmetry_across_multiple_runs() { + let client = MockClient::new("test-model"); + client.add_text_response("first"); + client.add_text_response("second"); + client.add_text_response("third"); + + let plugin = Arc::new(CountingObserver::new()); + let config = make_config(); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + agent.register_observer(plugin.clone()); + + for _ in 0..3 { + let _ = agent.run("Hi", &RunConfig::default()).await.unwrap(); + } + + assert_eq!( + plugin.run_starts.load(Ordering::SeqCst), + 3, + "on_run_start must fire once per run" + ); + assert_eq!( + plugin.run_ends.load(Ordering::SeqCst), + 3, + "on_run_end must fire once per run" + ); + } + #[tokio::test] async fn test_observer_tool_events() { let client = MockClient::new("test-model"); @@ -4450,7 +4482,7 @@ mod tests { #[cfg(feature = "hooks")] struct ReasonCaptureHook { - reason: Mutex>, + reason: Mutex>, } #[cfg(feature = "hooks")] @@ -4461,7 +4493,7 @@ mod tests { }) } - fn captured(&self) -> Option { + fn captured(&self) -> Option { *crate::error::recover_guard(self.reason.lock()) } } @@ -4472,7 +4504,7 @@ mod tests { "ReasonCaptureHook" } - fn on_session_end(&self, ctx: &HookSessionEndContext) { + fn on_run_end(&self, ctx: &HookRunEndContext) { *crate::error::recover_guard(self.reason.lock()) = Some(ctx.reason); } } @@ -4499,7 +4531,7 @@ mod tests { #[cfg(feature = "hooks")] #[tokio::test] - async fn session_end_reason_complete() { + async fn run_end_reason_complete() { let (mut loop_, hook) = loop_with_reason_hook(); // Normal completion: success true, not cancelled, under max_turns. loop_.current_run_mut().unwrap().turns = vec![ @@ -4521,18 +4553,18 @@ mod tests { }, ]; - loop_.notify_session_end( + loop_.notify_run_end( &loop_.current_run().unwrap().clone(), true, Duration::from_millis(100), ); - assert_eq!(hook.captured(), Some(SessionEndReason::Complete)); + assert_eq!(hook.captured(), Some(RunEndReason::Complete)); } #[cfg(feature = "hooks")] #[tokio::test] - async fn session_end_reason_cancelled() { + async fn run_end_reason_cancelled() { let (mut loop_, hook) = loop_with_reason_hook(); // Cancel signal fired — success is true (not Failed) but cancelled. loop_.current_run_mut().unwrap().turns = vec![ @@ -4555,18 +4587,18 @@ mod tests { ]; loop_.cancelled.cancel(); - loop_.notify_session_end( + loop_.notify_run_end( &loop_.current_run().unwrap().clone(), true, Duration::from_millis(100), ); - assert_eq!(hook.captured(), Some(SessionEndReason::Cancelled)); + assert_eq!(hook.captured(), Some(RunEndReason::Cancelled)); } #[cfg(feature = "hooks")] #[tokio::test] - async fn session_end_reason_max_turns() { + async fn run_end_reason_max_turns() { let (mut loop_, hook) = loop_with_reason_hook(); // Hit max_turns: turn count == max_turns, not cancelled, success true. loop_.current_run_mut().unwrap().turns = (0..5) @@ -4580,41 +4612,41 @@ mod tests { }) .collect(); - loop_.notify_session_end( + loop_.notify_run_end( &loop_.current_run().unwrap().clone(), true, Duration::from_millis(100), ); - assert_eq!(hook.captured(), Some(SessionEndReason::MaxTurns)); + assert_eq!(hook.captured(), Some(RunEndReason::MaxTurns)); } #[cfg(feature = "hooks")] #[tokio::test] - async fn session_end_reason_error() { + async fn run_end_reason_error() { let (loop_, hook) = loop_with_reason_hook(); - loop_.notify_session_end( + loop_.notify_run_end( &loop_.current_run().unwrap().clone(), false, Duration::from_millis(100), ); - assert_eq!(hook.captured(), Some(SessionEndReason::Error)); + assert_eq!(hook.captured(), Some(RunEndReason::Error)); } #[cfg(feature = "hooks")] #[tokio::test] - async fn session_end_reason_context_overflow() { + async fn run_end_reason_context_overflow() { let (loop_, hook) = loop_with_reason_hook(); - loop_.notify_session_end( + loop_.notify_run_end( &loop_.current_run().unwrap().clone(), false, Duration::from_millis(100), ); - assert_eq!(hook.captured(), Some(SessionEndReason::Error)); + assert_eq!(hook.captured(), Some(RunEndReason::Error)); } #[tokio::test] @@ -4727,9 +4759,9 @@ mod tests { "on_turn_end(false) must fire on cancel during dispatch", ); assert_eq!( - observer.session_ends.load(Ordering::SeqCst), + observer.run_ends.load(Ordering::SeqCst), 1, - "on_session_end must fire via finalize after cancel", + "on_run_end must fire via finalize after cancel", ); } @@ -5341,7 +5373,7 @@ mod tests { let config = contributor_config(); let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); // The first run() establishes the session (capturing the start time - // and firing on_session_start), moving the loop out of Idle. A + // and firing on_run_start), moving the loop out of Idle. A // subsequent add_contributor must panic in debug builds (matches // set_reflector's contract). // Box the future so we can drop it without awaiting; the session-init diff --git a/src/engine/bare/emission.rs b/src/engine/bare/emission.rs index 49d7e5e..e59a9cd 100644 --- a/src/engine/bare/emission.rs +++ b/src/engine/bare/emission.rs @@ -1,6 +1,6 @@ -//! Session lifecycle notifications — start and end events. +//! Run lifecycle notifications — start and end events. //! -//! Fires observer callbacks and hook notifications when a session begins and +//! Fires observer callbacks and hook notifications when a run begins and //! ends. Other observer events (`on_turn_start`, `on_response`, etc.) are fired //! directly at their call sites in the driver loop. @@ -9,53 +9,48 @@ use super::{ApiClient, BareLoop, Duration, Run}; use crate::capabilities::Hookable; #[cfg(feature = "hooks")] use crate::hooks::context::{ - SessionEndContext as HookSessionEndContext, SessionEndReason, - SessionStartContext as HookSessionStartContext, + RunEndContext as HookRunEndContext, RunEndReason, RunStartContext as HookRunStartContext, }; -use crate::observer::{SessionEndContext, SessionStartContext}; +use crate::observer::{RunEndContext, RunStartContext}; impl BareLoop { - /// Notify all observers and hooks that the session has started. + /// Notify all observers and hooks that a run has started. /// - /// Fires [`on_session_start`](crate::observer::LoopObserver::on_session_start) + /// Fires [`on_run_start`](crate::observer::LoopObserver::on_run_start) /// on every registered observer with the session id, then fires the - /// `on_session_start` hook (when the `hooks` feature is enabled and a + /// `on_run_start` hook (when the `hooks` feature is enabled and a /// hook executor is configured). Called once at the beginning of a /// run, before the first turn. - pub(super) fn notify_session_start(&self) { - self.managers - .observers() - .on_session_start(&SessionStartContext { - session_id: self.session.id, - }); - self.notify_session_start_hook(); + pub(super) fn notify_run_start(&self) { + self.managers.observers().on_run_start(&RunStartContext { + session_id: self.session.id, + }); + self.notify_run_start_hook(); } - /// Notify all observers and hooks that the session has ended. + /// Notify all observers and hooks that a run has ended. /// - /// Fires [`on_session_end`](crate::observer::LoopObserver::on_session_end) - /// on every registered observer, then fires the `on_session_end` hook + /// Fires [`on_run_end`](crate::observer::LoopObserver::on_run_end) + /// on every registered observer, then fires the `on_run_end` hook /// (when the `hooks` feature is enabled). The [`Run`] supplies the /// per-run totals (turn count, tokens); `success` distinguishes a /// normal exit from a failure so observers and hooks can branch on /// the outcome, and `duration` is the wall-clock run length. - pub(super) fn notify_session_end(&self, result: &Run, success: bool, duration: Duration) { - self.managers - .observers() - .on_session_end(&SessionEndContext { - success, - error: if success { - None - } else { - Some("run failed".to_string()) - }, - total_turns: result.turn_count(), - duration_ms: Self::millis_u64(duration), - }); - self.notify_session_end_hook(result, success, duration); + pub(super) fn notify_run_end(&self, result: &Run, success: bool, duration: Duration) { + self.managers.observers().on_run_end(&RunEndContext { + success, + error: if success { + None + } else { + Some("run failed".to_string()) + }, + total_turns: result.turn_count(), + duration_ms: Self::millis_u64(duration), + }); + self.notify_run_end_hook(result, success, duration); } - /// Derive the structured [`SessionEndReason`] from the loop's + /// Derive the structured [`RunEndReason`] from the loop's /// terminal state. /// /// Maps the boolean `success` plus the cancellation flag and the @@ -64,80 +59,80 @@ impl BareLoop { /// normal completion. Used only to populate the hook end-context — /// observers receive the simpler boolean `success` field. #[cfg(feature = "hooks")] - fn session_end_reason(&self, success: bool) -> SessionEndReason { + fn run_end_reason(&self, success: bool) -> RunEndReason { if self.is_cancelled() { - SessionEndReason::Cancelled + RunEndReason::Cancelled } else if !success { - SessionEndReason::Error + RunEndReason::Error } else if self.current_run().map_or(0, Run::turn_count) >= self.run_config().max_turns { - SessionEndReason::MaxTurns + RunEndReason::MaxTurns } else { - SessionEndReason::Complete + RunEndReason::Complete } } - /// Fire the `on_session_start` hook when a hook executor is + /// Fire the `on_run_start` hook when a hook executor is /// configured. /// - /// Builds a [`HookSessionStartContext`] from the session id, the + /// Builds a [`HookRunStartContext`] from the session id, the /// client's current model, and the process working directory, then - /// dispatches it to every registered session-start hook. No-op when + /// dispatches it to every registered run-start hook. No-op when /// no hook executor is set. #[cfg(feature = "hooks")] - fn notify_session_start_hook(&self) { + fn notify_run_start_hook(&self) { let Some(executor) = self.managers.hook_executor() else { return; }; - let ctx = HookSessionStartContext { + let ctx = HookRunStartContext { session_id: self.session.id, model: self.client.model(), working_directory: std::env::current_dir() .map(|p| p.to_string_lossy().into_owned()) .unwrap_or_default(), }; - executor.notify_session_start(&ctx); + executor.notify_run_start(&ctx); } - /// No-op session-start hook for builds without the `hooks` feature. + /// No-op run-start hook for builds without the `hooks` feature. /// /// Does nothing — there are no hooks to notify. Kept so - /// [`notify_session_start`](Self::notify_session_start) compiles + /// [`notify_run_start`](Self::notify_run_start) compiles /// identically with and without the feature. #[cfg(not(feature = "hooks"))] - fn notify_session_start_hook(&self) {} + fn notify_run_start_hook(&self) {} - /// Fire the `on_session_end` hook when a hook executor is + /// Fire the `on_run_end` hook when a hook executor is /// configured. /// - /// Derives the [`SessionEndReason`] via - /// [`session_end_reason`](Self::session_end_reason), then builds a - /// [`HookSessionEndContext`] from the session id, reason, turn + /// Derives the [`RunEndReason`] via + /// [`run_end_reason`](Self::run_end_reason), then builds a + /// [`HookRunEndContext`] from the session id, reason, turn /// count, total tokens, and run duration, and dispatches it to - /// every registered session-end hook. No-op when no hook executor + /// every registered run-end hook. No-op when no hook executor /// is set. #[cfg(feature = "hooks")] - fn notify_session_end_hook(&self, result: &Run, success: bool, duration: Duration) { + fn notify_run_end_hook(&self, result: &Run, success: bool, duration: Duration) { let Some(executor) = self.managers.hook_executor() else { return; }; - let reason = self.session_end_reason(success); - let ctx = HookSessionEndContext { + let reason = self.run_end_reason(success); + let ctx = HookRunEndContext { session_id: self.session.id, reason, total_turns: result.turn_count(), total_tokens: result.total_tokens(), duration_secs: duration.as_secs(), }; - executor.notify_session_end(&ctx); + executor.notify_run_end(&ctx); } - /// No-op session-end hook for builds without the `hooks` feature. + /// No-op run-end hook for builds without the `hooks` feature. /// /// Does nothing — there are no hooks to notify. Kept so - /// [`notify_session_end`](Self::notify_session_end) compiles + /// [`notify_run_end`](Self::notify_run_end) compiles /// identically with and without the feature. #[cfg(not(feature = "hooks"))] - fn notify_session_end_hook(&self, _result: &Run, _success: bool, _duration: Duration) {} + fn notify_run_end_hook(&self, _result: &Run, _success: bool, _duration: Duration) {} /// Convert a [`Duration`] to milliseconds as a `u64`. /// diff --git a/src/engine/core/lifecycle.rs b/src/engine/core/lifecycle.rs index 961e853..96149d4 100644 --- a/src/engine/core/lifecycle.rs +++ b/src/engine/core/lifecycle.rs @@ -61,10 +61,12 @@ fn instant_now() -> Instant { /// Per-run configuration. /// /// The slice of agent configuration that varies across `run()` calls on the -/// same agent (turn/token budgets, compaction policy, dispatch mode). It is -/// owned by the machine so that the machine's decisions are self-contained and -/// a serialized machine carries its configuration with it. +/// same agent (turn/token budgets, compaction policy, dispatch mode, manager +/// reset). It is owned by the machine so that the machine's decisions are +/// self-contained and a serialized machine carries its configuration with +/// it. #[derive(Debug, Clone, Serialize, Deserialize)] +#[non_exhaustive] pub struct RunConfig { /// Maximum number of turns before forcing completion. /// @@ -80,6 +82,20 @@ pub struct RunConfig { /// sequentially or in parallel, up to a concurrency limit. See /// [`ParallelDispatchConfig`]. pub parallel_tool_dispatch: ParallelDispatchConfig, + + /// Whether to reset all managers to their initial state before this run. + /// + /// When `true`, [`LoopManagers::reset_all`] is called at the top of + /// `run()`, clearing the fallback circuit breaker, loop/convergence + /// detection history, and per-observer accumulators. Use this when a + /// run is logically independent from the previous one (different task, + /// fresh context) and you do not want accumulated state to carry over. + /// Defaults to `false` — manager state persists across runs within a + /// session, which is usually the right behavior (e.g. a tripped + /// circuit breaker should stay tripped). + /// + /// [`LoopManagers::reset_all`]: crate::managers::LoopManagers::reset_all + pub reset_managers: bool, } impl Default for RunConfig { @@ -87,10 +103,29 @@ impl Default for RunConfig { Self { max_turns: 200, parallel_tool_dispatch: ParallelDispatchConfig::default(), + reset_managers: false, } } } +impl RunConfig { + /// Set the maximum number of turns for this run. + /// + /// Builder-style convenience for `#[non_exhaustive]` compliance. + /// + /// ``` + /// use loopctl::engine::RunConfig; + /// + /// let config = RunConfig::default().with_max_turns(50); + /// assert_eq!(config.max_turns, 50); + /// ``` + #[must_use] + pub fn with_max_turns(mut self, max_turns: usize) -> Self { + self.max_turns = max_turns; + self + } +} + /// Why the API stopped generating. /// /// Mirrors the stop reasons returned by LLM APIs. The framework uses this diff --git a/src/hooks.rs b/src/hooks.rs index f45a2f6..954e27a 100644 --- a/src/hooks.rs +++ b/src/hooks.rs @@ -49,7 +49,7 @@ use std::fmt; use context::{ CompactResult, PostCompactContext, PostToolUseContext, PreCompactContext, PreToolUseContext, - SessionEndContext, SessionStartContext, + RunEndContext, RunStartContext, }; /// Hook trait for bidirectional lifecycle control. @@ -136,17 +136,17 @@ pub trait Hook: Send + Sync { /// Notification-only; cannot alter the compaction result. fn on_post_compact(&self, _ctx: &PostCompactContext) {} - /// Called when an agent session starts. + /// Called when a run starts. /// - /// Fires once at session start. Use for initialization, + /// Fires at the start of each `run()` call. Use for initialization, /// credential setup, or state restoration. - fn on_session_start(&self, _ctx: &SessionStartContext) {} + fn on_run_start(&self, _ctx: &RunStartContext) {} - /// Called when an agent session ends. + /// Called when a run ends. /// - /// Fires once at session end. Use for cleanup, teardown, + /// Fires at the end of each `run()` call. Use for cleanup, teardown, /// or final metric emission. - fn on_session_end(&self, _ctx: &SessionEndContext) {} + fn on_run_end(&self, _ctx: &RunEndContext) {} } /// Whether the session can interact with a human operator. diff --git a/src/hooks/builtin.rs b/src/hooks/builtin.rs index 8453093..db2c68d 100644 --- a/src/hooks/builtin.rs +++ b/src/hooks/builtin.rs @@ -13,7 +13,7 @@ //! | [`LoggingHook`] | Logs pre/post tool use and compact events via `tracing`. | //! | [`BlocklistHook`] | Blocks or allows tools by name (denylist/allowlist). | //! | [`ConfirmationHook`] | Requires interactive confirmation for specific tools. | -//! | [`AutoCommitHook`] | Tracks file modifications and auto-commits at session end. | +//! | [`AutoCommitHook`] | Tracks file modifications and auto-commits at run end. | mod auto_commit; mod blocklist_hook; diff --git a/src/hooks/builtin/auto_commit.rs b/src/hooks/builtin/auto_commit.rs index 2f12c36..ded60bb 100644 --- a/src/hooks/builtin/auto_commit.rs +++ b/src/hooks/builtin/auto_commit.rs @@ -1,13 +1,13 @@ //! Auto-commit hook — automatically commits file changes during agent sessions. //! //! Tracks file modifications from tool calls (Write, Edit) and commits -//! them at session end. Useful for keeping a git trail of agent actions. +//! them at run end. Useful for keeping a git trail of agent actions. //! //! # How it fits together //! //! 1. [`AutoCommitHook`] observes post-tool-use events, recording each //! `file_path` touched by the configured tracking tools. -//! 2. At session end the hook hands the recorded paths (plus the +//! 2. At run end the hook hands the recorded paths (plus the //! [`AutoCommitConfig`]) to [`GitExecutor`], which shells out to //! `git` to stage, commit, and optionally push. //! 3. [`AutoCommitResult`] describes the outcome so callers can log @@ -18,7 +18,7 @@ use std::sync::Mutex; use std::time::Duration; use crate::hooks::Hook; -use crate::hooks::context::{PostToolUseContext, SessionEndContext, SessionStartContext}; +use crate::hooks::context::{PostToolUseContext, RunEndContext, RunStartContext}; /// Default timeout for git subprocess invocations. /// @@ -455,7 +455,7 @@ pub struct AutoCommitConfig { /// Commit message template (supports `{{tool}}`, `{{session}}` placeholders). /// /// Template string used to build the commit message; `{{tool}}` and - /// `{{session}}` are expanded at session end into the triggering tool + /// `{{session}}` are expanded at run end into the triggering tool /// name and the session identifier respectively. pub message_template: String, @@ -490,7 +490,7 @@ pub struct AutoCommitConfig { /// /// Tool names whose output should be watched for a `file_path`; when one /// of these tools runs, the touched file is recorded and staged at - /// session end. + /// run end. pub commit_on_tools: Vec, /// Files to add before commit. @@ -558,7 +558,7 @@ impl AutoCommitConfigBuilder { /// /// Accepts any `Into` so literals work directly. The /// template may include `{{tool}}` and `{{session}}` placeholders - /// expanded at session end. + /// expanded at run end. #[must_use] pub fn with_message_template(mut self, template: impl Into) -> Self { self.config.message_template = template.into(); @@ -602,7 +602,7 @@ impl Default for AutoCommitConfigBuilder { } /// A hook that tracks file modifications from tool calls and commits -/// them at session end. +/// them at run end. /// /// Useful for keeping a git trail of agent actions. Configure which /// tools trigger tracking via [`AutoCommitConfig::commit_on_tools`]. @@ -678,7 +678,7 @@ impl AutoCommitHook { /// Clear all recorded file modifications. /// - /// Called at session start so a fresh session does not inherit + /// Called at run start so a fresh run does not inherit /// modifications from the previous one. Lock failures are silently /// ignored. fn clear_modifications(&self) { @@ -711,11 +711,11 @@ impl Hook for AutoCommitHook { } } - fn on_session_start(&self, _ctx: &SessionStartContext) { + fn on_run_start(&self, _ctx: &RunStartContext) { self.clear_modifications(); } - fn on_session_end(&self, _ctx: &SessionEndContext) { + fn on_run_end(&self, _ctx: &RunEndContext) { let files = self.modified_files.lock().ok().filter(|f| !f.is_empty()); let result = GitExecutor::auto_commit_with_files( &self.config, diff --git a/src/hooks/builtin/logging_hook.rs b/src/hooks/builtin/logging_hook.rs index 843536d..e4aec0c 100644 --- a/src/hooks/builtin/logging_hook.rs +++ b/src/hooks/builtin/logging_hook.rs @@ -68,8 +68,8 @@ impl Hook for LoggingHook { mod tests { use super::*; use crate::hooks::context::{ - CompactTrigger, PostCompactContext, PostToolUseContext, SessionEndContext, - SessionEndReason, SessionStartContext, + CompactTrigger, PostCompactContext, PostToolUseContext, RunEndContext, RunEndReason, + RunStartContext, }; use serde_json::json; @@ -129,26 +129,26 @@ mod tests { } #[test] - fn logging_hook_session_start_no_panic() { + fn logging_hook_run_start_no_panic() { let hook = LoggingHook; - let ctx = SessionStartContext { + let ctx = RunStartContext { session_id: uuid::Uuid::nil(), model: "test".to_string(), working_directory: "/tmp".to_string(), }; - hook.on_session_start(&ctx); + hook.on_run_start(&ctx); } #[test] - fn logging_hook_session_end_no_panic() { + fn logging_hook_run_end_no_panic() { let hook = LoggingHook; - let ctx = SessionEndContext { + let ctx = RunEndContext { session_id: uuid::Uuid::nil(), - reason: SessionEndReason::Complete, + reason: RunEndReason::Complete, total_turns: 5, total_tokens: 1000, duration_secs: 30, }; - hook.on_session_end(&ctx); + hook.on_run_end(&ctx); } } diff --git a/src/hooks/context.rs b/src/hooks/context.rs index e47745e..0973764 100644 --- a/src/hooks/context.rs +++ b/src/hooks/context.rs @@ -7,13 +7,13 @@ //! //! # Lifecycle //! -//! Every session opens with [`SessionStartContext`], then interleaves +//! Every run opens with [`RunStartContext`], then interleaves //! any number of tool calls (`PreToolUse` → tool → `PostToolUse`) and //! compaction passes (`PreCompact` → compact → `PostCompact`), and -//! finally closes with [`SessionEndContext`]. Every callback receives a +//! finally closes with [`RunEndContext`]. Every callback receives a //! context whose fields mirror the moment it fires: pre-contexts carry //! the *about-to-happen* inputs, post-contexts carry the -//! *what-happened* results, and the session contexts bracket the whole +//! *what-happened* results, and the run contexts bracket the whole //! run. use serde::{Deserialize, Serialize}; @@ -339,15 +339,15 @@ impl CompactResult { } } -/// Why the session ended. +/// Why the run ended. /// -/// Carried by [`SessionEndContext::reason`] so end-hooks can branch on -/// the terminal condition. The variants cover the four ways a session +/// Carried by [`RunEndContext::reason`] so end-hooks can branch on +/// the terminal condition. The variants cover the four ways a run /// can stop — normal completion, cancellation, an unrecoverable error, /// hitting a turn cap — plus context overflow when compaction could not /// recover. Use this for outcome-specific logging and cleanup. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -pub enum SessionEndReason { +pub enum RunEndReason { /// The agent finished normally. /// /// The agent declared its task complete (for example the model @@ -355,10 +355,10 @@ pub enum SessionEndReason { /// This is the "happy path" outcome. Complete, - /// The user interrupted the session. + /// The user interrupted the run. /// /// A cancel signal — Ctrl+C in an interactive host, or a programmatic - /// cancellation — stopped the loop. The session may have partial + /// cancellation — stopped the loop. The run may have partial /// results available in its state at the time of interruption. Cancelled, @@ -387,19 +387,19 @@ pub enum SessionEndReason { ContextOverflow, } -/// Context provided to `on_session_start` hooks. +/// Context provided to `on_run_start` hooks. /// -/// Fires exactly once per session, before any turns run. Carries the -/// session identifier, the model that will handle the first request, -/// and the working directory the agent treats as its root. -/// Notification-only — use it to initialize per-session state (open +/// Fires at the start of each `run()` call, before any turns run. Carries +/// the session identifier, the model that will handle the first +/// request, and the working directory the agent treats as its root. +/// Notification-only — use it to initialize per-run state (open /// resources, set up tracking, log the start). #[derive(Debug, Clone)] -pub struct SessionStartContext { +pub struct RunStartContext { /// Session ID. /// - /// Unique identifier for the agent session that is starting, letting hooks - /// initialize any per-session state they track. + /// Unique identifier for the agent session that owns this run, letting + /// hooks initialize any per-session state they track. pub session_id: uuid::Uuid, /// Model identifier. @@ -416,43 +416,43 @@ pub struct SessionStartContext { pub working_directory: String, } -/// Context provided to `on_session_end` hooks. +/// Context provided to `on_run_end` hooks. /// -/// Fires exactly once per session, after the loop has terminated. +/// Fires at the end of each `run()` call, after the loop has terminated. /// Carries the session identifier, the terminal -/// [`SessionEndReason`], and aggregate counters (turns, tokens, +/// [`RunEndReason`], and aggregate counters (turns, tokens, /// duration) for bookkeeping and cost accounting. Notification-only — /// use it to flush resources, finalize tracking, and emit a summary log /// line. #[derive(Debug, Clone)] -pub struct SessionEndContext { +pub struct RunEndContext { /// Session ID. /// - /// Unique identifier for the session that is ending, matching the value - /// carried by [`SessionStartContext::session_id`]. + /// Unique identifier for the session that owns this run, matching the + /// value carried by [`RunStartContext::session_id`]. pub session_id: uuid::Uuid, - /// Why the session ended. + /// Why the run ended. /// /// The terminal condition — normal completion, cancellation, error, turn /// cap, or context overflow — so hooks can branch on the outcome. - pub reason: SessionEndReason, + pub reason: RunEndReason, /// Total turns executed. /// - /// How many turns ran to completion during the session, for bookkeeping + /// How many turns ran to completion during this run, for bookkeeping /// and budget reporting. pub total_turns: usize, /// Total tokens consumed (input + output across all turns). /// - /// Aggregate token usage over the whole session, combining input and + /// Aggregate token usage over the whole run, combining input and /// output tokens from every turn for cost accounting. pub total_tokens: u64, - /// Wall-clock session duration in seconds. + /// Wall-clock run duration in seconds. /// - /// Elapsed time from session start to session end, rounded to whole + /// Elapsed time from run start to run end, rounded to whole /// seconds, for latency and uptime reporting. pub duration_secs: u64, } @@ -498,17 +498,17 @@ mod tests { } #[test] - fn session_end_reason_serialization() { + fn run_end_reason_serialization() { let reasons = [ - SessionEndReason::Complete, - SessionEndReason::Cancelled, - SessionEndReason::Error, - SessionEndReason::MaxTurns, - SessionEndReason::ContextOverflow, + RunEndReason::Complete, + RunEndReason::Cancelled, + RunEndReason::Error, + RunEndReason::MaxTurns, + RunEndReason::ContextOverflow, ]; for reason in reasons { let json = serde_json::to_string(&reason).expect("serialize"); - let back: SessionEndReason = serde_json::from_str(&json).expect("deserialize"); + let back: RunEndReason = serde_json::from_str(&json).expect("deserialize"); assert_eq!(back, reason); } } diff --git a/src/hooks/executor.rs b/src/hooks/executor.rs index 33f1986..ceb332d 100644 --- a/src/hooks/executor.rs +++ b/src/hooks/executor.rs @@ -23,7 +23,7 @@ //! //! # Post-hook Execution Model //! -//! For `on_post_*` and session hooks, the executor calls **all** hooks. +//! For `on_post_*` and run hooks, the executor calls **all** hooks. //! There's no short-circuit because post-hooks are notification-only. //! //! # Thread Safety @@ -40,7 +40,7 @@ use crate::hooks::HookAction; use crate::hooks::Interactivity; use crate::hooks::context::{ CompactResult, PostCompactContext, PostToolUseContext, PreCompactContext, PreToolUseContext, - SessionEndContext, SessionStartContext, + RunEndContext, RunStartContext, }; /// Executes hooks in registration order with short-circuit semantics. @@ -268,28 +268,28 @@ impl HookExecutor { } } - /// Notify every registered session-start hook. + /// Notify every registered run-start hook. /// - /// Fires [`Hook::on_session_start`] on each hook in registration - /// order, once, at the beginning of a session. Notification-only — - /// every hook always runs. Use it to initialize per-session state + /// Fires [`Hook::on_run_start`] on each hook in registration + /// order, once, at the beginning of a run. Notification-only — + /// every hook always runs. Use it to initialize per-run state /// (open resources, reset counters, emit a start log line). - pub fn notify_session_start(&self, ctx: &SessionStartContext) { + pub fn notify_run_start(&self, ctx: &RunStartContext) { for hook in &self.hooks { - hook.on_session_start(ctx); + hook.on_run_start(ctx); } } - /// Notify every registered session-end hook. + /// Notify every registered run-end hook. /// - /// Fires [`Hook::on_session_end`] on each hook in registration + /// Fires [`Hook::on_run_end`] on each hook in registration /// order, once, after the loop has terminated. Notification-only — /// every hook always runs. Use it to flush resources, finalize /// tracking, and emit a summary log line keyed off - /// [`SessionEndContext::reason`]. - pub fn notify_session_end(&self, ctx: &SessionEndContext) { + /// [`RunEndContext::reason`]. + pub fn notify_run_end(&self, ctx: &RunEndContext) { for hook in &self.hooks { - hook.on_session_end(ctx); + hook.on_run_end(ctx); } } @@ -306,29 +306,29 @@ impl HookExecutor { Box::pin(async {}) } - /// Async wrapper for [`notify_session_start`](Self::notify_session_start). + /// Async wrapper for [`notify_run_start`](Self::notify_run_start). /// - /// Runs all session-start hooks synchronously and wraps completion + /// Runs all run-start hooks synchronously and wraps completion /// in a `Pin>` for async compatibility. #[must_use] - pub fn notify_session_start_async( + pub fn notify_run_start_async( &self, - ctx: &SessionStartContext, + ctx: &RunStartContext, ) -> Pin + Send + '_>> { - self.notify_session_start(ctx); + self.notify_run_start(ctx); Box::pin(async {}) } - /// Async wrapper for [`notify_session_end`](Self::notify_session_end). + /// Async wrapper for [`notify_run_end`](Self::notify_run_end). /// - /// Runs all session-end hooks synchronously and wraps completion + /// Runs all run-end hooks synchronously and wraps completion /// in a `Pin>` for async compatibility. #[must_use] - pub fn notify_session_end_async( + pub fn notify_run_end_async( &self, - ctx: &SessionEndContext, + ctx: &RunEndContext, ) -> Pin + Send + '_>> { - self.notify_session_end(ctx); + self.notify_run_end(ctx); Box::pin(async {}) } } @@ -336,7 +336,7 @@ impl HookExecutor { #[cfg(test)] mod tests { use super::*; - use crate::hooks::context::{CompactTrigger, SessionEndReason}; + use crate::hooks::context::{CompactTrigger, RunEndReason}; use serde_json::json; struct AllowHook; @@ -538,7 +538,7 @@ mod tests { } #[test] - fn session_start_end_notify_all() { + fn run_start_end_notify_all() { use std::sync::atomic::{AtomicUsize, Ordering}; struct CounterHook { starts: AtomicUsize, @@ -548,10 +548,10 @@ mod tests { fn name(&self) -> &'static str { "counter" } - fn on_session_start(&self, _ctx: &SessionStartContext) { + fn on_run_start(&self, _ctx: &RunStartContext) { self.starts.fetch_add(1, Ordering::Relaxed); } - fn on_session_end(&self, _ctx: &SessionEndContext) { + fn on_run_end(&self, _ctx: &RunEndContext) { self.ends.fetch_add(1, Ordering::Relaxed); } } @@ -562,15 +562,15 @@ mod tests { let executor = HookExecutor::new() .with_hook(counter.clone()) .with_hook(counter.clone()); - executor.notify_session_start(&SessionStartContext { + executor.notify_run_start(&RunStartContext { session_id: uuid::Uuid::nil(), model: "test".to_string(), working_directory: "/tmp".to_string(), }); assert_eq!(counter.starts.load(Ordering::Relaxed), 2); - executor.notify_session_end(&SessionEndContext { + executor.notify_run_end(&RunEndContext { session_id: uuid::Uuid::nil(), - reason: SessionEndReason::Complete, + reason: RunEndReason::Complete, total_turns: 5, total_tokens: 1000, duration_secs: 30, @@ -638,9 +638,9 @@ mod tests { } #[test] - fn notify_session_start_empty_executor() { + fn notify_run_start_empty_executor() { let executor = HookExecutor::new(); - executor.notify_session_start(&SessionStartContext { + executor.notify_run_start(&RunStartContext { session_id: uuid::Uuid::nil(), model: "test".to_string(), working_directory: "/tmp".to_string(), @@ -648,11 +648,11 @@ mod tests { } #[test] - fn notify_session_end_empty_executor() { + fn notify_run_end_empty_executor() { let executor = HookExecutor::new(); - executor.notify_session_end(&SessionEndContext { + executor.notify_run_end(&RunEndContext { session_id: uuid::Uuid::nil(), - reason: SessionEndReason::Complete, + reason: RunEndReason::Complete, total_turns: 0, total_tokens: 0, duration_secs: 0, diff --git a/src/managers.rs b/src/managers.rs index d195c12..6537caf 100644 --- a/src/managers.rs +++ b/src/managers.rs @@ -124,8 +124,8 @@ pub struct LoopManagers { /// Circuit breaker for API model fallback. /// /// Tracks consecutive API failures and, when the threshold is exceeded, - /// switches to a configured backup model. Reset at the start of each run - /// via [`reset_all`](Self::reset_all). + /// switches to a configured backup model. Fresh on construction; call + /// [`reset_all`](Self::reset_all) to reinitialise mid-session. fallback: FallbackManager, /// Loop and convergence detection orchestrator. @@ -399,16 +399,20 @@ impl LoopManagers { /// 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. + /// Clears the fallback circuit breaker, loop/convergence detection + /// history, and per-observer accumulators. Call this when you want a + /// clean slate mid-session — for example after a provider outage + /// resolves (so the circuit breaker does not stay tripped) or when + /// switching to an unrelated task (so stale detection history does not + /// skew the next run). The engine does not call this automatically; + /// manager state persists across `run()` calls within a session. /// /// # Example /// /// ``` /// # use loopctl::managers::LoopManagers; /// let managers = LoopManagers::new(); - /// // ... after a session ... + /// // ... after several runs ... /// managers.reset_all(); /// // All managers are back to their initial state /// ``` diff --git a/src/observer.rs b/src/observer.rs index 1cb0f50..f6e8946 100644 --- a/src/observer.rs +++ b/src/observer.rs @@ -8,7 +8,7 @@ //! //! Each callback receives a typed context struct with relevant fields: //! -//! - [`SessionStartContext`] / [`SessionEndContext`] — session boundaries +//! - [`RunStartContext`] / [`RunEndContext`] — run boundaries (one per `run()` call) //! - [`TurnStartContext`] / [`TurnEndContext`] — turn boundaries //! - [`StreamContext`] / [`StreamFailureContext`] — stream success/failure //! - [`ResponseContext`] — model response text and usage @@ -24,14 +24,14 @@ //! # Example //! //! ```rust,ignore -//! use loopctl::observer::{LoopObserver, SessionStartContext}; +//! use loopctl::observer::{LoopObserver, RunStartContext}; //! //! struct MetricsObserver; //! //! impl LoopObserver for MetricsObserver { //! fn name(&self) -> &str { "metrics" } //! -//! fn on_session_start(&self, ctx: &SessionStartContext) { +//! fn on_run_start(&self, ctx: &RunStartContext) { //! println!("session {} started", ctx.session_id); //! } //! } @@ -43,7 +43,7 @@ pub mod context; pub use context::{ CompactedContext, ConvergenceDetectedContext, FallbackContext, LoopDetectedContext, - ModelSwitchedContext, ResponseContext, SessionEndContext, SessionStartContext, StreamContext, + ModelSwitchedContext, ResponseContext, RunEndContext, RunStartContext, StreamContext, StreamFailureContext, TextDeltaContext, ThinkingDeltaContext, ToolCallReceivedContext, ToolPostContext, ToolPreContext, TurnEndContext, TurnStartContext, }; @@ -63,16 +63,18 @@ pub trait LoopObserver: Send + Sync { /// produced a side-effect. fn name(&self) -> &str; - /// Called when an agent session begins. + /// Called when a run begins. /// - /// Fired once per session, before the first turn starts. - fn on_session_start(&self, _ctx: &SessionStartContext) {} + /// Fired at the start of each `run()` call, before the first turn of + /// that run. A session may contain many runs. + fn on_run_start(&self, _ctx: &RunStartContext) {} - /// Called when an agent session ends. + /// Called when a run 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) {} + /// Fired after the run completes — whether normally, by error, or by + /// cancellation. Check [`RunEndContext::success`] to distinguish normal + /// exit from failure. + fn on_run_end(&self, _ctx: &RunEndContext) {} /// Called at the start of processing a turn. /// @@ -238,8 +240,9 @@ pub trait LoopObserver: Send + Sync { /// Reset observer state for a new session. /// - /// Called before [`on_session_start`](Self::on_session_start) to allow - /// observers to clear per-session accumulators. + /// Called once per session (on the first `run()`), before + /// [`on_run_start`](Self::on_run_start), to allow observers to clear + /// per-session accumulators. fn reset(&self) {} } @@ -323,23 +326,23 @@ impl ObserverHost { self.dispatch(|obs| obs.reset()); } - /// Dispatch [`LoopObserver::on_session_start`] to all observers. + /// Dispatch [`LoopObserver::on_run_start`] to all observers. /// - /// Fired once per session, before the first turn begins. Iterates - /// registered observers in registration order; use it to initialize - /// per-session observer state. - pub fn on_session_start(&self, ctx: &SessionStartContext) { - self.dispatch(|obs| obs.on_session_start(ctx)); + /// Fired at the start of each `run()` call, before the first turn + /// begins. Iterates registered observers in registration order; use it + /// to initialize per-run observer state. + pub fn on_run_start(&self, ctx: &RunStartContext) { + self.dispatch(|obs| obs.on_run_start(ctx)); } - /// Dispatch [`LoopObserver::on_session_end`] to all observers. + /// Dispatch [`LoopObserver::on_run_end`] to all observers. /// - /// Fired once per session, after the loop has terminated. Iterates - /// registered observers in registration order; check - /// [`SessionEndContext::success`] in each observer to distinguish a + /// Fired at the end of each `run()` call, after the loop has + /// terminated. Iterates registered observers in registration order; + /// check [`RunEndContext::success`] in each observer to distinguish a /// normal exit from a failure. - pub fn on_session_end(&self, ctx: &SessionEndContext) { - self.dispatch(|obs| obs.on_session_end(ctx)); + pub fn on_run_end(&self, ctx: &RunEndContext) { + self.dispatch(|obs| obs.on_run_end(ctx)); } /// Dispatch [`LoopObserver::on_turn_start`] to all observers. diff --git a/src/observer/context.rs b/src/observer/context.rs index 7034789..3faf117 100644 --- a/src/observer/context.rs +++ b/src/observer/context.rs @@ -4,44 +4,46 @@ //! Observers receive shared references (`&Context`) — the structs are //! notification-only data carriers. -/// Context for [`LoopObserver::on_session_start`](crate::observer::LoopObserver::on_session_start). +/// Context for [`LoopObserver::on_run_start`](crate::observer::LoopObserver::on_run_start). /// -/// Carries the session identifier so observers can correlate -/// lifecycle events with a specific agent run. +/// Carries the session identifier so observers can correlate lifecycle +/// events with a specific agent run. One run is one `run()` call on the +/// loop; a session may contain many runs. #[derive(Debug, Clone)] -pub struct SessionStartContext { +pub struct RunStartContext { /// Unique session identifier. /// - /// Correlates all lifecycle events belonging to the same agent run. + /// Correlates all lifecycle events belonging to the same agent session. + /// Stable across every `run()` call on the same loop. pub session_id: uuid::Uuid, } -/// Context for [`LoopObserver::on_session_end`](crate::observer::LoopObserver::on_session_end). +/// Context for [`LoopObserver::on_run_end`](crate::observer::LoopObserver::on_run_end). /// -/// Captures the session's completion status, optional error description, -/// total turns executed, and wall-clock duration in milliseconds. +/// Captures the run'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. +pub struct RunEndContext { + /// Whether the run completed successfully. /// - /// `true` when the loop exited normally, `false` on error or cancellation. + /// `true` when the run exited normally, `false` on error or cancellation. pub success: bool, - /// Error description, if the session ended due to an error. + /// Error description, if the run ended due to an error. /// /// `None` when [`success`](Self::success) is `true`. pub error: Option, - /// Total turns completed during the session. + /// Total turns completed during this run. /// /// 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. + /// Run 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). + /// Measured wall-clock from [`on_run_start`](crate::observer::LoopObserver::on_run_start) + /// to [`on_run_end`](crate::observer::LoopObserver::on_run_end). pub duration_ms: u64, } From 38a1803a1fb776fa1447bbf6a2e91fd31fe438d6 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 28 Jul 2026 21:01:13 +1200 Subject: [PATCH 18/21] fix: prevent infinite compation when no progress --- src/engine/core/machine.rs | 81 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/src/engine/core/machine.rs b/src/engine/core/machine.rs index 5f13572..84b0886 100644 --- a/src/engine/core/machine.rs +++ b/src/engine/core/machine.rs @@ -380,6 +380,17 @@ pub struct LoopMachine { /// policy to decide whether to emit a [`MachineStep::Compact`]. context_tokens: u64, + /// Context tokens at the time of the last compaction request. + /// + /// When the machine emits `Compact`, it records `context_tokens` + /// here. After `compaction_result`, if `tokens_after` has not + /// decreased below this value, the machine transitions to + /// [`MachineOutcome::Failed`] with + /// [`LoopError::ContextExceeded`] instead of requesting another + /// compaction — preventing an infinite compaction cycle when the + /// compactor cannot reduce the context. + last_compaction_tokens: Option, + /// Whether [`Self::cancel`] has been called. /// /// A cancellation request from the driver. Once set, the next @@ -414,6 +425,7 @@ impl LoopMachine { state: MachineState::Start, turns_taken: 0, context_tokens: 0, + last_compaction_tokens: None, cancelled: false, pending_tools: Vec::new(), } @@ -499,11 +511,13 @@ impl LoopMachine { } if self.is_emergency(policy) { let reason = CompactReason::Emergency; + self.last_compaction_tokens = Some(self.context_tokens); self.state = MachineState::AwaitingCompaction { reason }; return MachineStep::Compact { reason }; } if policy.auto_compact && self.should_compact(policy) { let reason = CompactReason::ThresholdExceeded; + self.last_compaction_tokens = Some(self.context_tokens); self.state = MachineState::AwaitingCompaction { reason }; return MachineStep::Compact { reason }; } @@ -633,6 +647,19 @@ impl LoopMachine { if self.is_terminal() { return; } + if let Some(before) = self.last_compaction_tokens + && tokens_after >= before + { + self.state = MachineState::Terminal(MachineOutcome::Failed { + error: LoopError::ContextExceeded { + used: tokens_after, + limit: before, + }, + }); + self.last_compaction_tokens = None; + return; + } + self.last_compaction_tokens = None; self.history = compacted; self.pending.clear(); self.context_tokens = tokens_after; @@ -1182,6 +1209,60 @@ mod tests { ); } + #[test] + fn compaction_no_progress_terminates_not_loops() { + let policy = MachinePolicy { + max_turns: 5, + context_window: 100, + compact_threshold: 50, + auto_compact: true, + }; + let mut machine = LoopMachine::from_history(vec![Message::user("hello")]); + let _ = machine.next_step(policy); + machine.model_response(tool_response("echo", &["echo"], 60)); + let _ = machine.next_step(policy); + machine.tool_results(vec![Message::user("tool-out")]); + assert!(matches!( + machine.next_step(policy), + MachineStep::Compact { .. } + )); + machine.compaction_result(vec![Message::user("compacted")], 70); + + match machine.next_step(policy) { + MachineStep::Done(MachineOutcome::Failed { + error: LoopError::ContextExceeded { .. }, + }) => {} + other => { + panic!("no-progress compaction must terminate with ContextExceeded, got {other:?}") + } + } + } + + #[test] + fn compaction_progress_continues_normally() { + let policy = MachinePolicy { + max_turns: 5, + context_window: 100, + compact_threshold: 50, + auto_compact: true, + }; + let mut machine = LoopMachine::from_history(vec![Message::user("hello")]); + let _ = machine.next_step(policy); + machine.model_response(tool_response("echo", &["echo"], 60)); + let _ = machine.next_step(policy); + machine.tool_results(vec![Message::user("tool-out")]); + assert!(matches!( + machine.next_step(policy), + MachineStep::Compact { .. } + )); + machine.compaction_result(vec![Message::user("compacted")], 30); + + assert!( + matches!(machine.next_step(policy), MachineStep::CallLLM { .. }), + "compaction that reduced tokens must continue to CallLLM" + ); + } + #[test] fn history_accumulates_user_assistant_tool_round() { let mut machine = small_machine(5); From 86e47753a0857683ef7eb06bfe16b30b13748fab Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 28 Jul 2026 21:17:23 +1200 Subject: [PATCH 19/21] fix: kill child git process on timeout --- src/hooks/builtin/auto_commit.rs | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/hooks/builtin/auto_commit.rs b/src/hooks/builtin/auto_commit.rs index ded60bb..f465d0d 100644 --- a/src/hooks/builtin/auto_commit.rs +++ b/src/hooks/builtin/auto_commit.rs @@ -170,6 +170,8 @@ impl GitExecutor { .spawn() .map_err(|e| GitExecutorError::ExecutionFailed(e.to_string()))?; + let pid = child.id(); + let (tx, rx) = std::sync::mpsc::channel::>(); std::thread::spawn(move || { let result = child.wait_with_output(); @@ -187,6 +189,13 @@ impl GitExecutor { } Ok(Err(e)) => Err(GitExecutorError::ExecutionFailed(e.to_string())), Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { + let _ = Command::new("kill") + .args(["-9", &pid.to_string()]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .ok(); + let _ = rx.recv().ok(); Err(GitExecutorError::Timeout(GIT_TIMEOUT)) } Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => Err( From 6873b62e1c8b0ce2732bf494a07a3280f8ac308e Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 28 Jul 2026 21:36:29 +1200 Subject: [PATCH 20/21] chore: recover guard for health state lock, cfg providers for structured, bytes optional dep --- CHANGELOG.md | 2 +- Cargo.toml | 4 +- README.md | 3 +- src/hooks/builtin/auto_commit.rs | 7 +- src/structured.rs | 116 +++++++++++++++++++++++++++++++ src/tool/health.rs | 29 +++----- 6 files changed, 134 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7f8663a..cf09438 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,7 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. `from_machine()` for inspection and serialize-and-resume. - `LoopMachine::inject(message)` — add an arbitrary message to the machine's history (host steering, or `ContextContributor` goal re-injection). -- `Session`/`Run`/`Turn`/`Run` lifetime types (`engine::core`) — +- `Session`/`Run`/`Turn`/`RunResult` lifetime types (`engine::core`) — the Session ⊃ [Run ⊃ [Turn]] hierarchy: one `Session` spans the process, one `Run` per `run()` prompt, one `Turn` per loop iteration. `Session` derives per-session totals (`total_turns`/`total_duration`/ diff --git a/Cargo.toml b/Cargo.toml index a4a9640..bb699f1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,7 +28,7 @@ uuid = { version = "1", features = ["v4", "serde"] } tracing = "0.1" reqwest = { version = "0.13", default-features = false, features = ["json", "stream", "rustls"], optional = true } -bytes = "1" +bytes = { version = "1", optional = true } async-stream = "0.3" httpdate = { version = "1", optional = true } jsonschema = { version = "0.48", optional = true } @@ -45,7 +45,7 @@ tool_health = [] tool_shield = ["tool_health"] # Providers -providers = ["dep:reqwest", "dep:httpdate"] +providers = ["dep:reqwest", "dep:httpdate", "dep:bytes"] openai = ["providers"] anthropic = ["providers"] ollama = ["providers", "openai"] diff --git a/README.md b/README.md index bca04b0..d9a57cc 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ and tool implementations; the framework handles the rest. | [`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) | `LoopManagers` — the default infrastructure bundle | +| [`managers`](https://docs.rs/loopctl/latest/loopctl/managers/index.html) | `LoopManagers` — 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.* | @@ -152,6 +152,7 @@ let agent = BareLoop::new( | `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) | +| `xai` | No | `grok` | Alias for `grok` (xAI API client) | `gemini` | No | `providers` | Google Gemini API client (`provider::gemini`) | | `zai` | No | `providers`, `anthropic` | Z.AI API client (Anthropic-compatible) | | `grammar` | No | `providers` | Tool-call grammar providers for grammar-aware samplers (vLLM `guided_json`); enables the `Grammar` mode of `ToolConstraint` | diff --git a/src/hooks/builtin/auto_commit.rs b/src/hooks/builtin/auto_commit.rs index f465d0d..a85c9eb 100644 --- a/src/hooks/builtin/auto_commit.rs +++ b/src/hooks/builtin/auto_commit.rs @@ -497,7 +497,7 @@ pub struct AutoCommitConfig { /// Tools that trigger file tracking (e.g., `"Write"`, `"Edit"`). /// - /// Tool names whose output should be watched for a `file_path`; when one + /// Tool names whose `input` is read for a `file_path`; when one /// of these tools runs, the touched file is recorded and staged at /// run end. pub commit_on_tools: Vec, @@ -586,8 +586,9 @@ impl AutoCommitConfigBuilder { /// Set the explicit file allow-list to stage before commit. /// - /// Mirrors [`AutoCommitConfig::files`]; an empty vector stages every - /// change in the working tree. + /// Mirrors [`AutoCommitConfig::files`]; an empty vector is rejected by + /// [`stage_files`](GitExecutor::stage_files) to avoid sweeping up + /// unrelated or secret files. #[must_use] pub fn with_files(mut self, files: Vec) -> Self { self.config.files = files; diff --git a/src/structured.rs b/src/structured.rs index 033447c..4a618f8 100644 --- a/src/structured.rs +++ b/src/structured.rs @@ -674,30 +674,40 @@ mod tests { assert_eq!(opts.response_format.as_ref().unwrap().name, "action"); } + #[cfg(any(feature = "openai", feature = "gemini"))] + #[cfg(any(feature = "openai", feature = "gemini"))] #[test] fn parse_json_lenient_plain_json() { let v = parse_json_lenient(r#"{"a": 1}"#).unwrap(); assert_eq!(v["a"], 1); } + #[cfg(any(feature = "openai", feature = "gemini"))] + #[cfg(any(feature = "openai", feature = "gemini"))] #[test] fn parse_json_lenient_with_prefix() { let v = parse_json_lenient(r#"Here is the JSON: {"a": 1}"#).unwrap(); assert_eq!(v["a"], 1); } + #[cfg(any(feature = "openai", feature = "gemini"))] + #[cfg(any(feature = "openai", feature = "gemini"))] #[test] fn parse_json_lenient_markdown_fences() { let v = parse_json_lenient("```json\n{\"a\": 1}\n```").unwrap(); assert_eq!(v["a"], 1); } + #[cfg(any(feature = "openai", feature = "gemini"))] + #[cfg(any(feature = "openai", feature = "gemini"))] #[test] fn parse_json_lenient_array() { let v = parse_json_lenient(r#"prefix [1, 2, 3] suffix"#).unwrap(); assert_eq!(v[0], 1); } + #[cfg(any(feature = "openai", feature = "gemini"))] + #[cfg(any(feature = "openai", feature = "gemini"))] #[test] fn parse_json_lenient_no_json() { let result = parse_json_lenient("just prose, nothing here"); @@ -711,30 +721,40 @@ mod tests { assert!(err.to_string().contains("schema")); } + #[cfg(any(feature = "openai", feature = "gemini"))] + #[cfg(any(feature = "openai", feature = "gemini"))] #[test] fn parse_json_lenient_brace_inside_string() { let v = parse_json_lenient(r#"prefix {"a": "}"} suffix"#).unwrap(); assert_eq!(v["a"], "}"); } + #[cfg(any(feature = "openai", feature = "gemini"))] + #[cfg(any(feature = "openai", feature = "gemini"))] #[test] fn parse_json_lenient_bracket_inside_string() { let v = parse_json_lenient(r#"before {"x": "]"} after"#).unwrap(); assert_eq!(v["x"], "]"); } + #[cfg(any(feature = "openai", feature = "gemini"))] + #[cfg(any(feature = "openai", feature = "gemini"))] #[test] fn parse_json_lenient_escaped_quote_in_string() { let v = parse_json_lenient(r#"here {"a": "he said \"hi\""} there"#).unwrap(); assert_eq!(v["a"], "he said \"hi\""); } + #[cfg(any(feature = "openai", feature = "gemini"))] + #[cfg(any(feature = "openai", feature = "gemini"))] #[test] fn parse_json_lenient_nested_objects_in_prose() { let v = parse_json_lenient(r#"result: {"outer": {"inner": 42}}"#).unwrap(); assert_eq!(v["outer"]["inner"], 42); } + #[cfg(any(feature = "openai", feature = "gemini"))] + #[cfg(any(feature = "openai", feature = "gemini"))] #[test] fn parse_json_lenient_mismatched_delimiter_then_valid() { let v = parse_json_lenient(r#"{oops] then {"a":1}"#).unwrap(); @@ -949,6 +969,18 @@ mod tests { assert!(matches!(cloned.tool_constraint, ToolConstraint::Strict)); } + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] #[test] fn tighten_sets_additional_properties_false() { let schema = serde_json::json!({ @@ -959,6 +991,18 @@ mod tests { assert_eq!(tightened["additionalProperties"], false); } + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] #[test] fn tighten_enumerates_required() { let schema = serde_json::json!({ @@ -976,6 +1020,18 @@ mod tests { assert!(keys.contains(&"b")); } + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] #[test] fn tighten_recurses_into_nested_objects() { let schema = serde_json::json!({ @@ -999,6 +1055,18 @@ mod tests { assert_eq!(inner_required[0], "x"); } + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] #[test] fn tighten_preserves_non_object_schemas() { let schema = serde_json::json!({"type": "string"}); @@ -1009,6 +1077,18 @@ mod tests { assert!(tightened.get("required").is_none()); } + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] #[test] fn tighten_idempotent_on_already_strict() { let schema = serde_json::json!({ @@ -1022,6 +1102,18 @@ mod tests { assert_eq!(once, twice); } + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] #[test] fn tighten_object_without_properties() { let schema = serde_json::json!({"type": "object"}); @@ -1031,6 +1123,18 @@ mod tests { assert_eq!(tightened["required"].as_array().unwrap().len(), 0); } + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] #[test] fn tighten_recurses_into_local_defs() { // A property that $refs a local definition: the reference itself @@ -1071,6 +1175,18 @@ mod tests { assert_eq!(tightened["properties"]["filter"]["$ref"], "#/$defs/Filter"); } + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] + #[cfg(any( + feature = "anthropic", + feature = "grammar", + feature = "openai", + feature = "gemini" + ))] #[test] fn tighten_recurses_into_legacy_definitions() { // The Draft 07 keyword `definitions` should be walked the same way diff --git a/src/tool/health.rs b/src/tool/health.rs index dd5cf94..870e798 100644 --- a/src/tool/health.rs +++ b/src/tool/health.rs @@ -570,9 +570,7 @@ impl ToolCircuitBreaker { /// (returns `false` to prevent thundering-herd). #[must_use] pub fn allow_request(&self) -> bool { - let Ok(mut state) = self.state.lock() else { - return false; - }; + let mut state = crate::error::recover_guard(self.state.lock()); match state.circuit { CircuitState::Closed => true, CircuitState::HalfOpen => false, @@ -595,10 +593,9 @@ impl ToolCircuitBreaker { /// Resets consecutive failures to zero and transitions the breaker /// to Closed. pub fn record_success(&self) { - if let Ok(mut state) = self.state.lock() { - state.consecutive_failures = 0; - state.circuit = CircuitState::Closed; - } + let mut state = crate::error::recover_guard(self.state.lock()); + state.consecutive_failures = 0; + state.circuit = CircuitState::Closed; } /// Record a failed call. @@ -607,9 +604,7 @@ impl ToolCircuitBreaker { /// `failure_threshold`, the breaker transitions to Open. In the /// `HalfOpen` state, a single failure reopens the breaker. pub fn record_failure(&self) { - let Ok(mut state) = self.state.lock() else { - return; - }; + let mut state = crate::error::recover_guard(self.state.lock()); state.consecutive_failures = state.consecutive_failures.saturating_add(1); state.last_failure_time = Some(Instant::now()); match state.circuit { @@ -649,7 +644,7 @@ impl ToolCircuitBreaker { /// trips the breaker. #[must_use] pub fn consecutive_failures(&self) -> u64 { - self.state.lock().map_or(0, |s| s.consecutive_failures) + crate::error::recover_guard(self.state.lock()).consecutive_failures } /// Whether the breaker is currently in the Closed (healthy) state. @@ -657,9 +652,7 @@ impl ToolCircuitBreaker { /// `true` when requests are allowed unconditionally. #[must_use] pub fn is_closed(&self) -> bool { - self.state - .lock() - .is_ok_and(|s| s.circuit == CircuitState::Closed) + crate::error::recover_guard(self.state.lock()).circuit == CircuitState::Closed } /// Whether the breaker is currently in the Open (blocking) state. @@ -669,9 +662,7 @@ impl ToolCircuitBreaker { /// [`allow_request`](Self::allow_request)). #[must_use] pub fn is_open(&self) -> bool { - self.state - .lock() - .is_ok_and(|s| s.circuit == CircuitState::Open) + crate::error::recover_guard(self.state.lock()).circuit == CircuitState::Open } /// Whether the breaker is currently in the `HalfOpen` (probing) @@ -681,9 +672,7 @@ impl ToolCircuitBreaker { /// probes are refused to avoid a thundering herd. #[must_use] pub fn is_half_open(&self) -> bool { - self.state - .lock() - .is_ok_and(|s| s.circuit == CircuitState::HalfOpen) + crate::error::recover_guard(self.state.lock()).circuit == CircuitState::HalfOpen } } From 4d7b738a92a965c2a0d6bba1cb5b473a8fcd520f Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 28 Jul 2026 23:32:29 +1200 Subject: [PATCH 21/21] fix: thread LoopError through run-end notifications, Gemini usage fix --- README.md | 4 ++-- src/engine/bare.rs | 37 ++++++++++++++++++++---------------- src/engine/bare/emission.rs | 35 +++++++++++++++++++--------------- src/engine/core/lifecycle.rs | 2 +- src/hooks/executor.rs | 17 ++++++++--------- src/observer.rs | 7 ++++++- src/provider/gemini.rs | 21 +++++++------------- src/tool/health.rs | 6 +----- 8 files changed, 66 insertions(+), 63 deletions(-) diff --git a/README.md b/README.md index d9a57cc..01bc4d8 100644 --- a/README.md +++ b/README.md @@ -89,7 +89,7 @@ use std::sync::Arc; # struct MyClient; # use loopctl::api::ApiClient; # impl ApiClient for MyClient { -# fn model(&self) -> &str { "llm-70b" } +# fn model(&self) -> String { "llm-70b".to_string() } # fn stream_messages(&self, _req: &loopctl::api::StreamRequest) # -> std::pin::Pin> + Send>> { # unimplemented!() @@ -152,7 +152,7 @@ let agent = BareLoop::new( | `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) | -| `xai` | No | `grok` | Alias for `grok` (xAI API client) +| `xai` | No | `grok` | Alias for `grok` (xAI API client) | | `gemini` | No | `providers` | Google Gemini API client (`provider::gemini`) | | `zai` | No | `providers`, `anthropic` | Z.AI API client (Anthropic-compatible) | | `grammar` | No | `providers` | Tool-call grammar providers for grammar-aware samplers (vLLM `guided_json`); enables the `Grammar` mode of `ToolConstraint` | diff --git a/src/engine/bare.rs b/src/engine/bare.rs index 69011a6..e22f82c 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -1710,7 +1710,7 @@ impl crate::engine::core::Loop for BareLoop { MachineStep::CallLLM { turn } => { if let Err(e) = self.handle_call_llm(turn).await { self.set_error_state(&e); - self.finalize(false).await?; + self.finalize(Some(&e)).await?; return Err(e); } } @@ -1718,14 +1718,14 @@ impl crate::engine::core::Loop for BareLoop { let turn = self.machine.turns_taken(); if let Err(e) = self.handle_call_tools(turn, &calls).await { self.set_error_state(&e); - self.finalize(false).await?; + self.finalize(Some(&e)).await?; return Err(e); } } MachineStep::Compact { reason } => { if let Err(e) = self.handle_compact(reason).await { self.set_error_state(&e); - self.finalize(false).await?; + self.finalize(Some(&e)).await?; return Err(e); } } @@ -1741,22 +1741,22 @@ impl crate::engine::core::Loop for BareLoop { let err = LoopError::MaxTurnsExceeded { max: self.run_config().max_turns, }; - self.finalize(false).await?; + self.finalize(Some(&err)).await?; return Err(err); } MachineOutcome::Cancelled => { - self.finalize(false).await?; + self.finalize(Some(&LoopError::Cancelled)).await?; return Err(LoopError::Cancelled); } MachineOutcome::Failed { error } => { - self.finalize(false).await?; + self.finalize(Some(&error)).await?; return Err(error); } }, } } - self.finalize(true).await + self.finalize(None).await }) } @@ -1776,14 +1776,14 @@ impl crate::engine::core::Loop for BareLoop { /// so the agent is never left permanently dead after one cancel. fn finalize<'a>( &'a mut self, - success: bool, + error: Option<&'a LoopError>, ) -> Pin + Send + 'a>> { Box::pin(async move { if let Some(run) = self.current_run_mut() { run.end = Some(Instant::now()); } - if success { + if error.is_none() { self.machine.commit_pending(); } else { self.machine.discard_pending(); @@ -1792,7 +1792,7 @@ impl crate::engine::core::Loop for BareLoop { let run = self.current_run().cloned().unwrap_or_default(); let duration = run.duration(); - self.notify_run_end(&run, success, duration); + self.notify_run_end(&run, duration, error); self.cancelled.reset(); Ok(run) @@ -4555,8 +4555,8 @@ mod tests { loop_.notify_run_end( &loop_.current_run().unwrap().clone(), - true, Duration::from_millis(100), + None, ); assert_eq!(hook.captured(), Some(RunEndReason::Complete)); @@ -4589,8 +4589,8 @@ mod tests { loop_.notify_run_end( &loop_.current_run().unwrap().clone(), - true, Duration::from_millis(100), + None, ); assert_eq!(hook.captured(), Some(RunEndReason::Cancelled)); @@ -4614,8 +4614,8 @@ mod tests { loop_.notify_run_end( &loop_.current_run().unwrap().clone(), - true, Duration::from_millis(100), + None, ); assert_eq!(hook.captured(), Some(RunEndReason::MaxTurns)); @@ -4625,11 +4625,12 @@ mod tests { #[tokio::test] async fn run_end_reason_error() { let (loop_, hook) = loop_with_reason_hook(); + let err = LoopError::Api("something went wrong".into()); loop_.notify_run_end( &loop_.current_run().unwrap().clone(), - false, Duration::from_millis(100), + Some(&err), ); assert_eq!(hook.captured(), Some(RunEndReason::Error)); @@ -4639,14 +4640,18 @@ mod tests { #[tokio::test] async fn run_end_reason_context_overflow() { let (loop_, hook) = loop_with_reason_hook(); + let err = LoopError::ContextExceeded { + used: 100_000, + limit: 50_000, + }; loop_.notify_run_end( &loop_.current_run().unwrap().clone(), - false, Duration::from_millis(100), + Some(&err), ); - assert_eq!(hook.captured(), Some(RunEndReason::Error)); + assert_eq!(hook.captured(), Some(RunEndReason::ContextOverflow)); } #[tokio::test] diff --git a/src/engine/bare/emission.rs b/src/engine/bare/emission.rs index e59a9cd..8f8f47b 100644 --- a/src/engine/bare/emission.rs +++ b/src/engine/bare/emission.rs @@ -4,7 +4,7 @@ //! ends. Other observer events (`on_turn_start`, `on_response`, etc.) are fired //! directly at their call sites in the driver loop. -use super::{ApiClient, BareLoop, Duration, Run}; +use super::{ApiClient, BareLoop, Duration, LoopError, Run}; #[cfg(feature = "hooks")] use crate::capabilities::Hookable; #[cfg(feature = "hooks")] @@ -36,18 +36,19 @@ impl BareLoop { /// per-run totals (turn count, tokens); `success` distinguishes a /// normal exit from a failure so observers and hooks can branch on /// the outcome, and `duration` is the wall-clock run length. - pub(super) fn notify_run_end(&self, result: &Run, success: bool, duration: Duration) { + pub(super) fn notify_run_end( + &self, + result: &Run, + duration: Duration, + error: Option<&LoopError>, + ) { + self.notify_run_end_hook(result, error, duration); self.managers.observers().on_run_end(&RunEndContext { - success, - error: if success { - None - } else { - Some("run failed".to_string()) - }, + success: error.is_none(), + error: error.map_or_else(|| None, |e| Some(e.to_string())), total_turns: result.turn_count(), duration_ms: Self::millis_u64(duration), }); - self.notify_run_end_hook(result, success, duration); } /// Derive the structured [`RunEndReason`] from the loop's @@ -59,11 +60,15 @@ impl BareLoop { /// normal completion. Used only to populate the hook end-context — /// observers receive the simpler boolean `success` field. #[cfg(feature = "hooks")] - fn run_end_reason(&self, success: bool) -> RunEndReason { + fn run_end_reason(&self, error: Option<&LoopError>) -> RunEndReason { if self.is_cancelled() { RunEndReason::Cancelled - } else if !success { - RunEndReason::Error + } else if let Some(e) = error { + if matches!(e, LoopError::ContextExceeded { .. }) { + RunEndReason::ContextOverflow + } else { + RunEndReason::Error + } } else if self.current_run().map_or(0, Run::turn_count) >= self.run_config().max_turns { RunEndReason::MaxTurns } else { @@ -111,11 +116,11 @@ impl BareLoop { /// every registered run-end hook. No-op when no hook executor /// is set. #[cfg(feature = "hooks")] - fn notify_run_end_hook(&self, result: &Run, success: bool, duration: Duration) { + fn notify_run_end_hook(&self, result: &Run, error: Option<&LoopError>, duration: Duration) { let Some(executor) = self.managers.hook_executor() else { return; }; - let reason = self.run_end_reason(success); + let reason = self.run_end_reason(error); let ctx = HookRunEndContext { session_id: self.session.id, reason, @@ -132,7 +137,7 @@ impl BareLoop { /// [`notify_run_end`](Self::notify_run_end) compiles /// identically with and without the feature. #[cfg(not(feature = "hooks"))] - fn notify_run_end_hook(&self, _result: &Run, _success: bool, _duration: Duration) {} + fn notify_run_end_hook(&self, _result: &Run, _error: Option<&LoopError>, _duration: Duration) {} /// Convert a [`Duration`] to milliseconds as a `u64`. /// diff --git a/src/engine/core/lifecycle.rs b/src/engine/core/lifecycle.rs index 96149d4..783ba09 100644 --- a/src/engine/core/lifecycle.rs +++ b/src/engine/core/lifecycle.rs @@ -685,7 +685,7 @@ pub trait Loop: Send + Sync { /// and produce a final [`Run`]. fn finalize<'a>( &'a mut self, - success: bool, + error: Option<&'a LoopError>, ) -> Pin + Send + 'a>>; /// Get the current state of the agent. diff --git a/src/hooks/executor.rs b/src/hooks/executor.rs index ceb332d..327f2e5 100644 --- a/src/hooks/executor.rs +++ b/src/hooks/executor.rs @@ -187,16 +187,15 @@ impl HookExecutor { Box::pin(async move { action }) } - /// Check pre-compact hooks, merging every hook's [`CompactResult`]. + /// Check pre-compact hooks, merging results until the first abort. /// - /// Unlike [`check_pre_tool_use`](Self::check_pre_tool_use), there is - /// no short-circuit on the first result — all hooks run, and their - /// results are combined: if any hook returns - /// [`abort: true`](CompactResult::abort), the merged result is that - /// abort (the first one wins, returned immediately); otherwise the - /// last hook's `new_instructions` overrides earlier ones, and every - /// hook's `additional_context` entries accumulate in registration - /// order. + /// Hooks run in registration order. Each hook's + /// [`CompactResult`] is merged + /// into the accumulated result: `new_instructions` overrides earlier + /// values, and `additional_context` entries accumulate. If any hook + /// returns [`abort: true`](CompactResult::abort), execution stops + /// immediately and that abort result is returned — remaining hooks + /// are not called. #[must_use] pub fn check_pre_compact(&self, ctx: &PreCompactContext) -> CompactResult { let mut result = CompactResult::allow(); diff --git a/src/observer.rs b/src/observer.rs index f6e8946..3af642c 100644 --- a/src/observer.rs +++ b/src/observer.rs @@ -289,9 +289,14 @@ impl ObserverHost { for obs in &self.observers { let obs: &dyn LoopObserver = obs.as_ref(); if let Err(payload) = catch_unwind(AssertUnwindSafe(|| f(obs))) { + let msg = payload + .downcast_ref::<&'static str>() + .map(|s| (*s).to_string()) + .or_else(|| payload.downcast_ref::().cloned()) + .unwrap_or_else(|| "non-string panic payload".to_string()); tracing::error!( observer = obs.name(), - payload_downcast = payload.is::<&str>(), + panic_message = %msg, "observer panicked; continuing with remaining observers" ); } diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs index 123e434..d265156 100644 --- a/src/provider/gemini.rs +++ b/src/provider/gemini.rs @@ -760,11 +760,6 @@ struct StreamEmitter { /// Whether a tool-call part is currently open. /// /// Each `functionCall` part emits a [`PartStart`]. Without tracking, - /// [`extract_finish_reason`](Self::extract_finish_reason) would never - /// close tool parts, leaving the downstream accumulator with unbalanced - /// `PartStart`/`PartStop` events. - tool_part_open: bool, - /// Next tool-call index for multiple function calls in one response. /// /// Gemini can emit several `functionCall` parts in a single chunk. @@ -950,7 +945,6 @@ impl StreamEmitter { let idx = self.next_tool_index; self.next_tool_index = self.next_tool_index.saturating_add(1); - self.tool_part_open = true; self.push(StreamEvent::PartStart(PartStart { index: idx, @@ -967,7 +961,6 @@ impl StreamEmitter { }, })); self.push(StreamEvent::PartStop); - self.tool_part_open = false; } } @@ -1000,11 +993,6 @@ impl StreamEmitter { self.push(StreamEvent::PartStop); } - if self.tool_part_open { - self.tool_part_open = false; - self.push(StreamEvent::PartStop); - } - if self.text_part_open { self.text_part_open = false; self.push(StreamEvent::PartStop); @@ -1019,12 +1007,17 @@ impl StreamEmitter { .get("candidatesTokenCount") .and_then(Value::as_u64) .unwrap_or(0); - if input == 0 && output == 0 { + let thoughts = u + .get("thoughtsTokenCount") + .and_then(Value::as_u64) + .unwrap_or(0); + let total_output = output.saturating_add(thoughts); + if input == 0 && total_output == 0 { None } else { Some(Usage::new( u32::try_from(input).unwrap_or(0), - u32::try_from(output).unwrap_or(0), + u32::try_from(total_output).unwrap_or(0), )) } }); diff --git a/src/tool/health.rs b/src/tool/health.rs index 870e798..82b1d78 100644 --- a/src/tool/health.rs +++ b/src/tool/health.rs @@ -627,11 +627,7 @@ impl ToolCircuitBreaker { /// numeric encoding. #[must_use] pub fn state_label(&self) -> &'static str { - match self - .state - .lock() - .map_or(CircuitState::Closed, |s| s.circuit) - { + match crate::error::recover_guard(self.state.lock()).circuit { CircuitState::Closed => "closed", CircuitState::Open => "open", CircuitState::HalfOpen => "half-open",