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,