Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

refactor: finalize interceptor lifecycle boundaries

Browse Source
This commit is contained in:
Hare
2026-09-03 15:33:15 +09:00
parent 68b1aa64e9
commit 0245980ea5
11 changed files with 587 additions and 289 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/examples/engine_cli.rs
+1 -1
View File
@@ -280,7 +280,7 @@ impl ToolResultPrinterPolicy {
#[async_trait]
impl Interceptor for ToolResultPrinterPolicy {
async fn post_tool_call(&self, info: &mut ToolResultInfo) -> InterceptorResult<PostToolAction> {
async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult<PostToolAction> {
let name = self
.call_names
.lock()
crates/agen/src/engine.rs
+100 -110
View File
@@ -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<Result<EngineResult, EngineError>> 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<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
self.active_run_turn_count.get_or_insert(0);
}
fn finish_logical_run(&mut self, result: &Result<EngineResult, EngineError>) {
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<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
request
}
/// Hooks: on_prompt_submit
///
async fn finalize_interruption<T>(
async fn finalize_run_exit(
&mut self,
result: Result<T, EngineError>,
) -> Result<T, EngineError> {
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<EngineResult, EngineError>,
) -> 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<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
}
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<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
"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<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
// 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<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
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<C: LlmClient, A> Engine<C, Locked, A> {
user_input: impl Into<String>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
) -> 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<C: LlmClient, A> Engine<C, Locked, A> {
// 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<C: LlmClient, A> Engine<C, Locked, A> {
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<C: LlmClient, A> Engine<C, Locked, A> {
history: &mut History<A>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
) -> 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<C: LlmClient, A> Engine<C, Locked, A> {
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
) -> Result<EngineResult, EngineError> {
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
crates/agen/src/interceptor.rs
+52 -21
View File
@@ -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<T> = Result<T, InterceptorError>;
// =============================================================================
// 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<Item>,
}
/// 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<Item>),
/// 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<PromptAction> {
async fn on_prompt_submit(
&self,
_context: PromptSubmitContext<'_>,
) -> InterceptorResult<PromptAction> {
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<Item>,
_context: PreLlmRequestContext<'_>,
) -> InterceptorResult<PreRequestAction> {
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<PostToolAction> {
/// Called after each tool reaches one terminal result and that result is committed.
async fn post_tool_call(&self, _info: &ToolResultInfo) -> InterceptorResult<PostToolAction> {
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<TurnEndAction> {
/// 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<TurnEndAction> {
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(())
}
}
crates/agen/src/lib.rs
+2 -1
View File
@@ -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::{
crates/agen/tests/engine_state_test.rs
+267 -46
View File
@@ -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<Item>,
_context: PreLlmRequestContext<'_>,
) -> InterceptorResult<PreRequestAction> {
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<TurnEndAction> {
async fn on_assistant_turn_end(
&self,
_context: AssistantTurnEndContext<'_>,
) -> InterceptorResult<TurnEndAction> {
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<PromptAction> {
async fn on_prompt_submit(
&self,
_context: PromptSubmitContext<'_>,
) -> InterceptorResult<PromptAction> {
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<Item>,
_context: PreLlmRequestContext<'_>,
) -> InterceptorResult<PreRequestAction> {
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<PreToolAction> {
@@ -719,76 +719,83 @@ impl Interceptor for FailingLifecycleInterceptor {
self.record(InterceptorPoint::PreToolCall, PreToolAction::Continue)
}
async fn post_tool_call(
&self,
_info: &mut ToolResultInfo,
) -> InterceptorResult<PostToolAction> {
async fn post_tool_call(&self, _info: &ToolResultInfo) -> InterceptorResult<PostToolAction> {
tokio::task::yield_now().await;
self.record(InterceptorPoint::PostToolCall, PostToolAction::Continue)
}
async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult<TurnEndAction> {
async fn on_assistant_turn_end(
&self,
context: AssistantTurnEndContext<'_>,
) -> InterceptorResult<TurnEndAction> {
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<InterceptorPoint> {
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<AtomicUsize>,
exits: Arc<Mutex<Vec<&'static str>>>,
}
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<PreRequestAction> {
Ok(if self.mode == TerminalMode::Yield {
PreRequestAction::Yield
} else {
PreRequestAction::Continue
})
}
async fn on_assistant_turn_end(
&self,
context: AssistantTurnEndContext<'_>,
) -> InterceptorResult<TurnEndAction> {
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<ResponseStream, ClientError> {
Err(ClientError::ContextWindowExceeded)
}
fn clone_boxed(&self) -> Box<dyn LlmClient> {
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();
crates/agen/tests/parallel_execution_test.rs
+30 -32
View File
@@ -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<PostToolAction> {
async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult<PostToolAction> {
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<std::sync::Mutex<Option<String>>>,
// Policy to observe the committed terminal result.
struct ObservingPolicy {
observed_content: Arc<std::sync::Mutex<Option<String>>>,
}
#[async_trait]
impl Interceptor for ModifyingPolicy {
async fn post_tool_call(
&self,
info: &mut ToolResultInfo,
) -> InterceptorResult<PostToolAction> {
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<PostToolAction> {
*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<Mutex<Vec<&'static str>>>,
}
#[async_trait]
impl Interceptor for AbortAfterResult {
async fn post_tool_call(
&self,
_info: &mut ToolResultInfo,
_info: &ToolResultInfo,
) -> InterceptorResult<PostToolAction> {
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,
crates/session-store/tests/session_test.rs
+8 -4
View File
@@ -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<TurnEndAction> {
async fn on_assistant_turn_end(
&self,
_context: AssistantTurnEndContext<'_>,
) -> InterceptorResult<TurnEndAction> {
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]
crates/worker/src/compact/worker.rs
+46 -9
View File
@@ -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<Item>,
context: PreLlmRequestContext<'_>,
) -> InterceptorResult<PreRequestAction> {
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")
));
}
crates/worker/src/hook.rs
+1 -21
View File
@@ -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<Box<dyn Hook<PreToolCall>>>,
post_tool_call: Vec<Box<dyn Hook<PostToolCall>>>,
on_turn_end: Vec<Box<dyn Hook<OnTurnEnd>>>,
on_abort: Vec<Box<dyn Hook<OnAbort>>>,
}
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<OnAbort> + '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<Box<dyn Hook<PreToolCall>>>,
pub(crate) post_tool_call: Vec<Box<dyn Hook<PostToolCall>>>,
pub(crate) on_turn_end: Vec<Box<dyn Hook<OnTurnEnd>>>,
pub(crate) on_abort: Vec<Box<dyn Hook<OnAbort>>>,
}
#[cfg(test)]
crates/worker/src/ipc/interceptor.rs
+79 -33
View File
@@ -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<PromptAction> {
async fn on_prompt_submit(
&self,
context: PromptSubmitContext<'_>,
) -> InterceptorResult<PromptAction> {
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<Item>,
context: PreLlmRequestContext<'_>,
) -> InterceptorResult<PreRequestAction> {
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<PostToolAction> {
async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult<PostToolAction> {
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<TurnEndAction> {
async fn on_assistant_turn_end(
&self,
context: AssistantTurnEndContext<'_>,
) -> InterceptorResult<TurnEndAction> {
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<Item> = 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<Item> = 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<Item> = 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<Item> = 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<Item> = 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));
crates/worker/src/worker.rs
+1 -11
View File
@@ -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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<OnAbort> + '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<C: LlmClient + 'static, St: Store> Worker<C, St> {
| 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}"),
}
}