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/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}"),
}
}