feat: type async interceptor failures

This commit is contained in:
2026-09-03 14:46:07 +09:00
parent 74bfbe941e
commit 68b1aa64e9
10 changed files with 505 additions and 123 deletions
+3 -3
View File
@@ -40,7 +40,7 @@ use tracing_subscriber::EnvFilter;
use agen::{ use agen::{
Engine, EngineRunExit, RunInterruptionReason, Engine, EngineRunExit, RunInterruptionReason,
interceptor::{Interceptor, PostToolAction, ToolResultInfo}, interceptor::{Interceptor, InterceptorResult, PostToolAction, ToolResultInfo},
llm_client::{ llm_client::{
LlmClient, LlmClient,
capability::{CacheStrategy, ModelCapability, StructuredOutput, ToolCallingSupport}, capability::{CacheStrategy, ModelCapability, StructuredOutput, ToolCallingSupport},
@@ -280,7 +280,7 @@ impl ToolResultPrinterPolicy {
#[async_trait] #[async_trait]
impl Interceptor for ToolResultPrinterPolicy { impl Interceptor for ToolResultPrinterPolicy {
async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction { async fn post_tool_call(&self, info: &mut ToolResultInfo) -> InterceptorResult<PostToolAction> {
let name = self let name = self
.call_names .call_names
.lock() .lock()
@@ -294,7 +294,7 @@ impl Interceptor for ToolResultPrinterPolicy {
println!(" Result ({}): ✅ {}", name, info.result.summary); println!(" Result ({}): ✅ {}", name, info.result.summary);
} }
PostToolAction::Continue Ok(PostToolAction::Continue)
} }
} }
+68 -9
View File
@@ -15,8 +15,8 @@ use crate::{
}, },
handler::{ErrorKind, StatusKind, ToolUseBlockStart, UsageKind}, handler::{ErrorKind, StatusKind, ToolUseBlockStart, UsageKind},
interceptor::{ interceptor::{
DefaultInterceptor, Interceptor, PostToolAction, PreRequestAction, PreToolAction, DefaultInterceptor, Interceptor, InterceptorFailure, InterceptorPoint, PostToolAction,
PromptAction, ToolCallInfo, ToolResultInfo, TurnEndAction, PreRequestAction, PreToolAction, PromptAction, ToolCallInfo, ToolResultInfo, TurnEndAction,
}, },
llm_client::{ llm_client::{
ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream, ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream,
@@ -58,6 +58,9 @@ pub enum EngineError {
/// A durable-history observer rejected an item before it entered history. /// A durable-history observer rejected an item before it entered history.
#[error("History append failed: {0}")] #[error("History append failed: {0}")]
HistoryAppend(String), HistoryAppend(String),
/// A trusted host interceptor callback failed.
#[error(transparent)]
Interceptor(#[from] InterceptorFailure),
/// Tool terminalization lost its execution-attempt compare-and-set fence. /// Tool terminalization lost its execution-attempt compare-and-set fence.
#[error("Tool execution attempt fence failed: {0}")] #[error("Tool execution attempt fence failed: {0}")]
ToolAttemptFence(String), ToolAttemptFence(String),
@@ -153,6 +156,8 @@ pub enum EngineRunExit {
/// A typed reason why an engine run could not finish normally. /// A typed reason why an engine run could not finish normally.
#[derive(Debug)] #[derive(Debug)]
pub enum RunInterruptionReason { pub enum RunInterruptionReason {
/// A trusted host interceptor callback failed at a typed lifecycle point.
Interceptor(InterceptorFailure),
LimitReached, LimitReached,
ContextWindowExceeded, ContextWindowExceeded,
Cancelled, Cancelled,
@@ -173,6 +178,9 @@ impl From<Result<EngineResult, EngineError>> for EngineRunExit {
} }
Err(EngineError::Cancelled) => Self::Interrupted(RunInterruptionReason::Cancelled), Err(EngineError::Cancelled) => Self::Interrupted(RunInterruptionReason::Cancelled),
Err(EngineError::PauseRequested) => Self::Paused, Err(EngineError::PauseRequested) => Self::Paused,
Err(EngineError::Interceptor(failure)) => {
Self::Interrupted(RunInterruptionReason::Interceptor(failure))
}
Err(error) => Self::Interrupted(RunInterruptionReason::Unexpected(error)), Err(error) => Self::Interrupted(RunInterruptionReason::Unexpected(error)),
} }
} }
@@ -1092,7 +1100,9 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
EngineError::Cancelled => "Cancelled".to_string(), EngineError::Cancelled => "Cancelled".to_string(),
_ => err.to_string(), _ => err.to_string(),
}; };
self.interceptor.on_abort(&reason).await; if let Err(error) = self.interceptor.on_abort(&reason).await {
return Err(InterceptorFailure::new(InterceptorPoint::Abort, error).into());
}
Err(err) Err(err)
} }
} }
@@ -1175,7 +1185,17 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
context, context,
}; };
match self.interceptor.pre_tool_call(&mut info).await { let pre_tool_action =
self.interceptor
.pre_tool_call(&mut info)
.await
.map_err(|error| {
EngineError::from(InterceptorFailure::new(
InterceptorPoint::PreToolCall,
error,
))
})?;
match pre_tool_action {
PreToolAction::Continue => {} PreToolAction::Continue => {}
PreToolAction::Skip => { PreToolAction::Skip => {
continue; continue;
@@ -1476,7 +1496,17 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
context: context.clone(), context: context.clone(),
}; };
match self.interceptor.post_tool_call(&mut info).await { 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::Continue => {}
PostToolAction::Abort(reason) => { PostToolAction::Abort(reason) => {
abort_reason = Some(reason); abort_reason = Some(reason);
@@ -1612,7 +1642,12 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
.interceptor .interceptor
.pending_history_appends() .pending_history_appends()
.await .await
.map_err(EngineError::HistoryAppend)?; .map_err(|error| {
EngineError::from(InterceptorFailure::new(
InterceptorPoint::PendingHistoryAppends,
error,
))
})?;
if !pending.is_empty() { if !pending.is_empty() {
self.append_history_items(history, pending, annotate)?; self.append_history_items(history, pending, annotate)?;
} }
@@ -1679,7 +1714,17 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
} }
// Interceptor: pre_llm_request // Interceptor: pre_llm_request
match self.interceptor.pre_llm_request(&mut request_context).await { let pre_request_action = self
.interceptor
.pre_llm_request(&mut request_context)
.await
.map_err(|error| {
EngineError::from(InterceptorFailure::new(
InterceptorPoint::PreLlmRequest,
error,
))
})?;
match pre_request_action {
PreRequestAction::Cancel(reason) => { PreRequestAction::Cancel(reason) => {
info!(reason = %reason, "Aborted by interceptor"); info!(reason = %reason, "Aborted by interceptor");
for cb in &self.turn_end_cbs { for cb in &self.turn_end_cbs {
@@ -1795,7 +1840,14 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
if tool_calls.is_empty() { if tool_calls.is_empty() {
let turn_end_context = history.items_cloned(); let turn_end_context = history.items_cloned();
match self.interceptor.on_turn_end(&turn_end_context).await { 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 => { TurnEndAction::Finish => {
return Ok(EngineResult::Finished); return Ok(EngineResult::Finished);
} }
@@ -2502,7 +2554,14 @@ impl<C: LlmClient, A> Engine<C, Locked, A> {
// Supplying new user input abandons any paused/yielded logical run. // Supplying new user input abandons any paused/yielded logical run.
self.active_run_turn_count = None; self.active_run_turn_count = None;
let mut user_item = Item::user_message(user_input); let mut user_item = Item::user_message(user_input);
let extras = match self.interceptor.on_prompt_submit(&mut user_item).await { 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 extras = match prompt_action {
PromptAction::Cancel(reason) => { PromptAction::Cancel(reason) => {
return self return self
.finalize_interruption(Err(EngineError::Aborted(reason))) .finalize_interruption(Err(EngineError::Aborted(reason)))
+130 -19
View File
@@ -11,6 +11,101 @@ use async_trait::async_trait;
use crate::Item; use crate::Item;
use crate::tool::{Tool, ToolCall, ToolExecutionContext, ToolMeta, ToolResult}; use crate::tool::{Tool, ToolCall, ToolExecutionContext, ToolMeta, ToolResult};
// =============================================================================
// Failure Types
// =============================================================================
/// A typed failure returned by an [`Interceptor`] implementation.
///
/// The Engine attaches the exact [`InterceptorPoint`] at which the failure was
/// observed before exposing it through the run termination boundary.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("{message}")]
pub struct InterceptorError {
message: String,
}
impl InterceptorError {
/// Create an interceptor failure with a caller-defined message.
pub fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
/// Return the failure message supplied by the interceptor.
pub fn message(&self) -> &str {
&self.message
}
}
impl From<String> for InterceptorError {
fn from(message: String) -> Self {
Self::new(message)
}
}
impl From<&str> for InterceptorError {
fn from(message: &str) -> Self {
Self::new(message)
}
}
/// The Engine lifecycle point at which an interceptor failed.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum InterceptorPoint {
PromptSubmit,
PendingHistoryAppends,
PreLlmRequest,
PreToolCall,
PostToolCall,
TurnEnd,
Abort,
}
impl std::fmt::Display for InterceptorPoint {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let name = match self {
Self::PromptSubmit => "prompt_submit",
Self::PendingHistoryAppends => "pending_history_appends",
Self::PreLlmRequest => "pre_llm_request",
Self::PreToolCall => "pre_tool_call",
Self::PostToolCall => "post_tool_call",
Self::TurnEnd => "turn_end",
Self::Abort => "abort",
};
formatter.write_str(name)
}
}
/// An interceptor failure bound to the exact Engine lifecycle point that ran it.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("{point} interceptor failed: {error}")]
pub struct InterceptorFailure {
point: InterceptorPoint,
#[source]
error: InterceptorError,
}
impl InterceptorFailure {
pub(crate) fn new(point: InterceptorPoint, error: InterceptorError) -> Self {
Self { point, error }
}
/// The lifecycle point that returned the failure.
pub fn point(&self) -> InterceptorPoint {
self.point
}
/// The typed error returned by the interceptor.
pub fn error(&self) -> &InterceptorError {
&self.error
}
}
/// Result returned by asynchronous interceptor lifecycle methods.
pub type InterceptorResult<T> = Result<T, InterceptorError>;
// ============================================================================= // =============================================================================
// Action Enums // Action Enums
// ============================================================================= // =============================================================================
@@ -130,14 +225,19 @@ pub struct ToolResultInfo {
/// Intercepts the Engine execution loop at key decision points. /// Intercepts the Engine execution loop at key decision points.
/// ///
/// All methods have default implementations that let the Engine /// Every lifecycle method is asynchronous and returns [`InterceptorResult`],
/// proceed without intervention. Callers provide richer implementations for /// keeping implementation failure separate from the method's control-flow
/// approval flows, permission checks, etc. /// action. The Engine reports a failure as a typed run interruption annotated
/// with the exact [`InterceptorPoint`] that failed.
///
/// All methods have default implementations that let the Engine proceed
/// without intervention. Callers provide richer implementations for approval
/// flows, permission checks, and other trusted host adaptation.
#[async_trait] #[async_trait]
pub trait Interceptor: Send + Sync { pub trait Interceptor: Send + Sync {
/// Called after receiving user input, before adding to history. /// Called after receiving user input, before adding it to Engine history.
async fn on_prompt_submit(&self, _item: &mut Item) -> PromptAction { async fn on_prompt_submit(&self, _item: &mut Item) -> InterceptorResult<PromptAction> {
PromptAction::Continue Ok(PromptAction::Continue)
} }
/// Items that should be **committed to `engine.history`** just /// Items that should be **committed to `engine.history`** just
@@ -158,7 +258,7 @@ pub trait Interceptor: Send + Sync {
/// reproducible per-request transformations (pruning, content /// reproducible per-request transformations (pruning, content
/// trimming, cache anchors) that depend only on the existing /// trimming, cache anchors) that depend only on the existing
/// history. /// history.
async fn pending_history_appends(&self) -> Result<Vec<Item>, String> { async fn pending_history_appends(&self) -> InterceptorResult<Vec<Item>> {
Ok(Vec::new()) Ok(Vec::new())
} }
@@ -170,27 +270,38 @@ pub trait Interceptor: Send + Sync {
/// If an interceptor derives a human/model-visible nudge from the current /// If an interceptor derives a human/model-visible nudge from the current
/// request context, return [`PreRequestAction::ContinueWith`] so the Engine /// request context, return [`PreRequestAction::ContinueWith`] so the Engine
/// commits it to history before the request is sent. /// commits it to history before the request is sent.
async fn pre_llm_request(&self, _context: &mut Vec<Item>) -> PreRequestAction { async fn pre_llm_request(
PreRequestAction::Continue &self,
_context: &mut Vec<Item>,
) -> InterceptorResult<PreRequestAction> {
Ok(PreRequestAction::Continue)
} }
/// Called before each tool is executed. /// Called before each tool is executed.
async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> PreToolAction { async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> InterceptorResult<PreToolAction> {
PreToolAction::Continue Ok(PreToolAction::Continue)
} }
/// Called after each tool completes. /// Called after each tool reaches one terminal result.
async fn post_tool_call(&self, _info: &mut ToolResultInfo) -> PostToolAction { async fn post_tool_call(
PostToolAction::Continue &self,
_info: &mut ToolResultInfo,
) -> InterceptorResult<PostToolAction> {
Ok(PostToolAction::Continue)
} }
/// Called when a turn ends with no tool calls. /// Called at the assistant boundary when a completed response has no tool calls.
async fn on_turn_end(&self, _history: &[Item]) -> TurnEndAction { ///
TurnEndAction::Finish /// 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> {
Ok(TurnEndAction::Finish)
} }
/// Called when execution is interrupted (abort or cancel). /// Called once when execution is interrupted (abort, cancellation, or failure).
async fn on_abort(&self, _reason: &str) {} async fn on_abort(&self, _reason: &str) -> InterceptorResult<()> {
Ok(())
}
} }
/// Default interceptor: no intervention. Engine proceeds through the loop /// Default interceptor: no intervention. Engine proceeds through the loop
+3 -1
View File
@@ -26,7 +26,9 @@ pub use engine::{
}; };
pub use handler::ToolUseBlockStart; pub use handler::ToolUseBlockStart;
pub use history::{History, HistoryEntry}; pub use history::{History, HistoryEntry};
pub use interceptor::Interceptor; pub use interceptor::{
Interceptor, InterceptorError, InterceptorFailure, InterceptorPoint, InterceptorResult,
};
pub use message::{ContentPart, Item, Message, Role}; pub use message::{ContentPart, Item, Message, Role};
pub use tool::{ pub use tool::{
ToolCall, ToolExecutionContext, ToolExecutionHandle, ToolExecutionPolicy, ToolCall, ToolExecutionContext, ToolExecutionHandle, ToolExecutionPolicy,
+193 -9
View File
@@ -10,7 +10,8 @@ use std::sync::{Arc, Mutex};
use agen::Item; use agen::Item;
use agen::interceptor::{ use agen::interceptor::{
Interceptor, PreRequestAction, PreToolAction, ToolCallInfo, TurnEndAction, Interceptor, InterceptorError, InterceptorPoint, InterceptorResult, PostToolAction,
PreRequestAction, PreToolAction, PromptAction, ToolCallInfo, ToolResultInfo, TurnEndAction,
}; };
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
@@ -613,12 +614,15 @@ struct YieldOnce {
#[async_trait] #[async_trait]
impl Interceptor for YieldOnce { impl Interceptor for YieldOnce {
async fn pre_llm_request(&self, _context: &mut Vec<Item>) -> PreRequestAction { async fn pre_llm_request(
if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { &self,
_context: &mut Vec<Item>,
) -> InterceptorResult<PreRequestAction> {
Ok(if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
PreRequestAction::Yield PreRequestAction::Yield
} else { } else {
PreRequestAction::Continue PreRequestAction::Continue
} })
} }
} }
@@ -628,12 +632,12 @@ struct PauseToolOnce {
#[async_trait] #[async_trait]
impl Interceptor for PauseToolOnce { impl Interceptor for PauseToolOnce {
async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> PreToolAction { async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> InterceptorResult<PreToolAction> {
if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { Ok(if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
PreToolAction::Pause PreToolAction::Pause
} else { } else {
PreToolAction::Continue PreToolAction::Continue
} })
} }
} }
@@ -643,12 +647,192 @@ struct ContinueTurnOnce {
#[async_trait] #[async_trait]
impl Interceptor for ContinueTurnOnce { impl Interceptor for ContinueTurnOnce {
async fn on_turn_end(&self, _history: &[Item]) -> TurnEndAction { async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult<TurnEndAction> {
if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { Ok(if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
TurnEndAction::ContinueWithMessages(vec![Item::system_message("continue")]) TurnEndAction::ContinueWithMessages(vec![Item::system_message("continue")])
} else { } else {
TurnEndAction::Finish TurnEndAction::Finish
})
} }
}
#[derive(Debug, Clone)]
struct FailingLifecycleInterceptor {
failure: InterceptorPoint,
calls: Arc<Mutex<Vec<InterceptorPoint>>>,
}
impl FailingLifecycleInterceptor {
fn new(failure: InterceptorPoint) -> Self {
Self {
failure,
calls: Arc::new(Mutex::new(Vec::new())),
}
}
fn record<T>(&self, point: InterceptorPoint, action: T) -> InterceptorResult<T> {
self.calls.lock().unwrap().push(point);
if self.failure == point {
Err(InterceptorError::new(format!("{point} rejected")))
} else {
Ok(action)
}
}
fn calls(&self) -> Vec<InterceptorPoint> {
self.calls.lock().unwrap().clone()
}
}
#[async_trait]
impl Interceptor for FailingLifecycleInterceptor {
async fn on_prompt_submit(&self, _item: &mut Item) -> InterceptorResult<PromptAction> {
tokio::task::yield_now().await;
self.record(InterceptorPoint::PromptSubmit, PromptAction::Continue)
}
async fn pending_history_appends(&self) -> InterceptorResult<Vec<Item>> {
tokio::task::yield_now().await;
self.record(InterceptorPoint::PendingHistoryAppends, Vec::new())
}
async fn pre_llm_request(
&self,
_context: &mut Vec<Item>,
) -> 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)
}
}
async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> InterceptorResult<PreToolAction> {
tokio::task::yield_now().await;
self.record(InterceptorPoint::PreToolCall, PreToolAction::Continue)
}
async fn post_tool_call(
&self,
_info: &mut ToolResultInfo,
) -> InterceptorResult<PostToolAction> {
tokio::task::yield_now().await;
self.record(InterceptorPoint::PostToolCall, PostToolAction::Continue)
}
async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult<TurnEndAction> {
tokio::task::yield_now().await;
self.record(InterceptorPoint::TurnEnd, TurnEndAction::Finish)
}
async fn on_abort(&self, _reason: &str) -> InterceptorResult<()> {
tokio::task::yield_now().await;
self.record(InterceptorPoint::Abort, ())
}
}
fn expected_interceptor_calls(failure: InterceptorPoint) -> Vec<InterceptorPoint> {
use InterceptorPoint as Point;
match failure {
Point::PromptSubmit => vec![Point::PromptSubmit, Point::Abort],
Point::PendingHistoryAppends => {
vec![
Point::PromptSubmit,
Point::PendingHistoryAppends,
Point::Abort,
]
}
Point::PreLlmRequest => vec![
Point::PromptSubmit,
Point::PendingHistoryAppends,
Point::PreLlmRequest,
Point::Abort,
],
Point::PreToolCall => vec![
Point::PromptSubmit,
Point::PendingHistoryAppends,
Point::PreLlmRequest,
Point::PreToolCall,
Point::Abort,
],
Point::PostToolCall => vec![
Point::PromptSubmit,
Point::PendingHistoryAppends,
Point::PreLlmRequest,
Point::PreToolCall,
Point::PostToolCall,
Point::Abort,
],
Point::TurnEnd => vec![
Point::PromptSubmit,
Point::PendingHistoryAppends,
Point::PreLlmRequest,
Point::TurnEnd,
Point::Abort,
],
Point::Abort => vec![
Point::PromptSubmit,
Point::PendingHistoryAppends,
Point::PreLlmRequest,
Point::Abort,
],
}
}
#[tokio::test]
async fn interceptor_failures_are_typed_and_each_lifecycle_point_runs_once() {
use InterceptorPoint as Point;
for failure_point in [
Point::PromptSubmit,
Point::PendingHistoryAppends,
Point::PreLlmRequest,
Point::PreToolCall,
Point::PostToolCall,
Point::TurnEnd,
Point::Abort,
] {
let interceptor = FailingLifecycleInterceptor::new(failure_point);
let needs_tool = matches!(failure_point, Point::PreToolCall | Point::PostToolCall);
let events = if needs_tool {
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,
}),
]
} else {
completed_text_events()
};
let mut engine = Engine::new(MockLlmClient::new(events));
engine.register_tool(CountingTool::new("count_tool").definition());
engine.set_interceptor(interceptor.clone());
let mut history = History::new();
let mut engine = engine.lock(&history);
let exit = engine.run(&mut history, "test").await;
let EngineRunExit::Interrupted(RunInterruptionReason::Interceptor(failure)) = exit else {
panic!("expected typed interceptor interruption at {failure_point}, got {exit:?}");
};
assert_eq!(failure.point(), failure_point);
assert_eq!(
failure.error().message(),
format!("{failure_point} rejected")
);
assert_eq!(
interceptor.calls(),
expected_interceptor_calls(failure_point)
);
} }
} }
+27 -16
View File
@@ -6,7 +6,9 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use agen::interceptor::{Interceptor, PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo}; use agen::interceptor::{
Interceptor, InterceptorResult, PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo,
};
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::tool::{ use agen::tool::{
Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult, Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult,
@@ -905,24 +907,27 @@ async fn test_tool_execution_context_for_skipped_and_synthetic_paths() {
#[async_trait] #[async_trait]
impl Interceptor for ContextPolicy { impl Interceptor for ContextPolicy {
async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction { async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult<PreToolAction> {
self.pre_contexts.lock().unwrap().push(info.context.clone()); self.pre_contexts.lock().unwrap().push(info.context.clone());
match info.call.name.as_str() { Ok(match info.call.name.as_str() {
"skip_tool" => PreToolAction::Skip, "skip_tool" => PreToolAction::Skip,
"synthetic_tool" => PreToolAction::SyntheticResult(ToolResult::from_output( "synthetic_tool" => PreToolAction::SyntheticResult(ToolResult::from_output(
&info.call.id, &info.call.id,
ToolOutput::from("synthetic result".to_string()), ToolOutput::from("synthetic result".to_string()),
)), )),
_ => PreToolAction::Continue, _ => PreToolAction::Continue,
} })
} }
async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction { async fn post_tool_call(
&self,
info: &mut ToolResultInfo,
) -> InterceptorResult<PostToolAction> {
self.post_contexts self.post_contexts
.lock() .lock()
.unwrap() .unwrap()
.push(info.context.clone()); .push(info.context.clone());
PostToolAction::Continue Ok(PostToolAction::Continue)
} }
} }
@@ -994,12 +999,12 @@ async fn test_before_tool_call_skip() {
#[async_trait] #[async_trait]
impl Interceptor for BlockingPolicy { impl Interceptor for BlockingPolicy {
async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction { async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult<PreToolAction> {
if info.call.name == "blocked_tool" { Ok(if info.call.name == "blocked_tool" {
PreToolAction::Skip PreToolAction::Skip
} else { } else {
PreToolAction::Continue PreToolAction::Continue
} })
} }
} }
@@ -1081,10 +1086,13 @@ async fn test_post_tool_call_modification() {
#[async_trait] #[async_trait]
impl Interceptor for ModifyingPolicy { impl Interceptor for ModifyingPolicy {
async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction { async fn post_tool_call(
&self,
info: &mut ToolResultInfo,
) -> InterceptorResult<PostToolAction> {
info.result.summary = format!("[Modified] {}", info.result.summary); info.result.summary = format!("[Modified] {}", info.result.summary);
*self.modified_content.lock().unwrap() = Some(info.result.summary.clone()); *self.modified_content.lock().unwrap() = Some(info.result.summary.clone());
PostToolAction::Continue Ok(PostToolAction::Continue)
} }
} }
@@ -1143,11 +1151,11 @@ async fn test_before_tool_call_synthetic_result_committed() {
#[async_trait] #[async_trait]
impl Interceptor for SyntheticPolicy { impl Interceptor for SyntheticPolicy {
async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction { async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult<PreToolAction> {
PreToolAction::SyntheticResult(ToolResult::error( Ok(PreToolAction::SyntheticResult(ToolResult::error(
info.call.id.clone(), info.call.id.clone(),
"permission denied", "permission denied",
)) )))
} }
} }
@@ -1184,8 +1192,11 @@ async fn post_tool_abort_commits_confirmed_result_before_stopping_run() {
struct AbortAfterResult; struct AbortAfterResult;
#[async_trait] #[async_trait]
impl Interceptor for AbortAfterResult { impl Interceptor for AbortAfterResult {
async fn post_tool_call(&self, _info: &mut ToolResultInfo) -> PostToolAction { async fn post_tool_call(
PostToolAction::Abort("policy stopped the run".to_string()) &self,
_info: &mut ToolResultInfo,
) -> InterceptorResult<PostToolAction> {
Ok(PostToolAction::Abort("policy stopped the run".to_string()))
} }
} }
engine.set_interceptor(AbortAfterResult); engine.set_interceptor(AbortAfterResult);
+3 -3
View File
@@ -3,7 +3,7 @@ mod common;
use std::ops::{Deref, DerefMut}; use std::ops::{Deref, DerefMut};
use std::sync::Arc; use std::sync::Arc;
use agen::interceptor::{Interceptor, TurnEndAction}; use agen::interceptor::{Interceptor, InterceptorResult, TurnEndAction};
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::llm_client::types::{Item, RequestConfig}; use agen::llm_client::types::{Item, RequestConfig};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
@@ -100,8 +100,8 @@ struct PausePolicy;
#[async_trait] #[async_trait]
impl Interceptor for PausePolicy { impl Interceptor for PausePolicy {
async fn on_turn_end(&self, _history: &[Item]) -> TurnEndAction { async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult<TurnEndAction> {
TurnEndAction::Pause Ok(TurnEndAction::Pause)
} }
} }
+23 -18
View File
@@ -22,7 +22,9 @@ use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use agen::Item; use agen::Item;
use agen::interceptor::{Interceptor, PreRequestAction, PreToolAction, ToolCallInfo}; use agen::interceptor::{
Interceptor, InterceptorResult, PreRequestAction, PreToolAction, ToolCallInfo,
};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput, ToolResult}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput, ToolResult};
use async_trait::async_trait; use async_trait::async_trait;
use serde::Deserialize; use serde::Deserialize;
@@ -398,14 +400,17 @@ impl CompactWorkerInterceptor {
#[async_trait] #[async_trait]
impl Interceptor for CompactWorkerInterceptor { impl Interceptor for CompactWorkerInterceptor {
async fn pre_llm_request(&self, context: &mut Vec<Item>) -> PreRequestAction { async fn pre_llm_request(
&self,
context: &mut Vec<Item>,
) -> InterceptorResult<PreRequestAction> {
let records = self.usage_tracker.records(); let records = self.usage_tracker.records();
let estimate = agen::token_counter::total_tokens(context, &records); let estimate = agen::token_counter::total_tokens(context, &records);
if estimate.tokens > self.max_input_tokens { if estimate.tokens > self.max_input_tokens {
return PreRequestAction::Cancel(format!( return Ok(PreRequestAction::Cancel(format!(
"compact worker input occupancy exceeded {} tokens", "compact worker input occupancy exceeded {} tokens",
self.max_input_tokens self.max_input_tokens
)); )));
} }
let remaining = self.max_input_tokens.saturating_sub(estimate.tokens); let remaining = self.max_input_tokens.saturating_sub(estimate.tokens);
@@ -413,25 +418,25 @@ impl Interceptor for CompactWorkerInterceptor {
.store(remaining, Ordering::Release); .store(remaining, Ordering::Release);
if let Some(item) = self.maybe_emit_warning(remaining) { if let Some(item) = self.maybe_emit_warning(remaining) {
self.usage_tracker.note_request(context.len() + 1); self.usage_tracker.note_request(context.len() + 1);
return PreRequestAction::ContinueWith(vec![item]); return Ok(PreRequestAction::ContinueWith(vec![item]));
} }
self.usage_tracker.note_request(context.len()); self.usage_tracker.note_request(context.len());
PreRequestAction::Continue Ok(PreRequestAction::Continue)
} }
async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction { async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult<PreToolAction> {
if self.final_reserve_tokens == 0 || info.call.name == "write_summary" { if self.final_reserve_tokens == 0 || info.call.name == "write_summary" {
return PreToolAction::Continue; return Ok(PreToolAction::Continue);
} }
let remaining = self.last_remaining_tokens.load(Ordering::Acquire); let remaining = self.last_remaining_tokens.load(Ordering::Acquire);
if remaining > self.final_reserve_tokens { if remaining > self.final_reserve_tokens {
return PreToolAction::Continue; return Ok(PreToolAction::Continue);
} }
PreToolAction::SyntheticResult(ToolResult::error( Ok(PreToolAction::SyntheticResult(ToolResult::error(
info.call.id.clone(), info.call.id.clone(),
"compact worker final reserve reached; do not perform more exploratory tool reads. Call `write_summary` now.", "compact worker final reserve reached; do not perform more exploratory tool reads. Call `write_summary` now.",
)) )))
} }
} }
@@ -467,13 +472,13 @@ mod tests {
let mut context = vec![Item::user_message("hello")]; let mut context = vec![Item::user_message("hello")];
assert!(matches!( assert!(matches!(
interceptor.pre_llm_request(&mut context).await, interceptor.pre_llm_request(&mut context).await.unwrap(),
PreRequestAction::Continue PreRequestAction::Continue
)); ));
tracker.record_usage(&make_usage(100)); tracker.record_usage(&make_usage(100));
assert!(matches!( assert!(matches!(
interceptor.pre_llm_request(&mut context).await, interceptor.pre_llm_request(&mut context).await.unwrap(),
PreRequestAction::Continue PreRequestAction::Continue
)); ));
tracker.record_usage(&make_usage(100)); tracker.record_usage(&make_usage(100));
@@ -481,7 +486,7 @@ mod tests {
// Two 100-token requests would exceed a cumulative 150-token cap, but // Two 100-token requests would exceed a cumulative 150-token cap, but
// current occupancy is still the latest 100-token measurement. // current occupancy is still the latest 100-token measurement.
assert!(matches!( assert!(matches!(
interceptor.pre_llm_request(&mut context).await, interceptor.pre_llm_request(&mut context).await.unwrap(),
PreRequestAction::Continue PreRequestAction::Continue
)); ));
} }
@@ -503,13 +508,13 @@ mod tests {
let mut context = vec![Item::user_message("hello")]; let mut context = vec![Item::user_message("hello")];
assert!(matches!( assert!(matches!(
interceptor.pre_llm_request(&mut context).await, interceptor.pre_llm_request(&mut context).await.unwrap(),
PreRequestAction::Continue PreRequestAction::Continue
)); ));
tracker.record_usage(&make_usage(100)); tracker.record_usage(&make_usage(100));
assert!(matches!( assert!(matches!(
interceptor.pre_llm_request(&mut context).await, interceptor.pre_llm_request(&mut context).await.unwrap(),
PreRequestAction::ContinueWith(items) PreRequestAction::ContinueWith(items)
if items.len() == 1 && items[0].as_text().unwrap_or_default().contains("write_summary") if items.len() == 1 && items[0].as_text().unwrap_or_default().contains("write_summary")
)); ));
@@ -523,13 +528,13 @@ mod tests {
let mut context = vec![Item::user_message("hello")]; let mut context = vec![Item::user_message("hello")];
assert!(matches!( assert!(matches!(
interceptor.pre_llm_request(&mut context).await, interceptor.pre_llm_request(&mut context).await.unwrap(),
PreRequestAction::Continue PreRequestAction::Continue
)); ));
tracker.record_usage(&make_usage(100)); tracker.record_usage(&make_usage(100));
assert!(matches!( assert!(matches!(
interceptor.pre_llm_request(&mut context).await, interceptor.pre_llm_request(&mut context).await.unwrap(),
PreRequestAction::Cancel(message) if message.contains("occupancy") PreRequestAction::Cancel(message) if message.contains("occupancy")
)); ));
} }
+51 -45
View File
@@ -15,8 +15,8 @@ use std::sync::{Arc, Mutex};
use agen::Item; use agen::Item;
use agen::UsageRecord; use agen::UsageRecord;
use agen::interceptor::{ use agen::interceptor::{
Interceptor, PostToolAction, PreRequestAction, PreToolAction, PromptAction, ToolCallInfo, Interceptor, InterceptorResult, PostToolAction, PreRequestAction, PreToolAction, PromptAction,
ToolResultInfo, TurnEndAction, ToolCallInfo, ToolResultInfo, TurnEndAction,
}; };
use agen::tool::ToolOutput; use agen::tool::ToolOutput;
use arc_swap::ArcSwap; use arc_swap::ArcSwap;
@@ -232,7 +232,7 @@ impl WorkerInterceptor {
#[async_trait] #[async_trait]
impl Interceptor for WorkerInterceptor { impl Interceptor for WorkerInterceptor {
async fn on_prompt_submit(&self, item: &mut Item) -> PromptAction { async fn on_prompt_submit(&self, item: &mut Item) -> InterceptorResult<PromptAction> {
let turn_index = self.next_turn_index.fetch_add(1, Ordering::Relaxed); let turn_index = self.next_turn_index.fetch_add(1, Ordering::Relaxed);
self.tool_calls_this_turn.store(0, Ordering::Relaxed); self.tool_calls_this_turn.store(0, Ordering::Relaxed);
@@ -243,7 +243,7 @@ impl Interceptor for WorkerInterceptor {
for hook in &self.registry.on_prompt_submit { for hook in &self.registry.on_prompt_submit {
let action = hook.call(&info).await; let action = hook.call(&info).await;
if !matches!(action, HookPromptAction::Continue) { if !matches!(action, HookPromptAction::Continue) {
return action.into(); return Ok(action.into());
} }
} }
let mut extras: Vec<SystemItem> = std::mem::take( let mut extras: Vec<SystemItem> = std::mem::take(
@@ -252,7 +252,7 @@ impl Interceptor for WorkerInterceptor {
.lock() .lock()
.expect("pending_attachments poisoned"), .expect("pending_attachments poisoned"),
); );
if extras.is_empty() { Ok(if extras.is_empty() {
PromptAction::Continue PromptAction::Continue
} else { } else {
// Commit the typed system items first, then hand the // Commit the typed system items first, then hand the
@@ -266,10 +266,10 @@ impl Interceptor for WorkerInterceptor {
Ok(()) => PromptAction::ContinueWith(items), Ok(()) => PromptAction::ContinueWith(items),
Err(error) => PromptAction::Cancel(format!("session persistence failed: {error}")), Err(error) => PromptAction::Cancel(format!("session persistence failed: {error}")),
} }
} })
} }
async fn pending_history_appends(&self) -> Result<Vec<Item>, String> { async fn pending_history_appends(&self) -> InterceptorResult<Vec<Item>> {
let drained = self.pending_notifies.drain(); let drained = self.pending_notifies.drain();
if drained.is_empty() { if drained.is_empty() {
return Ok(Vec::new()); return Ok(Vec::new());
@@ -295,7 +295,7 @@ impl Interceptor for WorkerInterceptor {
Ok(system_item) => system_item, Ok(system_item) => system_item,
Err(error) => { Err(error) => {
self.pending_notifies.requeue_front(drained); self.pending_notifies.requeue_front(drained);
return Err(format!("failed to render notify_wrapper: {error}")); return Err(format!("failed to render notify_wrapper: {error}").into());
} }
}; };
items.push(system_item.to_history_item()); items.push(system_item.to_history_item());
@@ -303,15 +303,18 @@ impl Interceptor for WorkerInterceptor {
} }
if let Err(error) = self.commit_system_items(&system_items) { if let Err(error) = self.commit_system_items(&system_items) {
self.pending_notifies.requeue_front(drained); self.pending_notifies.requeue_front(drained);
return Err(format!("session persistence failed: {error}")); return Err(format!("session persistence failed: {error}").into());
} }
Ok(items) Ok(items)
} }
async fn pre_llm_request(&self, context: &mut Vec<Item>) -> PreRequestAction { async fn pre_llm_request(
&self,
context: &mut Vec<Item>,
) -> InterceptorResult<PreRequestAction> {
let initial_tokens = self.estimated_tokens(context); let initial_tokens = self.estimated_tokens(context);
if self.request_threshold_exceeded(initial_tokens, context) { if self.request_threshold_exceeded(initial_tokens, context) {
return PreRequestAction::Yield; return Ok(PreRequestAction::Yield);
} }
let info = PreRequestInfo { let info = PreRequestInfo {
item_count: context.len(), item_count: context.len(),
@@ -328,7 +331,7 @@ impl Interceptor for WorkerInterceptor {
for hook in &self.registry.pre_llm_request { for hook in &self.registry.pre_llm_request {
let action = hook.call(&hook_context).await; let action = hook.call(&hook_context).await;
if !matches!(action, HookPreRequestAction::Continue) { if !matches!(action, HookPreRequestAction::Continue) {
return action.into(); return Ok(action.into());
} }
} }
@@ -353,28 +356,30 @@ impl Interceptor for WorkerInterceptor {
if self.request_threshold_exceeded(current_tokens, effective_context.as_ref()) { if self.request_threshold_exceeded(current_tokens, effective_context.as_ref()) {
if let Err(error) = self.commit_system_items(&system_items) { if let Err(error) = self.commit_system_items(&system_items) {
return PreRequestAction::Cancel(format!("session persistence failed: {error}")); return Ok(PreRequestAction::Cancel(format!(
"session persistence failed: {error}"
)));
} }
return if appended_items.is_empty() { return Ok(if appended_items.is_empty() {
PreRequestAction::Yield PreRequestAction::Yield
} else { } else {
PreRequestAction::YieldWith(appended_items) PreRequestAction::YieldWith(appended_items)
}; });
} }
if let Some(usage_tracker) = self.usage_tracker.as_ref() { if let Some(usage_tracker) = self.usage_tracker.as_ref() {
usage_tracker.note_request(effective_context.len()); usage_tracker.note_request(effective_context.len());
} }
if system_items.is_empty() { if system_items.is_empty() {
return PreRequestAction::Continue; return Ok(PreRequestAction::Continue);
} }
match self.commit_system_items(&system_items) { Ok(match self.commit_system_items(&system_items) {
Ok(()) => PreRequestAction::ContinueWith(appended_items), Ok(()) => PreRequestAction::ContinueWith(appended_items),
Err(error) => PreRequestAction::Cancel(format!("session persistence failed: {error}")), Err(error) => PreRequestAction::Cancel(format!("session persistence failed: {error}")),
} })
} }
async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction { async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult<PreToolAction> {
let summary = ToolCallSummary { let summary = ToolCallSummary {
call_id: info.call.id.clone(), call_id: info.call.id.clone(),
tool_name: info.call.name.clone(), tool_name: info.call.name.clone(),
@@ -383,14 +388,14 @@ impl Interceptor for WorkerInterceptor {
for hook in &self.registry.pre_tool_call { for hook in &self.registry.pre_tool_call {
let action = hook.call(&summary).await; let action = hook.call(&summary).await;
if !matches!(action, HookPreToolAction::Continue) { if !matches!(action, HookPreToolAction::Continue) {
return action.into_worker_action(summary.call_id.clone()); return Ok(action.into_worker_action(summary.call_id.clone()));
} }
} }
self.tool_calls_this_turn.fetch_add(1, Ordering::Relaxed); self.tool_calls_this_turn.fetch_add(1, Ordering::Relaxed);
PreToolAction::Continue Ok(PreToolAction::Continue)
} }
async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction { async fn post_tool_call(&self, info: &mut ToolResultInfo) -> InterceptorResult<PostToolAction> {
let summary = ToolResultSummary { let summary = ToolResultSummary {
call_id: info.result.tool_use_id.clone(), call_id: info.result.tool_use_id.clone(),
tool_name: info.call.name.clone(), tool_name: info.call.name.clone(),
@@ -405,13 +410,13 @@ impl Interceptor for WorkerInterceptor {
for hook in &self.registry.post_tool_call { for hook in &self.registry.post_tool_call {
let action = hook.call(&summary).await; let action = hook.call(&summary).await;
if !matches!(action, HookPostToolAction::Continue) { if !matches!(action, HookPostToolAction::Continue) {
return action.into(); return Ok(action.into());
} }
} }
PostToolAction::Continue Ok(PostToolAction::Continue)
} }
async fn on_turn_end(&self, history: &[Item]) -> TurnEndAction { async fn on_turn_end(&self, history: &[Item]) -> InterceptorResult<TurnEndAction> {
let final_text_preview = history let final_text_preview = history
.iter() .iter()
.rev() .rev()
@@ -427,19 +432,20 @@ impl Interceptor for WorkerInterceptor {
for hook in &self.registry.on_turn_end { for hook in &self.registry.on_turn_end {
let action = hook.call(&info).await; let action = hook.call(&info).await;
if !matches!(action, HookTurnEndAction::Finish) { if !matches!(action, HookTurnEndAction::Finish) {
return action.into(); return Ok(action.into());
} }
} }
TurnEndAction::Finish Ok(TurnEndAction::Finish)
} }
async fn on_abort(&self, reason: &str) { async fn on_abort(&self, reason: &str) -> InterceptorResult<()> {
let info = AbortInfo { let info = AbortInfo {
reason: reason.to_string(), reason: reason.to_string(),
}; };
for hook in &self.registry.on_abort { for hook in &self.registry.on_abort {
hook.call(&info).await; hook.call(&info).await;
} }
Ok(())
} }
} }
@@ -623,7 +629,7 @@ mod tests {
None, None,
); );
let mut ctx = ctx_items; let mut ctx = ctx_items;
let action = interceptor.pre_llm_request(&mut ctx).await; let action = interceptor.pre_llm_request(&mut ctx).await.unwrap();
assert!(matches!(action, PreRequestAction::Yield)); assert!(matches!(action, PreRequestAction::Yield));
// Hook must not run when an internal mechanism short-circuits first. // Hook must not run when an internal mechanism short-circuits first.
@@ -655,7 +661,7 @@ mod tests {
})), })),
); );
let mut ctx = ctx_items; let mut ctx = ctx_items;
let action = interceptor.pre_llm_request(&mut ctx).await; let action = interceptor.pre_llm_request(&mut ctx).await.unwrap();
match action { match action {
PreRequestAction::YieldWith(items) => assert_eq!(items.len(), 1), PreRequestAction::YieldWith(items) => assert_eq!(items.len(), 1),
@@ -692,7 +698,7 @@ mod tests {
) )
.with_usage_tracker(usage_tracker); .with_usage_tracker(usage_tracker);
let mut ctx = ctx_items; let mut ctx = ctx_items;
let action = interceptor.pre_llm_request(&mut ctx).await; let action = interceptor.pre_llm_request(&mut ctx).await.unwrap();
assert!(matches!(action, PreRequestAction::Yield)); assert!(matches!(action, PreRequestAction::Yield));
} }
@@ -716,7 +722,7 @@ mod tests {
None, None,
); );
let mut ctx = ctx_items; let mut ctx = ctx_items;
let action = interceptor.pre_llm_request(&mut ctx).await; let action = interceptor.pre_llm_request(&mut ctx).await.unwrap();
assert!(matches!(action, PreRequestAction::Continue)); assert!(matches!(action, PreRequestAction::Continue));
assert_eq!(count.load(Ordering::Relaxed), 1); assert_eq!(count.load(Ordering::Relaxed), 1);
@@ -757,7 +763,7 @@ mod tests {
None, None,
); );
let mut ctx = ctx_items; let mut ctx = ctx_items;
let action = interceptor.pre_llm_request(&mut ctx).await; let action = interceptor.pre_llm_request(&mut ctx).await.unwrap();
assert!(matches!(action, PreRequestAction::Continue)); assert!(matches!(action, PreRequestAction::Continue));
assert_eq!(count.load(Ordering::Relaxed), 1); assert_eq!(count.load(Ordering::Relaxed), 1);
@@ -784,7 +790,7 @@ mod tests {
None, None,
); );
let mut ctx = ctx_items; let mut ctx = ctx_items;
let action = interceptor.pre_llm_request(&mut ctx).await; let action = interceptor.pre_llm_request(&mut ctx).await.unwrap();
assert!(matches!(action, PreRequestAction::Continue)); assert!(matches!(action, PreRequestAction::Continue));
assert_eq!(count.load(Ordering::Relaxed), 1); assert_eq!(count.load(Ordering::Relaxed), 1);
@@ -805,7 +811,7 @@ mod tests {
None, None,
); );
let mut ctx: Vec<Item> = Vec::new(); let mut ctx: Vec<Item> = Vec::new();
let action = interceptor.pre_llm_request(&mut ctx).await; let action = interceptor.pre_llm_request(&mut ctx).await.unwrap();
assert!(matches!(action, PreRequestAction::Continue)); assert!(matches!(action, PreRequestAction::Continue));
assert_eq!(count.load(Ordering::Relaxed), 1); assert_eq!(count.load(Ordering::Relaxed), 1);
@@ -834,7 +840,7 @@ mod tests {
); );
let mut ctx: Vec<Item> = Vec::new(); let mut ctx: Vec<Item> = Vec::new();
let action = interceptor.pre_llm_request(&mut ctx).await; let action = interceptor.pre_llm_request(&mut ctx).await.unwrap();
assert!(saw_handle.load(Ordering::Relaxed)); assert!(saw_handle.load(Ordering::Relaxed));
let PreRequestAction::ContinueWith(items) = action else { let PreRequestAction::ContinueWith(items) = action else {
@@ -881,7 +887,7 @@ mod tests {
); );
let mut ctx: Vec<Item> = Vec::new(); let mut ctx: Vec<Item> = Vec::new();
let action = interceptor.pre_llm_request(&mut ctx).await; let action = interceptor.pre_llm_request(&mut ctx).await.unwrap();
assert!(!saw_handle.load(Ordering::Relaxed)); assert!(!saw_handle.load(Ordering::Relaxed));
assert!(matches!(action, PreRequestAction::Continue)); assert!(matches!(action, PreRequestAction::Continue));
@@ -938,7 +944,7 @@ mod tests {
); );
let mut info = task_tool_call_info("TaskList", serde_json::json!({"scope": "all"})); let mut info = task_tool_call_info("TaskList", serde_json::json!({"scope": "all"}));
let action = interceptor.pre_tool_call(&mut info).await; let action = interceptor.pre_tool_call(&mut info).await.unwrap();
match action { match action {
PreToolAction::SyntheticResult(result) => { PreToolAction::SyntheticResult(result) => {
@@ -1000,7 +1006,7 @@ mod tests {
context: info.context, context: info.context,
}; };
let action = interceptor.post_tool_call(&mut result_info).await; let action = interceptor.post_tool_call(&mut result_info).await.unwrap();
assert_eq!(action, PostToolAction::Abort("post tool abort".to_string())); assert_eq!(action, PostToolAction::Abort("post tool abort".to_string()));
assert_eq!(count.load(Ordering::Relaxed), 1); assert_eq!(count.load(Ordering::Relaxed), 1);
@@ -1036,7 +1042,7 @@ mod tests {
); );
let history = vec![Item::user_message("hi"), Item::assistant_message("done")]; let history = vec![Item::user_message("hi"), Item::assistant_message("done")];
let action = interceptor.on_turn_end(&history).await; let action = interceptor.on_turn_end(&history).await.unwrap();
assert!(matches!(action, TurnEndAction::Pause)); assert!(matches!(action, TurnEndAction::Pause));
assert_eq!(count.load(Ordering::Relaxed), 1); assert_eq!(count.load(Ordering::Relaxed), 1);
@@ -1073,7 +1079,7 @@ mod tests {
let ctx_items = vec![Item::user_message("hi")]; let ctx_items = vec![Item::user_message("hi")];
for _ in 0..23 { for _ in 0..23 {
let mut ctx = ctx_items.clone(); let mut ctx = ctx_items.clone();
let action = interceptor.pre_llm_request(&mut ctx).await; let action = interceptor.pre_llm_request(&mut ctx).await.unwrap();
assert!(matches!(action, PreRequestAction::Continue)); assert!(matches!(action, PreRequestAction::Continue));
usage_tracker.record_usage(&agen::event::UsageEvent { usage_tracker.record_usage(&agen::event::UsageEvent {
input_tokens: Some(10), input_tokens: Some(10),
@@ -1085,7 +1091,7 @@ mod tests {
} }
let mut ctx = ctx_items.clone(); let mut ctx = ctx_items.clone();
let action = interceptor.pre_llm_request(&mut ctx).await; let action = interceptor.pre_llm_request(&mut ctx).await.unwrap();
let appended_len = match action { let appended_len = match action {
PreRequestAction::ContinueWith(items) => items.len(), PreRequestAction::ContinueWith(items) => items.len(),
other => panic!("expected reminder append, got {other:?}"), other => panic!("expected reminder append, got {other:?}"),
@@ -1210,7 +1216,7 @@ mod tests {
let error = interceptor.pending_history_appends().await.unwrap_err(); let error = interceptor.pending_history_appends().await.unwrap_err();
assert!(error.contains("failed to render notify_wrapper")); assert!(error.message().contains("failed to render notify_wrapper"));
let requeued = buffer.drain(); let requeued = buffer.drain();
assert_eq!(requeued.len(), 1); assert_eq!(requeued.len(), 1);
} }
@@ -1269,7 +1275,7 @@ mod tests {
None, None,
); );
let mut ctx: Vec<Item> = vec![Item::user_message("hi")]; let mut ctx: Vec<Item> = vec![Item::user_message("hi")];
let action = interceptor.pre_llm_request(&mut ctx).await; let action = interceptor.pre_llm_request(&mut ctx).await.unwrap();
assert!(matches!(action, PreRequestAction::Continue)); assert!(matches!(action, PreRequestAction::Continue));
assert_eq!(ctx.len(), 1, "pre_llm_request must not append notifies"); assert_eq!(ctx.len(), 1, "pre_llm_request must not append notifies");
@@ -1299,7 +1305,7 @@ mod tests {
None, None,
); );
let mut ctx: Vec<Item> = Vec::new(); let mut ctx: Vec<Item> = Vec::new();
let action = interceptor.pre_llm_request(&mut ctx).await; let action = interceptor.pre_llm_request(&mut ctx).await.unwrap();
assert!(matches!(action, PreRequestAction::Cancel(_))); assert!(matches!(action, PreRequestAction::Cancel(_)));
assert!(first_called.load(Ordering::Relaxed)); assert!(first_called.load(Ordering::Relaxed));
+4
View File
@@ -3733,6 +3733,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
| EngineRunExit::Yielded | EngineRunExit::Yielded
| EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) | EngineRunExit::Interrupted(RunInterruptionReason::Cancelled)
| EngineRunExit::Interrupted(RunInterruptionReason::ContextWindowExceeded) | EngineRunExit::Interrupted(RunInterruptionReason::ContextWindowExceeded)
| EngineRunExit::Interrupted(RunInterruptionReason::Interceptor(_))
| EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(_)) | EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(_))
); );
let active_run_turn_count = self.engine.as_ref().unwrap().active_run_turn_count(); let active_run_turn_count = self.engine.as_ref().unwrap().active_run_turn_count();
@@ -6129,12 +6130,14 @@ fn run_interruption_reason_error_code(reason: &RunInterruptionReason) -> ErrorCo
RunInterruptionReason::Unexpected(EngineError::Tool(_)) => ErrorCode::ToolError, RunInterruptionReason::Unexpected(EngineError::Tool(_)) => ErrorCode::ToolError,
RunInterruptionReason::LimitReached RunInterruptionReason::LimitReached
| RunInterruptionReason::Cancelled | RunInterruptionReason::Cancelled
| RunInterruptionReason::Interceptor(_)
| RunInterruptionReason::Unexpected( | RunInterruptionReason::Unexpected(
EngineError::Aborted(_) EngineError::Aborted(_)
| EngineError::Cancelled | EngineError::Cancelled
| EngineError::PauseRequested | EngineError::PauseRequested
| EngineError::ConfigWarnings(_) | EngineError::ConfigWarnings(_)
| EngineError::HistoryAppend(_) | EngineError::HistoryAppend(_)
| EngineError::Interceptor(_)
| EngineError::ToolAttemptFence(_), | EngineError::ToolAttemptFence(_),
) => ErrorCode::Internal, ) => ErrorCode::Internal,
} }
@@ -6145,6 +6148,7 @@ fn run_interruption_reason_message(reason: &RunInterruptionReason) -> String {
RunInterruptionReason::LimitReached => "engine turn limit reached".to_string(), RunInterruptionReason::LimitReached => "engine turn limit reached".to_string(),
RunInterruptionReason::ContextWindowExceeded => "model context window reached".to_string(), RunInterruptionReason::ContextWindowExceeded => "model context window reached".to_string(),
RunInterruptionReason::Cancelled => "engine run cancelled".to_string(), RunInterruptionReason::Cancelled => "engine run cancelled".to_string(),
RunInterruptionReason::Interceptor(failure) => failure.to_string(),
RunInterruptionReason::Unexpected(error) => format!("unexpected engine failure: {error}"), RunInterruptionReason::Unexpected(error) => format!("unexpected engine failure: {error}"),
} }
} }