diff --git a/crates/worker/src/feature.rs b/crates/worker/src/feature.rs index c13bb603..e91d44a6 100644 --- a/crates/worker/src/feature.rs +++ b/crates/worker/src/feature.rs @@ -24,8 +24,8 @@ use serde::{Deserialize, Serialize}; use thiserror::Error; use crate::hook::{ - BeforeSessionRewrite, Hook, HookFailurePolicy, HookRegistryBuilder, OnPromptSubmit, OnTurnEnd, - PostToolCall, PreLlmRequest, PreToolCall, RunCommitted, RunExit, WorkerStopping, + BeforeSessionRewrite, Hook, HookExecutionPolicy, HookRegistryBuilder, OnPromptSubmit, + OnTurnEnd, PostToolCall, PreLlmRequest, PreToolCall, RunCommitted, RunExit, WorkerStopping, }; use background::{ BackgroundTaskSpec, FeatureBackgroundTask, FeatureBackgroundTaskRegistry, @@ -903,46 +903,6 @@ fn reject_undeclared_contribution( error } -/// Model-visible durable notification sink skeleton. The first slice exposes -/// the boundary without implementing a new event channel. -pub struct FeatureNotificationSink<'a> { - report: &'a mut FeatureInstallReport, -} - -impl FeatureNotificationSink<'_> { - pub fn notify_model(&mut self, message: impl Into) -> Result<(), FeatureInstallError> { - let message = message.into(); - self.report.diagnostics.push(FeatureDiagnostic::warning(format!( - "model notification requested during feature installation but no durable Notify host is attached: {message}" - ))); - self.report.mark_skipped( - FeatureContributionKind::Notification, - "notify_model", - "durable Notify/SystemItem host is not connected during feature installation", - ); - Ok(()) - } -} - -/// Transient human-facing alert sink skeleton. -pub struct FeatureAlertSink<'a> { - report: &'a mut FeatureInstallReport, -} - -impl FeatureAlertSink<'_> { - pub fn alert(&mut self, message: impl Into) { - let message = message.into(); - self.report - .diagnostics - .push(FeatureDiagnostic::info(format!("feature alert: {message}"))); - self.report.mark_skipped( - FeatureContributionKind::Alert, - "alert", - "transient alert host is not connected during feature installation", - ); - } -} - /// Diagnostic sink available to feature installers. pub struct FeatureDiagnosticSink<'a> { report: &'a mut FeatureInstallReport, @@ -1073,16 +1033,18 @@ impl HookContributionRegistrar<'_> { pub fn add_prompt_submit( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::PromptSubmit); self.require_declared(&declaration)?; - self.hook_builder.add_named_on_prompt_submit( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_on_prompt_submit( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1090,16 +1052,18 @@ impl HookContributionRegistrar<'_> { pub fn add_pre_llm_request( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::PreLlmRequest); self.require_declared(&declaration)?; - self.hook_builder.add_named_pre_llm_request( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_pre_llm_request( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1109,22 +1073,24 @@ impl HookContributionRegistrar<'_> { name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - self.add_pre_llm_request(name, HookFailurePolicy::FailClosed, hook) + self.add_pre_llm_request(name, HookExecutionPolicy::fail_closed(), hook) } pub fn add_pre_tool_call_with_policy( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::PreToolCall); self.require_declared(&declaration)?; - self.hook_builder.add_named_pre_tool_call( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_pre_tool_call( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1134,22 +1100,24 @@ impl HookContributionRegistrar<'_> { name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - self.add_pre_tool_call_with_policy(name, HookFailurePolicy::FailClosed, hook) + self.add_pre_tool_call_with_policy(name, HookExecutionPolicy::fail_closed(), hook) } pub fn add_post_tool_call( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::PostToolCall); self.require_declared(&declaration)?; - self.hook_builder.add_named_post_tool_call( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_post_tool_call( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1159,22 +1127,24 @@ impl HookContributionRegistrar<'_> { name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - self.add_post_tool_call(name, HookFailurePolicy::FailClosed, hook) + self.add_post_tool_call(name, HookExecutionPolicy::fail_closed(), hook) } pub fn add_assistant_turn_end( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::AssistantTurnEnd); self.require_declared(&declaration)?; - self.hook_builder.add_named_on_turn_end( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_on_turn_end( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1184,22 +1154,24 @@ impl HookContributionRegistrar<'_> { name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - self.add_assistant_turn_end(name, HookFailurePolicy::FailClosed, hook) + self.add_assistant_turn_end(name, HookExecutionPolicy::fail_closed(), hook) } pub fn add_run_exit( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::RunExit); self.require_declared(&declaration)?; - self.hook_builder.add_named_run_exit( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_run_exit( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1207,16 +1179,18 @@ impl HookContributionRegistrar<'_> { pub fn add_run_committed( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::RunCommitted); self.require_declared(&declaration)?; - self.hook_builder.add_named_run_committed( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_run_committed( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1224,16 +1198,18 @@ impl HookContributionRegistrar<'_> { pub fn add_before_session_rewrite( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::BeforeSessionRewrite); self.require_declared(&declaration)?; - self.hook_builder.add_named_before_session_rewrite( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_before_session_rewrite( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1241,16 +1217,18 @@ impl HookContributionRegistrar<'_> { pub fn add_worker_stopping( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::WorkerStopping); self.require_declared(&declaration)?; - self.hook_builder.add_named_worker_stopping( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_worker_stopping( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1592,18 +1570,6 @@ impl FeatureInstallContext<'_> { } } - pub fn notifications(&mut self) -> FeatureNotificationSink<'_> { - FeatureNotificationSink { - report: self.report, - } - } - - pub fn alerts(&mut self) -> FeatureAlertSink<'_> { - FeatureAlertSink { - report: self.report, - } - } - pub fn diagnostics(&mut self) -> FeatureDiagnosticSink<'_> { FeatureDiagnosticSink { report: self.report, diff --git a/crates/worker/src/feature/background.rs b/crates/worker/src/feature/background.rs index 76254e5e..695e45d3 100644 --- a/crates/worker/src/feature/background.rs +++ b/crates/worker/src/feature/background.rs @@ -12,7 +12,9 @@ use std::time::Duration; use async_trait::async_trait; use serde::{Deserialize, Serialize}; +#[cfg(test)] use tokio::sync::Notify; +use tokio::sync::watch; use tokio::task::JoinHandle; use super::{BackgroundTaskDeclaration, FeatureId, FeatureInstallError}; @@ -22,6 +24,7 @@ const MAX_TASK_CONCURRENCY: u16 = 64; const MAX_TASK_ATTEMPTS: u16 = 16; const MAX_TASK_TIMEOUT_MS: u64 = 24 * 60 * 60 * 1_000; const MAX_RETAINED_DIAGNOSTICS: usize = 128; +const TASK_SETTLE_TIMEOUT_MS: u64 = 30_000; #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] @@ -107,31 +110,43 @@ pub struct BackgroundTaskContext { pub feature_id: FeatureId, pub task_name: String, pub execution_id: u64, + pub session_generation: u64, pub attempt: u16, } -#[derive(Clone, Default)] +#[derive(Clone)] pub struct BackgroundTaskCancellation { - cancelled: Arc, - notify: Arc, + sender: Arc>, +} + +impl Default for BackgroundTaskCancellation { + fn default() -> Self { + let (sender, _receiver) = watch::channel(false); + Self { + sender: Arc::new(sender), + } + } } impl BackgroundTaskCancellation { pub fn is_cancelled(&self) -> bool { - self.cancelled.load(Ordering::Acquire) + *self.sender.borrow() } pub async fn cancelled(&self) { - if self.is_cancelled() { + let mut receiver = self.sender.subscribe(); + if *receiver.borrow_and_update() { return; } - self.notify.notified().await; + while receiver.changed().await.is_ok() { + if *receiver.borrow_and_update() { + return; + } + } } fn cancel(&self) { - if !self.cancelled.swap(true, Ordering::AcqRel) { - self.notify.notify_one(); - } + self.sender.send_replace(true); } } @@ -200,6 +215,7 @@ impl FeatureBackgroundTaskRegistryBuilder { running: Mutex::new(BTreeMap::new()), diagnostics: Mutex::new(Vec::new()), next_execution_id: AtomicU64::new(1), + session_generation: AtomicU64::new(1), accepting: AtomicBool::new(true), }), } @@ -218,6 +234,7 @@ struct RegistryInner { running: Mutex>, diagnostics: Mutex>, next_execution_id: AtomicU64, + session_generation: AtomicU64, accepting: AtomicBool, } @@ -280,6 +297,7 @@ pub enum BackgroundTaskOutcome { TimedOut, Failed(HookError), JoinFailed, + StaleGenerationDiscarded, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -330,6 +348,7 @@ impl FeatureBackgroundTaskRegistry { } let execution_id = self.inner.next_execution_id.fetch_add(1, Ordering::Relaxed); + let session_generation = self.inner.session_generation.load(Ordering::Acquire); let cancellation = BackgroundTaskCancellation::default(); let task_cancellation = cancellation.clone(); let task = Arc::clone(®istration.task); @@ -347,10 +366,17 @@ impl FeatureBackgroundTaskRegistry { task_feature_id.clone(), task_task_name.clone(), execution_id, + session_generation, task_cancellation, ) .await; if let Some(inner) = weak_inner.upgrade() { + let outcome = + if inner.session_generation.load(Ordering::Acquire) == session_generation { + outcome + } else { + BackgroundTaskOutcome::StaleGenerationDiscarded + }; let mut diagnostics = inner.diagnostics.lock().expect("diagnostics poisoned"); diagnostics.push(BackgroundTaskDiagnostic { execution_id, @@ -405,7 +431,9 @@ impl FeatureBackgroundTaskRegistry { } pub async fn before_session_rewrite(&self) -> Result<(), HookError> { - self.settle(false).await + self.settle(false).await?; + self.inner.session_generation.fetch_add(1, Ordering::AcqRel); + Ok(()) } pub async fn shutdown(&self) -> Result<(), HookError> { @@ -469,17 +497,31 @@ impl FeatureBackgroundTaskRegistry { } } + let settle_deadline = + tokio::time::Instant::now() + Duration::from_millis(TASK_SETTLE_TIMEOUT_MS); let mut join_failed = false; - for (execution_id, feature_id, task_name, handle) in waiting { - if handle.await.is_err() { - join_failed = true; + for (execution_id, feature_id, task_name, mut handle) in waiting { + let outcome = match tokio::time::timeout_at(settle_deadline, &mut handle).await { + Ok(Ok(())) => None, + Ok(Err(_)) => { + join_failed = true; + Some(BackgroundTaskOutcome::JoinFailed) + } + Err(_) => { + join_failed = true; + handle.abort(); + let _ = handle.await; + Some(BackgroundTaskOutcome::TimedOut) + } + }; + if let Some(outcome) = outcome { let mut diagnostics = self.inner.diagnostics.lock().expect("diagnostics poisoned"); diagnostics.push(BackgroundTaskDiagnostic { execution_id, feature_id, task_name, attempts: 0, - outcome: BackgroundTaskOutcome::JoinFailed, + outcome, }); if diagnostics.len() > MAX_RETAINED_DIAGNOSTICS { let remove = diagnostics.len() - MAX_RETAINED_DIAGNOSTICS; @@ -490,7 +532,7 @@ impl FeatureBackgroundTaskRegistry { if join_failed { return Err(HookError::new( HookErrorCategory::Internal, - "feature background task join failed", + "feature background task settlement failed or exceeded its host deadline", )); } Ok(()) @@ -504,6 +546,7 @@ async fn execute_task( feature_id: FeatureId, task_name: String, execution_id: u64, + session_generation: u64, cancellation: BackgroundTaskCancellation, ) -> (u16, BackgroundTaskOutcome) { let (max_attempts, delay_ms) = match spec.retry { @@ -513,22 +556,24 @@ async fn execute_task( delay_ms, } => (max_attempts, delay_ms), }; + let deadline = tokio::time::Instant::now() + Duration::from_millis(spec.timeout_ms); for attempt in 1..=max_attempts { if cancellation.is_cancelled() { return (attempt, BackgroundTaskOutcome::Cancelled); } + if tokio::time::Instant::now() >= deadline { + return (attempt, BackgroundTaskOutcome::TimedOut); + } let context = BackgroundTaskContext { invocation: invocation.clone(), feature_id: feature_id.clone(), task_name: task_name.clone(), execution_id, + session_generation, attempt, }; - let result = tokio::time::timeout( - Duration::from_millis(spec.timeout_ms), - task.run(context, cancellation.clone()), - ) - .await; + let result = + tokio::time::timeout_at(deadline, task.run(context, cancellation.clone())).await; match result { Ok(Ok(())) => return (attempt, BackgroundTaskOutcome::Completed), Ok(Err(_error)) if cancellation.is_cancelled() => { @@ -541,8 +586,16 @@ async fn execute_task( Err(_) => return (attempt, BackgroundTaskOutcome::TimedOut), } if delay_ms > 0 { + let retry_at = std::cmp::min( + deadline, + tokio::time::Instant::now() + Duration::from_millis(delay_ms), + ); tokio::select! { - () = tokio::time::sleep(Duration::from_millis(delay_ms)) => {} + () = tokio::time::sleep_until(retry_at) => { + if tokio::time::Instant::now() >= deadline { + return (attempt, BackgroundTaskOutcome::TimedOut); + } + } () = cancellation.cancelled() => { return (attempt, BackgroundTaskOutcome::Cancelled); } @@ -684,6 +737,98 @@ mod tests { ); } + #[tokio::test] + async fn cancellation_observed_when_cancel_races_with_wait_registration() { + for _ in 0..100 { + let cancellation = BackgroundTaskCancellation::default(); + let waiter = cancellation.clone(); + let task = tokio::spawn(async move { + waiter.cancelled().await; + }); + cancellation.cancel(); + tokio::time::timeout(Duration::from_millis(100), task) + .await + .expect("watch-backed cancellation must not lose the transition") + .unwrap(); + } + } + + struct AlwaysFails { + calls: Arc, + } + + #[async_trait] + impl FeatureBackgroundTask for AlwaysFails { + async fn run( + &self, + _context: BackgroundTaskContext, + _cancellation: BackgroundTaskCancellation, + ) -> Result<(), HookError> { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(HookError::new(HookErrorCategory::Dependency, "retry")) + } + } + + #[tokio::test] + async fn retry_delay_consumes_one_total_execution_deadline() { + let feature = FeatureId::builtin("deadline-test"); + let declaration = BackgroundTaskDeclaration::worker_managed("deadline", "deadline"); + let calls = Arc::new(AtomicUsize::new(0)); + let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); + let mut spec = BackgroundTaskSpec::single_flight(declaration, Duration::from_millis(20)); + spec.shutdown = BackgroundTaskShutdownPolicy::Wait; + spec.retry = BackgroundTaskRetryPolicy::Bounded { + max_attempts: 3, + delay_ms: 100, + }; + builder + .register( + feature.clone(), + spec, + AlwaysFails { + calls: Arc::clone(&calls), + }, + ) + .unwrap(); + let registry = builder.build(); + registry.start(&feature, "deadline", invocation()).unwrap(); + registry.shutdown().await.unwrap(); + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!( + registry.diagnostics()[0].outcome, + BackgroundTaskOutcome::TimedOut + ); + } + + #[tokio::test] + async fn rewrite_joins_old_tasks_before_advancing_session_generation() { + let feature = FeatureId::builtin("generation-test"); + let declaration = BackgroundTaskDeclaration::worker_managed("generation", "generation"); + let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); + builder + .register( + feature.clone(), + BackgroundTaskSpec::single_flight(declaration, Duration::from_secs(1)), + WaitForCancellation, + ) + .unwrap(); + let registry = builder.build(); + registry + .start(&feature, "generation", invocation()) + .unwrap(); + assert_eq!(registry.inner.session_generation.load(Ordering::Acquire), 1); + + registry.before_session_rewrite().await.unwrap(); + + assert_eq!(registry.inner.session_generation.load(Ordering::Acquire), 2); + assert_eq!(registry.diagnostics().len(), 1); + assert_eq!( + registry.diagnostics()[0].outcome, + BackgroundTaskOutcome::Cancelled + ); + } + #[tokio::test] async fn block_policy_fences_rewrite_without_detaching_the_task() { let feature = FeatureId::builtin("rewrite-test"); diff --git a/crates/worker/src/hook.rs b/crates/worker/src/hook.rs index b326697a..1bebebfb 100644 --- a/crates/worker/src/hook.rs +++ b/crates/worker/src/hook.rs @@ -18,7 +18,6 @@ use std::ops::Deref; use std::sync::{Arc, Mutex}; -use agen::HistoryEntry; use agen::interceptor::{ PostToolAction, PreRequestAction, PreToolAction, PromptAction, TurnEndAction, }; @@ -29,7 +28,7 @@ use serde_json::Value; use session_store::{SystemItem, SystemReminder}; use thiserror::Error; -use crate::session_history::SessionHistoryMetadata; +use crate::SessionEntryRef; const HOOK_DIAGNOSTIC_MAX_BYTES: usize = 1_024; @@ -79,6 +78,39 @@ pub enum HookFailurePolicy { AttentionRequired, } +const DEFAULT_HOOK_TIMEOUT_MS: u64 = 30_000; +const MAX_HOOK_TIMEOUT_MS: u64 = 120_000; + +/// Host-enforced execution and failure budget for one hook registration. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct HookExecutionPolicy { + pub failure: HookFailurePolicy, + pub timeout_ms: u64, +} + +impl HookExecutionPolicy { + pub const fn new(failure: HookFailurePolicy, timeout_ms: u64) -> Self { + Self { + failure, + timeout_ms, + } + } + + pub const fn fail_closed() -> Self { + Self::new(HookFailurePolicy::FailClosed, DEFAULT_HOOK_TIMEOUT_MS) + } + + fn validate(self) -> Result { + if self.timeout_ms == 0 || self.timeout_ms > MAX_HOOK_TIMEOUT_MS { + return Err(HookError::new( + HookErrorCategory::InvalidInput, + format!("hook timeout_ms must be within 1..={MAX_HOOK_TIMEOUT_MS}"), + )); + } + Ok(self) + } +} + fn bounded_utf8(mut value: String, max_bytes: usize) -> String { if value.len() <= max_bytes { return value; @@ -421,6 +453,16 @@ pub trait Hook: Send + Sync { // Hook Registry // ============================================================================= +/// Bounded identity-only view of the committed history visible at a lifecycle +/// boundary. Hook context never carries history payloads; a Feature that needs +/// content must use a separately granted session-exploration service. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct HookHistoryRange { + pub first_entry_ref: Option, + pub last_entry_ref: Option, + pub entry_count: usize, +} + /// Stable provenance attached to every Worker lifecycle callback. #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct HookInvocationContext { @@ -452,8 +494,7 @@ pub struct RunExitContext { pub struct RunCommittedContext { pub invocation: HookInvocationContext, pub exit: RunCommittedExit, - pub committed_history: Vec>, - pub committed_history_len: usize, + pub committed_history: HookHistoryRange, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -468,8 +509,7 @@ pub enum SessionRewriteKind { pub struct BeforeSessionRewriteContext { pub invocation: HookInvocationContext, pub kind: SessionRewriteKind, - pub current_history: Vec>, - pub current_history_len: usize, + pub current_history: HookHistoryRange, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -511,20 +551,28 @@ impl HookEventKind for WorkerStopping { pub(crate) struct RegisteredHook { owner: String, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: Box>, } impl RegisteredHook { pub(crate) async fn call(&self, input: &E::Input) -> Result { - self.hook - .call(input) - .await - .map_err(|source| HookExecutionError { - owner: self.owner.clone(), - policy: self.policy, - source, - }) + let result = tokio::time::timeout( + std::time::Duration::from_millis(self.policy.timeout_ms), + self.hook.call(input), + ) + .await + .unwrap_or_else(|_| { + Err(HookError::new( + HookErrorCategory::Timeout, + "hook exceeded its host-enforced execution deadline", + )) + }); + result.map_err(|source| HookExecutionError { + owner: self.owner.clone(), + policy: self.policy.failure, + source, + }) } pub(crate) async fn call_optional( @@ -567,20 +615,23 @@ pub struct HookRegistryBuilder { macro_rules! add_hook_methods { ($default:ident, $named:ident, $field:ident, $event:ty) => { pub fn $default(&mut self, hook: impl Hook<$event> + 'static) { - self.$named("worker.host", HookFailurePolicy::FailClosed, hook); + self.$named("worker.host", HookExecutionPolicy::fail_closed(), hook) + .expect("host hook policy is valid"); } pub fn $named( &mut self, owner: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook<$event> + 'static, - ) { + ) -> Result<(), HookError> { + let policy = policy.validate()?; self.$field.push(RegisteredHook { owner: owner.into(), policy, hook: Box::new(hook), }); + Ok(()) } }; } @@ -706,6 +757,17 @@ impl HookRegistry { } } + pub(crate) fn record_chain_timeout(&self, lifecycle: &str) { + self.record_diagnostic(HookExecutionError { + owner: format!("worker.{lifecycle}"), + policy: HookFailurePolicy::AttentionRequired, + source: HookError::new( + HookErrorCategory::Timeout, + "hook chain exceeded its host-enforced lifecycle deadline", + ), + }); + } + pub fn diagnostics(&self) -> Vec { self.diagnostics .lock() @@ -899,28 +961,31 @@ mod tests { call_id: None, }, kind: SessionRewriteKind::Compact, - current_history: Vec::new(), - current_history_len: 8, + current_history: HookHistoryRange::default(), } } #[tokio::test] async fn rewrite_denials_are_resolved_by_owner_not_registration_order() { let mut builder = HookRegistryBuilder::new(); - builder.add_named_before_session_rewrite( - "z-feature", - HookFailurePolicy::FailClosed, - RewriteHook { - action: BeforeSessionRewriteAction::Deny("z denied".into()), - }, - ); - builder.add_named_before_session_rewrite( - "a-feature", - HookFailurePolicy::FailClosed, - RewriteHook { - action: BeforeSessionRewriteAction::Deny("a denied".into()), - }, - ); + builder + .add_named_before_session_rewrite( + "z-feature", + HookExecutionPolicy::fail_closed(), + RewriteHook { + action: BeforeSessionRewriteAction::Deny("z denied".into()), + }, + ) + .unwrap(); + builder + .add_named_before_session_rewrite( + "a-feature", + HookExecutionPolicy::fail_closed(), + RewriteHook { + action: BeforeSessionRewriteAction::Deny("a denied".into()), + }, + ) + .unwrap(); assert_eq!( builder @@ -935,11 +1000,13 @@ mod tests { #[tokio::test] async fn hook_failure_policy_is_applied_at_the_registry_boundary() { let mut fail_open = HookRegistryBuilder::new(); - fail_open.add_named_before_session_rewrite( - "feature", - HookFailurePolicy::FailOpenWithDiagnostic, - FailingRewriteHook, - ); + fail_open + .add_named_before_session_rewrite( + "feature", + HookExecutionPolicy::new(HookFailurePolicy::FailOpenWithDiagnostic, 30_000), + FailingRewriteHook, + ) + .unwrap(); let fail_open = fail_open.build(); assert_eq!( fail_open @@ -951,11 +1018,13 @@ mod tests { assert_eq!(fail_open.diagnostics().len(), 1); let mut fail_closed = HookRegistryBuilder::new(); - fail_closed.add_named_before_session_rewrite( - "feature", - HookFailurePolicy::FailClosed, - FailingRewriteHook, - ); + fail_closed + .add_named_before_session_rewrite( + "feature", + HookExecutionPolicy::fail_closed(), + FailingRewriteHook, + ) + .unwrap(); let error = fail_closed .build() .before_session_rewrite(&rewrite_context()) @@ -964,6 +1033,54 @@ mod tests { assert_eq!(error.source.category, HookErrorCategory::Dependency); } + struct NeverReturns; + + #[async_trait] + impl Hook for NeverReturns { + async fn call( + &self, + _input: &BeforeSessionRewriteContext, + ) -> Result { + std::future::pending().await + } + } + + #[tokio::test] + async fn hook_execution_deadline_cancels_a_non_cooperative_callback() { + let mut builder = HookRegistryBuilder::new(); + builder + .add_named_before_session_rewrite( + "feature", + HookExecutionPolicy::new(HookFailurePolicy::FailClosed, 10), + NeverReturns, + ) + .unwrap(); + let registry = builder.build(); + + let error = tokio::time::timeout( + std::time::Duration::from_millis(100), + registry.before_session_rewrite(&rewrite_context()), + ) + .await + .expect("host hook deadline must terminate the callback") + .unwrap_err(); + assert_eq!(error.source.category, HookErrorCategory::Timeout); + assert_eq!(registry.diagnostics().len(), 1); + } + + #[test] + fn invalid_hook_execution_budget_is_rejected_before_registration() { + let mut builder = HookRegistryBuilder::new(); + let error = builder + .add_named_before_session_rewrite( + "feature", + HookExecutionPolicy::new(HookFailurePolicy::FailClosed, 0), + NeverReturns, + ) + .unwrap_err(); + assert_eq!(error.category, HookErrorCategory::InvalidInput); + } + #[test] fn hook_diagnostics_are_utf8_bounded() { let error = HookError::new(HookErrorCategory::Internal, "界".repeat(1_000)); diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index 0e5c5b40..30bf4e12 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -30,9 +30,9 @@ use crate::compact::usage_tracker::UsageTracker; use session_store::SystemItem; use crate::hook::{ - HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookPromptAction, HookRegistry, - HookTurnEndAction, PreRequestContext, PreRequestInfo, PromptSubmitInfo, SystemItemAppendHandle, - ToolCallSummary, ToolResultSummary, TurnEndInfo, + HookEventKind, HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookPromptAction, + HookRegistry, HookTurnEndAction, PreRequestContext, PreRequestInfo, PromptSubmitInfo, + RegisteredHook, SystemItemAppendHandle, ToolCallSummary, ToolResultSummary, TurnEndInfo, }; use crate::ipc::notify_buffer::{NotifyBuffer, build_system_item_with_provenance}; use crate::prompt::catalog::PromptCatalog; @@ -43,6 +43,32 @@ use agen::token_counter::total_tokens; /// Maximum number of bytes copied into `TurnEndInfo::final_text_preview`. const FINAL_TEXT_PREVIEW_LIMIT: usize = 512; +const INLINE_HOOK_CHAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); + +async fn call_hook_before_deadline( + hook: &RegisteredHook, + input: &E::Input, + deadline: tokio::time::Instant, +) -> Result, InterceptorError> { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + return Err(InterceptorError::new( + InterceptorErrorCategory::Policy, + "Worker hook chain exceeded its host-enforced lifecycle deadline", + )); + } + tokio::time::timeout(remaining, hook.call_optional(input)) + .await + .map_err(|_| { + InterceptorError::new( + InterceptorErrorCategory::Policy, + "Worker hook chain exceeded its host-enforced lifecycle deadline", + ) + })? + .map_err(|error| { + InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) + }) +} pub(crate) struct WorkerInterceptor { registry: Arc, @@ -246,12 +272,10 @@ impl Interceptor for WorkerInterceptor { input_text: extract_message_text(item).unwrap_or_default(), turn_index, }; + let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT; let mut cancellations = Vec::new(); for hook in &self.registry.on_prompt_submit { - let Some(action) = hook.call_optional(&info).await.map_err(|error| { - InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) - })? - else { + let Some(action) = call_hook_before_deadline(hook, &info, deadline).await? else { continue; }; if let HookPromptAction::Cancel(reason) = action { @@ -354,12 +378,11 @@ impl Interceptor for WorkerInterceptor { .as_ref() .map(|_| SystemItemAppendHandle::new(Arc::clone(&pending_hook_system_items))); let hook_context = PreRequestContext::new(info, system_item_sink); + let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT; let mut cancellations = Vec::new(); let mut should_yield = false; for hook in &self.registry.pre_llm_request { - let Some(action) = hook.call_optional(&hook_context).await.map_err(|error| { - InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) - })? + let Some(action) = call_hook_before_deadline(hook, &hook_context, deadline).await? else { continue; }; @@ -431,14 +454,12 @@ impl Interceptor for WorkerInterceptor { tool_name: info.call.name.clone(), arguments: info.call.input.clone(), }; + let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT; let mut aborts = Vec::new(); let mut should_pause = false; let mut denials = Vec::new(); for hook in &self.registry.pre_tool_call { - let Some(action) = hook.call_optional(&summary).await.map_err(|error| { - InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) - })? - else { + let Some(action) = call_hook_before_deadline(hook, &summary, deadline).await? else { continue; }; @@ -479,12 +500,10 @@ impl Interceptor for WorkerInterceptor { attachments: Vec::new(), }, }; + let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT; let mut aborts = Vec::new(); for hook in &self.registry.post_tool_call { - let Some(action) = hook.call_optional(&summary).await.map_err(|error| { - InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) - })? - else { + let Some(action) = call_hook_before_deadline(hook, &summary, deadline).await? else { continue; }; @@ -516,12 +535,10 @@ impl Interceptor for WorkerInterceptor { tool_calls_count: self.tool_calls_this_turn.load(Ordering::Relaxed), final_text_preview, }; + let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT; let mut should_pause = false; for hook in &self.registry.on_turn_end { - let Some(action) = hook.call_optional(&info).await.map_err(|error| { - InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) - })? - else { + let Some(action) = call_hook_before_deadline(hook, &info, deadline).await? else { continue; }; if matches!(action, HookTurnEndAction::Pause) { diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 7db1f37d..d5f7f8d1 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -50,6 +50,7 @@ pub use prompt::system::{SystemPromptContext, SystemPromptError, SystemPromptTem pub use protocol::{ErrorCode, Event, Method, TurnResult, WorkerStatus}; pub use runtime::dir::RuntimeDir; pub use segment_log_sink::SegmentLogSink; +pub use session_capture::SessionEntryRef; pub use session_history::{ SessionHistoryDerivation, SessionHistoryEntryId, SessionHistoryMetadata, WorkerHistoryProvenance, WorkerSubjectSnapshot, diff --git a/crates/worker/src/session_capture.rs b/crates/worker/src/session_capture.rs index 19026fb8..92f14518 100644 --- a/crates/worker/src/session_capture.rs +++ b/crates/worker/src/session_capture.rs @@ -23,14 +23,14 @@ const OVERVIEW_ANCHOR_STRIDE: usize = 8; #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(transparent)] -pub(crate) struct SessionEntryRef(String); +pub struct SessionEntryRef(String); impl SessionEntryRef { - pub(crate) fn from_history_entry_id(entry_id: &crate::SessionHistoryEntryId) -> Self { + pub fn from_history_entry_id(entry_id: &crate::SessionHistoryEntryId) -> Self { Self(format!("E{}", entry_id.0)) } - pub(crate) fn parse(value: &str) -> Option { + pub fn parse(value: &str) -> Option { let suffix = value.strip_prefix('E')?; if suffix.is_empty() || suffix.len() > 64 @@ -43,7 +43,7 @@ impl SessionEntryRef { Some(Self(value.to_string())) } - pub(crate) fn as_str(&self) -> &str { + pub fn as_str(&self) -> &str { &self.0 } diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 582a6a81..6b944d39 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -51,10 +51,10 @@ use crate::feature::{ FeatureRegistryInstallReport, dedupe_instruction_contributions, }; use crate::hook::{ - BeforeSessionRewriteAction, BeforeSessionRewriteContext, Hook, HookInvocationContext, - HookRegistry, HookRegistryBuilder, OnPromptSubmit, OnTurnEnd, PostToolCall, PreLlmRequest, - PreToolCall, RunCommittedContext, RunCommittedExit, RunExitContext, SessionRewriteKind, - WorkerStoppingContext, + BeforeSessionRewriteAction, BeforeSessionRewriteContext, Hook, HookHistoryRange, + HookInvocationContext, HookRegistry, HookRegistryBuilder, OnPromptSubmit, OnTurnEnd, + PostToolCall, PreLlmRequest, PreToolCall, RunCommittedContext, RunCommittedExit, + RunExitContext, SessionRewriteKind, WorkerStoppingContext, }; use crate::in_flight::InFlightEvents; use crate::internal_worker::{ @@ -66,6 +66,7 @@ const COMPACTION_EXTENSION_DOMAIN: &str = "yoi.compaction"; const LARGE_PASTE_INLINE_MAX_BYTES: usize = 32 * 1024; const WORKER_ORCHESTRATION_INSTRUCTION_ID: &str = "worker.orchestration"; const WORKER_ORCHESTRATION_PROMPT_REF: &str = "common.worker_orchestration"; +const FEATURE_HOOK_CHAIN_TIMEOUT: Duration = Duration::from_secs(30); fn hook_run_exit(exit: &EngineRunExit) -> RunCommittedExit { match exit { @@ -1281,8 +1282,20 @@ impl Worker { invocation: self.hook_invocation_context(None), reason: reason.into(), }; - if let Err(error) = hooks.on_worker_stopping(&context).await { - tracing::warn!(error = %error, "worker-stopping hook requires attention"); + match tokio::time::timeout( + FEATURE_HOOK_CHAIN_TIMEOUT, + hooks.on_worker_stopping(&context), + ) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + tracing::warn!(error = %error, "worker-stopping hook requires attention"); + } + Err(_) => { + hooks.record_chain_timeout("worker-stopping"); + tracing::warn!("worker-stopping hook chain exceeded its host deadline"); + } } } if let Err(error) = self.feature_background_tasks.shutdown().await { @@ -1860,6 +1873,19 @@ impl Worker { self.engine.as_ref().expect("worker taken during run") } + fn hook_history_range(&self) -> HookHistoryRange { + let entries = self.session.history().entries(); + HookHistoryRange { + first_entry_ref: entries.first().map(|entry| { + crate::SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id) + }), + last_entry_ref: entries.last().map(|entry| { + crate::SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id) + }), + entry_count: entries.len(), + } + } + fn hook_invocation_context(&self, run_id: Option) -> HookInvocationContext { HookInvocationContext { workspace_id: self @@ -3444,8 +3470,17 @@ impl Worker { exit: hook_exit, history_len: self.session.history().len(), }; - if let Err(error) = hooks.on_run_exit(&context).await { - tracing::warn!(error = %error, "run-exit hook failed; preserving terminal commit"); + match tokio::time::timeout(FEATURE_HOOK_CHAIN_TIMEOUT, hooks.on_run_exit(&context)) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + tracing::warn!(error = %error, "run-exit hook failed; preserving terminal commit"); + } + Err(_) => { + hooks.record_chain_timeout("run-exit"); + tracing::warn!("run-exit hook chain exceeded its host deadline"); + } } } if matches!(&result, EngineRunExit::Interrupted(_)) { @@ -3457,11 +3492,19 @@ impl Worker { let context = RunCommittedContext { invocation: committed_invocation.clone(), exit: hook_exit, - committed_history: self.session.history().entries().to_vec(), - committed_history_len: self.session.history().len(), + committed_history: self.hook_history_range(), }; - if let Err(error) = hooks.on_run_committed(&context).await { - tracing::warn!(error = %error, "run-committed hook requires attention"); + match tokio::time::timeout(FEATURE_HOOK_CHAIN_TIMEOUT, hooks.on_run_committed(&context)) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + tracing::warn!(error = %error, "run-committed hook requires attention"); + } + Err(_) => { + hooks.record_chain_timeout("run-committed"); + tracing::warn!("run-committed hook chain exceeded its host deadline"); + } } } if let Err(error) = self @@ -3659,14 +3702,21 @@ impl Worker { let context = BeforeSessionRewriteContext { invocation: self.hook_invocation_context(None), kind, - current_history: self.session.history().entries().to_vec(), - current_history_len: self.session.history().len(), + current_history: self.hook_history_range(), }; - match hooks - .before_session_rewrite(&context) - .await - .map_err(|error| WorkerError::FeatureLifecycle(error.to_string()))? - { + let action = tokio::time::timeout( + FEATURE_HOOK_CHAIN_TIMEOUT, + hooks.before_session_rewrite(&context), + ) + .await + .map_err(|_| { + hooks.record_chain_timeout("before-session-rewrite"); + WorkerError::FeatureLifecycle( + "session-rewrite hook chain exceeded its host deadline".to_string(), + ) + })? + .map_err(|error| WorkerError::FeatureLifecycle(error.to_string()))?; + match action { BeforeSessionRewriteAction::Continue => {} BeforeSessionRewriteAction::Deny(reason) => { return Err(WorkerError::FeatureLifecycle(reason));