interceptorの修正

This commit is contained in:
2026-04-18 17:19:59 +09:00
parent 79f342ca60
commit 84a8bd099b
13 changed files with 952 additions and 448 deletions
-76
View File
@@ -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
View File
@@ -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;
}
// =============================================================================
-87
View File
@@ -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;
}
}
}
+1 -2
View File
@@ -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
View File
@@ -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;
}
}
+294
View File
@@ -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);
}
}