diff --git a/src/compact.rs b/src/compact.rs index ab976d0..9e9f2aa 100644 --- a/src/compact.rs +++ b/src/compact.rs @@ -6,29 +6,12 @@ //! //! # Architecture //! -//! The design separates **when to compact** from **how to compact**: +//! The design separates when to compact from how to compact: //! -//! - [`ContextManager`] — a concrete struct that monitors token usage, -//! checks thresholds, and decides when to trigger compaction. -//! - [`ContextCompactor`] — a trait that defines the compaction strategy. -//! Plug in truncation, summarization, Q&A extraction, or any custom -//! approach. -//! -//! ```text -//! ┌────────────────────────────┐ -//! │ ContextManager │ -//! │ │ -//! │ estimate_tokens() │ -//! │ should_compact() │ -//! │ ensure_context_fits() │ -//! │ │ │ -//! │ ▼ │ -//! │ ┌──────────────────────┐ │ -//! │ │ dyn ContextCompactor│ │ -//! │ │ .compact() │ │ -//! │ └──────────────────────┘ │ -//! └────────────────────────────┘ -//! ``` +//! 1. [`ContextManager`] monitors token usage, checks thresholds, and +//! decides when to trigger compaction. +//! 2. [`ContextCompactor`] is the trait that defines the compaction +//! strategy. Plug in truncation, summarization, or any custom approach. //! //! # Provided Compactors //! diff --git a/src/engine/bare.rs b/src/engine/bare.rs index 440f050..eb0d2f8 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -1,7 +1,6 @@ //! `BareLoop` — the framework's default agent loop implementation. //! -//! [`BareLoop`] — a generic, framework-level agent -//! loop that orchestrates the full lifecycle of an LLM-based agent session: +//! [`BareLoop`] orchestrates the full lifecycle of an LLM-based agent session: //! sending messages to an LLM API, accumulating streaming responses, //! dispatching tool calls, and feeding results back into the conversation //! until the model ends its turn or a configured limit is reached. @@ -24,8 +23,7 @@ //! [`ApiClient`](crate::api::ApiClient) type parameter `C`, //! avoiding `dyn` overhead for the hot path. //! - **Sequential tool dispatch** — tools within a single turn are -//! executed one after another so cancellation is checked between each. -//! Parallel execution may be added in a future release. +//! executed one after another. //! - **Soft tool errors** — when a tool is not found or returns an error, //! the loop records the error as a tool result and continues, letting //! the model decide how to recover. Only hard errors (API failures, @@ -843,11 +841,6 @@ impl BareLoop { Ok((msg, usage, stop)) } Err(e) => { - if matches!(e, LoopError::Cancelled) { - self.state = LoopState::Cancelled; - return Err(e); - } - let tripped = self.managers.fallback.record_api_failure(); if tripped { let from = self.client.model(); @@ -1057,15 +1050,13 @@ impl BareLoop { /// Set the terminal state for a propagated error. /// - /// Cancellation is a clean termination ([`LoopError::Cancelled`]), not a - /// failure; all other errors are recorded as [`LoopState::Failed`]. + /// 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. fn set_error_state(&mut self, e: &LoopError) { - self.state = if matches!(e, LoopError::Cancelled) { - LoopState::Cancelled - } else { - LoopState::Failed { - error: e.to_string(), - } + self.state = LoopState::Failed { + error: e.to_string(), }; } } @@ -1171,6 +1162,87 @@ impl ModelSwitch<'_, C> { } } +impl BareLoop { + /// Execute the turn body without cancellation awareness. + /// + /// Cancellation is handled by the `select!` in `process_turn`, which + /// drops this future if `cancel.notified()` fires. + /// + /// # 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()); + + 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); + + if let Some(e) = self.apply_loop_detection(current_turn, &pattern) { + return Err(e); + } + + let tool_calls = Self::extract_tool_calls(&msg); + self.conversation.push(msg); + self.budget.total_turns = self.budget.total_turns.saturating_add(1); + + if tool_calls.is_empty() { + return Ok(self.complete_session(text, turn_in, turn_out, turn_start)); + } + + 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); + } + + self.try_compact_context().await; + self.state = LoopState::Processing { + turn: self.budget.total_turns, + }; + + Ok(TurnResult { + text, + tool_calls, + tool_results: Vec::new(), + input_tokens: turn_in, + output_tokens: turn_out, + duration: turn_start.elapsed(), + is_complete: false, + stop_reason: StopReason::ToolCall, + }) + } +} + impl crate::engine::loop_core::Loop for BareLoop { fn initialize<'a>( &'a mut self, @@ -1200,70 +1272,23 @@ impl crate::engine::loop_core::Loop for BareLoop { self.record_user_input(input); self.fire_turn_start(current_turn, input); - let (msg, usage, _stream_stop) = self.do_stream().await?; - self.accumulate_usage(usage.as_ref()); - - 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); - - if let Some(e) = self.apply_loop_detection(current_turn, &pattern) { - return Err(e); - } - - let tool_calls = Self::extract_tool_calls(&msg); - self.conversation.push(msg); - self.budget.total_turns = self.budget.total_turns.saturating_add(1); - - if tool_calls.is_empty() { - return Ok(self.complete_session(text, turn_in, turn_out, turn_start)); - } - - 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(), - }; - - // Temporarily extract budget to avoid double mutable borrow. - 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); + tokio::select! { + biased; + () = cancel.notified() => { + self.state = LoopState::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) + } + result = self.run_turn_body(current_turn, turn_start) => result, } - - self.try_compact_context().await; - self.state = LoopState::Processing { - turn: self.budget.total_turns, - }; - - Ok(TurnResult { - text, - tool_calls, - tool_results: Vec::new(), - input_tokens: turn_in, - output_tokens: turn_out, - duration: turn_start.elapsed(), - is_complete: false, - stop_reason: StopReason::ToolCall, - }) }) } @@ -3471,4 +3496,189 @@ mod tests { assert_eq!(hook.captured(), Some(SessionEndReason::ContextOverflow)); } + + #[tokio::test] + async fn process_turn_cancel_during_streaming_returns_fast() { + let (client, tx) = StreamingMockClient::new("test-model"); + tx.send(Ok(StreamEvent::MessageStart(MessageStart { + message: MessageMetadata { + id: "msg-1".into(), + role: "assistant".into(), + model: "test-model".into(), + }, + }))) + .await + .unwrap(); + + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + let observer = Arc::new(CountingObserver::new()); + agent.register_observer(observer.clone()); + let signal = agent.cancel_signal(); + + let handle = tokio::spawn(async move { agent.run("Hi").await }); + + for _ in 0..5 { + tokio::task::yield_now().await; + } + let start = Instant::now(); + signal.cancel(); + + let result = handle.await.unwrap(); + let elapsed = start.elapsed(); + + match result { + Err(LoopError::Cancelled) => {} + other => panic!("expected Err(LoopError::Cancelled), got {other:?}"), + } + assert!( + elapsed < Duration::from_secs(5), + "cancel during streaming should return fast; elapsed {elapsed:?}", + ); + assert_eq!( + observer.turn_ends.load(Ordering::SeqCst), + 1, + "on_turn_end should fire once on cancel", + ); + } + + #[tokio::test] + async fn process_turn_cancel_during_dispatch_fires_turn_end() { + struct SlowTool { + notify: Arc, + } + impl Tool for SlowTool { + fn name(&self) -> &'static str { + "slow" + } + fn description(&self) -> &'static str { + "Blocks until notified" + } + fn schema(&self) -> ToolSchema { + ToolSchema { + tool: "slow".into(), + description: "Blocks until notified".into(), + input_schema: json!({"type": "object", "properties": {}}), + } + } + fn call( + &self, + _input: Value, + _ctx: &ToolContext, + ) -> Pin> + Send + '_>> + { + let notify = self.notify.clone(); + Box::pin(async move { + notify.notified().await; + Ok(ToolOutput::text("done")) + }) + } + } + + let notify = Arc::new(tokio::sync::Notify::new()); + let mut registry = ToolRegistry::new(); + registry.register(SlowTool { + notify: notify.clone(), + }); + + let client = MockClient::new("test"); + client.add_tool_only_response("tc-1", "slow", json!({})); + + let mut agent = BareLoop::new(Arc::new(client), registry, make_config()); + let observer = Arc::new(CountingObserver::new()); + agent.register_observer(observer.clone()); + let signal = agent.cancel_signal(); + + let handle = tokio::spawn(async move { agent.run("Use slow tool").await }); + + for _ in 0..10 { + tokio::task::yield_now().await; + } + signal.cancel(); + + let result = handle.await.unwrap(); + match result { + Err(LoopError::Cancelled) => {} + other => panic!("expected Err(LoopError::Cancelled), got {other:?}"), + } + assert_eq!( + observer.turn_ends.load(Ordering::SeqCst), + 1, + "on_turn_end(false) must fire on cancel during dispatch", + ); + assert_eq!( + observer.session_ends.load(Ordering::SeqCst), + 1, + "on_session_end must fire via finalize after cancel", + ); + } + + #[tokio::test] + async fn process_turn_cancel_during_recovery_backoff_returns_fast() { + struct AlwaysRecoverable; + impl crate::reflection::Reflector for AlwaysRecoverable { + fn analyze( + &self, + error: &str, + tool_name: &str, + _tool_input: &serde_json::Value, + _context: &crate::reflection::ReflectionContext, + ) -> Pin< + Box< + dyn Future< + Output = Result< + crate::reflection::FailureAnalysis, + crate::reflection::ReflectionError, + >, + > + Send + + '_, + >, + > { + let error = error.to_string(); + let tool_name = tool_name.to_string(); + Box::pin(async move { + Ok(crate::reflection::FailureAnalysis { + is_recoverable: true, + root_cause: error, + severity: crate::reflection::FailureSeverity::Medium, + correction: None, + context: format!("tool: {tool_name}"), + }) + }) + } + } + + let client = MockClient::new("test"); + client.add_tool_only_response("tc-1", "fail", json!({})); + + let mut registry = ToolRegistry::new(); + registry.register(FailingTool); + + let mut agent = BareLoop::new(Arc::new(client), registry, make_config()); + agent.set_reflector(Arc::new(AlwaysRecoverable)); + agent.set_recovery_strategy(Arc::new( + crate::reflection::ExponentialBackoffRecovery::new(5) + .with_base_delay(Duration::from_secs(60)), + )); + let signal = agent.cancel_signal(); + + let handle = tokio::spawn(async move { agent.run("Use failing tool").await }); + + for _ in 0..10 { + tokio::task::yield_now().await; + } + let start = Instant::now(); + signal.cancel(); + + let result = handle.await.unwrap(); + let elapsed = start.elapsed(); + + match result { + Err(LoopError::Cancelled) => {} + other => panic!("expected Err(LoopError::Cancelled), got {other:?}"), + } + assert!( + elapsed < Duration::from_secs(5), + "cancel during recovery backoff should return fast, not wait 60s; elapsed {elapsed:?}", + ); + } } diff --git a/src/engine/bare/compact.rs b/src/engine/bare/compact.rs index 5ce5007..08f13d5 100644 --- a/src/engine/bare/compact.rs +++ b/src/engine/bare/compact.rs @@ -1,8 +1,8 @@ -//! Context compaction phase — check and compact the conversation when needed. +//! Context compaction for agent conversations. //! -//! Extracted from [`BareLoop`] to isolate the compaction concern. -//! When a [`ContextManager`] is configured, checks token usage after each -//! tool dispatch and triggers compaction if usage exceeds the threshold. +//! 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. use super::{ApiClient, BareLoop, Instant, LoopError}; #[cfg(feature = "hooks")] @@ -34,66 +34,32 @@ impl BareLoop { return Ok(()); }; - #[cfg(feature = "hooks")] - let messages_before = self.conversation.len(); - - // Pre-compact hook check - #[cfg(feature = "hooks")] - if let Some(executor) = self.managers.hook_executor() { - let tokens_before = - crate::compact::CompactionOutcome::estimate_tokens(&self.conversation); - let ctx = PreCompactContext { - trigger: CompactTrigger::Auto, - custom_instructions: None, - message_count: messages_before, - tokens_before, - context_window: self.config.context_window, - session_id: self.config.session_id, - }; - let hook_result = executor.check_pre_compact(&ctx); - if hook_result.abort { - // Hook aborted compaction — return Ok, conversation unchanged. - return Ok(()); - } - // Note: hook_result.new_instructions and hook_result.additional_context - // are available for future use with a hook-aware compactor. + if self.pre_compact_hook_aborts() { + return Ok(()); } + let messages_before = self.conversation.len(); let compact_start = Instant::now(); - let result = ctx_manager - .ensure_context_fits(std::mem::take(&mut self.conversation), turn) - .await; - #[cfg(feature = "hooks")] - let compact_duration_ms = u64::try_from(compact_start.elapsed().as_millis()).unwrap_or(0); - #[cfg(not(feature = "hooks"))] - let _ = compact_start; + let conversation = self.conversation.clone(); + let result = ctx_manager.ensure_context_fits(conversation, turn).await; + match result { Ok(EnsureContextResult::Compacted(outcome)) => { + let tokens_after = outcome.tokens_after; + let tokens_saved = outcome.tokens_saved; self.conversation = outcome.messages; - #[cfg(feature = "hooks")] - let messages_after = self.conversation.len(); - let tokens_before = outcome.tokens_after.saturating_add(outcome.tokens_saved); + let tokens_before = tokens_after.saturating_add(tokens_saved); self.managers.observers().on_compaction(&CompactedContext { tokens_before, - tokens_after: outcome.tokens_after, - tokens_saved: outcome.tokens_saved, + tokens_after, + tokens_saved, }); - - // Post-compact hook notification - #[cfg(feature = "hooks")] - if let Some(executor) = self.managers.hook_executor() { - let messages_compacted = messages_before.saturating_sub(messages_after); - let ctx = PostCompactContext { - trigger: CompactTrigger::Auto, - messages_compacted, - tokens_saved: outcome.tokens_saved, - tokens_after: outcome.tokens_after, - duration_ms: compact_duration_ms, - session_id: self.config.session_id, - }; - executor.notify_post_compact(&ctx); - } - + self.notify_post_compact_hook( + messages_before, + tokens_after, + tokens_saved, + compact_start.elapsed(), + ); Ok(()) } Ok(EnsureContextResult::NoAction(messages)) => { @@ -106,4 +72,61 @@ impl BareLoop { }), } } + + /// Run the pre-compact hook. Returns `true` if the hook aborts compaction. + #[cfg(feature = "hooks")] + fn pre_compact_hook_aborts(&self) -> bool { + let Some(executor) = self.managers.hook_executor() else { + return false; + }; + let tokens_before = crate::compact::CompactionOutcome::estimate_tokens(&self.conversation); + let ctx = PreCompactContext { + trigger: CompactTrigger::Auto, + custom_instructions: None, + message_count: self.conversation.len(), + tokens_before, + context_window: self.config.context_window, + session_id: self.config.session_id, + }; + executor.check_pre_compact(&ctx).abort + } + + #[cfg(not(feature = "hooks"))] + fn pre_compact_hook_aborts(&self) -> bool { + false + } + + /// Notify the post-compact hook that compaction completed. + #[cfg(feature = "hooks")] + fn notify_post_compact_hook( + &self, + messages_before: usize, + tokens_after: u64, + tokens_saved: u64, + duration: std::time::Duration, + ) { + 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, + }; + executor.notify_post_compact(&ctx); + } + + #[cfg(not(feature = "hooks"))] + fn notify_post_compact_hook( + &self, + _messages_before: 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 ee67155..112b93a 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -1,8 +1,7 @@ -//! Tool dispatch phase — execute tool calls requested by the model. +//! Tool dispatch — execute tool calls requested by the model. //! -//! Extracted from [`BareLoop`] to isolate the tool dispatch concern. -//! Handles sequential tool execution, reflection/recovery on errors, -//! hook interception, health recording, and middleware pipeline dispatch. +//! Sequential tool execution with reflection and recovery on errors, hook +//! interception, health recording, and middleware pipeline support. #[cfg(feature = "hooks")] use super::HookAction; @@ -33,35 +32,22 @@ use std::panic::AssertUnwindSafe; enum RecoveryOutcome { /// Return this soft-error result to the caller. SoftError(ToolDispatchResult), - /// The session was cancelled during the recovery wait. - Cancelled, } impl BareLoop { /// Execute tool calls and return results. /// - /// Iterates over each [`ToolCall`] extracted from the assistant - /// message, looks up the corresponding tool in the [`ToolRegistry`](crate::tool::ToolRegistry), - /// and invokes it. Each result is wrapped in a [`ToolDispatchResult`]. + /// Dispatch a batch of tool calls sequentially. /// - /// Tool execution is **sequential** so that cancellation can be - /// checked between invocations. A tool that is not found in the - /// registry produces a soft error result (not a hard [`LoopError`]), - /// allowing the model to recover. - /// - /// When a tool returns an error (execution failure or not-found), - /// the framework consults the [`Reflector`](crate::reflection::Reflector) and [`RecoveryStrategy`](crate::reflection::RecoveryStrategy) - /// to decide whether to retry, skip, ask user, or fail. Retry - /// attempts use the delay specified by the [`RecoveryAction`]. - /// - /// Observers are notified before and after each tool invocation via - /// [`LoopObserver::on_tool_pre`](crate::observer::LoopObserver::on_tool_pre) and - /// [`LoopObserver::on_tool_post`](crate::observer::LoopObserver::on_tool_post). + /// Each [`ToolCall`] is executed one at a time via + /// [`dispatch_tool_with_recovery`](Self::dispatch_tool_with_recovery). + /// A tool that is not found in the registry produces a soft error result + /// (not a hard [`LoopError`]), allowing the model to recover. /// /// # Errors /// - /// Returns [`LoopError::Cancelled`] if the cancellation flag is set - /// between tool invocations. + /// Returns [`LoopError`] if any tool dispatch fails with a hard error + /// (e.g. loop detection forces an abort). pub(super) async fn dispatch_tools( &self, tool_calls: &[ToolCall], @@ -69,34 +55,27 @@ impl BareLoop { ) -> Result, LoopError> { let mut results = Vec::with_capacity(tool_calls.len()); for tc in tool_calls { - if self.is_cancelled() { - return Err(LoopError::Cancelled); - } let result = self.dispatch_tool_with_recovery(tc, turn_idx).await?; results.push(result); } Ok(results) } - /// Dispatch a single tool call, using reflector + recovery on errors. + /// Dispatch a single tool call with reflection and recovery on errors. /// - /// If the tool call succeeds, returns the result immediately. If it - /// fails, calls [`Reflector::analyze`](crate::reflection::Reflector::analyze) and [`RecoveryStrategy::decide`](crate::reflection::RecoveryStrategy::decide) - /// to determine the next action: + /// 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. /// - /// - [`Retry`](RecoveryAction::Retry) — re-dispatch the tool after the - /// specified delay, up to the recovery strategy's retry limit. - /// - [`Skip`](RecoveryAction::Skip) — produce a soft error result and - /// continue to the next tool. - /// - [`Fail`](RecoveryAction::Fail) — produce a soft error result (the - /// model sees the failure and can decide how to respond). - /// - [`AskUser`](RecoveryAction::AskUser) — treated as `Skip` (interactive - /// recovery not yet supported in `BareLoop`). + /// 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. /// /// # Errors /// - /// Returns [`LoopError::Cancelled`] if the cancellation signal fires - /// during tool execution or between retry attempts. + /// Returns [`LoopError`] if the detection manager signals a hard stop + /// (e.g. [`LoopError::LoopDetected`]). async fn dispatch_tool_with_recovery( &self, tc: &ToolCall, @@ -107,37 +86,15 @@ impl BareLoop { let mut tc = tc.clone(); loop { - if self.is_cancelled() { - return Err(LoopError::Cancelled); - } - - self.managers.observers().on_tool_pre(&ToolPreContext { - turn: turn_idx, - tool: tc.tool.clone(), - tool_call_id: tc.id.clone(), - }); + self.fire_tool_pre(turn_idx, &tc); if let Some(blocked) = self.check_pre_tool_use_hooks(&tc, turn_idx) { - // Pair on_tool_pre with on_tool_post so observers see a - // complete lifecycle even when a hook blocks the call. - self.managers.observers().on_tool_post(&ToolPostContext { - turn: turn_idx, - tool: tc.tool.clone(), - result_hash: loop_detector::hash_result(&blocked.output.to_string()), - is_error: blocked.is_error, - duration: Duration::ZERO, - }); + self.fire_tool_post(turn_idx, &tc, &blocked); return Ok(blocked); } if let Some(blocked) = self.pre_detection(&tc, turn_idx)? { - self.managers.observers().on_tool_post(&ToolPostContext { - turn: turn_idx, - tool: tc.tool.clone(), - result_hash: loop_detector::hash_result(&blocked.output.to_string()), - is_error: blocked.is_error, - duration: Duration::ZERO, - }); + self.fire_tool_post(turn_idx, &tc, &blocked); return Ok(blocked); } @@ -147,13 +104,7 @@ impl BareLoop { .await?; self.post_detection(&tc, &tool_result); - self.managers.observers().on_tool_post(&ToolPostContext { - turn: turn_idx, - tool: tc.tool.clone(), - result_hash: loop_detector::hash_result(&tool_result.output.to_string()), - is_error: tool_result.is_error, - duration: tool_result.duration, - }); + self.fire_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); @@ -167,29 +118,66 @@ impl BareLoop { { Ok((next_attempt, correction)) => { attempt = next_attempt; - if let Some(ref correction) = correction { - let correction_result = tc.apply_correction(correction, &tool_result); - if let CorrectionResult::Failed(msg) = &correction_result { - tracing::warn!( - tool = %tc.tool, - error = %msg, - "correction failed to produce a usable retry" - ); - } - } + Self::apply_correction_if_present(&mut tc, correction, &tool_result); } Err(RecoveryOutcome::SoftError(returned_result)) => return Ok(returned_result), - Err(RecoveryOutcome::Cancelled) => return Err(LoopError::Cancelled), } } } - /// Check for a loop pattern before executing the tool. + /// Fire [`on_tool_pre`](crate::observer::LoopObserver::on_tool_pre) for + /// the tool call about to be dispatched. + fn fire_tool_pre(&self, turn_idx: usize, tc: &ToolCall) { + self.managers.observers().on_tool_pre(&ToolPreContext { + turn: turn_idx, + tool: tc.tool.clone(), + tool_call_id: tc.id.clone(), + }); + } + + /// Fire [`on_tool_post`](crate::observer::LoopObserver::on_tool_post). + /// + /// Called for both successful dispatches and blocked paths (hooks, + /// detection) so observers always see a complete `pre` / `post` pair. + fn fire_tool_post(&self, turn_idx: usize, tc: &ToolCall, result: &ToolDispatchResult) { + self.managers.observers().on_tool_post(&ToolPostContext { + turn: turn_idx, + tool: tc.tool.clone(), + result_hash: loop_detector::hash_result(&result.output.to_string()), + is_error: result.is_error, + duration: result.duration, + }); + } + + /// Apply a correction produced by the recovery strategy, if any. + /// + /// 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, + ) { + let Some(correction) = correction else { return }; + if let CorrectionResult::Failed(msg) = tc.apply_correction(&correction, tool_result) { + tracing::warn!( + tool = %tc.tool, + error = %msg, + "correction failed to produce a usable retry" + ); + } + } + + /// Record the tool call's input signature and check for loop patterns. + /// + /// Returns `Some(blocked_result)` if loop detection blocks the call, + /// or `None` if dispatch should proceed. /// /// # Errors /// - /// Returns [`LoopError`] when the detection manager signals a hard - /// stop (e.g. [`LoopError::LoopDetected`]). + /// Returns [`LoopError::LoopDetected`] if the detection manager signals + /// a hard stop. fn pre_detection( &self, tc: &ToolCall, @@ -211,11 +199,10 @@ impl BareLoop { } } - /// Record the tool result with the detection manager (post-execution). + /// Record the tool result's output hash for loop detection. /// - /// Constructs an [`Operation`] with the result hash and records it with - /// the detection manager. This lets the detector distinguish "same input, - /// same output" (stuck) from "same input, different output" (progress). + /// Lets the detector distinguish "same input, same output" (stuck) from + /// "same input, different output" (progress). fn post_detection(&self, tc: &ToolCall, tool_result: &ToolDispatchResult) { let result_hash = match &tool_result.output { ToolContent::Text(t) => loop_detector::hash_result(t), @@ -230,17 +217,18 @@ impl BareLoop { self.managers.detection.record_operation(operation); } - /// Execute a single tool call through the pipeline or registry. + /// Execute a single tool call. + /// + /// Tries the middleware pipeline first (if configured), then falls back + /// to a direct registry lookup. Tool panics are caught and converted to + /// error results. A tool not in the registry produces a soft error. /// - /// Tries the middleware pipeline first, then a direct registry lookup, - /// then produces a not-found error result. Handles cancellation during - /// execution. Observer notification is handled by the caller - /// (`dispatch_tool_with_recovery`). + /// Observer notifications are handled by the caller + /// ([`dispatch_tool_with_recovery`](Self::dispatch_tool_with_recovery)). /// /// # Errors /// - /// Returns [`LoopError::Cancelled`] if the cancel signal fires - /// during tool execution. + /// Returns [`LoopError`] if loop detection forces a hard stop. async fn dispatch_tool( &self, tc: &ToolCall, @@ -255,16 +243,9 @@ impl BareLoop { } let tool_result = if let Some(tool) = self.tools.get(&tc.tool) { - let cancel = Arc::clone(&self.cancelled); - // Wrap the tool call in `catch_unwind` so a panicking tool - // implementation produces an error result instead of unwinding - // through and aborting the entire agent loop. - let call_result = tokio::select! { - r = AssertUnwindSafe(tool.call(tc.input.clone(), tool_context)).catch_unwind() => r, - () = cancel.notified() => { - return Err(LoopError::Cancelled); - } - }; + let call_result = AssertUnwindSafe(tool.call(tc.input.clone(), tool_context)) + .catch_unwind() + .await; match call_result { Ok(Ok(result)) => { let duration = start.elapsed(); @@ -317,10 +298,9 @@ impl BareLoop { Ok(tool_result) } - /// Build a soft-error result for a tool that isn't in the registry. + /// Build a soft-error result for a tool that is not in the registry. /// - /// Notifies observers with the error message - /// that lists available tool names to help the model recover. + /// The error message lists available tool names to help the model recover. fn tool_not_found(&self, tc: &ToolCall) -> ToolDispatchResult { let available: Vec = self.tools.tool_names(); let available_refs: Vec<&str> = available.iter().map(String::as_str).collect(); @@ -335,18 +315,19 @@ impl BareLoop { } } - /// Decide whether to retry a failed tool or return the error result. + /// Decide whether to retry a failed tool or return the error as a soft result. /// - /// Consults the reflector and recovery strategy. On `Retry`, sleeps for - /// the prescribed delay (cancellation-aware) and returns the updated - /// attempt count and the [`Correction`] (if any) via `Ok`. On all other - /// recovery actions, returns the original error result via `Err` (which - /// ends the retry loop). + /// Consults the [`Reflector`](crate::reflection::Reflector) and + /// [`RecoveryStrategy`](crate::reflection::RecoveryStrategy). On + /// [`Retry`](RecoveryAction::Retry), sleeps for the prescribed delay and + /// returns the updated attempt count and optional [`Correction`]. On all + /// other actions (`Skip`, `Fail`, `AskUser`), returns the original error + /// result as a soft error. /// /// # Errors /// - /// Returns `Err(RecoveryOutcome)` when the recovery strategy decides - /// not to retry — the caller should return this as a soft error. + /// Returns `Err(RecoveryOutcome::SoftError)` when the recovery strategy + /// decides not to retry. async fn recovery_wait_or_return( &self, tc: &ToolCall, @@ -357,12 +338,7 @@ impl BareLoop { match recovery_action { RecoveryAction::Retry { delay } => { let next_attempt = attempt.saturating_add(1); - tokio::select! { - () = tokio::time::sleep(delay) => {}, - () = self.cancelled.notified() => { - return Err(RecoveryOutcome::Cancelled); - } - } + tokio::time::sleep(delay).await; Ok((next_attempt, correction)) } RecoveryAction::Skip(_) | RecoveryAction::AskUser(_) | RecoveryAction::Fail(_) => { @@ -376,120 +352,112 @@ impl BareLoop { /// /// Returns `Some(ToolDispatchResult)` with an error result if a hook /// blocked the call, or `None` if the call should proceed. - /// - /// *Requires `hooks` feature; returns `None` otherwise.* - #[allow(clippy::unused_self)] + #[cfg(feature = "hooks")] fn check_pre_tool_use_hooks( &self, tc: &ToolCall, turn_idx: usize, ) -> Option { - #[cfg(feature = "hooks")] - if let Some(executor) = self.managers.hook_executor() { - let ctx = PreToolUseContext { - tool_name: tc.tool.clone(), - input: tc.input.clone(), - session_id: self.config.session_id, - turn_number: turn_idx, - }; - match executor.check_pre_tool_use(&ctx) { - HookAction::Allow => None, - HookAction::Block { reason } => Some(ToolDispatchResult { - tool_call_id: tc.id.clone(), - output: ToolContent::Text(reason), - is_error: true, - duration: Duration::ZERO, - resolved_tool_name: tc.tool.clone(), - }), - HookAction::Ask { message } => { - // In Headless mode (the default) the executor already - // downgrades Ask → Block. If we reach this arm the - // executor is Interactive, but BareLoop has no UI to - // show a prompt, so we still treat it as Block. - Some(ToolDispatchResult { - tool_call_id: tc.id.clone(), - output: ToolContent::Text(message), - is_error: true, - duration: Duration::ZERO, - resolved_tool_name: tc.tool.clone(), - }) - } - } - } else { - None - } - #[cfg(not(feature = "hooks"))] - { - let _ = (tc, turn_idx); - None + let executor = self.managers.hook_executor()?; + let ctx = PreToolUseContext { + tool_name: tc.tool.clone(), + input: tc.input.clone(), + session_id: self.config.session_id, + turn_number: turn_idx, + }; + match executor.check_pre_tool_use(&ctx) { + HookAction::Allow => None, + HookAction::Block { reason } => Some(ToolDispatchResult { + tool_call_id: tc.id.clone(), + output: ToolContent::Text(reason), + is_error: true, + duration: Duration::ZERO, + resolved_tool_name: tc.tool.clone(), + }), + HookAction::Ask { message } => Some(ToolDispatchResult { + tool_call_id: tc.id.clone(), + output: ToolContent::Text(message), + is_error: true, + duration: Duration::ZERO, + resolved_tool_name: tc.tool.clone(), + }), } } + #[cfg(not(feature = "hooks"))] + fn check_pre_tool_use_hooks( + &self, + _tc: &ToolCall, + _turn_idx: usize, + ) -> Option { + None + } + /// Notify post-tool-use hooks with the execution result. - /// - /// *Requires `hooks` feature; no-op otherwise.* - #[allow(clippy::unused_self)] + #[cfg(feature = "hooks")] fn notify_post_tool_use_hooks( &self, tc: &ToolCall, tool_result: &ToolDispatchResult, turn_idx: usize, ) { - #[cfg(feature = "hooks")] - if let Some(executor) = self.managers.hook_executor() { - let output_text = tool_result.output.to_string(); - let ctx = PostToolUseContext { - tool_name: tc.tool.clone(), - input: tc.input.clone(), - output: output_text, - is_error: tool_result.is_error, - duration_ms: tool_result - .duration - .as_millis() - .try_into() - .unwrap_or(u64::MAX), - session_id: self.config.session_id, - turn_number: turn_idx, - }; - executor.notify_post_tool_use(&ctx); - } - #[cfg(not(feature = "hooks"))] - { - let _ = (tc, tool_result, turn_idx); - } + let Some(executor) = self.managers.hook_executor() else { + return; + }; + let output_text = tool_result.output.to_string(); + let ctx = PostToolUseContext { + tool_name: tc.tool.clone(), + input: tc.input.clone(), + output: output_text, + is_error: tool_result.is_error, + duration_ms: tool_result + .duration + .as_millis() + .try_into() + .unwrap_or(u64::MAX), + session_id: self.config.session_id, + turn_number: turn_idx, + }; + executor.notify_post_tool_use(&ctx); + } + + #[cfg(not(feature = "hooks"))] + fn notify_post_tool_use_hooks( + &self, + _tc: &ToolCall, + _tool_result: &ToolDispatchResult, + _turn_idx: usize, + ) { } /// Record tool health (success or failure) in the health registry. - /// - /// *Requires `tool_health` feature; no-op otherwise.* - #[allow(clippy::unused_self)] + #[cfg(feature = "tool_health")] fn record_tool_health(&self, tool_name: &str, tool_result: &ToolDispatchResult) { - #[cfg(feature = "tool_health")] - if let Some(health) = self.managers.health_registry() { - if tool_result.is_error { - health.record_failure(tool_name, tool_result.duration); - } else { - health.record_success(tool_name, tool_result.duration); - } - } - #[cfg(not(feature = "tool_health"))] - { - let _ = (tool_name, tool_result); + let Some(health) = self.managers.health_registry() else { + return; + }; + if tool_result.is_error { + health.record_failure(tool_name, tool_result.duration); + } else { + health.record_success(tool_name, tool_result.duration); } } + #[cfg(not(feature = "tool_health"))] + fn record_tool_health(&self, _tool_name: &str, _tool_result: &ToolDispatchResult) {} + /// Dispatch a tool call through the middleware pipeline. /// - /// Builds a [`ToolDispatchContext`] from the tool call info, delegates - /// to the pipeline's middleware chain, and converts the - /// [`ToolDispatchResult`] back to a [`ToolDispatchResult`]. - /// Observer notification is handled by the caller - /// ([`dispatch_tool_with_recovery`]). + /// Builds a [`ToolDispatchContext`] and delegates to the pipeline's + /// middleware chain (timeout, permissions, output limits, etc.). + /// + /// Observer notifications are handled by the caller + /// ([`dispatch_tool_with_recovery`](Self::dispatch_tool_with_recovery)). /// /// # Errors /// - /// Returns [`LoopError::Cancelled`] if the cancel signal fires - /// during pipeline dispatch. + /// Never returns an error — pipeline dispatch always produces a result + /// (soft errors are returned as `Ok` with `is_error: true`). async fn dispatch_via_pipeline( &self, pipeline: &ToolPipeline, @@ -506,20 +474,7 @@ impl BareLoop { permission: PermissionCheck::Allow, tool_context: tool_context.clone(), }; - let cancel = Arc::clone(&self.cancelled); - let dispatch_result = tokio::select! { - r = pipeline.invoke(ctx) => r, - () = cancel.notified() => { - return Err(LoopError::Cancelled); - } - }; - // Guard against a late-arriving cancellation that races with the - // pipeline future resolving first. Without this check the result - // would be treated as a soft tool error by ToolCallMiddleware - // instead of a hard cancellation. - if cancel.is_cancelled() { - return Err(LoopError::Cancelled); - } + let dispatch_result = pipeline.invoke(ctx).await; Ok(ToolDispatchResult { tool_call_id: if dispatch_result.tool_call_id.is_empty() { tc.id.clone() @@ -535,13 +490,14 @@ impl BareLoop { /// Analyse a tool error and decide on a recovery action. /// - /// Calls [`Reflector::analyze()`] and then [`RecoveryStrategy::decide()`]. - /// If the reflector itself fails, logs the error and returns - /// [`RecoveryAction::Fail`] (conservative default). + /// Calls [`Reflector::analyze`](crate::reflection::Reflector::analyze) to + /// classify the failure, then + /// [`RecoveryStrategy::decide`](crate::reflection::RecoveryStrategy::decide) + /// to choose the action. If the reflector itself fails, conservatively + /// returns [`RecoveryAction::Fail`]. /// - /// Returns the [`RecoveryAction`] alongside the [`Correction`] (if any) - /// produced by the reflector. The correction is threaded through so the - /// retry loop can apply it before re-dispatching. + /// Returns the [`RecoveryAction`] and an optional [`Correction`] that the + /// retry loop applies to the tool input before re-dispatching. async fn recover_tool_error( &self, tc: &ToolCall, diff --git a/src/engine/bare/emission.rs b/src/engine/bare/emission.rs index 50ae8cc..65ef617 100644 --- a/src/engine/bare/emission.rs +++ b/src/engine/bare/emission.rs @@ -1,12 +1,8 @@ -//! Session lifecycle notifications. +//! Session lifecycle notifications — start and end events. //! -//! Split from [`BareLoop`] for clarity — these methods dispatch to -//! the [`ObserverHost`](crate::observer::ObserverHost) and the hook executor. -//! -//! Only session start/end live here because they do *two* things: -//! observer notification + hook dispatch. All other observer notifications -//! are called directly at their call sites via -//! `self.managers.observers().on_*()`. +//! 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`. use super::{ApiClient, BareLoop, Duration, SessionResult}; #[cfg(feature = "hooks")] @@ -30,18 +26,7 @@ impl BareLoop { .on_session_start(&SessionStartContext { session_id: self.config.session_id, }); - - #[cfg(feature = "hooks")] - if let Some(executor) = self.managers.hook_executor() { - let ctx = HookSessionStartContext { - session_id: self.config.session_id, - model: self.config.model.clone(), - working_directory: std::env::current_dir() - .map(|p| p.to_string_lossy().into_owned()) - .unwrap_or_default(), - }; - executor.notify_session_start(&ctx); - } + self.notify_session_start_hook(); } /// Notify all observers and hooks that the session has ended. @@ -54,19 +39,7 @@ impl BareLoop { total_turns: result.total_turns, duration_ms: Self::millis_u64(duration), }); - - #[cfg(feature = "hooks")] - if let Some(executor) = self.managers.hook_executor() { - let reason = self.session_end_reason(result.success); - 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); - } + self.notify_session_end_hook(result, duration); } /// Derive the structured [`SessionEndReason`] from the loop's @@ -76,9 +49,9 @@ impl BareLoop { /// cancellation, max-turns exhaustion, and context overflow. #[cfg(feature = "hooks")] fn session_end_reason(&self, success: bool) -> SessionEndReason { - if !success { - // Context overflow is a specific failure mode distinguishable - // from a generic error by its message. + if self.is_cancelled() { + SessionEndReason::Cancelled + } else if !success { if self .budget .error @@ -89,8 +62,6 @@ impl BareLoop { } else { SessionEndReason::Error } - } else if self.is_cancelled() { - SessionEndReason::Cancelled } else if self.budget.total_turns >= self.config.max_turns { SessionEndReason::MaxTurns } else { @@ -98,6 +69,43 @@ impl BareLoop { } } + #[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(), + working_directory: std::env::current_dir() + .map(|p| p.to_string_lossy().into_owned()) + .unwrap_or_default(), + }; + executor.notify_session_start(&ctx); + } + + #[cfg(not(feature = "hooks"))] + fn notify_session_start_hook(&self) {} + + #[cfg(feature = "hooks")] + fn notify_session_end_hook(&self, result: &SessionResult, duration: Duration) { + let Some(executor) = self.managers.hook_executor() else { + return; + }; + let reason = self.session_end_reason(result.success); + 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); + } + + #[cfg(not(feature = "hooks"))] + fn notify_session_end_hook(&self, _result: &SessionResult, _duration: Duration) {} + /// Convert a [`Duration`] to milliseconds as `u64`. pub(super) fn millis_u64(duration: Duration) -> u64 { u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) diff --git a/src/engine/bare/message.rs b/src/engine/bare/message.rs index 1ce6053..d9e5187 100644 --- a/src/engine/bare/message.rs +++ b/src/engine/bare/message.rs @@ -1,8 +1,8 @@ -//! Message construction and extraction helpers. +//! Message construction and extraction. //! -//! Pure functions and `&self` methods that build or extract data from -//! [`Message`] instances. Extracted from [`BareLoop`] so the main loop -//! file focuses on orchestration rather than message wrangling. +//! Pure functions that build or extract data from [`Message`] instances — +//! assembling tool-result messages, extracting text or tool calls from an +//! assistant response, and computing token counts. use super::{ ApiClient, BareLoop, Message, MessagePart, Role, ToolCall, ToolContext, ToolDispatchResult, diff --git a/src/engine/bare/stream.rs b/src/engine/bare/stream.rs index 6ccf333..8883f3d 100644 --- a/src/engine/bare/stream.rs +++ b/src/engine/bare/stream.rs @@ -1,8 +1,8 @@ -//! Streaming phase — send conversation to the LLM API and accumulate the response. +//! Streaming — send the conversation to the LLM API and accumulate the response. //! -//! Extracted from [`BareLoop`] to isolate the streaming concern. When a -//! [`StreamHandler`] is configured, delegates to it for resilient streaming -//! (retry, timeout, fallback). Otherwise, uses basic inline logic. +//! When a [`StreamHandler`](crate::stream::handler::StreamHandler) is configured, +//! delegates to it for resilient streaming (retry, timeout, fallback). Otherwise, +//! uses basic inline logic. use super::{ ApiClient, BareLoop, LoopError, Message, StreamAccumulator, StreamEvent, StreamStopReason, @@ -40,8 +40,12 @@ impl BareLoop { /// /// # Errors /// - /// Returns [`LoopError::Api`] if any stream event is an error. - /// Returns [`LoopError::Cancelled`] if the cancellation signal fires mid-stream. + /// Returns [`LoopError::Api`] if any stream event is an error. When a + /// [`StreamHandler`](crate::stream::handler::StreamHandler) is configured, + /// may also return [`LoopError::Cancelled`] if the handler's cancel-aware + /// `select!` fires mid-stream. The inline path does not check cancellation + /// itself — that is handled by the `select!` in `process_turn`, which drops + /// this future if cancelled. pub(super) async fn stream_turn( &self, ) -> Result<(Message, Option, StreamStopReason), LoopError> { @@ -66,12 +70,7 @@ impl BareLoop { let mut accumulator = StreamAccumulator::new(); let mut stop_reason = StreamStopReason::EndTurn; loop { - let event_result = tokio::select! { - event = stream.next() => event, - () = self.cancelled.notified() => { - return Err(LoopError::Cancelled); - } - }; + let event_result = stream.next().await; match event_result { Some(Ok(event)) => {