interceptorの修正
This commit is contained in:
@@ -1,76 +0,0 @@
|
||||
//! CompactInterceptor — wraps HookInterceptor with urgent compaction check.
|
||||
//!
|
||||
//! Decorator that delegates all [`Interceptor`] methods to the inner
|
||||
//! `HookInterceptor`, then adds a token-count check in `pre_llm_request`.
|
||||
//! When `last_input_tokens` exceeds the turn threshold, returns
|
||||
//! `PreRequestAction::Yield` so the Worker exits the turn loop cleanly
|
||||
//! with `WorkerResult::Yielded` and Pod can perform compaction.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_worker::Item;
|
||||
use llm_worker::interceptor::{
|
||||
Interceptor, PostToolAction, PreRequestAction, PreToolAction, PromptAction, ToolCallInfo,
|
||||
ToolResultInfo, TurnEndAction,
|
||||
};
|
||||
use tracing::info;
|
||||
|
||||
use crate::compact_state::CompactState;
|
||||
use crate::hook_interceptor::HookInterceptor;
|
||||
|
||||
/// Interceptor that wraps HookInterceptor and adds between-turns
|
||||
/// compaction threshold check.
|
||||
pub(crate) struct CompactInterceptor {
|
||||
inner: HookInterceptor,
|
||||
state: Arc<CompactState>,
|
||||
}
|
||||
|
||||
impl CompactInterceptor {
|
||||
pub(crate) fn new(inner: HookInterceptor, state: Arc<CompactState>) -> Self {
|
||||
Self { inner, state }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Interceptor for CompactInterceptor {
|
||||
async fn on_prompt_submit(&self, item: &mut Item) -> PromptAction {
|
||||
self.inner.on_prompt_submit(item).await
|
||||
}
|
||||
|
||||
async fn pre_llm_request(&self, context: &mut Vec<Item>) -> PreRequestAction {
|
||||
// Step 1: Delegate to inner hooks first.
|
||||
let inner_action = self.inner.pre_llm_request(context).await;
|
||||
if !matches!(inner_action, PreRequestAction::Continue) {
|
||||
return inner_action;
|
||||
}
|
||||
|
||||
// Step 2: Check between-turns compaction threshold.
|
||||
if !self.state.is_disabled() && self.state.exceeds_turn() {
|
||||
info!(
|
||||
input_tokens = self.state.last_input_tokens(),
|
||||
threshold = self.state.turn_threshold(),
|
||||
"Between-turns compaction threshold exceeded, yielding"
|
||||
);
|
||||
return PreRequestAction::Yield;
|
||||
}
|
||||
|
||||
PreRequestAction::Continue
|
||||
}
|
||||
|
||||
async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction {
|
||||
self.inner.pre_tool_call(info).await
|
||||
}
|
||||
|
||||
async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction {
|
||||
self.inner.post_tool_call(info).await
|
||||
}
|
||||
|
||||
async fn on_turn_end(&self, history: &[Item]) -> TurnEndAction {
|
||||
self.inner.on_turn_end(history).await
|
||||
}
|
||||
|
||||
async fn on_abort(&self, reason: &str) {
|
||||
self.inner.on_abort(reason).await;
|
||||
}
|
||||
}
|
||||
+100
-25
@@ -1,18 +1,95 @@
|
||||
//! Pod-layer hook infrastructure
|
||||
//!
|
||||
//! Provides the `Hook<E>` trait and `HookRegistry` for orchestration hooks
|
||||
//! that govern control-flow decisions in the Worker execution loop.
|
||||
//! Hooks are the **public** orchestration extension point. They receive
|
||||
//! read-only summary information about each event in the Worker
|
||||
//! execution loop and return a control-flow action
|
||||
//! (continue / skip / abort / pause).
|
||||
//!
|
||||
//! The type system (`HookEventKind` / `Hook<E>`) mirrors the pattern
|
||||
//! originally in llm-worker, now at the insomnia layer where orchestration
|
||||
//! concerns belong.
|
||||
//! Hooks intentionally cannot mutate the Worker's context, history, tool
|
||||
//! call, or tool result. Internal mechanisms that need such access (e.g.
|
||||
//! compaction, notification injection, output truncation) implement
|
||||
//! `llm_worker::Interceptor` directly inside Pod, never via this trait.
|
||||
//!
|
||||
//! This separation lets Hooks be exposed safely to user-facing
|
||||
//! extension surfaces (scripting, plugins) in the future without
|
||||
//! exposing the underlying mutable state.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_worker::Item;
|
||||
use llm_worker::tool::ToolOutput;
|
||||
use llm_worker::interceptor::{
|
||||
PostToolAction, PreRequestAction, PreToolAction, PromptAction, ToolCallInfo, ToolResultInfo,
|
||||
TurnEndAction,
|
||||
PostToolAction, PreRequestAction, PreToolAction, PromptAction, TurnEndAction,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
// =============================================================================
|
||||
// Hook input summary types (read-only)
|
||||
// =============================================================================
|
||||
|
||||
/// Information passed to `OnPromptSubmit` hooks.
|
||||
pub struct PromptSubmitInfo {
|
||||
/// Concatenated text content of the user's input message.
|
||||
pub input_text: String,
|
||||
/// 0-based turn index this prompt opens.
|
||||
pub turn_index: usize,
|
||||
}
|
||||
|
||||
/// Information passed to `PreLlmRequest` hooks.
|
||||
pub struct PreRequestInfo {
|
||||
/// Number of items currently in the Worker context.
|
||||
pub item_count: usize,
|
||||
/// Most recently observed `input_tokens` from the LLM provider.
|
||||
/// `None` when the Pod has no compaction state attached, or when
|
||||
/// no LLM call has completed yet.
|
||||
pub estimated_tokens: Option<u64>,
|
||||
/// Current turn index (0-based).
|
||||
pub turn_index: usize,
|
||||
/// Tool calls already executed in this turn.
|
||||
pub tool_calls_this_turn: usize,
|
||||
}
|
||||
|
||||
/// Information passed to `PreToolCall` hooks.
|
||||
pub struct ToolCallSummary {
|
||||
/// Provider-assigned tool call id.
|
||||
pub call_id: String,
|
||||
/// Registered tool name.
|
||||
pub tool_name: String,
|
||||
/// Tool arguments as a JSON value (cloned).
|
||||
///
|
||||
/// LLM-generated arguments are bounded by max_tokens, so cloning
|
||||
/// is cheap relative to tool execution. Structural access is
|
||||
/// required for permission decisions (e.g. inspecting a `path`
|
||||
/// field), which a stringified preview would not support.
|
||||
pub arguments: Value,
|
||||
}
|
||||
|
||||
/// Information passed to `PostToolCall` hooks.
|
||||
pub struct ToolResultSummary {
|
||||
/// Provider-assigned tool call id this result corresponds to.
|
||||
pub call_id: String,
|
||||
/// Registered tool name.
|
||||
pub tool_name: String,
|
||||
/// Whether the tool reported an error.
|
||||
pub is_error: bool,
|
||||
/// Tool output (`summary` always present, `content` may be `None`).
|
||||
pub output: ToolOutput,
|
||||
}
|
||||
|
||||
/// Information passed to `OnTurnEnd` hooks.
|
||||
pub struct TurnEndInfo {
|
||||
/// Turn that just ended (0-based).
|
||||
pub turn_index: usize,
|
||||
/// Tool calls executed in this turn.
|
||||
pub tool_calls_count: usize,
|
||||
/// Preview of the assistant's final text in this turn.
|
||||
/// Truncated at a UTF-8 boundary; empty when no assistant text exists.
|
||||
pub final_text_preview: String,
|
||||
}
|
||||
|
||||
/// Information passed to `OnAbort` hooks.
|
||||
pub struct AbortInfo {
|
||||
/// Reason supplied by the aborter.
|
||||
pub reason: String,
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Hook Event Kinds
|
||||
@@ -20,17 +97,15 @@ use llm_worker::interceptor::{
|
||||
|
||||
/// Marker trait for hook event kinds.
|
||||
///
|
||||
/// Each event kind specifies its input (passed mutably to hooks) and
|
||||
/// output (the control-flow action returned by hooks).
|
||||
/// Each event kind specifies its read-only input and the control-flow
|
||||
/// action returned by hooks.
|
||||
pub trait HookEventKind: Send + Sync + 'static {
|
||||
/// Mutable input passed to the hook.
|
||||
type Input;
|
||||
/// Read-only input passed to the hook.
|
||||
type Input: Send + Sync;
|
||||
/// Control-flow action returned by the hook.
|
||||
type Output;
|
||||
}
|
||||
|
||||
// --- Event kind markers ---
|
||||
|
||||
/// After receiving user input, before adding to history.
|
||||
pub struct OnPromptSubmit;
|
||||
/// Before each LLM request.
|
||||
@@ -45,32 +120,32 @@ pub struct OnTurnEnd;
|
||||
pub struct OnAbort;
|
||||
|
||||
impl HookEventKind for OnPromptSubmit {
|
||||
type Input = Item;
|
||||
type Input = PromptSubmitInfo;
|
||||
type Output = PromptAction;
|
||||
}
|
||||
|
||||
impl HookEventKind for PreLlmRequest {
|
||||
type Input = Vec<Item>;
|
||||
type Input = PreRequestInfo;
|
||||
type Output = PreRequestAction;
|
||||
}
|
||||
|
||||
impl HookEventKind for PreToolCall {
|
||||
type Input = ToolCallInfo;
|
||||
type Input = ToolCallSummary;
|
||||
type Output = PreToolAction;
|
||||
}
|
||||
|
||||
impl HookEventKind for PostToolCall {
|
||||
type Input = ToolResultInfo;
|
||||
type Input = ToolResultSummary;
|
||||
type Output = PostToolAction;
|
||||
}
|
||||
|
||||
impl HookEventKind for OnTurnEnd {
|
||||
type Input = Vec<Item>;
|
||||
type Input = TurnEndInfo;
|
||||
type Output = TurnEndAction;
|
||||
}
|
||||
|
||||
impl HookEventKind for OnAbort {
|
||||
type Input = String;
|
||||
type Input = AbortInfo;
|
||||
type Output = ();
|
||||
}
|
||||
|
||||
@@ -80,13 +155,13 @@ impl HookEventKind for OnAbort {
|
||||
|
||||
/// Async hook for a specific event kind.
|
||||
///
|
||||
/// Hooks receive mutable access to the event's input and return a
|
||||
/// control-flow action. Multiple hooks can be registered per event;
|
||||
/// they are evaluated in registration order and short-circuit on the
|
||||
/// first non-Continue result.
|
||||
/// Hooks receive a shared reference to the event's read-only input
|
||||
/// and return a control-flow action. Multiple hooks can be registered
|
||||
/// per event; they are evaluated in registration order and
|
||||
/// short-circuit on the first non-Continue (or non-Finish) result.
|
||||
#[async_trait]
|
||||
pub trait Hook<E: HookEventKind>: Send + Sync {
|
||||
async fn call(&self, input: &mut E::Input) -> E::Output;
|
||||
async fn call(&self, input: &E::Input) -> E::Output;
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
|
||||
@@ -1,87 +0,0 @@
|
||||
//! HookInterceptor — bridges Pod-layer hooks to Worker's Interceptor trait.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_worker::Item;
|
||||
use llm_worker::interceptor::{
|
||||
Interceptor, PostToolAction, PreRequestAction, PreToolAction, PromptAction, ToolCallInfo,
|
||||
ToolResultInfo, TurnEndAction,
|
||||
};
|
||||
|
||||
use crate::hook::HookRegistry;
|
||||
|
||||
/// An `Interceptor` implementation that delegates to a `HookRegistry`.
|
||||
///
|
||||
/// Each method iterates the registered hooks in order and short-circuits
|
||||
/// on the first non-Continue (or non-Finish) result.
|
||||
pub(crate) struct HookInterceptor {
|
||||
registry: Arc<HookRegistry>,
|
||||
}
|
||||
|
||||
impl HookInterceptor {
|
||||
pub(crate) fn new(registry: Arc<HookRegistry>) -> Self {
|
||||
Self { registry }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Interceptor for HookInterceptor {
|
||||
async fn on_prompt_submit(&self, item: &mut Item) -> PromptAction {
|
||||
for hook in &self.registry.on_prompt_submit {
|
||||
let action = hook.call(item).await;
|
||||
if !matches!(action, PromptAction::Continue) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
PromptAction::Continue
|
||||
}
|
||||
|
||||
async fn pre_llm_request(&self, context: &mut Vec<Item>) -> PreRequestAction {
|
||||
for hook in &self.registry.pre_llm_request {
|
||||
let action = hook.call(context).await;
|
||||
if !matches!(action, PreRequestAction::Continue) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
PreRequestAction::Continue
|
||||
}
|
||||
|
||||
async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction {
|
||||
for hook in &self.registry.pre_tool_call {
|
||||
let action = hook.call(info).await;
|
||||
if !matches!(action, PreToolAction::Continue) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
PreToolAction::Continue
|
||||
}
|
||||
|
||||
async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction {
|
||||
for hook in &self.registry.post_tool_call {
|
||||
let action = hook.call(info).await;
|
||||
if !matches!(action, PostToolAction::Continue) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
PostToolAction::Continue
|
||||
}
|
||||
|
||||
async fn on_turn_end(&self, history: &[Item]) -> TurnEndAction {
|
||||
let mut history_vec = history.to_vec();
|
||||
for hook in &self.registry.on_turn_end {
|
||||
let action = hook.call(&mut history_vec).await;
|
||||
if !matches!(action, TurnEndAction::Finish) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
TurnEndAction::Finish
|
||||
}
|
||||
|
||||
async fn on_abort(&self, reason: &str) {
|
||||
let mut reason_string = reason.to_string();
|
||||
for hook in &self.registry.on_abort {
|
||||
hook.call(&mut reason_string).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -6,11 +6,10 @@ pub mod shared_state;
|
||||
pub mod socket_server;
|
||||
|
||||
mod agents_md;
|
||||
mod compact_interceptor;
|
||||
mod compact_state;
|
||||
mod factory;
|
||||
mod hook_interceptor;
|
||||
mod pod;
|
||||
mod pod_interceptor;
|
||||
mod prompt_loader;
|
||||
mod prune;
|
||||
mod system_prompt;
|
||||
|
||||
+14
-22
@@ -14,14 +14,13 @@ use tracing::{info, warn};
|
||||
use manifest::{PodManifest, PodManifestConfig, ResolveError, Scope, ScopeError, WorkerManifest};
|
||||
|
||||
use crate::agents_md::read_agents_md;
|
||||
use crate::compact_interceptor::CompactInterceptor;
|
||||
use crate::compact_state::CompactState;
|
||||
use crate::hook::{
|
||||
Hook, HookRegistryBuilder, OnAbort, OnPromptSubmit, OnTurnEnd, PostToolCall, PreLlmRequest,
|
||||
PreToolCall,
|
||||
PreRequestInfo, PreToolCall,
|
||||
};
|
||||
use crate::hook_interceptor::HookInterceptor;
|
||||
use crate::notifier::Notifier;
|
||||
use crate::pod_interceptor::PodInterceptor;
|
||||
use crate::prompt_loader::PromptLoader;
|
||||
use crate::system_prompt::{SystemPromptContext, SystemPromptError, SystemPromptTemplate};
|
||||
use crate::usage_tracker::UsageTracker;
|
||||
@@ -38,8 +37,8 @@ struct UsageTrackingHook {
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreLlmRequest> for UsageTrackingHook {
|
||||
async fn call(&self, context: &mut Vec<Item>) -> PreRequestAction {
|
||||
self.tracker.note_request(context.len());
|
||||
async fn call(&self, info: &PreRequestInfo) -> PreRequestAction {
|
||||
self.tracker.note_request(info.item_count);
|
||||
PreRequestAction::Continue
|
||||
}
|
||||
}
|
||||
@@ -346,16 +345,15 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
|
||||
/// `on_usage` callback to track `input_tokens`.
|
||||
fn ensure_interceptor_installed(&mut self) {
|
||||
if !self.interceptor_installed {
|
||||
// Pre-LLM-request hook: capture history.len() into the
|
||||
// UsageTracker so the upcoming on_usage callback can pair
|
||||
// it with the measured input_tokens.
|
||||
// Pre-LLM-request hook: record the item count at send time
|
||||
// so the on_usage callback can pair it with the measured
|
||||
// input_tokens.
|
||||
self.hook_builder.add_pre_llm_request(UsageTrackingHook {
|
||||
tracker: self.usage_tracker.clone(),
|
||||
});
|
||||
|
||||
let builder = std::mem::take(&mut self.hook_builder);
|
||||
let registry = Arc::new(builder.build());
|
||||
let hook_interceptor = HookInterceptor::new(registry);
|
||||
|
||||
let compact_threshold = self
|
||||
.manifest
|
||||
@@ -363,12 +361,9 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
|
||||
.as_ref()
|
||||
.and_then(|c| c.compact_threshold);
|
||||
|
||||
// Usage tracking via on_usage callback. Independent of
|
||||
// compact_threshold so that LlmUsage entries are persisted
|
||||
// unconditionally.
|
||||
let tracker_for_usage = self.usage_tracker.clone();
|
||||
|
||||
if let Some(threshold) = compact_threshold {
|
||||
let compact_state = if let Some(threshold) = compact_threshold {
|
||||
let retained = self
|
||||
.manifest
|
||||
.compaction
|
||||
@@ -377,9 +372,6 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
|
||||
.unwrap_or(2);
|
||||
|
||||
let state = Arc::new(CompactState::new(threshold, retained));
|
||||
|
||||
// Combined on_usage: feed both the legacy compact threshold
|
||||
// tracker and the new UsageTracker.
|
||||
let state_for_usage = state.clone();
|
||||
self.worker_mut().on_usage(move |event| {
|
||||
if let Some(tokens) = event.input_tokens {
|
||||
@@ -387,17 +379,17 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
|
||||
}
|
||||
tracker_for_usage.record_usage(event);
|
||||
});
|
||||
|
||||
let interceptor = CompactInterceptor::new(hook_interceptor, state.clone());
|
||||
self.worker_mut().set_interceptor(interceptor);
|
||||
self.compact_state = Some(state);
|
||||
self.compact_state = Some(state.clone());
|
||||
Some(state)
|
||||
} else {
|
||||
self.worker_mut().on_usage(move |event| {
|
||||
tracker_for_usage.record_usage(event);
|
||||
});
|
||||
self.worker_mut().set_interceptor(hook_interceptor);
|
||||
}
|
||||
None
|
||||
};
|
||||
|
||||
let interceptor = PodInterceptor::new(registry, compact_state);
|
||||
self.worker_mut().set_interceptor(interceptor);
|
||||
self.interceptor_installed = true;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,294 @@
|
||||
//! Pod-owned `Interceptor` implementation.
|
||||
//!
|
||||
//! Bridges Pod's internal mechanisms (compaction trigger today;
|
||||
//! notification injection / output truncation in the future) and the
|
||||
//! public `HookRegistry`. Internal mechanisms run first and have full
|
||||
//! mutable access via the `Interceptor` trait. Hooks then receive
|
||||
//! read-only summary information and only return control-flow
|
||||
//! decisions (continue / skip / abort / pause).
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_worker::Item;
|
||||
use llm_worker::interceptor::{
|
||||
Interceptor, PostToolAction, PreRequestAction, PreToolAction, PromptAction, ToolCallInfo,
|
||||
ToolResultInfo, TurnEndAction,
|
||||
};
|
||||
use llm_worker::tool::ToolOutput;
|
||||
use tracing::info;
|
||||
|
||||
use crate::compact_state::CompactState;
|
||||
use crate::hook::{
|
||||
AbortInfo, HookRegistry, PreRequestInfo, PromptSubmitInfo, ToolCallSummary, ToolResultSummary,
|
||||
TurnEndInfo,
|
||||
};
|
||||
|
||||
/// Maximum number of bytes copied into `TurnEndInfo::final_text_preview`.
|
||||
const FINAL_TEXT_PREVIEW_LIMIT: usize = 512;
|
||||
|
||||
pub(crate) struct PodInterceptor {
|
||||
registry: Arc<HookRegistry>,
|
||||
compact_state: Option<Arc<CompactState>>,
|
||||
/// Next turn index assigned by `on_prompt_submit`.
|
||||
next_turn_index: AtomicUsize,
|
||||
/// Tool calls observed in the current turn (reset on each new prompt).
|
||||
tool_calls_this_turn: AtomicUsize,
|
||||
}
|
||||
|
||||
impl PodInterceptor {
|
||||
pub(crate) fn new(
|
||||
registry: Arc<HookRegistry>,
|
||||
compact_state: Option<Arc<CompactState>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
registry,
|
||||
compact_state,
|
||||
next_turn_index: AtomicUsize::new(0),
|
||||
tool_calls_this_turn: AtomicUsize::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
fn current_turn_index(&self) -> usize {
|
||||
self.next_turn_index
|
||||
.load(Ordering::Relaxed)
|
||||
.saturating_sub(1)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Interceptor for PodInterceptor {
|
||||
async fn on_prompt_submit(&self, item: &mut Item) -> PromptAction {
|
||||
let turn_index = self.next_turn_index.fetch_add(1, Ordering::Relaxed);
|
||||
self.tool_calls_this_turn.store(0, Ordering::Relaxed);
|
||||
|
||||
let info = PromptSubmitInfo {
|
||||
input_text: extract_message_text(item).unwrap_or_default(),
|
||||
turn_index,
|
||||
};
|
||||
for hook in &self.registry.on_prompt_submit {
|
||||
let action = hook.call(&info).await;
|
||||
if !matches!(action, PromptAction::Continue) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
PromptAction::Continue
|
||||
}
|
||||
|
||||
async fn pre_llm_request(&self, context: &mut Vec<Item>) -> PreRequestAction {
|
||||
// Internal mechanism: between-turns compaction trigger.
|
||||
if let Some(state) = self.compact_state.as_ref() {
|
||||
if !state.is_disabled() && state.exceeds_turn() {
|
||||
info!(
|
||||
input_tokens = state.last_input_tokens(),
|
||||
threshold = state.turn_threshold(),
|
||||
"Between-turns compaction threshold exceeded, yielding"
|
||||
);
|
||||
return PreRequestAction::Yield;
|
||||
}
|
||||
}
|
||||
|
||||
let info = PreRequestInfo {
|
||||
item_count: context.len(),
|
||||
estimated_tokens: self.compact_state.as_ref().map(|s| s.last_input_tokens()),
|
||||
turn_index: self.current_turn_index(),
|
||||
tool_calls_this_turn: self.tool_calls_this_turn.load(Ordering::Relaxed),
|
||||
};
|
||||
for hook in &self.registry.pre_llm_request {
|
||||
let action = hook.call(&info).await;
|
||||
if !matches!(action, PreRequestAction::Continue) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
PreRequestAction::Continue
|
||||
}
|
||||
|
||||
async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction {
|
||||
let summary = ToolCallSummary {
|
||||
call_id: info.call.id.clone(),
|
||||
tool_name: info.call.name.clone(),
|
||||
arguments: info.call.input.clone(),
|
||||
};
|
||||
for hook in &self.registry.pre_tool_call {
|
||||
let action = hook.call(&summary).await;
|
||||
if !matches!(action, PreToolAction::Continue) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
self.tool_calls_this_turn.fetch_add(1, Ordering::Relaxed);
|
||||
PreToolAction::Continue
|
||||
}
|
||||
|
||||
async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction {
|
||||
let summary = ToolResultSummary {
|
||||
call_id: info.result.tool_use_id.clone(),
|
||||
tool_name: info.call.name.clone(),
|
||||
is_error: info.result.is_error,
|
||||
output: ToolOutput {
|
||||
summary: info.result.summary.clone(),
|
||||
content: info.result.content.clone(),
|
||||
},
|
||||
};
|
||||
for hook in &self.registry.post_tool_call {
|
||||
let action = hook.call(&summary).await;
|
||||
if !matches!(action, PostToolAction::Continue) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
PostToolAction::Continue
|
||||
}
|
||||
|
||||
async fn on_turn_end(&self, history: &[Item]) -> TurnEndAction {
|
||||
let final_text_preview = history
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|i| i.is_assistant_message())
|
||||
.and_then(extract_message_text)
|
||||
.map(|t| preview(&t, FINAL_TEXT_PREVIEW_LIMIT))
|
||||
.unwrap_or_default();
|
||||
let info = TurnEndInfo {
|
||||
turn_index: self.current_turn_index(),
|
||||
tool_calls_count: self.tool_calls_this_turn.load(Ordering::Relaxed),
|
||||
final_text_preview,
|
||||
};
|
||||
for hook in &self.registry.on_turn_end {
|
||||
let action = hook.call(&info).await;
|
||||
if !matches!(action, TurnEndAction::Finish) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
TurnEndAction::Finish
|
||||
}
|
||||
|
||||
async fn on_abort(&self, reason: &str) {
|
||||
let info = AbortInfo {
|
||||
reason: reason.to_string(),
|
||||
};
|
||||
for hook in &self.registry.on_abort {
|
||||
hook.call(&info).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_message_text(item: &Item) -> Option<String> {
|
||||
match item {
|
||||
Item::Message { content, .. } => Some(
|
||||
content
|
||||
.iter()
|
||||
.map(|p| p.as_text())
|
||||
.collect::<Vec<_>>()
|
||||
.join(""),
|
||||
),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn preview(text: &str, limit: usize) -> String {
|
||||
if text.len() <= limit {
|
||||
return text.to_string();
|
||||
}
|
||||
let mut end = limit;
|
||||
while end > 0 && !text.is_char_boundary(end) {
|
||||
end -= 1;
|
||||
}
|
||||
text[..end].to_string()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize};
|
||||
|
||||
use super::*;
|
||||
use crate::hook::{Hook, HookRegistryBuilder, PreLlmRequest};
|
||||
|
||||
struct CountingHook(Arc<AtomicUsize>);
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreLlmRequest> for CountingHook {
|
||||
async fn call(&self, _info: &PreRequestInfo) -> PreRequestAction {
|
||||
self.0.fetch_add(1, Ordering::Relaxed);
|
||||
PreRequestAction::Continue
|
||||
}
|
||||
}
|
||||
|
||||
fn registry_with_pre_llm_hook(counter: Arc<AtomicUsize>) -> Arc<HookRegistry> {
|
||||
let mut builder = HookRegistryBuilder::new();
|
||||
builder.add_pre_llm_request(CountingHook(counter));
|
||||
Arc::new(builder.build())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_llm_request_yields_and_skips_hooks_when_compact_threshold_exceeded() {
|
||||
let count = Arc::new(AtomicUsize::new(0));
|
||||
let registry = registry_with_pre_llm_hook(count.clone());
|
||||
|
||||
let state = Arc::new(CompactState::new(100, 2));
|
||||
state.update_input_tokens(200); // exceeds turn threshold
|
||||
|
||||
let interceptor = PodInterceptor::new(registry, Some(state));
|
||||
let mut ctx: Vec<Item> = vec![Item::user_message("hi")];
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Yield));
|
||||
// Hook must not run when an internal mechanism short-circuits first.
|
||||
assert_eq!(count.load(Ordering::Relaxed), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_llm_request_runs_hooks_when_under_threshold() {
|
||||
let count = Arc::new(AtomicUsize::new(0));
|
||||
let registry = registry_with_pre_llm_hook(count.clone());
|
||||
|
||||
let state = Arc::new(CompactState::new(100, 2));
|
||||
// last_input_tokens stays at 0, well below threshold.
|
||||
|
||||
let interceptor = PodInterceptor::new(registry, Some(state));
|
||||
let mut ctx: Vec<Item> = vec![Item::user_message("hi")];
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Continue));
|
||||
assert_eq!(count.load(Ordering::Relaxed), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_llm_request_runs_hooks_when_no_compact_state() {
|
||||
let count = Arc::new(AtomicUsize::new(0));
|
||||
let registry = registry_with_pre_llm_hook(count.clone());
|
||||
|
||||
let interceptor = PodInterceptor::new(registry, None);
|
||||
let mut ctx: Vec<Item> = Vec::new();
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Continue));
|
||||
assert_eq!(count.load(Ordering::Relaxed), 1);
|
||||
}
|
||||
|
||||
struct AbortingHook(Arc<AtomicBool>);
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreLlmRequest> for AbortingHook {
|
||||
async fn call(&self, _info: &PreRequestInfo) -> PreRequestAction {
|
||||
self.0.store(true, Ordering::Relaxed);
|
||||
PreRequestAction::Cancel("nope".into())
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_llm_request_short_circuits_on_first_non_continue() {
|
||||
let first_called = Arc::new(AtomicBool::new(false));
|
||||
let second_count = Arc::new(AtomicUsize::new(0));
|
||||
let mut builder = HookRegistryBuilder::new();
|
||||
builder.add_pre_llm_request(AbortingHook(first_called.clone()));
|
||||
builder.add_pre_llm_request(CountingHook(second_count.clone()));
|
||||
let registry = Arc::new(builder.build());
|
||||
|
||||
let interceptor = PodInterceptor::new(registry, None);
|
||||
let mut ctx: Vec<Item> = Vec::new();
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Cancel(_)));
|
||||
assert!(first_called.load(Ordering::Relaxed));
|
||||
assert_eq!(second_count.load(Ordering::Relaxed), 0);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user