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/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,