diff --git a/crates/agen/examples/engine_cli.rs b/crates/agen/examples/engine_cli.rs index 2c50d354..10eb28ef 100644 --- a/crates/agen/examples/engine_cli.rs +++ b/crates/agen/examples/engine_cli.rs @@ -280,7 +280,7 @@ impl ToolResultPrinterPolicy { #[async_trait] impl Interceptor for ToolResultPrinterPolicy { - async fn post_tool_call(&self, info: &mut ToolResultInfo) -> InterceptorResult { + async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult { let name = self .call_names .lock() diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index ae3279b2..276ee3a4 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -15,8 +15,10 @@ use crate::{ }, handler::{ErrorKind, StatusKind, ToolUseBlockStart, UsageKind}, interceptor::{ - DefaultInterceptor, Interceptor, InterceptorFailure, InterceptorPoint, PostToolAction, - PreRequestAction, PreToolAction, PromptAction, ToolCallInfo, ToolResultInfo, TurnEndAction, + AssistantTurnEndContext, DefaultInterceptor, Interceptor, InterceptorFailure, + InterceptorPoint, PostToolAction, PreLlmRequestContext, PreRequestAction, PreToolAction, + PromptAction, PromptSubmitContext, RunExitContext, ToolCallInfo, ToolResultInfo, + TurnEndAction, }, llm_client::{ ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream, @@ -156,8 +158,6 @@ pub enum EngineRunExit { /// A typed reason why an engine run could not finish normally. #[derive(Debug)] pub enum RunInterruptionReason { - /// A trusted host interceptor callback failed at a typed lifecycle point. - Interceptor(InterceptorFailure), LimitReached, ContextWindowExceeded, Cancelled, @@ -178,9 +178,6 @@ impl From> for EngineRunExit { } Err(EngineError::Cancelled) => Self::Interrupted(RunInterruptionReason::Cancelled), Err(EngineError::PauseRequested) => Self::Paused, - Err(EngineError::Interceptor(failure)) => { - Self::Interrupted(RunInterruptionReason::Interceptor(failure)) - } Err(error) => Self::Interrupted(RunInterruptionReason::Unexpected(error)), } } @@ -438,11 +435,8 @@ impl Engine { self.active_run_turn_count.get_or_insert(0); } - fn finish_logical_run(&mut self, result: &Result) { - if !matches!( - result, - Ok(EngineResult::Paused | EngineResult::Yielded) | Err(EngineError::PauseRequested) - ) { + fn finish_logical_run(&mut self, exit: &EngineRunExit) { + if !matches!(exit, EngineRunExit::Paused | EngineRunExit::Yielded) { self.active_run_turn_count = None; } } @@ -1086,26 +1080,23 @@ impl Engine { request } - /// Hooks: on_prompt_submit - /// - async fn finalize_interruption( + async fn finalize_run_exit( &mut self, - result: Result, - ) -> Result { - match result { - Ok(value) => Ok(value), - Err(err) => { - let reason = match &err { - EngineError::Aborted(reason) => reason.clone(), - EngineError::Cancelled => "Cancelled".to_string(), - _ => err.to_string(), - }; - if let Err(error) = self.interceptor.on_abort(&reason).await { - return Err(InterceptorFailure::new(InterceptorPoint::Abort, error).into()); - } - Err(err) - } - } + result: Result, + ) -> EngineRunExit { + let exit = EngineRunExit::from(result); + let exit = match self + .interceptor + .on_run_exit(RunExitContext { exit: &exit }) + .await + { + Ok(()) => exit, + Err(error) => EngineRunExit::Interrupted(RunInterruptionReason::Unexpected( + InterceptorFailure::new(InterceptorPoint::RunExit, error).into(), + )), + }; + self.finish_logical_run(&exit); + exit } /// Check for pending tool calls (for resuming from Pause) @@ -1486,41 +1477,13 @@ impl Engine { } let call_info = call_info_map.get(&tool_result.tool_use_id); - let mut abort_reason = None; - if let Some((tool_call, meta, tool, context)) = call_info { - let mut info = ToolResultInfo { - call: tool_call.clone(), - result: tool_result, - meta: meta.clone(), - tool: tool.clone(), - context: context.clone(), - }; - - let post_tool_action = - self.interceptor - .post_tool_call(&mut info) - .await - .map_err(|error| { - EngineError::from(InterceptorFailure::new( - InterceptorPoint::PostToolCall, - error, - )) - })?; - match post_tool_action { - PostToolAction::Continue => {} - PostToolAction::Abort(reason) => { - abort_reason = Some(reason); - } - } - tool_result = info.result; - } if tool_result.is_error && tool_result.disposition.is_success() { tool_result.disposition = ToolResultDisposition::Error; } tool_result.is_error = !tool_result.disposition.is_success(); - // Cap content only after post_tool_call so interceptors still observe - // the full payload and any content they inject is bounded too. + // Bound the terminal payload before committing it so the post-tool + // interceptor observes exactly the model-visible durable result. if let (Some(limits), Some((tool_call, _, _, _)), Some(content)) = ( self.tool_output_limits.as_ref(), call_info, @@ -1573,9 +1536,30 @@ impl Engine { "Tool execution terminalized" ); self.emit_tool_result(&tool_result); - if let Some(reason) = abort_reason { - return Err(EngineError::Aborted(reason)); + + if let Some((tool_call, meta, tool, context)) = call_info { + let info = ToolResultInfo { + call: tool_call.clone(), + result: tool_result, + meta: meta.clone(), + tool: tool.clone(), + context: context.clone(), + }; + let post_tool_action = + self.interceptor + .post_tool_call(&info) + .await + .map_err(|error| { + EngineError::from(InterceptorFailure::new( + InterceptorPoint::PostToolCall, + error, + )) + })?; + if let PostToolAction::Abort(reason) = post_tool_action { + return Err(EngineError::Aborted(reason)); + } } + Ok(true) } @@ -1716,7 +1700,9 @@ impl Engine { // Interceptor: pre_llm_request let pre_request_action = self .interceptor - .pre_llm_request(&mut request_context) + .pre_llm_request(PreLlmRequestContext { + items: &mut request_context, + }) .await .map_err(|error| { EngineError::from(InterceptorFailure::new( @@ -1836,28 +1822,37 @@ impl Engine { let tool_calls = self.tool_call_collector.take_collected(); let assistant_items = self.build_assistant_items(&reasoning_items, &text_blocks, &tool_calls); + let committed_assistant_items = assistant_items.clone(); self.append_history_items(history, assistant_items, annotate)?; - if tool_calls.is_empty() { - let turn_end_context = history.items_cloned(); - let turn_end_action = self - .interceptor - .on_turn_end(&turn_end_context) - .await - .map_err(|error| { - EngineError::from(InterceptorFailure::new(InterceptorPoint::TurnEnd, error)) - })?; - match turn_end_action { - TurnEndAction::Finish => { - return Ok(EngineResult::Finished); - } - TurnEndAction::ContinueWithMessages(additional) => { - self.append_history_items(history, additional, annotate)?; + let assistant_turn_history = history.items_cloned(); + let assistant_turn_action = self + .interceptor + .on_assistant_turn_end(AssistantTurnEndContext { + assistant_items: &committed_assistant_items, + history: &assistant_turn_history, + tool_calls: &tool_calls, + }) + .await + .map_err(|error| { + EngineError::from(InterceptorFailure::new( + InterceptorPoint::AssistantTurnEnd, + error, + )) + })?; + match assistant_turn_action { + TurnEndAction::Finish if tool_calls.is_empty() => { + return Ok(EngineResult::Finished); + } + TurnEndAction::Finish => {} + TurnEndAction::ContinueWithMessages(additional) => { + self.append_history_items(history, additional, annotate)?; + if tool_calls.is_empty() { continue; } - TurnEndAction::Pause => { - return Ok(EngineResult::Paused); - } + } + TurnEndAction::Pause => { + return Ok(EngineResult::Paused); } } @@ -2540,9 +2535,10 @@ impl Engine { user_input: impl Into, annotate: &mut impl FnMut(&Item) -> Result, ) -> EngineRunExit { - self.run_result_with_annotation(history, user_input.into(), annotate) - .await - .into() + let result = self + .run_result_with_annotation(history, user_input.into(), annotate) + .await; + self.finalize_run_exit(result).await } async fn run_result_with_annotation( @@ -2554,19 +2550,20 @@ impl Engine { // Supplying new user input abandons any paused/yielded logical run. self.active_run_turn_count = None; let mut user_item = Item::user_message(user_input); - let prompt_action = match self.interceptor.on_prompt_submit(&mut user_item).await { - Ok(action) => action, - Err(error) => { - let error = InterceptorFailure::new(InterceptorPoint::PromptSubmit, error).into(); - return self.finalize_interruption(Err(error)).await; - } - }; + let prompt_action = self + .interceptor + .on_prompt_submit(PromptSubmitContext { + item: &mut user_item, + }) + .await + .map_err(|error| { + EngineError::from(InterceptorFailure::new( + InterceptorPoint::PromptSubmit, + error, + )) + })?; let extras = match prompt_action { - PromptAction::Cancel(reason) => { - return self - .finalize_interruption(Err(EngineError::Aborted(reason))) - .await; - } + PromptAction::Cancel(reason) => return Err(EngineError::Aborted(reason)), PromptAction::Continue => Vec::new(), PromptAction::ContinueWith(items) => items, }; @@ -2575,13 +2572,10 @@ impl Engine { self.append_history_items(history, extras, annotate)?; } self.start_logical_run(); - let result = match self.run_turn_loop(history, annotate).await { + match self.run_turn_loop(history, annotate).await { Err(EngineError::PauseRequested) => Ok(EngineResult::Paused), other => other, - }; - let result = self.finalize_interruption(result).await; - self.finish_logical_run(&result); - result + } } /// Resume execution (from Paused state). @@ -2590,9 +2584,8 @@ impl Engine { history: &mut History, annotate: &mut impl FnMut(&Item) -> Result, ) -> EngineRunExit { - self.resume_result_with_annotation(history, annotate) - .await - .into() + let result = self.resume_result_with_annotation(history, annotate).await; + self.finalize_run_exit(result).await } async fn resume_result_with_annotation( @@ -2601,13 +2594,10 @@ impl Engine { annotate: &mut impl FnMut(&Item) -> Result, ) -> Result { self.ensure_logical_run(); - let result = match self.run_turn_loop(history, annotate).await { + match self.run_turn_loop(history, annotate).await { Err(EngineError::PauseRequested) => Ok(EngineResult::Paused), other => other, - }; - let result = self.finalize_interruption(result).await; - self.finish_logical_run(&result); - result + } } /// Get the prefix length at lock time diff --git a/crates/agen/src/interceptor.rs b/crates/agen/src/interceptor.rs index ecb1fd8f..b85e0d9e 100644 --- a/crates/agen/src/interceptor.rs +++ b/crates/agen/src/interceptor.rs @@ -9,6 +9,7 @@ use std::sync::Arc; use async_trait::async_trait; use crate::Item; +use crate::engine::EngineRunExit; use crate::tool::{Tool, ToolCall, ToolExecutionContext, ToolMeta, ToolResult}; // ============================================================================= @@ -59,8 +60,8 @@ pub enum InterceptorPoint { PreLlmRequest, PreToolCall, PostToolCall, - TurnEnd, - Abort, + AssistantTurnEnd, + RunExit, } impl std::fmt::Display for InterceptorPoint { @@ -71,8 +72,8 @@ impl std::fmt::Display for InterceptorPoint { Self::PreLlmRequest => "pre_llm_request", Self::PreToolCall => "pre_tool_call", Self::PostToolCall => "post_tool_call", - Self::TurnEnd => "turn_end", - Self::Abort => "abort", + Self::AssistantTurnEnd => "assistant_turn_end", + Self::RunExit => "run_exit", }; formatter.write_str(name) } @@ -106,6 +107,35 @@ impl InterceptorFailure { /// Result returned by asynchronous interceptor lifecycle methods. pub type InterceptorResult = Result; +// ============================================================================= +// Lifecycle Contexts +// ============================================================================= + +/// Mutable prompt input presented before it is committed to Engine history. +pub struct PromptSubmitContext<'a> { + pub item: &'a mut Item, +} + +/// Mutable provider-visible item projection presented before an LLM request. +pub struct PreLlmRequestContext<'a> { + pub items: &'a mut Vec, +} + +/// A terminalized and committed assistant response at the next-phase boundary. +pub struct AssistantTurnEndContext<'a> { + /// The exact assistant items committed for this response. + pub assistant_items: &'a [Item], + /// The committed Engine history after the assistant items were appended. + pub history: &'a [Item], + /// Terminal tool calls collected from the response, if any. + pub tool_calls: &'a [ToolCall], +} + +/// The one terminal outcome produced by a public Engine run or resume call. +pub struct RunExitContext<'a> { + pub exit: &'a EngineRunExit, +} + // ============================================================================= // Action Enums // ============================================================================= @@ -181,9 +211,9 @@ pub enum PostToolAction { /// Action at the end of a turn (when LLM produces no tool calls). #[derive(Debug, Clone)] pub enum TurnEndAction { - /// Turn is finished, return to caller. + /// Accept the Engine's natural next phase: execute tools, or finish when none exist. Finish, - /// Continue with additional messages injected into history. + /// Commit additional messages, then continue through the natural next phase. ContinueWithMessages(Vec), /// Pause execution (can be resumed later). Pause, @@ -209,7 +239,7 @@ pub struct ToolCallInfo { pub struct ToolResultInfo { /// Original tool call. pub call: ToolCall, - /// Tool execution result (modifiable). + /// Committed terminal tool execution result. pub result: ToolResult, /// Tool meta information. pub meta: ToolMeta, @@ -236,7 +266,10 @@ pub struct ToolResultInfo { #[async_trait] pub trait Interceptor: Send + Sync { /// Called after receiving user input, before adding it to Engine history. - async fn on_prompt_submit(&self, _item: &mut Item) -> InterceptorResult { + async fn on_prompt_submit( + &self, + _context: PromptSubmitContext<'_>, + ) -> InterceptorResult { Ok(PromptAction::Continue) } @@ -272,7 +305,7 @@ pub trait Interceptor: Send + Sync { /// commits it to history before the request is sent. async fn pre_llm_request( &self, - _context: &mut Vec, + _context: PreLlmRequestContext<'_>, ) -> InterceptorResult { Ok(PreRequestAction::Continue) } @@ -282,24 +315,22 @@ pub trait Interceptor: Send + Sync { Ok(PreToolAction::Continue) } - /// Called after each tool reaches one terminal result. - async fn post_tool_call( - &self, - _info: &mut ToolResultInfo, - ) -> InterceptorResult { + /// Called after each tool reaches one terminal result and that result is committed. + async fn post_tool_call(&self, _info: &ToolResultInfo) -> InterceptorResult { Ok(PostToolAction::Continue) } - /// Called at the assistant boundary when a completed response has no tool calls. - /// - /// This is not the logical run termination observer. A host that needs that - /// boundary must inspect the returned [`crate::EngineRunExit`]. - async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult { + /// Called after every terminal assistant response is committed and before + /// the Engine decides whether to execute tools, continue, or finish. + async fn on_assistant_turn_end( + &self, + _context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { Ok(TurnEndAction::Finish) } - /// Called once when execution is interrupted (abort, cancellation, or failure). - async fn on_abort(&self, _reason: &str) -> InterceptorResult<()> { + /// Called once for the terminal outcome of each public run or resume call. + async fn on_run_exit(&self, _context: RunExitContext<'_>) -> InterceptorResult<()> { Ok(()) } } diff --git a/crates/agen/src/lib.rs b/crates/agen/src/lib.rs index 52a7a246..27bed76e 100644 --- a/crates/agen/src/lib.rs +++ b/crates/agen/src/lib.rs @@ -27,7 +27,8 @@ pub use engine::{ pub use handler::ToolUseBlockStart; pub use history::{History, HistoryEntry}; pub use interceptor::{ - Interceptor, InterceptorError, InterceptorFailure, InterceptorPoint, InterceptorResult, + AssistantTurnEndContext, Interceptor, InterceptorError, InterceptorFailure, InterceptorPoint, + InterceptorResult, PreLlmRequestContext, PromptSubmitContext, RunExitContext, }; pub use message::{ContentPart, Item, Message, Role}; pub use tool::{ diff --git a/crates/agen/tests/engine_state_test.rs b/crates/agen/tests/engine_state_test.rs index c4d25428..ab810e9e 100644 --- a/crates/agen/tests/engine_state_test.rs +++ b/crates/agen/tests/engine_state_test.rs @@ -10,10 +10,14 @@ use std::sync::{Arc, Mutex}; use agen::Item; use agen::interceptor::{ - Interceptor, InterceptorError, InterceptorPoint, InterceptorResult, PostToolAction, - PreRequestAction, PreToolAction, PromptAction, ToolCallInfo, ToolResultInfo, TurnEndAction, + AssistantTurnEndContext, Interceptor, InterceptorError, InterceptorPoint, InterceptorResult, + PostToolAction, PreLlmRequestContext, PreRequestAction, PreToolAction, PromptAction, + PromptSubmitContext, RunExitContext, ToolCallInfo, ToolResultInfo, TurnEndAction, +}; +use agen::llm_client::{ + ClientError, LlmClient, Request, ResponseStream, + event::{Event, ResponseStatus, StatusEvent}, }; -use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::{Engine, EngineError, EngineRunExit, History, RunInterruptionReason}; use async_trait::async_trait; @@ -616,7 +620,7 @@ struct YieldOnce { impl Interceptor for YieldOnce { async fn pre_llm_request( &self, - _context: &mut Vec, + _context: PreLlmRequestContext<'_>, ) -> InterceptorResult { Ok(if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { PreRequestAction::Yield @@ -647,7 +651,10 @@ struct ContinueTurnOnce { #[async_trait] impl Interceptor for ContinueTurnOnce { - async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult { + async fn on_assistant_turn_end( + &self, + _context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { Ok(if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { TurnEndAction::ContinueWithMessages(vec![Item::system_message("continue")]) } else { @@ -686,7 +693,10 @@ impl FailingLifecycleInterceptor { #[async_trait] impl Interceptor for FailingLifecycleInterceptor { - async fn on_prompt_submit(&self, _item: &mut Item) -> InterceptorResult { + async fn on_prompt_submit( + &self, + _context: PromptSubmitContext<'_>, + ) -> InterceptorResult { tokio::task::yield_now().await; self.record(InterceptorPoint::PromptSubmit, PromptAction::Continue) } @@ -698,20 +708,10 @@ impl Interceptor for FailingLifecycleInterceptor { async fn pre_llm_request( &self, - _context: &mut Vec, + _context: PreLlmRequestContext<'_>, ) -> InterceptorResult { tokio::task::yield_now().await; - if self.failure == InterceptorPoint::Abort { - self.calls - .lock() - .unwrap() - .push(InterceptorPoint::PreLlmRequest); - Ok(PreRequestAction::Cancel( - "trigger abort callback".to_string(), - )) - } else { - self.record(InterceptorPoint::PreLlmRequest, PreRequestAction::Continue) - } + self.record(InterceptorPoint::PreLlmRequest, PreRequestAction::Continue) } async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> InterceptorResult { @@ -719,76 +719,83 @@ impl Interceptor for FailingLifecycleInterceptor { self.record(InterceptorPoint::PreToolCall, PreToolAction::Continue) } - async fn post_tool_call( - &self, - _info: &mut ToolResultInfo, - ) -> InterceptorResult { + async fn post_tool_call(&self, _info: &ToolResultInfo) -> InterceptorResult { tokio::task::yield_now().await; self.record(InterceptorPoint::PostToolCall, PostToolAction::Continue) } - async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult { + async fn on_assistant_turn_end( + &self, + context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { tokio::task::yield_now().await; - self.record(InterceptorPoint::TurnEnd, TurnEndAction::Finish) + assert!(context.history.ends_with(context.assistant_items)); + if !context.tool_calls.is_empty() { + assert_eq!( + context + .assistant_items + .iter() + .filter(|item| matches!(item, Item::ToolCall { .. })) + .count(), + context.tool_calls.len() + ); + } + self.record(InterceptorPoint::AssistantTurnEnd, TurnEndAction::Finish) } - async fn on_abort(&self, _reason: &str) -> InterceptorResult<()> { + async fn on_run_exit(&self, _context: RunExitContext<'_>) -> InterceptorResult<()> { tokio::task::yield_now().await; - self.record(InterceptorPoint::Abort, ()) + self.record(InterceptorPoint::RunExit, ()) } } fn expected_interceptor_calls(failure: InterceptorPoint) -> Vec { use InterceptorPoint as Point; - match failure { - Point::PromptSubmit => vec![Point::PromptSubmit, Point::Abort], + let mut calls = match failure { + Point::PromptSubmit => vec![Point::PromptSubmit], Point::PendingHistoryAppends => { - vec![ - Point::PromptSubmit, - Point::PendingHistoryAppends, - Point::Abort, - ] + vec![Point::PromptSubmit, Point::PendingHistoryAppends] } Point::PreLlmRequest => vec![ Point::PromptSubmit, Point::PendingHistoryAppends, Point::PreLlmRequest, - Point::Abort, ], Point::PreToolCall => vec![ Point::PromptSubmit, Point::PendingHistoryAppends, Point::PreLlmRequest, + Point::AssistantTurnEnd, Point::PreToolCall, - Point::Abort, ], Point::PostToolCall => vec![ Point::PromptSubmit, Point::PendingHistoryAppends, Point::PreLlmRequest, + Point::AssistantTurnEnd, Point::PreToolCall, Point::PostToolCall, - Point::Abort, ], - Point::TurnEnd => vec![ + Point::AssistantTurnEnd => vec![ Point::PromptSubmit, Point::PendingHistoryAppends, Point::PreLlmRequest, - Point::TurnEnd, - Point::Abort, + Point::AssistantTurnEnd, ], - Point::Abort => vec![ + Point::RunExit => vec![ Point::PromptSubmit, Point::PendingHistoryAppends, Point::PreLlmRequest, - Point::Abort, + Point::AssistantTurnEnd, ], - } + }; + calls.push(Point::RunExit); + calls } #[tokio::test] -async fn interceptor_failures_are_typed_and_each_lifecycle_point_runs_once() { +async fn interceptor_failures_are_typed_unexpected_and_each_lifecycle_point_runs_once() { use InterceptorPoint as Point; for failure_point in [ @@ -797,8 +804,8 @@ async fn interceptor_failures_are_typed_and_each_lifecycle_point_runs_once() { Point::PreLlmRequest, Point::PreToolCall, Point::PostToolCall, - Point::TurnEnd, - Point::Abort, + Point::AssistantTurnEnd, + Point::RunExit, ] { let interceptor = FailingLifecycleInterceptor::new(failure_point); let needs_tool = matches!(failure_point, Point::PreToolCall | Point::PostToolCall); @@ -821,7 +828,10 @@ async fn interceptor_failures_are_typed_and_each_lifecycle_point_runs_once() { let mut engine = engine.lock(&history); let exit = engine.run(&mut history, "test").await; - let EngineRunExit::Interrupted(RunInterruptionReason::Interceptor(failure)) = exit else { + let EngineRunExit::Interrupted(RunInterruptionReason::Unexpected( + EngineError::Interceptor(failure), + )) = exit + else { panic!("expected typed interceptor interruption at {failure_point}, got {exit:?}"); }; assert_eq!(failure.point(), failure_point); @@ -833,9 +843,220 @@ async fn interceptor_failures_are_typed_and_each_lifecycle_point_runs_once() { interceptor.calls(), expected_interceptor_calls(failure_point) ); + if failure_point == Point::PostToolCall { + assert!( + history + .items() + .any(|item| matches!(item, Item::ToolResult { .. })), + "post-tool failure must not precede terminal output commit" + ); + } } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TerminalMode { + Finish, + PauseOnce, + Yield, +} + +#[derive(Debug, Clone)] +struct RecordingTerminalInterceptor { + mode: TerminalMode, + assistant_turns: Arc, + exits: Arc>>, +} + +impl RecordingTerminalInterceptor { + fn new(mode: TerminalMode) -> Self { + Self { + mode, + assistant_turns: Arc::new(AtomicUsize::new(0)), + exits: Arc::new(Mutex::new(Vec::new())), + } + } + + fn exits(&self) -> Vec<&'static str> { + self.exits.lock().unwrap().clone() + } +} + +#[async_trait] +impl Interceptor for RecordingTerminalInterceptor { + async fn pre_llm_request( + &self, + _context: PreLlmRequestContext<'_>, + ) -> InterceptorResult { + Ok(if self.mode == TerminalMode::Yield { + PreRequestAction::Yield + } else { + PreRequestAction::Continue + }) + } + + async fn on_assistant_turn_end( + &self, + context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { + assert!(!context.assistant_items.is_empty()); + assert!( + context.history.ends_with(context.assistant_items), + "assistant-turn callback must observe committed terminal items" + ); + let turn = self.assistant_turns.fetch_add(1, Ordering::SeqCst); + Ok(if self.mode == TerminalMode::PauseOnce && turn == 0 { + TurnEndAction::Pause + } else { + TurnEndAction::Finish + }) + } + + async fn on_run_exit(&self, context: RunExitContext<'_>) -> InterceptorResult<()> { + let kind = match context.exit { + EngineRunExit::Finished => "finished", + EngineRunExit::Paused => "paused", + EngineRunExit::Yielded => "yielded", + EngineRunExit::Interrupted(RunInterruptionReason::LimitReached) => "limit", + EngineRunExit::Interrupted(RunInterruptionReason::ContextWindowExceeded) => "context", + EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) => "cancelled", + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(_)) => "unexpected", + }; + self.exits.lock().unwrap().push(kind); + Ok(()) + } +} + +#[derive(Clone)] +struct ContextWindowClient; + +#[async_trait] +impl LlmClient for ContextWindowClient { + async fn stream(&self, _request: Request) -> Result { + Err(ClientError::ContextWindowExceeded) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +#[tokio::test] +async fn terminal_observer_runs_once_for_every_exit_and_interruption_kind() { + let finished = RecordingTerminalInterceptor::new(TerminalMode::Finish); + let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); + engine.set_interceptor(finished.clone()); + let mut history = History::new(); + assert!(matches!( + engine.lock(&history).run(&mut history, "finish").await, + EngineRunExit::Finished + )); + assert_eq!(finished.exits(), ["finished"]); + + let yielded = RecordingTerminalInterceptor::new(TerminalMode::Yield); + let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); + engine.set_interceptor(yielded.clone()); + let mut history = History::new(); + assert!(matches!( + engine.lock(&history).run(&mut history, "yield").await, + EngineRunExit::Yielded + )); + assert_eq!(yielded.exits(), ["yielded"]); + + let limited = RecordingTerminalInterceptor::new(TerminalMode::Finish); + let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); + engine.set_max_turns(Some(0)); + engine.set_interceptor(limited.clone()); + let mut history = History::new(); + assert!(matches!( + engine.lock(&history).run(&mut history, "limit").await, + EngineRunExit::Interrupted(RunInterruptionReason::LimitReached) + )); + assert_eq!(limited.exits(), ["limit"]); + + let cancelled = RecordingTerminalInterceptor::new(TerminalMode::Finish); + let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); + engine.set_interceptor(cancelled.clone()); + engine.cancel(); + let mut history = History::new(); + assert!(matches!( + engine.lock(&history).run(&mut history, "cancel").await, + EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) + )); + assert_eq!(cancelled.exits(), ["cancelled"]); + + let context = RecordingTerminalInterceptor::new(TerminalMode::Finish); + let mut engine = Engine::new(ContextWindowClient); + engine.set_interceptor(context.clone()); + let mut history = History::new(); + assert!(matches!( + engine.lock(&history).run(&mut history, "context").await, + EngineRunExit::Interrupted(RunInterruptionReason::ContextWindowExceeded) + )); + assert_eq!(context.exits(), ["context"]); + + let unexpected = FailingLifecycleInterceptor::new(InterceptorPoint::PromptSubmit); + let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); + engine.set_interceptor(unexpected.clone()); + let mut history = History::new(); + assert!(matches!( + engine.lock(&history).run(&mut history, "fail").await, + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(EngineError::Interceptor( + _ + ))) + )); + assert_eq!( + unexpected + .calls() + .iter() + .filter(|point| **point == InterceptorPoint::RunExit) + .count(), + 1 + ); +} + +#[tokio::test] +async fn terminal_observer_does_not_duplicate_on_resume() { + let interceptor = RecordingTerminalInterceptor::new(TerminalMode::PauseOnce); + let first_response = vec![ + Event::tool_use_start(0, "call-1", "count_tool"), + Event::tool_input_delta(0, "{}"), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]; + let client = MockLlmClient::with_responses(vec![first_response, completed_text_events()]); + let tool = CountingTool::new("count_tool"); + let mut engine = Engine::new(client); + engine.register_tool(tool.definition()); + engine.set_interceptor(interceptor.clone()); + let mut history = History::new(); + let mut engine = engine.lock(&history); + + assert!(matches!( + engine.run(&mut history, "pause").await, + EngineRunExit::Paused + )); + assert_eq!(interceptor.exits(), ["paused"]); + assert_eq!( + tool.call_count(), + 0, + "pause must retain the pending tool phase" + ); + + assert!(matches!( + engine.resume(&mut history).await, + EngineRunExit::Finished + )); + assert_eq!(interceptor.exits(), ["paused", "finished"]); + assert_eq!( + tool.call_count(), + 1, + "resume must execute the retained tool once" + ); +} + #[tokio::test] async fn max_turns_is_scoped_to_each_fresh_run() { let mut history: History = History::new(); diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index ead143b4..f63b20a1 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -919,10 +919,7 @@ async fn test_tool_execution_context_for_skipped_and_synthetic_paths() { }) } - async fn post_tool_call( - &self, - info: &mut ToolResultInfo, - ) -> InterceptorResult { + async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult { self.post_contexts .lock() .unwrap() @@ -1026,9 +1023,9 @@ async fn test_before_tool_call_skip() { ); } -/// Hook: post_tool_call - verify that results can be modified +/// Hook: post_tool_call - verify that the committed terminal result is observed. #[tokio::test] -async fn test_post_tool_call_modification() { +async fn test_post_tool_call_observes_committed_result() { // Prepare responses for multiple requests let client = MockLlmClient::with_responses(vec![ // First request: tool call @@ -1079,43 +1076,39 @@ async fn test_post_tool_call_modification() { engine.register_tool(simple_tool_definition()); - // Policy to modify results - struct ModifyingPolicy { - modified_content: Arc>>, + // Policy to observe the committed terminal result. + struct ObservingPolicy { + observed_content: Arc>>, } #[async_trait] - impl Interceptor for ModifyingPolicy { - async fn post_tool_call( - &self, - info: &mut ToolResultInfo, - ) -> InterceptorResult { - info.result.summary = format!("[Modified] {}", info.result.summary); - *self.modified_content.lock().unwrap() = Some(info.result.summary.clone()); + impl Interceptor for ObservingPolicy { + async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult { + *self.observed_content.lock().unwrap() = Some(info.result.summary.clone()); Ok(PostToolAction::Continue) } } - let modified_content = Arc::new(std::sync::Mutex::new(None)); - engine.set_interceptor(ModifyingPolicy { - modified_content: modified_content.clone(), + let observed_content = Arc::new(std::sync::Mutex::new(None)); + engine.set_interceptor(ObservingPolicy { + observed_content: observed_content.clone(), }); // Mutable::run consumes self, returns (Locked, EngineResult) - let result = engine.run(&mut history, "Test modification").await; + let result = engine.run(&mut history, "Test observation").await; assert!( matches!(result.result, agen::EngineRunExit::Finished), "Engine should complete" ); - // Verify hook was called and content was modified - let content = modified_content.lock().unwrap().clone(); - assert!(content.is_some(), "Hook should have been called"); - assert!( - content.unwrap().contains("[Modified]"), - "Result should be modified" - ); + // Verify the interceptor observed the exact committed result. + let observed = observed_content.lock().unwrap().clone(); + assert_eq!(observed.as_deref(), Some("Original Result")); + assert!(history.items().any(|item| matches!( + item, + Item::ToolResult { summary, .. } if summary == "Original Result" + ))); } /// Hook: pre_tool_call synthetic result - skipped tool gets an error result in history. @@ -1189,19 +1182,24 @@ async fn post_tool_abort_commits_confirmed_result_before_stopping_run() { let tool = SlowTool::new("confirmed", 1); engine.register_tool(tool.definition()); - struct AbortAfterResult; + let observed = Arc::new(Mutex::new(Vec::<&'static str>::new())); + struct AbortAfterResult { + lifecycle: Arc>>, + } #[async_trait] impl Interceptor for AbortAfterResult { async fn post_tool_call( &self, - _info: &mut ToolResultInfo, + _info: &ToolResultInfo, ) -> InterceptorResult { + self.lifecycle.lock().unwrap().push("post_tool_call"); Ok(PostToolAction::Abort("policy stopped the run".to_string())) } } - engine.set_interceptor(AbortAfterResult); + engine.set_interceptor(AbortAfterResult { + lifecycle: observed.clone(), + }); - let observed = Arc::new(Mutex::new(Vec::<&'static str>::new())); let published = observed.clone(); engine.on_tool_result(move |_| published.lock().unwrap().push("published")); let committed = observed.clone(); @@ -1221,7 +1219,7 @@ async fn post_tool_abort_commits_confirmed_result_before_stopping_run() { assert_eq!(tool.call_count(), 1); assert_eq!( observed.lock().unwrap().as_slice(), - ["committed", "published", "run-returned"] + ["committed", "published", "post_tool_call", "run-returned"] ); assert!(matches!( output.result, diff --git a/crates/session-store/tests/session_test.rs b/crates/session-store/tests/session_test.rs index 2f0ffbeb..dced4ca9 100644 --- a/crates/session-store/tests/session_test.rs +++ b/crates/session-store/tests/session_test.rs @@ -3,7 +3,7 @@ mod common; use std::ops::{Deref, DerefMut}; use std::sync::Arc; -use agen::interceptor::{Interceptor, InterceptorResult, TurnEndAction}; +use agen::interceptor::{AssistantTurnEndContext, Interceptor, InterceptorResult, TurnEndAction}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::types::{Item, RequestConfig}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; @@ -100,7 +100,10 @@ struct PausePolicy; #[async_trait] impl Interceptor for PausePolicy { - async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult { + async fn on_assistant_turn_end( + &self, + _context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { Ok(TurnEndAction::Pause) } } @@ -350,7 +353,8 @@ async fn session_run_with_tool_call() { async fn session_resume_after_pause() { let (_dir, store) = make_store(); - // First run: tool call with pause policy → Paused + // First terminal assistant response requests a tool; the assistant-turn + // interceptor pauses before the Engine enters the tool phase. let client = MockLlmClient::with_responses(tool_call_events()); let mut worker = TestWorker::new(Engine::new(client)); worker.register_tool(weather_tool_definition()); @@ -386,7 +390,7 @@ async fn session_resume_after_pause() { // Restore state and verify let state = session_store::restore(&store, sid, segid).unwrap(); assert!(state.last_run_interrupted); - assert_eq!(state.active_run_turn_count, Some(2)); + assert_eq!(state.active_run_turn_count, Some(1)); } #[tokio::test] diff --git a/crates/worker/src/compact/worker.rs b/crates/worker/src/compact/worker.rs index cd218ac0..158c68ff 100644 --- a/crates/worker/src/compact/worker.rs +++ b/crates/worker/src/compact/worker.rs @@ -23,7 +23,8 @@ use std::sync::{Arc, Mutex}; use agen::Item; use agen::interceptor::{ - Interceptor, InterceptorResult, PreRequestAction, PreToolAction, ToolCallInfo, + Interceptor, InterceptorResult, PreLlmRequestContext, PreRequestAction, PreToolAction, + ToolCallInfo, }; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput, ToolResult}; use async_trait::async_trait; @@ -402,8 +403,9 @@ impl CompactWorkerInterceptor { impl Interceptor for CompactWorkerInterceptor { async fn pre_llm_request( &self, - context: &mut Vec, + context: PreLlmRequestContext<'_>, ) -> InterceptorResult { + let context = context.items; let records = self.usage_tracker.records(); let estimate = agen::token_counter::total_tokens(context, &records); if estimate.tokens > self.max_input_tokens { @@ -472,13 +474,23 @@ mod tests { let mut context = vec![Item::user_message("hello")]; assert!(matches!( - interceptor.pre_llm_request(&mut context).await.unwrap(), + interceptor + .pre_llm_request(PreLlmRequestContext { + items: &mut context, + }) + .await + .unwrap(), PreRequestAction::Continue )); tracker.record_usage(&make_usage(100)); assert!(matches!( - interceptor.pre_llm_request(&mut context).await.unwrap(), + interceptor + .pre_llm_request(PreLlmRequestContext { + items: &mut context, + }) + .await + .unwrap(), PreRequestAction::Continue )); tracker.record_usage(&make_usage(100)); @@ -486,7 +498,12 @@ mod tests { // Two 100-token requests would exceed a cumulative 150-token cap, but // current occupancy is still the latest 100-token measurement. assert!(matches!( - interceptor.pre_llm_request(&mut context).await.unwrap(), + interceptor + .pre_llm_request(PreLlmRequestContext { + items: &mut context, + }) + .await + .unwrap(), PreRequestAction::Continue )); } @@ -508,13 +525,23 @@ mod tests { let mut context = vec![Item::user_message("hello")]; assert!(matches!( - interceptor.pre_llm_request(&mut context).await.unwrap(), + interceptor + .pre_llm_request(PreLlmRequestContext { + items: &mut context, + }) + .await + .unwrap(), PreRequestAction::Continue )); tracker.record_usage(&make_usage(100)); assert!(matches!( - interceptor.pre_llm_request(&mut context).await.unwrap(), + interceptor + .pre_llm_request(PreLlmRequestContext { + items: &mut context, + }) + .await + .unwrap(), PreRequestAction::ContinueWith(items) if items.len() == 1 && items[0].as_text().unwrap_or_default().contains("write_summary") )); @@ -528,13 +555,23 @@ mod tests { let mut context = vec![Item::user_message("hello")]; assert!(matches!( - interceptor.pre_llm_request(&mut context).await.unwrap(), + interceptor + .pre_llm_request(PreLlmRequestContext { + items: &mut context, + }) + .await + .unwrap(), PreRequestAction::Continue )); tracker.record_usage(&make_usage(100)); assert!(matches!( - interceptor.pre_llm_request(&mut context).await.unwrap(), + interceptor + .pre_llm_request(PreLlmRequestContext { + items: &mut context, + }) + .await + .unwrap(), PreRequestAction::Cancel(message) if message.contains("occupancy") )); } diff --git a/crates/worker/src/hook.rs b/crates/worker/src/hook.rs index d34ec5cc..4318f9c3 100644 --- a/crates/worker/src/hook.rs +++ b/crates/worker/src/hook.rs @@ -285,12 +285,6 @@ pub struct TurnEndInfo { pub final_text_preview: String, } -/// Information passed to `OnAbort` hooks. -pub struct AbortInfo { - /// Reason supplied by the aborter. - pub reason: String, -} - // ============================================================================= // Hook Event Kinds // ============================================================================= @@ -315,10 +309,8 @@ pub struct PreLlmRequest; pub struct PreToolCall; /// After each tool completes; observational except it may abort the run. pub struct PostToolCall; -/// When a turn ends with no tool calls; observational except it may pause. +/// After every terminal assistant response is committed; observational except it may pause. pub struct OnTurnEnd; -/// When execution is interrupted; observational only. -pub struct OnAbort; impl HookEventKind for OnPromptSubmit { type Input = PromptSubmitInfo; @@ -345,11 +337,6 @@ impl HookEventKind for OnTurnEnd { type Output = HookTurnEndAction; } -impl HookEventKind for OnAbort { - type Input = AbortInfo; - type Output = (); -} - // ============================================================================= // Hook Trait // ============================================================================= @@ -380,7 +367,6 @@ pub struct HookRegistryBuilder { pre_tool_call: Vec>>, post_tool_call: Vec>>, on_turn_end: Vec>>, - on_abort: Vec>>, } impl HookRegistryBuilder { @@ -408,10 +394,6 @@ impl HookRegistryBuilder { self.on_turn_end.push(Box::new(hook)); } - pub fn add_on_abort(&mut self, hook: impl Hook + 'static) { - self.on_abort.push(Box::new(hook)); - } - /// Freeze the builder into an immutable registry. pub fn build(self) -> HookRegistry { HookRegistry { @@ -420,7 +402,6 @@ impl HookRegistryBuilder { pre_tool_call: self.pre_tool_call, post_tool_call: self.post_tool_call, on_turn_end: self.on_turn_end, - on_abort: self.on_abort, } } } @@ -432,7 +413,6 @@ pub struct HookRegistry { pub(crate) pre_tool_call: Vec>>, pub(crate) post_tool_call: Vec>>, pub(crate) on_turn_end: Vec>>, - pub(crate) on_abort: Vec>>, } #[cfg(test)] diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index a1468443..8d380ea7 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -15,8 +15,9 @@ use std::sync::{Arc, Mutex}; use agen::Item; use agen::UsageRecord; use agen::interceptor::{ - Interceptor, InterceptorResult, PostToolAction, PreRequestAction, PreToolAction, PromptAction, - ToolCallInfo, ToolResultInfo, TurnEndAction, + AssistantTurnEndContext, Interceptor, InterceptorResult, PostToolAction, PreLlmRequestContext, + PreRequestAction, PreToolAction, PromptAction, PromptSubmitContext, ToolCallInfo, + ToolResultInfo, TurnEndAction, }; use agen::tool::ToolOutput; use arc_swap::ArcSwap; @@ -28,9 +29,9 @@ use crate::compact::usage_tracker::UsageTracker; use session_store::SystemItem; use crate::hook::{ - AbortInfo, HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookPromptAction, - HookRegistry, HookTurnEndAction, PreRequestContext, PreRequestInfo, PromptSubmitInfo, - SystemItemAppendHandle, ToolCallSummary, ToolResultSummary, TurnEndInfo, + HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookPromptAction, HookRegistry, + HookTurnEndAction, PreRequestContext, PreRequestInfo, PromptSubmitInfo, SystemItemAppendHandle, + ToolCallSummary, ToolResultSummary, TurnEndInfo, }; use crate::ipc::notify_buffer::{NotifyBuffer, build_system_item_with_provenance}; use crate::prompt::catalog::PromptCatalog; @@ -232,7 +233,11 @@ impl WorkerInterceptor { #[async_trait] impl Interceptor for WorkerInterceptor { - async fn on_prompt_submit(&self, item: &mut Item) -> InterceptorResult { + async fn on_prompt_submit( + &self, + context: PromptSubmitContext<'_>, + ) -> InterceptorResult { + let item = context.item; let turn_index = self.next_turn_index.fetch_add(1, Ordering::Relaxed); self.tool_calls_this_turn.store(0, Ordering::Relaxed); @@ -310,8 +315,9 @@ impl Interceptor for WorkerInterceptor { async fn pre_llm_request( &self, - context: &mut Vec, + context: PreLlmRequestContext<'_>, ) -> InterceptorResult { + let context = context.items; let initial_tokens = self.estimated_tokens(context); if self.request_threshold_exceeded(initial_tokens, context) { return Ok(PreRequestAction::Yield); @@ -395,7 +401,7 @@ impl Interceptor for WorkerInterceptor { Ok(PreToolAction::Continue) } - async fn post_tool_call(&self, info: &mut ToolResultInfo) -> InterceptorResult { + async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult { let summary = ToolResultSummary { call_id: info.result.tool_use_id.clone(), tool_name: info.call.name.clone(), @@ -416,7 +422,11 @@ impl Interceptor for WorkerInterceptor { Ok(PostToolAction::Continue) } - async fn on_turn_end(&self, history: &[Item]) -> InterceptorResult { + async fn on_assistant_turn_end( + &self, + context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { + let history = context.history; let final_text_preview = history .iter() .rev() @@ -437,16 +447,6 @@ impl Interceptor for WorkerInterceptor { } Ok(TurnEndAction::Finish) } - - async fn on_abort(&self, reason: &str) -> InterceptorResult<()> { - let info = AbortInfo { - reason: reason.to_string(), - }; - for hook in &self.registry.on_abort { - hook.call(&info).await; - } - Ok(()) - } } struct ContextShape { @@ -629,7 +629,10 @@ mod tests { None, ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Yield)); // Hook must not run when an internal mechanism short-circuits first. @@ -661,7 +664,10 @@ mod tests { })), ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); match action { PreRequestAction::YieldWith(items) => assert_eq!(items.len(), 1), @@ -698,7 +704,10 @@ mod tests { ) .with_usage_tracker(usage_tracker); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Yield)); } @@ -722,7 +731,10 @@ mod tests { None, ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -763,7 +775,10 @@ mod tests { None, ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -790,7 +805,10 @@ mod tests { None, ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -811,7 +829,10 @@ mod tests { None, ); let mut ctx: Vec = Vec::new(); - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -840,7 +861,10 @@ mod tests { ); let mut ctx: Vec = Vec::new(); - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(saw_handle.load(Ordering::Relaxed)); let PreRequestAction::ContinueWith(items) = action else { @@ -887,7 +911,10 @@ mod tests { ); let mut ctx: Vec = Vec::new(); - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(!saw_handle.load(Ordering::Relaxed)); assert!(matches!(action, PreRequestAction::Continue)); @@ -1042,7 +1069,14 @@ mod tests { ); let history = vec![Item::user_message("hi"), Item::assistant_message("done")]; - let action = interceptor.on_turn_end(&history).await.unwrap(); + let action = interceptor + .on_assistant_turn_end(AssistantTurnEndContext { + assistant_items: &[], + history: &history, + tool_calls: &[], + }) + .await + .unwrap(); assert!(matches!(action, TurnEndAction::Pause)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -1079,7 +1113,10 @@ mod tests { let ctx_items = vec![Item::user_message("hi")]; for _ in 0..23 { let mut ctx = ctx_items.clone(); - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Continue)); usage_tracker.record_usage(&agen::event::UsageEvent { input_tokens: Some(10), @@ -1091,7 +1128,10 @@ mod tests { } let mut ctx = ctx_items.clone(); - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); let appended_len = match action { PreRequestAction::ContinueWith(items) => items.len(), other => panic!("expected reminder append, got {other:?}"), @@ -1275,7 +1315,10 @@ mod tests { None, ); let mut ctx: Vec = vec![Item::user_message("hi")]; - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(ctx.len(), 1, "pre_llm_request must not append notifies"); @@ -1305,7 +1348,10 @@ mod tests { None, ); let mut ctx: Vec = Vec::new(); - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Cancel(_))); assert!(first_called.load(Ordering::Relaxed)); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index f98a9b5a..332ab2e5 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -50,8 +50,7 @@ use crate::feature::{ FeatureRegistryInstallReport, dedupe_instruction_contributions, }; use crate::hook::{ - Hook, HookRegistryBuilder, OnAbort, OnPromptSubmit, OnTurnEnd, PostToolCall, PreLlmRequest, - PreToolCall, + Hook, HookRegistryBuilder, OnPromptSubmit, OnTurnEnd, PostToolCall, PreLlmRequest, PreToolCall, }; use crate::in_flight::InFlightEvents; use crate::internal_worker::{ @@ -2328,12 +2327,6 @@ impl Worker { self.hook_builder.add_on_turn_end(hook); } - /// Register a hook that runs when execution is aborted. - pub fn add_on_abort_hook(&mut self, hook: impl Hook + 'static) { - self.assert_hooks_open(); - self.hook_builder.add_on_abort(hook); - } - /// Install the hook-based interceptor on the Engine if not already done. /// /// When either compaction threshold (`threshold` or @@ -3733,7 +3726,6 @@ impl Worker { | EngineRunExit::Yielded | EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) | EngineRunExit::Interrupted(RunInterruptionReason::ContextWindowExceeded) - | EngineRunExit::Interrupted(RunInterruptionReason::Interceptor(_)) | EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(_)) ); let active_run_turn_count = self.engine.as_ref().unwrap().active_run_turn_count(); @@ -6130,7 +6122,6 @@ fn run_interruption_reason_error_code(reason: &RunInterruptionReason) -> ErrorCo RunInterruptionReason::Unexpected(EngineError::Tool(_)) => ErrorCode::ToolError, RunInterruptionReason::LimitReached | RunInterruptionReason::Cancelled - | RunInterruptionReason::Interceptor(_) | RunInterruptionReason::Unexpected( EngineError::Aborted(_) | EngineError::Cancelled @@ -6148,7 +6139,6 @@ fn run_interruption_reason_message(reason: &RunInterruptionReason) -> String { RunInterruptionReason::LimitReached => "engine turn limit reached".to_string(), RunInterruptionReason::ContextWindowExceeded => "model context window reached".to_string(), RunInterruptionReason::Cancelled => "engine run cancelled".to_string(), - RunInterruptionReason::Interceptor(failure) => failure.to_string(), RunInterruptionReason::Unexpected(error) => format!("unexpected engine failure: {error}"), } }