diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 5d7ae971..5e848f2c 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -1590,7 +1590,9 @@ async fn controller_loop( &working_event_tx, target, expected_head_entries, - ) { + ) + .await + { worker.clear_in_flight_events(); shared_state.set_status(WorkerStatus::Idle); let _ = working_event_tx.send(Event::Status { @@ -1711,6 +1713,10 @@ async fn controller_loop( tracing::warn!(%error, "Worker runtime socket cleanup failed"); } + // Feature callbacks and tasks share the Worker scope. Stop them before + // Memory/Workdir teardown so they cannot observe a partially closed Worker. + worker.stop_feature_runtime("controller shutdown").await; + // Background memory jobs own extract/consolidate workers after a // turn completes. Join them before closing the Workdir session so no // Worker-owned task can outlive its operation attachment. @@ -1993,7 +1999,7 @@ where } } -fn apply_rewind( +async fn apply_rewind( worker: &mut Worker, working_event_tx: &broadcast::Sender, target: RewindTargetId, @@ -2003,7 +2009,7 @@ where C: LlmClient + 'static, St: Store, { - match worker.rewind_to(target, expected_head_entries) { + match worker.rewind_to(target, expected_head_entries).await { Ok(applied) => { let session = session_store::public_snapshot::project_current_session_snapshot(&applied.entries); diff --git a/crates/worker/src/feature.rs b/crates/worker/src/feature.rs index 5698e84b..c13bb603 100644 --- a/crates/worker/src/feature.rs +++ b/crates/worker/src/feature.rs @@ -23,7 +23,14 @@ use agen::tool::ToolDefinition; use serde::{Deserialize, Serialize}; use thiserror::Error; -use crate::hook::{Hook, HookRegistryBuilder, OnTurnEnd, PostToolCall, PreLlmRequest, PreToolCall}; +use crate::hook::{ + BeforeSessionRewrite, Hook, HookFailurePolicy, HookRegistryBuilder, OnPromptSubmit, OnTurnEnd, + PostToolCall, PreLlmRequest, PreToolCall, RunCommitted, RunExit, WorkerStopping, +}; +use background::{ + BackgroundTaskSpec, FeatureBackgroundTask, FeatureBackgroundTaskRegistry, + FeatureBackgroundTaskRegistryBuilder, +}; /// Stable source-qualified identifier for a feature module. #[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)] @@ -253,10 +260,15 @@ pub enum FeatureRuntimeKind { #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum FeatureHookPoint { - PreRequest, + PromptSubmit, + PreLlmRequest, PreToolCall, - ToolResult, - TurnEnd, + PostToolCall, + AssistantTurnEnd, + RunExit, + RunCommitted, + BeforeSessionRewrite, + WorkerStopping, } /// Serializable declaration of a tool contribution. The executable factory is @@ -379,16 +391,17 @@ impl FeatureInstructionContribution { } } -/// Background task lifecycle phase represented by this registry slice. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +/// Background tasks are always Worker-managed and execute inside the owning +/// feature scope. Report-only and detached host-managed declarations are not +/// accepted by the current contract. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum BackgroundTaskLifecycle { - DescriptorOnly, - HostManaged, + WorkerManaged, } -/// Declaration for a feature-provided background task. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +/// Declaration for a feature-provided executable background task. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] pub struct BackgroundTaskDeclaration { pub name: String, pub description: String, @@ -396,11 +409,11 @@ pub struct BackgroundTaskDeclaration { } impl BackgroundTaskDeclaration { - pub fn descriptor_only(name: impl Into, description: impl Into) -> Self { + pub fn worker_managed(name: impl Into, description: impl Into) -> Self { Self { name: name.into(), description: description.into(), - lifecycle: BackgroundTaskLifecycle::DescriptorOnly, + lifecycle: BackgroundTaskLifecycle::WorkerManaged, } } } @@ -772,6 +785,15 @@ impl FeatureInstallReport { } } + fn clear_installed_contributions(&mut self) { + self.installed = false; + self.installed_tools.clear(); + self.installed_hooks.clear(); + self.installed_instructions.clear(); + self.declared_background_tasks.clear(); + self.provided_services.clear(); + } + fn mark_skipped( &mut self, kind: FeatureContributionKind, @@ -1042,15 +1064,68 @@ impl HookContributionRegistrar<'_> { )) } + fn record(&mut self, declaration: HookDeclaration) { + if !self.report.installed_hooks.contains(&declaration) { + self.report.installed_hooks.push(declaration); + } + } + + pub fn add_prompt_submit( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + 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.record(declaration); + Ok(()) + } + + pub fn add_pre_llm_request( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + 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.record(declaration); + Ok(()) + } + pub fn add_pre_request( &mut self, name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - let declaration = HookDeclaration::new(name, FeatureHookPoint::PreRequest); + self.add_pre_llm_request(name, HookFailurePolicy::FailClosed, hook) + } + + pub fn add_pre_tool_call_with_policy( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + hook: impl Hook + 'static, + ) -> Result<(), FeatureInstallError> { + let declaration = HookDeclaration::new(name, FeatureHookPoint::PreToolCall); self.require_declared(&declaration)?; - self.hook_builder.add_pre_llm_request(hook); - self.report.installed_hooks.push(declaration); + self.hook_builder.add_named_pre_tool_call( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ); + self.record(declaration); Ok(()) } @@ -1059,10 +1134,23 @@ impl HookContributionRegistrar<'_> { name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - let declaration = HookDeclaration::new(name, FeatureHookPoint::PreToolCall); + self.add_pre_tool_call_with_policy(name, HookFailurePolicy::FailClosed, hook) + } + + pub fn add_post_tool_call( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + hook: impl Hook + 'static, + ) -> Result<(), FeatureInstallError> { + let declaration = HookDeclaration::new(name, FeatureHookPoint::PostToolCall); self.require_declared(&declaration)?; - self.hook_builder.add_pre_tool_call(hook); - self.report.installed_hooks.push(declaration); + self.hook_builder.add_named_post_tool_call( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ); + self.record(declaration); Ok(()) } @@ -1071,10 +1159,23 @@ impl HookContributionRegistrar<'_> { name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - let declaration = HookDeclaration::new(name, FeatureHookPoint::ToolResult); + self.add_post_tool_call(name, HookFailurePolicy::FailClosed, hook) + } + + pub fn add_assistant_turn_end( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + hook: impl Hook + 'static, + ) -> Result<(), FeatureInstallError> { + let declaration = HookDeclaration::new(name, FeatureHookPoint::AssistantTurnEnd); self.require_declared(&declaration)?; - self.hook_builder.add_post_tool_call(hook); - self.report.installed_hooks.push(declaration); + self.hook_builder.add_named_on_turn_end( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ); + self.record(declaration); Ok(()) } @@ -1083,10 +1184,74 @@ impl HookContributionRegistrar<'_> { name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - let declaration = HookDeclaration::new(name, FeatureHookPoint::TurnEnd); + self.add_assistant_turn_end(name, HookFailurePolicy::FailClosed, hook) + } + + pub fn add_run_exit( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + hook: impl Hook + 'static, + ) -> Result<(), FeatureInstallError> { + let declaration = HookDeclaration::new(name, FeatureHookPoint::RunExit); self.require_declared(&declaration)?; - self.hook_builder.add_on_turn_end(hook); - self.report.installed_hooks.push(declaration); + self.hook_builder.add_named_run_exit( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ); + self.record(declaration); + Ok(()) + } + + pub fn add_run_committed( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + 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.record(declaration); + Ok(()) + } + + pub fn add_before_session_rewrite( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + 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.record(declaration); + Ok(()) + } + + pub fn add_worker_stopping( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + 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.record(declaration); Ok(()) } } @@ -1124,33 +1289,40 @@ impl FeatureInstructionRegistrar<'_> { } } -/// Background task registrar for descriptor/report-only contributions. +/// Registrar for executable, Worker-managed background task contributions. pub struct BackgroundTaskRegistrar<'a> { feature_id: &'a FeatureId, declarations: &'a FeatureContributionDeclarations, + registry: &'a mut FeatureBackgroundTaskRegistryBuilder, report: &'a mut FeatureInstallReport, } impl BackgroundTaskRegistrar<'_> { - pub fn declare( + pub fn register( &mut self, - declaration: BackgroundTaskDeclaration, + spec: BackgroundTaskSpec, + task: impl FeatureBackgroundTask + 'static, ) -> Result<(), FeatureInstallError> { - if !self.declarations.contains_background_task(&declaration) { + if !self + .declarations + .contains_background_task(&spec.declaration) + { return Err(reject_undeclared_contribution( self.feature_id, self.report, FeatureContributionKind::BackgroundTask, - declaration.name, + spec.declaration.name, )); } + self.registry + .register(self.feature_id.clone(), spec.clone(), task)?; if !self .report .declared_background_tasks .iter() - .any(|task| task.name == declaration.name) + .any(|task| task.name == spec.declaration.name) { - self.report.declared_background_tasks.push(declaration); + self.report.declared_background_tasks.push(spec.declaration); } Ok(()) } @@ -1330,15 +1502,17 @@ impl ProtocolProviderRegistrar<'_> { } } - for task in background_tasks { - if !self - .report - .declared_background_tasks - .iter() - .any(|declared| declared.name == task.name) - { - self.report.declared_background_tasks.push(task); - } + if let Some(task) = background_tasks.first() { + let reason = format!( + "protocol provider background task `{}` has no executable Worker-managed handler", + task.name + ); + self.report.mark_skipped( + FeatureContributionKind::BackgroundTask, + task.name.clone(), + reason.clone(), + ); + return Err(FeatureInstallError::InvalidDescriptor(reason)); } Ok(()) @@ -1352,6 +1526,7 @@ pub struct FeatureInstallContext<'a> { pending_tools: &'a mut Vec, installed_tool_names: &'a mut HashMap, hook_builder: &'a mut HookRegistryBuilder, + background_task_builder: &'a mut FeatureBackgroundTaskRegistryBuilder, service_registry: &'a mut FeatureServiceRegistry, report: &'a mut FeatureInstallReport, } @@ -1392,6 +1567,7 @@ impl FeatureInstallContext<'_> { BackgroundTaskRegistrar { feature_id: self.feature_id, declarations: self.declarations, + registry: self.background_task_builder, report: self.report, } } @@ -1440,6 +1616,7 @@ impl FeatureInstallContext<'_> { pub struct FeatureRegistryInstallReport { pub reports: Vec, pub services: FeatureServiceRegistry, + pub background_tasks: FeatureBackgroundTaskRegistry, pub plan_error: Option, } @@ -1861,12 +2038,16 @@ impl FeatureRegistryBuilder { return FeatureRegistryInstallReport { reports, services: FeatureServiceRegistry::default(), + background_tasks: FeatureBackgroundTaskRegistry::default(), plan_error: Some(error), }; } }; let mut service_registry = FeatureServiceRegistry::default(); + let mut background_task_builder = FeatureBackgroundTaskRegistryBuilder::default(); let mut reports = Vec::with_capacity(plan.ordered_indices.len()); + let install_hook_checkpoint = hook_builder.checkpoint(); + let install_tool_checkpoint = pending_tools.len(); let mut modules = self.modules.into_iter().map(Some).collect::>(); let ordered_modules = plan .ordered_indices @@ -1884,6 +2065,11 @@ impl FeatureRegistryBuilder { for (module, descriptor) in ordered_modules { let declarations = FeatureContributionDeclarations::from_descriptor(&descriptor); let mut report = FeatureInstallReport::new(&descriptor); + let hook_checkpoint = hook_builder.checkpoint(); + let background_checkpoint = background_task_builder.checkpoint(); + let service_checkpoint = service_registry.clone(); + let tool_checkpoint = pending_tools.len(); + let installed_tool_checkpoint = installed_tool_names.clone(); let mut required_service_failed = false; for requirement in descriptor.requires_services.iter().cloned() { @@ -1920,10 +2106,6 @@ impl FeatureRegistryBuilder { continue; } - for background_task in descriptor.background_tasks.iter().cloned() { - report.declared_background_tasks.push(background_task); - } - let install_result = { let mut context = FeatureInstallContext { feature_id: &descriptor.id, @@ -1931,6 +2113,7 @@ impl FeatureRegistryBuilder { pending_tools, installed_tool_names: &mut installed_tool_names, hook_builder, + background_task_builder: &mut background_task_builder, service_registry: &mut service_registry, report: &mut report, }; @@ -1940,18 +2123,81 @@ impl FeatureRegistryBuilder { match install_result { Ok(()) => report.installed = true, Err(error) => { + hook_builder.rollback_to(hook_checkpoint); + background_task_builder.rollback_to(&background_checkpoint); + service_registry = service_checkpoint.clone(); + pending_tools.truncate(tool_checkpoint); + installed_tool_names = installed_tool_checkpoint.clone(); + report.clear_installed_contributions(); report .diagnostics .push(FeatureDiagnostic::error(error.to_string())); } } + if report.installed { + for hook in &descriptor.hooks { + if !report.installed_hooks.contains(hook) { + report.diagnostics.push(FeatureDiagnostic::error(format!( + "feature `{}` declared hook `{}` at {:?} but did not register it", + descriptor.id, hook.name, hook.point + ))); + } + } + for task in &descriptor.background_tasks { + if !report.declared_background_tasks.contains(task) { + report.diagnostics.push(FeatureDiagnostic::error(format!( + "feature `{}` declared background task `{}` but did not register an executable handler", + descriptor.id, task.name + ))); + } + } + if report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.severity == FeatureDiagnosticSeverity::Error) + { + hook_builder.rollback_to(hook_checkpoint); + background_task_builder.rollback_to(&background_checkpoint); + service_registry = service_checkpoint.clone(); + pending_tools.truncate(tool_checkpoint); + installed_tool_names = installed_tool_checkpoint.clone(); + report.clear_installed_contributions(); + report.clear_installed_contributions(); + } + } reports.push(report); } - FeatureRegistryInstallReport { - reports, - services: service_registry, - plan_error: None, + let failed = reports.iter().any(|report| { + report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.severity == FeatureDiagnosticSeverity::Error) + }); + if failed { + hook_builder.rollback_to(install_hook_checkpoint); + pending_tools.truncate(install_tool_checkpoint); + for report in &mut reports { + if report.installed { + report.clear_installed_contributions(); + report.diagnostics.push(FeatureDiagnostic::warning( + "feature scope rolled back because another contribution failed", + )); + } + } + FeatureRegistryInstallReport { + reports, + services: FeatureServiceRegistry::default(), + background_tasks: FeatureBackgroundTaskRegistry::default(), + plan_error: None, + } + } else { + FeatureRegistryInstallReport { + reports, + services: service_registry, + background_tasks: background_task_builder.build(), + plan_error: None, + } } } } @@ -1996,6 +2242,7 @@ pub enum FeatureInstallError { Install(String), } +pub mod background; pub mod builtin; pub mod mcp; pub mod plugin; @@ -2398,13 +2645,9 @@ mod tests { } #[test] - fn descriptor_contributions_are_recorded() { + fn executable_contributions_are_recorded() { let descriptor = FeatureDescriptor::builtin("dummy", "Dummy") - .with_tool(ToolDeclaration::new("Dummy", "dummy tool")) - .with_background_task(BackgroundTaskDeclaration::descriptor_only( - "daily", - "descriptor-only background task", - )); + .with_tool(ToolDeclaration::new("Dummy", "dummy tool")); let mut hook_builder = HookRegistryBuilder::default(); let mut pending_tools = Vec::new(); let report = FeatureRegistryBuilder::new() @@ -2420,7 +2663,7 @@ mod tests { let feature_report = &report.reports[0]; assert!(feature_report.installed); assert_eq!(feature_report.installed_tools, vec!["Dummy"]); - assert_eq!(feature_report.declared_background_tasks[0].name, "daily"); + assert!(feature_report.declared_background_tasks.is_empty()); } #[test] @@ -2480,8 +2723,9 @@ mod tests { }) .install_into_pending(&mut pending_tools, &mut hook_builder); - assert_eq!(pending_tools.len(), 1); - assert!(report.reports[0].installed); + assert!(pending_tools.is_empty()); + assert!(!report.reports[0].installed); + assert!(report.reports[0].installed_tools.is_empty()); assert!(!report.reports[1].installed); assert!( report.reports[1] @@ -2558,7 +2802,7 @@ mod tests { "1.0.0", "startup-discovered service", )) - .with_background_task(BackgroundTaskDeclaration::descriptor_only( + .with_background_task(BackgroundTaskDeclaration::worker_managed( "provider-poller", "provider lifecycle poller", )) @@ -2568,7 +2812,7 @@ mod tests { } #[test] - fn protocol_provider_registers_startup_discovered_contributions_through_worker_path() { + fn protocol_provider_report_only_background_task_is_rejected_atomically() { let provider = ProtocolProviderDeclaration::new( ProviderId::builtin("dynamic-provider"), "test-protocol", @@ -2599,30 +2843,18 @@ mod tests { .collect(); let feature_report = &report.reports[0]; - assert!(feature_report.installed); - assert_eq!(feature_report.installed_tools, vec!["DynamicTool"]); - assert_eq!(tool_names, vec!["DynamicTool"]); + assert!(!feature_report.installed); + assert!(feature_report.installed_tools.is_empty()); + assert!(tool_names.is_empty()); assert_eq!(calls.load(Ordering::SeqCst), 1); - assert_eq!(feature_report.provided_services.len(), 1); - assert_eq!( - feature_report.provided_services[0].id, - ServiceId::builtin("dynamic-service") - ); - assert_eq!( - feature_report.declared_background_tasks[0].name, - "provider-poller" - ); + assert!(feature_report.provided_services.is_empty()); + assert!(feature_report.declared_background_tasks.is_empty()); assert_eq!(feature_report.protocol_providers.len(), 1); - assert_eq!( - feature_report.protocol_providers[0].state, - ProtocolProviderLifecycleState::Ready - ); - assert!( - feature_report - .diagnostics - .iter() - .any(|diagnostic| diagnostic.message.contains("startup discovery completed")) - ); + assert!(feature_report.diagnostics.iter().any(|diagnostic| { + diagnostic + .message + .contains("has no executable Worker-managed handler") + })); } #[test] @@ -2779,8 +3011,8 @@ mod tests { async fn call( &self, _input: &crate::hook::ToolCallSummary, - ) -> crate::hook::HookPreToolAction { - crate::hook::HookPreToolAction::Continue + ) -> Result { + Ok(crate::hook::HookPreToolAction::Continue) } } @@ -2804,6 +3036,19 @@ mod tests { } } + struct NoopBackgroundTask; + + #[async_trait] + impl FeatureBackgroundTask for NoopBackgroundTask { + async fn run( + &self, + _context: background::BackgroundTaskContext, + _cancellation: background::BackgroundTaskCancellation, + ) -> Result<(), crate::hook::HookError> { + Ok(()) + } + } + struct BackgroundFeature { descriptor: FeatureDescriptor, task_name: &'static str, @@ -2818,12 +3063,22 @@ mod tests { &self, context: &mut FeatureInstallContext<'_>, ) -> Result<(), FeatureInstallError> { - context - .background_tasks() - .declare(BackgroundTaskDeclaration::descriptor_only( - self.task_name, - "runtime background task", - )) + let declaration = self + .descriptor + .background_tasks + .iter() + .find(|task| task.name == self.task_name) + .cloned() + .unwrap_or_else(|| { + BackgroundTaskDeclaration::worker_managed( + self.task_name, + "undeclared background task", + ) + }); + context.background_tasks().register( + BackgroundTaskSpec::single_flight(declaration, std::time::Duration::from_secs(1)), + NoopBackgroundTask, + ) } } @@ -2985,25 +3240,44 @@ mod tests { ); } - #[test] - fn background_task_declaration_is_descriptor_contribution() { + #[tokio::test] + async fn executable_background_task_is_registered_in_worker_scope() { let descriptor = FeatureDescriptor::builtin("background", "Background") - .with_background_task(BackgroundTaskDeclaration::descriptor_only( + .with_background_task(BackgroundTaskDeclaration::worker_managed( "declared-task", "descriptor contribution", )); let mut hook_builder = HookRegistryBuilder::default(); let mut pending_tools = Vec::new(); let report = FeatureRegistryBuilder::new() - .with_module(ServiceFeature { descriptor }) + .with_module(BackgroundFeature { + descriptor, + task_name: "declared-task", + }) .install_into_pending(&mut pending_tools, &mut hook_builder); - assert!(report.reports[0].installed); assert_eq!( report.reports[0].declared_background_tasks[0].name, "declared-task" ); assert!(report.reports[0].skipped.is_empty()); + assert!(matches!( + report + .background_tasks + .start( + &FeatureId::builtin("background"), + "declared-task", + crate::hook::HookInvocationContext::default(), + ) + .unwrap(), + background::BackgroundTaskStart::Started { .. } + )); + report.background_tasks.shutdown().await.unwrap(); + assert!(matches!( + report.background_tasks.diagnostics()[0].outcome, + background::BackgroundTaskOutcome::Completed + | background::BackgroundTaskOutcome::Cancelled + )); } #[test] @@ -3118,7 +3392,10 @@ mod tests { assert_eq!(descriptor.runtime, FeatureRuntimeKind::Builtin); assert_eq!( hook_points, - vec![FeatureHookPoint::PreRequest, FeatureHookPoint::PreToolCall] + vec![ + FeatureHookPoint::PreLlmRequest, + FeatureHookPoint::PreToolCall + ] ); assert!(descriptor.background_tasks.is_empty()); assert!(descriptor.provides_services.is_empty()); diff --git a/crates/worker/src/feature/background.rs b/crates/worker/src/feature/background.rs new file mode 100644 index 00000000..76254e5e --- /dev/null +++ b/crates/worker/src/feature/background.rs @@ -0,0 +1,705 @@ +//! Executable, scope-owned background tasks contributed by Worker features. +//! +//! Tasks never receive provider handles, credentials, or raw Workdir paths from +//! this registry. Callers pass only stable Worker/session provenance. The task +//! implementation obtains any additional authority through the services its +//! feature was explicitly granted at install time. + +use std::collections::BTreeMap; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use tokio::sync::Notify; +use tokio::task::JoinHandle; + +use super::{BackgroundTaskDeclaration, FeatureId, FeatureInstallError}; +use crate::hook::{HookError, HookErrorCategory, HookInvocationContext}; + +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; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BackgroundTaskRewritePolicy { + CancelAndWait, + Wait, + Block, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BackgroundTaskShutdownPolicy { + CancelAndWait, + Wait, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BackgroundTaskRetryPolicy { + Never, + Bounded { max_attempts: u16, delay_ms: u64 }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BackgroundTaskTrigger { + Manual, + RunCommitted, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct BackgroundTaskSpec { + pub declaration: BackgroundTaskDeclaration, + pub trigger: BackgroundTaskTrigger, + pub max_concurrency: u16, + pub timeout_ms: u64, + pub retry: BackgroundTaskRetryPolicy, + pub rewrite: BackgroundTaskRewritePolicy, + pub shutdown: BackgroundTaskShutdownPolicy, +} + +impl BackgroundTaskSpec { + pub fn single_flight(declaration: BackgroundTaskDeclaration, timeout: Duration) -> Self { + Self { + declaration, + trigger: BackgroundTaskTrigger::Manual, + max_concurrency: 1, + timeout_ms: u64::try_from(timeout.as_millis()).unwrap_or(u64::MAX), + retry: BackgroundTaskRetryPolicy::Never, + rewrite: BackgroundTaskRewritePolicy::CancelAndWait, + shutdown: BackgroundTaskShutdownPolicy::CancelAndWait, + } + } + + fn validate(&self) -> Result<(), FeatureInstallError> { + if self.max_concurrency == 0 || self.max_concurrency > MAX_TASK_CONCURRENCY { + return Err(FeatureInstallError::InvalidDescriptor(format!( + "background task `{}` max_concurrency must be within 1..={MAX_TASK_CONCURRENCY}", + self.declaration.name + ))); + } + if self.timeout_ms == 0 || self.timeout_ms > MAX_TASK_TIMEOUT_MS { + return Err(FeatureInstallError::InvalidDescriptor(format!( + "background task `{}` timeout_ms must be within 1..={MAX_TASK_TIMEOUT_MS}", + self.declaration.name + ))); + } + if let BackgroundTaskRetryPolicy::Bounded { max_attempts, .. } = self.retry + && (max_attempts == 0 || max_attempts > MAX_TASK_ATTEMPTS) + { + return Err(FeatureInstallError::InvalidDescriptor(format!( + "background task `{}` max_attempts must be within 1..={MAX_TASK_ATTEMPTS}", + self.declaration.name + ))); + } + Ok(()) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct BackgroundTaskContext { + pub invocation: HookInvocationContext, + pub feature_id: FeatureId, + pub task_name: String, + pub execution_id: u64, + pub attempt: u16, +} + +#[derive(Clone, Default)] +pub struct BackgroundTaskCancellation { + cancelled: Arc, + notify: Arc, +} + +impl BackgroundTaskCancellation { + pub fn is_cancelled(&self) -> bool { + self.cancelled.load(Ordering::Acquire) + } + + pub async fn cancelled(&self) { + if self.is_cancelled() { + return; + } + self.notify.notified().await; + } + + fn cancel(&self) { + if !self.cancelled.swap(true, Ordering::AcqRel) { + self.notify.notify_one(); + } + } +} + +#[async_trait] +pub trait FeatureBackgroundTask: Send + Sync { + async fn run( + &self, + context: BackgroundTaskContext, + cancellation: BackgroundTaskCancellation, + ) -> Result<(), HookError>; +} + +struct Registration { + feature_id: FeatureId, + spec: BackgroundTaskSpec, + task: Arc, +} + +#[derive(Default)] +pub struct FeatureBackgroundTaskRegistryBuilder { + registrations: BTreeMap<(FeatureId, String), Registration>, +} + +impl FeatureBackgroundTaskRegistryBuilder { + pub fn register( + &mut self, + feature_id: FeatureId, + spec: BackgroundTaskSpec, + task: impl FeatureBackgroundTask + 'static, + ) -> Result<(), FeatureInstallError> { + spec.validate()?; + let key = (feature_id.clone(), spec.declaration.name.clone()); + if self.registrations.contains_key(&key) { + return Err(FeatureInstallError::InvalidDescriptor(format!( + "feature `{feature_id}` registered background task `{}` more than once", + spec.declaration.name + ))); + } + self.registrations.insert( + key, + Registration { + feature_id, + spec, + task: Arc::new(task), + }, + ); + Ok(()) + } + + pub(crate) fn checkpoint(&self) -> Vec<(FeatureId, String)> { + self.registrations.keys().cloned().collect() + } + + pub(crate) fn rollback_to(&mut self, checkpoint: &[(FeatureId, String)]) { + let retained = checkpoint + .iter() + .cloned() + .collect::>(); + self.registrations.retain(|key, _| retained.contains(key)); + } + + pub fn build(self) -> FeatureBackgroundTaskRegistry { + FeatureBackgroundTaskRegistry { + inner: Arc::new(RegistryInner { + registrations: self.registrations, + running: Mutex::new(BTreeMap::new()), + diagnostics: Mutex::new(Vec::new()), + next_execution_id: AtomicU64::new(1), + accepting: AtomicBool::new(true), + }), + } + } +} + +struct RunningTask { + feature_id: FeatureId, + task_name: String, + cancellation: BackgroundTaskCancellation, + handle: JoinHandle<()>, +} + +struct RegistryInner { + registrations: BTreeMap<(FeatureId, String), Registration>, + running: Mutex>, + diagnostics: Mutex>, + next_execution_id: AtomicU64, + accepting: AtomicBool, +} + +impl Drop for RegistryInner { + fn drop(&mut self) { + if let Ok(running) = self.running.get_mut() { + for task in running.values() { + task.cancellation.cancel(); + task.handle.abort(); + } + running.clear(); + } + } +} + +#[derive(Clone)] +pub struct FeatureBackgroundTaskRegistry { + inner: Arc, +} + +impl Default for FeatureBackgroundTaskRegistry { + fn default() -> Self { + FeatureBackgroundTaskRegistryBuilder::default().build() + } +} + +impl std::fmt::Debug for FeatureBackgroundTaskRegistry { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("FeatureBackgroundTaskRegistry") + .field( + "registrations", + &self.inner.registrations.keys().collect::>(), + ) + .finish_non_exhaustive() + } +} + +impl PartialEq for FeatureBackgroundTaskRegistry { + fn eq(&self, other: &Self) -> bool { + self.inner + .registrations + .keys() + .eq(other.inner.registrations.keys()) + } +} + +impl Eq for FeatureBackgroundTaskRegistry {} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum BackgroundTaskStart { + Started { execution_id: u64 }, + AtCapacity, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum BackgroundTaskOutcome { + Completed, + Cancelled, + TimedOut, + Failed(HookError), + JoinFailed, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct BackgroundTaskDiagnostic { + pub execution_id: u64, + pub feature_id: FeatureId, + pub task_name: String, + pub attempts: u16, + pub outcome: BackgroundTaskOutcome, +} + +impl FeatureBackgroundTaskRegistry { + /// Starts one execution immediately. The Worker scope intentionally owns no + /// hidden queue: reaching `max_concurrency` returns `AtCapacity`, so callers + /// must retain durable work in their domain authority and retry explicitly. + pub fn start( + &self, + feature_id: &FeatureId, + task_name: &str, + invocation: HookInvocationContext, + ) -> Result { + if !self.inner.accepting.load(Ordering::Acquire) { + return Err(HookError::new( + HookErrorCategory::ScopeDisposed, + "background task scope is stopping", + )); + } + let key = (feature_id.clone(), task_name.to_string()); + let registration = self.inner.registrations.get(&key).ok_or_else(|| { + HookError::new( + HookErrorCategory::InvalidInput, + format!("unknown background task `{feature_id}/{task_name}`"), + ) + })?; + + let mut running = self + .inner + .running + .lock() + .expect("background tasks poisoned"); + running.retain(|_, running| !running.handle.is_finished()); + let active = running + .values() + .filter(|running| running.feature_id == *feature_id && running.task_name == task_name) + .count(); + if active >= usize::from(registration.spec.max_concurrency) { + return Ok(BackgroundTaskStart::AtCapacity); + } + + let execution_id = self.inner.next_execution_id.fetch_add(1, Ordering::Relaxed); + let cancellation = BackgroundTaskCancellation::default(); + let task_cancellation = cancellation.clone(); + let task = Arc::clone(®istration.task); + let spec = registration.spec.clone(); + let feature_id = registration.feature_id.clone(); + let task_name = task_name.to_string(); + let weak_inner = Arc::downgrade(&self.inner); + let task_feature_id = feature_id.clone(); + let task_task_name = task_name.clone(); + let handle = tokio::spawn(async move { + let (attempts, outcome) = execute_task( + task, + spec, + invocation, + task_feature_id.clone(), + task_task_name.clone(), + execution_id, + task_cancellation, + ) + .await; + if let Some(inner) = weak_inner.upgrade() { + let mut diagnostics = inner.diagnostics.lock().expect("diagnostics poisoned"); + diagnostics.push(BackgroundTaskDiagnostic { + execution_id, + feature_id: task_feature_id, + task_name: task_task_name, + attempts, + outcome, + }); + if diagnostics.len() > MAX_RETAINED_DIAGNOSTICS { + let remove = diagnostics.len() - MAX_RETAINED_DIAGNOSTICS; + diagnostics.drain(..remove); + } + } + }); + running.insert( + execution_id, + RunningTask { + feature_id, + task_name, + cancellation, + handle, + }, + ); + Ok(BackgroundTaskStart::Started { execution_id }) + } + + /// Starts every task explicitly bound to the committed-run boundary in + /// deterministic `(FeatureId, task name)` order. Capacity is an expected + /// single-flight outcome and leaves the already-running execution intact. + pub fn start_run_committed(&self, invocation: HookInvocationContext) -> Result<(), HookError> { + let tasks = self + .inner + .registrations + .iter() + .filter(|(_, registration)| { + registration.spec.trigger == BackgroundTaskTrigger::RunCommitted + }) + .map(|((feature_id, task_name), _)| (feature_id.clone(), task_name.clone())) + .collect::>(); + for (feature_id, task_name) in tasks { + let _ = self.start(&feature_id, &task_name, invocation.clone())?; + } + Ok(()) + } + + pub fn diagnostics(&self) -> Vec { + self.inner + .diagnostics + .lock() + .expect("diagnostics poisoned") + .clone() + } + + pub async fn before_session_rewrite(&self) -> Result<(), HookError> { + self.settle(false).await + } + + pub async fn shutdown(&self) -> Result<(), HookError> { + self.inner.accepting.store(false, Ordering::Release); + self.settle(true).await + } + + async fn settle(&self, shutdown: bool) -> Result<(), HookError> { + let mut waiting = Vec::new(); + { + let mut running = self + .inner + .running + .lock() + .expect("background tasks poisoned"); + if !shutdown { + for task in running.values() { + let registration = self + .inner + .registrations + .get(&(task.feature_id.clone(), task.task_name.clone())) + .expect("running background task must retain registration"); + if registration.spec.rewrite == BackgroundTaskRewritePolicy::Block { + return Err(HookError::new( + HookErrorCategory::Dependency, + "session rewrite blocked by a running feature background task", + )); + } + } + } + let ids = running.keys().copied().collect::>(); + for id in ids { + let Some(task) = running.remove(&id) else { + continue; + }; + let registration = self + .inner + .registrations + .get(&(task.feature_id.clone(), task.task_name.clone())) + .expect("running background task must retain registration"); + let cancel = if shutdown { + registration.spec.shutdown == BackgroundTaskShutdownPolicy::CancelAndWait + } else { + match registration.spec.rewrite { + BackgroundTaskRewritePolicy::CancelAndWait => true, + BackgroundTaskRewritePolicy::Wait => false, + BackgroundTaskRewritePolicy::Block => unreachable!( + "blocking rewrite policies are rejected before task handles are drained" + ), + } + }; + if cancel { + task.cancellation.cancel(); + } + waiting.push(( + id, + task.feature_id.clone(), + task.task_name.clone(), + task.handle, + )); + } + } + + let mut join_failed = false; + for (execution_id, feature_id, task_name, handle) in waiting { + if handle.await.is_err() { + join_failed = true; + let mut diagnostics = self.inner.diagnostics.lock().expect("diagnostics poisoned"); + diagnostics.push(BackgroundTaskDiagnostic { + execution_id, + feature_id, + task_name, + attempts: 0, + outcome: BackgroundTaskOutcome::JoinFailed, + }); + if diagnostics.len() > MAX_RETAINED_DIAGNOSTICS { + let remove = diagnostics.len() - MAX_RETAINED_DIAGNOSTICS; + diagnostics.drain(..remove); + } + } + } + if join_failed { + return Err(HookError::new( + HookErrorCategory::Internal, + "feature background task join failed", + )); + } + Ok(()) + } +} + +async fn execute_task( + task: Arc, + spec: BackgroundTaskSpec, + invocation: HookInvocationContext, + feature_id: FeatureId, + task_name: String, + execution_id: u64, + cancellation: BackgroundTaskCancellation, +) -> (u16, BackgroundTaskOutcome) { + let (max_attempts, delay_ms) = match spec.retry { + BackgroundTaskRetryPolicy::Never => (1, 0), + BackgroundTaskRetryPolicy::Bounded { + max_attempts, + delay_ms, + } => (max_attempts, delay_ms), + }; + for attempt in 1..=max_attempts { + if cancellation.is_cancelled() { + return (attempt, BackgroundTaskOutcome::Cancelled); + } + let context = BackgroundTaskContext { + invocation: invocation.clone(), + feature_id: feature_id.clone(), + task_name: task_name.clone(), + execution_id, + attempt, + }; + let result = tokio::time::timeout( + Duration::from_millis(spec.timeout_ms), + task.run(context, cancellation.clone()), + ) + .await; + match result { + Ok(Ok(())) => return (attempt, BackgroundTaskOutcome::Completed), + Ok(Err(_error)) if cancellation.is_cancelled() => { + return (attempt, BackgroundTaskOutcome::Cancelled); + } + Ok(Err(error)) if attempt == max_attempts => { + return (attempt, BackgroundTaskOutcome::Failed(error)); + } + Ok(Err(_)) => {} + Err(_) => return (attempt, BackgroundTaskOutcome::TimedOut), + } + if delay_ms > 0 { + tokio::select! { + () = tokio::time::sleep(Duration::from_millis(delay_ms)) => {} + () = cancellation.cancelled() => { + return (attempt, BackgroundTaskOutcome::Cancelled); + } + } + } + } + ( + max_attempts, + BackgroundTaskOutcome::Failed(HookError::new( + HookErrorCategory::Internal, + "background task exhausted retry policy", + )), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + fn invocation() -> HookInvocationContext { + HookInvocationContext { + workspace_id: Some("workspace".into()), + worker_id: "worker".into(), + session_id: "session".into(), + session_revision: 3, + run_id: Some("run".into()), + turn_index: Some(2), + call_id: None, + } + } + + struct WaitForCancellation; + + #[async_trait] + impl FeatureBackgroundTask for WaitForCancellation { + async fn run( + &self, + _context: BackgroundTaskContext, + cancellation: BackgroundTaskCancellation, + ) -> Result<(), HookError> { + cancellation.cancelled().await; + Err(HookError::new(HookErrorCategory::Cancelled, "cancelled")) + } + } + + #[tokio::test] + async fn single_flight_rejects_overlap_and_shutdown_joins_the_task() { + let feature = FeatureId::builtin("background-test"); + let declaration = BackgroundTaskDeclaration::worker_managed("extract", "extract"); + let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); + builder + .register( + feature.clone(), + BackgroundTaskSpec::single_flight(declaration, Duration::from_secs(1)), + WaitForCancellation, + ) + .unwrap(); + let registry = builder.build(); + + assert!(matches!( + registry.start(&feature, "extract", invocation()).unwrap(), + BackgroundTaskStart::Started { .. } + )); + assert_eq!( + registry.start(&feature, "extract", invocation()).unwrap(), + BackgroundTaskStart::AtCapacity + ); + + registry.shutdown().await.unwrap(); + assert_eq!(registry.diagnostics().len(), 1); + assert_eq!( + registry.diagnostics()[0].outcome, + BackgroundTaskOutcome::Cancelled + ); + assert!(matches!( + registry.start(&feature, "extract", invocation()), + Err(HookError { + category: HookErrorCategory::ScopeDisposed, + .. + }) + )); + } + + struct FailTwice { + calls: Arc, + completed: Arc, + } + + #[async_trait] + impl FeatureBackgroundTask for FailTwice { + async fn run( + &self, + _context: BackgroundTaskContext, + _cancellation: BackgroundTaskCancellation, + ) -> Result<(), HookError> { + let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1; + if call < 3 { + return Err(HookError::new(HookErrorCategory::Dependency, "retry")); + } + self.completed.notify_one(); + Ok(()) + } + } + + #[tokio::test] + async fn bounded_retry_records_attempts_without_task_output() { + let feature = FeatureId::builtin("retry-test"); + let declaration = BackgroundTaskDeclaration::worker_managed("retry", "retry"); + let calls = Arc::new(AtomicUsize::new(0)); + let completed = Arc::new(Notify::new()); + let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); + let mut spec = BackgroundTaskSpec::single_flight(declaration, Duration::from_secs(1)); + spec.trigger = BackgroundTaskTrigger::RunCommitted; + spec.retry = BackgroundTaskRetryPolicy::Bounded { + max_attempts: 3, + delay_ms: 0, + }; + builder + .register( + feature.clone(), + spec, + FailTwice { + calls: Arc::clone(&calls), + completed: Arc::clone(&completed), + }, + ) + .unwrap(); + let registry = builder.build(); + registry.start_run_committed(invocation()).unwrap(); + completed.notified().await; + registry.shutdown().await.unwrap(); + + assert_eq!(calls.load(Ordering::SeqCst), 3); + assert_eq!(registry.diagnostics()[0].attempts, 3); + assert_eq!( + registry.diagnostics()[0].outcome, + BackgroundTaskOutcome::Completed + ); + } + + #[tokio::test] + async fn block_policy_fences_rewrite_without_detaching_the_task() { + let feature = FeatureId::builtin("rewrite-test"); + let declaration = BackgroundTaskDeclaration::worker_managed("rewrite", "rewrite"); + let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); + let mut spec = BackgroundTaskSpec::single_flight(declaration, Duration::from_secs(1)); + spec.rewrite = BackgroundTaskRewritePolicy::Block; + builder + .register(feature.clone(), spec, WaitForCancellation) + .unwrap(); + let registry = builder.build(); + registry.start(&feature, "rewrite", invocation()).unwrap(); + + let error = registry.before_session_rewrite().await.unwrap_err(); + assert_eq!(error.category, HookErrorCategory::Dependency); + registry.shutdown().await.unwrap(); + assert_eq!(registry.diagnostics().len(), 1); + } +} diff --git a/crates/worker/src/feature/builtin/task/mod.rs b/crates/worker/src/feature/builtin/task/mod.rs index 9df849bf..9ab38be3 100644 --- a/crates/worker/src/feature/builtin/task/mod.rs +++ b/crates/worker/src/feature/builtin/task/mod.rs @@ -125,7 +125,7 @@ impl FeatureModule for TaskFeature { )) .with_hook(HookDeclaration::new( "task-reminder-pre-request", - FeatureHookPoint::PreRequest, + FeatureHookPoint::PreLlmRequest, )) .with_hook(HookDeclaration::new( "task-reminder-tool-usage", @@ -209,24 +209,27 @@ struct TaskReminderPreRequestHook { #[async_trait] impl Hook for TaskReminderPreRequestHook { - async fn call(&self, input: &PreRequestContext) -> HookPreRequestAction { + async fn call( + &self, + input: &PreRequestContext, + ) -> Result { let tasks = self.state.task_store.list(); if tasks.is_empty() { - return HookPreRequestAction::Continue; + return Ok(HookPreRequestAction::Continue); } let (since_task_management, since_reminder) = self.state.reminder_state.note_request(); if since_task_management < TASK_REMINDER_REQUEST_THRESHOLD || since_reminder < TASK_REMINDER_COOLDOWN_REQUESTS { - return HookPreRequestAction::Continue; + return Ok(HookPreRequestAction::Continue); } if let Some(system_items) = input.system_items() { self.state.reminder_state.note_reminder(); system_items.append_task_reminder(render_task_reminder_body(&tasks)); } - HookPreRequestAction::Continue + Ok(HookPreRequestAction::Continue) } } @@ -236,11 +239,14 @@ struct TaskReminderToolUsageHook { #[async_trait] impl Hook for TaskReminderToolUsageHook { - async fn call(&self, input: &ToolCallSummary) -> HookPreToolAction { + async fn call( + &self, + input: &ToolCallSummary, + ) -> Result { if is_task_management_tool(&input.tool_name) { self.state.reminder_state.note_task_management(); } - HookPreToolAction::Continue + Ok(HookPreToolAction::Continue) } } diff --git a/crates/worker/src/feature/mcp.rs b/crates/worker/src/feature/mcp.rs index cfdd8f27..ca13ae13 100644 --- a/crates/worker/src/feature/mcp.rs +++ b/crates/worker/src/feature/mcp.rs @@ -2115,8 +2115,7 @@ read exit_notification || true meta.name }) .collect(); - assert!(!names.iter().any(|name| name == "Mcp_demo_search_files")); - assert!(names.iter().any(|name| name == "Mcp_demo_unique")); + assert!(names.is_empty()); } fn shell_tool_server(response: &str) -> McpStdioServerSpec { diff --git a/crates/worker/src/feature/plugin.rs b/crates/worker/src/feature/plugin.rs index 143d8e9f..39dbfecf 100644 --- a/crates/worker/src/feature/plugin.rs +++ b/crates/worker/src/feature/plugin.rs @@ -7621,7 +7621,7 @@ mod tests { ))) .install_into_pending(&mut pending, &mut hooks); - assert_eq!(pending.len(), 1); + assert!(pending.is_empty()); assert_eq!(skipped_count(&report), 1); assert!(has_diagnostic(&report, "duplicate tool contribution")); } diff --git a/crates/worker/src/hook.rs b/crates/worker/src/hook.rs index 4318f9c3..b326697a 100644 --- a/crates/worker/src/hook.rs +++ b/crates/worker/src/hook.rs @@ -18,13 +18,78 @@ use std::ops::Deref; use std::sync::{Arc, Mutex}; +use agen::HistoryEntry; use agen::interceptor::{ PostToolAction, PreRequestAction, PreToolAction, PromptAction, TurnEndAction, }; use agen::tool::{ToolOutput, ToolResult}; use async_trait::async_trait; +use serde::{Deserialize, Serialize}; use serde_json::Value; use session_store::{SystemItem, SystemReminder}; +use thiserror::Error; + +use crate::session_history::SessionHistoryMetadata; + +const HOOK_DIAGNOSTIC_MAX_BYTES: usize = 1_024; + +/// Failure category exposed by the safe Worker hook boundary. +/// +/// Categories are intentionally closed and payload-free so extensions cannot +/// smuggle provider, credential, or history data into diagnostics. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum HookErrorCategory { + InvalidInput, + Dependency, + Timeout, + Cancelled, + Trap, + ScopeDisposed, + Internal, +} + +/// Bounded hook callback failure. Raw tool arguments, output, prompts, and +/// credentials must never be placed in `diagnostic`. +#[derive(Clone, Debug, PartialEq, Eq, Error)] +#[error("{category:?}: {diagnostic}")] +pub struct HookError { + pub category: HookErrorCategory, + pub diagnostic: String, +} + +impl HookError { + pub fn new(category: HookErrorCategory, diagnostic: impl Into) -> Self { + Self { + category, + diagnostic: bounded_utf8(diagnostic.into(), HOOK_DIAGNOSTIC_MAX_BYTES), + } + } +} + +/// Failure behavior declared when a hook is registered. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum HookFailurePolicy { + /// Gate the current operation when the hook cannot decide safely. + FailClosed, + /// Keep the already-authorized operation moving and emit a diagnostic. + FailOpenWithDiagnostic, + /// Keep committed state intact and mark the failure for operator attention. + AttentionRequired, +} + +fn bounded_utf8(mut value: String, max_bytes: usize) -> String { + if value.len() <= max_bytes { + return value; + } + let mut end = max_bytes; + while end > 0 && !value.is_char_boundary(end) { + end -= 1; + } + value.truncate(end); + value +} /// Hook-facing prompt-submit action. /// @@ -349,24 +414,175 @@ impl HookEventKind for OnTurnEnd { /// short-circuit on the first non-continue action. #[async_trait] pub trait Hook: Send + Sync { - async fn call(&self, input: &E::Input) -> E::Output; + async fn call(&self, input: &E::Input) -> Result; } // ============================================================================= // Hook Registry // ============================================================================= +/// Stable provenance attached to every Worker lifecycle callback. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct HookInvocationContext { + pub workspace_id: Option, + pub worker_id: String, + pub session_id: String, + pub session_revision: u64, + pub run_id: Option, + pub turn_index: Option, + pub call_id: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RunCommittedExit { + Finished, + Paused, + Yielded, + Interrupted, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RunExitContext { + pub invocation: HookInvocationContext, + pub exit: RunCommittedExit, + pub history_len: usize, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct RunCommittedContext { + pub invocation: HookInvocationContext, + pub exit: RunCommittedExit, + pub committed_history: Vec>, + pub committed_history_len: usize, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SessionRewriteKind { + Rewind, + Compact, + Fork, + Restore, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct BeforeSessionRewriteContext { + pub invocation: HookInvocationContext, + pub kind: SessionRewriteKind, + pub current_history: Vec>, + pub current_history_len: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum BeforeSessionRewriteAction { + Continue, + Deny(String), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct WorkerStoppingContext { + pub invocation: HookInvocationContext, + pub reason: String, +} + +pub struct RunExit; +pub struct RunCommitted; +pub struct BeforeSessionRewrite; +pub struct WorkerStopping; + +impl HookEventKind for RunExit { + type Input = RunExitContext; + type Output = (); +} + +impl HookEventKind for RunCommitted { + type Input = RunCommittedContext; + type Output = (); +} + +impl HookEventKind for BeforeSessionRewrite { + type Input = BeforeSessionRewriteContext; + type Output = BeforeSessionRewriteAction; +} + +impl HookEventKind for WorkerStopping { + type Input = WorkerStoppingContext; + type Output = (); +} + +pub(crate) struct RegisteredHook { + owner: String, + policy: HookFailurePolicy, + 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, + }) + } + + pub(crate) async fn call_optional( + &self, + input: &E::Input, + ) -> Result, HookExecutionError> { + match self.call(input).await { + Ok(output) => Ok(Some(output)), + Err(error) if error.policy == HookFailurePolicy::FailOpenWithDiagnostic => { + tracing::warn!(owner = %error.owner, error = %error.source, "inline hook failed open"); + Ok(None) + } + Err(error) => Err(error), + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Error)] +#[error("hook `{owner}` failed under {policy:?}: {source}")] +pub struct HookExecutionError { + pub owner: String, + pub policy: HookFailurePolicy, + pub source: HookError, +} + /// Builder for constructing a frozen `HookRegistry`. -/// -/// Hooks are added during setup, then `build()` produces an immutable -/// registry that can be shared via `Arc`. #[derive(Default)] pub struct HookRegistryBuilder { - on_prompt_submit: Vec>>, - pre_llm_request: Vec>>, - pre_tool_call: Vec>>, - post_tool_call: Vec>>, - on_turn_end: Vec>>, + on_prompt_submit: Vec>, + pre_llm_request: Vec>, + pre_tool_call: Vec>, + post_tool_call: Vec>, + on_turn_end: Vec>, + run_exit: Vec>, + run_committed: Vec>, + before_session_rewrite: Vec>, + worker_stopping: Vec>, +} + +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); + } + + pub fn $named( + &mut self, + owner: impl Into, + policy: HookFailurePolicy, + hook: impl Hook<$event> + 'static, + ) { + self.$field.push(RegisteredHook { + owner: owner.into(), + policy, + hook: Box::new(hook), + }); + } + }; } impl HookRegistryBuilder { @@ -374,27 +590,82 @@ impl HookRegistryBuilder { Self::default() } - pub fn add_on_prompt_submit(&mut self, hook: impl Hook + 'static) { - self.on_prompt_submit.push(Box::new(hook)); + add_hook_methods!( + add_on_prompt_submit, + add_named_on_prompt_submit, + on_prompt_submit, + OnPromptSubmit + ); + add_hook_methods!( + add_pre_llm_request, + add_named_pre_llm_request, + pre_llm_request, + PreLlmRequest + ); + add_hook_methods!( + add_pre_tool_call, + add_named_pre_tool_call, + pre_tool_call, + PreToolCall + ); + add_hook_methods!( + add_post_tool_call, + add_named_post_tool_call, + post_tool_call, + PostToolCall + ); + add_hook_methods!( + add_on_turn_end, + add_named_on_turn_end, + on_turn_end, + OnTurnEnd + ); + add_hook_methods!(add_run_exit, add_named_run_exit, run_exit, RunExit); + add_hook_methods!( + add_run_committed, + add_named_run_committed, + run_committed, + RunCommitted + ); + add_hook_methods!( + add_before_session_rewrite, + add_named_before_session_rewrite, + before_session_rewrite, + BeforeSessionRewrite + ); + add_hook_methods!( + add_worker_stopping, + add_named_worker_stopping, + worker_stopping, + WorkerStopping + ); + + pub(crate) fn checkpoint(&self) -> [usize; 9] { + [ + self.on_prompt_submit.len(), + self.pre_llm_request.len(), + self.pre_tool_call.len(), + self.post_tool_call.len(), + self.on_turn_end.len(), + self.run_exit.len(), + self.run_committed.len(), + self.before_session_rewrite.len(), + self.worker_stopping.len(), + ] } - pub fn add_pre_llm_request(&mut self, hook: impl Hook + 'static) { - self.pre_llm_request.push(Box::new(hook)); + pub(crate) fn rollback_to(&mut self, checkpoint: [usize; 9]) { + self.on_prompt_submit.truncate(checkpoint[0]); + self.pre_llm_request.truncate(checkpoint[1]); + self.pre_tool_call.truncate(checkpoint[2]); + self.post_tool_call.truncate(checkpoint[3]); + self.on_turn_end.truncate(checkpoint[4]); + self.run_exit.truncate(checkpoint[5]); + self.run_committed.truncate(checkpoint[6]); + self.before_session_rewrite.truncate(checkpoint[7]); + self.worker_stopping.truncate(checkpoint[8]); } - pub fn add_pre_tool_call(&mut self, hook: impl Hook + 'static) { - self.pre_tool_call.push(Box::new(hook)); - } - - pub fn add_post_tool_call(&mut self, hook: impl Hook + 'static) { - self.post_tool_call.push(Box::new(hook)); - } - - pub fn add_on_turn_end(&mut self, hook: impl Hook + 'static) { - self.on_turn_end.push(Box::new(hook)); - } - - /// Freeze the builder into an immutable registry. pub fn build(self) -> HookRegistry { HookRegistry { on_prompt_submit: self.on_prompt_submit, @@ -402,17 +673,129 @@ impl HookRegistryBuilder { pre_tool_call: self.pre_tool_call, post_tool_call: self.post_tool_call, on_turn_end: self.on_turn_end, + run_exit: self.run_exit, + run_committed: self.run_committed, + before_session_rewrite: self.before_session_rewrite, + worker_stopping: self.worker_stopping, + diagnostics: std::sync::Mutex::new(Vec::new()), } } } /// Frozen registry of hooks. Constructed via [`HookRegistryBuilder::build()`]. pub struct HookRegistry { - pub(crate) on_prompt_submit: Vec>>, - pub(crate) pre_llm_request: Vec>>, - pub(crate) pre_tool_call: Vec>>, - pub(crate) post_tool_call: Vec>>, - pub(crate) on_turn_end: Vec>>, + pub(crate) on_prompt_submit: Vec>, + pub(crate) pre_llm_request: Vec>, + pub(crate) pre_tool_call: Vec>, + pub(crate) post_tool_call: Vec>, + pub(crate) on_turn_end: Vec>, + run_exit: Vec>, + run_committed: Vec>, + before_session_rewrite: Vec>, + worker_stopping: Vec>, + diagnostics: std::sync::Mutex>, +} + +impl HookRegistry { + fn record_diagnostic(&self, error: HookExecutionError) { + let mut diagnostics = self.diagnostics.lock().expect("hook diagnostics poisoned"); + diagnostics.push(error); + if diagnostics.len() > 128 { + let remove = diagnostics.len() - 128; + diagnostics.drain(..remove); + } + } + + pub fn diagnostics(&self) -> Vec { + self.diagnostics + .lock() + .expect("hook diagnostics poisoned") + .clone() + } + pub async fn on_run_exit(&self, context: &RunExitContext) -> Result<(), HookExecutionError> { + for registration in &self.run_exit { + if let Err(error) = registration.call(context).await { + self.record_diagnostic(error.clone()); + match error.policy { + HookFailurePolicy::FailOpenWithDiagnostic => { + tracing::warn!(owner = %error.owner, error = %error.source, "run-exit hook failed open"); + } + HookFailurePolicy::FailClosed | HookFailurePolicy::AttentionRequired => { + return Err(error); + } + } + } + } + Ok(()) + } + + pub async fn on_run_committed( + &self, + context: &RunCommittedContext, + ) -> Result<(), HookExecutionError> { + for registration in &self.run_committed { + if let Err(error) = registration.call(context).await { + self.record_diagnostic(error.clone()); + match error.policy { + HookFailurePolicy::FailOpenWithDiagnostic => { + tracing::warn!(owner = %error.owner, error = %error.source, "run-committed hook failed open"); + } + HookFailurePolicy::FailClosed | HookFailurePolicy::AttentionRequired => { + return Err(error); + } + } + } + } + Ok(()) + } + + pub async fn before_session_rewrite( + &self, + context: &BeforeSessionRewriteContext, + ) -> Result { + let mut denials = Vec::new(); + for registration in &self.before_session_rewrite { + match registration.call(context).await { + Ok(BeforeSessionRewriteAction::Continue) => {} + Ok(BeforeSessionRewriteAction::Deny(reason)) => { + denials.push((registration.owner.clone(), reason)); + } + Err(error) if error.policy == HookFailurePolicy::FailOpenWithDiagnostic => { + self.record_diagnostic(error.clone()); + tracing::warn!(owner = %error.owner, error = %error.source, "session-rewrite hook failed open"); + } + Err(error) => { + self.record_diagnostic(error.clone()); + return Err(error); + } + } + } + denials.sort_by(|left, right| left.0.cmp(&right.0)); + Ok(denials + .into_iter() + .next() + .map(|(_, reason)| BeforeSessionRewriteAction::Deny(reason)) + .unwrap_or(BeforeSessionRewriteAction::Continue)) + } + + pub async fn on_worker_stopping( + &self, + context: &WorkerStoppingContext, + ) -> Result<(), HookExecutionError> { + for registration in &self.worker_stopping { + if let Err(error) = registration.call(context).await { + self.record_diagnostic(error.clone()); + match error.policy { + HookFailurePolicy::FailOpenWithDiagnostic + | HookFailurePolicy::AttentionRequired => { + tracing::warn!(owner = %error.owner, error = %error.source, "worker-stopping hook requires attention"); + } + HookFailurePolicy::FailClosed => return Err(error), + } + } + } + Ok(()) + } } #[cfg(test)] @@ -474,4 +857,117 @@ mod tests { let pause_action = HookPreToolAction::Pause.into_worker_action("call_4".into()); assert!(matches!(pause_action, PreToolAction::Pause)); } + + struct RewriteHook { + action: BeforeSessionRewriteAction, + } + + #[async_trait] + impl Hook for RewriteHook { + async fn call( + &self, + _input: &BeforeSessionRewriteContext, + ) -> Result { + Ok(self.action.clone()) + } + } + + struct FailingRewriteHook; + + #[async_trait] + impl Hook for FailingRewriteHook { + async fn call( + &self, + _input: &BeforeSessionRewriteContext, + ) -> Result { + Err(HookError::new( + HookErrorCategory::Dependency, + "provider unavailable", + )) + } + } + + fn rewrite_context() -> BeforeSessionRewriteContext { + BeforeSessionRewriteContext { + invocation: HookInvocationContext { + workspace_id: Some("workspace".into()), + worker_id: "worker".into(), + session_id: "session".into(), + session_revision: 4, + run_id: None, + turn_index: None, + call_id: None, + }, + kind: SessionRewriteKind::Compact, + current_history: Vec::new(), + current_history_len: 8, + } + } + + #[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()), + }, + ); + + assert_eq!( + builder + .build() + .before_session_rewrite(&rewrite_context()) + .await + .unwrap(), + BeforeSessionRewriteAction::Deny("a denied".into()) + ); + } + + #[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, + ); + let fail_open = fail_open.build(); + assert_eq!( + fail_open + .before_session_rewrite(&rewrite_context()) + .await + .unwrap(), + BeforeSessionRewriteAction::Continue + ); + assert_eq!(fail_open.diagnostics().len(), 1); + + let mut fail_closed = HookRegistryBuilder::new(); + fail_closed.add_named_before_session_rewrite( + "feature", + HookFailurePolicy::FailClosed, + FailingRewriteHook, + ); + let error = fail_closed + .build() + .before_session_rewrite(&rewrite_context()) + .await + .unwrap_err(); + assert_eq!(error.source.category, HookErrorCategory::Dependency); + } + + #[test] + fn hook_diagnostics_are_utf8_bounded() { + let error = HookError::new(HookErrorCategory::Internal, "界".repeat(1_000)); + assert!(error.diagnostic.len() <= HOOK_DIAGNOSTIC_MAX_BYTES); + assert!(error.diagnostic.is_char_boundary(error.diagnostic.len())); + } } diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index 6dc23c51..b69ac42c 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -210,7 +210,7 @@ where let segment_id = worker.segment_id(); on_cancel_sender(worker.engine_mut().cancel_sender()); - match worker.run_text(&input).await { + let outcome = match worker.run_text(&input).await { Ok(lifecycle @ WorkerRunResult::Finished) | Ok(lifecycle @ WorkerRunResult::Paused) | Ok(lifecycle @ WorkerRunResult::RolledBack) => Ok(InternalWorkerResult { @@ -239,7 +239,11 @@ where identity, history_entries: store.entries_count(session_id, segment_id), }), - } + }; + worker + .stop_feature_runtime("internal Worker terminal outcome") + .await; + outcome } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -782,7 +786,8 @@ pub(crate) async fn prepare_internal_worker_session( }; tokio::spawn(async move { - while let Some(command) = command_rx.recv().await { + let mut stop_done = None; + 'actor: while let Some(command) = command_rx.recv().await { match command { InternalWorkerSessionCommand::Run(input) => { actor_in_flight.clear(); @@ -825,21 +830,16 @@ pub(crate) async fn prepare_internal_worker_session( Some(InternalWorkerSessionCommand::Stop(done)) => { let _ = cancel_sender.send(()).await; let _ = (&mut run).await; - actor_in_flight.clear(); - status.store(InternalWorkerSessionStatus::Stopped.encode(), std::sync::atomic::Ordering::Release); - let _ = event_tx.send(Event::Status { status: WorkerStatus::Stopped }); - let _ = event_tx.send(Event::Shutdown); - state_changed.notify_waiters(); - let _ = done.send(()); - return; + stop_done = Some(done); + break 'actor; } Some(InternalWorkerSessionCommand::Run(_)) => { // `send` reserves Running atomically, so a second Run cannot be enqueued. } None => { let _ = cancel_sender.send(()).await; - actor_in_flight.clear(); - return; + let _ = (&mut run).await; + break 'actor; } } } @@ -847,22 +847,27 @@ pub(crate) async fn prepare_internal_worker_session( } } InternalWorkerSessionCommand::Stop(done) => { - actor_in_flight.clear(); - status.store( - InternalWorkerSessionStatus::Stopped.encode(), - std::sync::atomic::Ordering::Release, - ); - let _ = event_tx.send(Event::Status { - status: WorkerStatus::Stopped, - }); - let _ = event_tx.send(Event::Shutdown); - state_changed.notify_waiters(); - let _ = done.send(()); - return; + stop_done = Some(done); + break; } } } + worker + .stop_feature_runtime("internal Worker session stopped") + .await; actor_in_flight.clear(); + status.store( + InternalWorkerSessionStatus::Stopped.encode(), + std::sync::atomic::Ordering::Release, + ); + let _ = event_tx.send(Event::Status { + status: WorkerStatus::Stopped, + }); + let _ = event_tx.send(Event::Shutdown); + state_changed.notify_waiters(); + if let Some(done) = stop_done { + let _ = done.send(()); + } }); Ok(handle) diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index 2706f715..0e5c5b40 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -246,12 +246,22 @@ impl Interceptor for WorkerInterceptor { input_text: extract_message_text(item).unwrap_or_default(), turn_index, }; + let mut cancellations = Vec::new(); for hook in &self.registry.on_prompt_submit { - let action = hook.call(&info).await; - if !matches!(action, HookPromptAction::Continue) { - return Ok(action.into()); + let Some(action) = hook.call_optional(&info).await.map_err(|error| { + InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) + })? + else { + continue; + }; + if let HookPromptAction::Cancel(reason) = action { + cancellations.push(reason); } } + cancellations.sort(); + if let Some(reason) = cancellations.into_iter().next() { + return Ok(PromptAction::Cancel(reason)); + } let mut extras: Vec = std::mem::take( &mut *self .pending_attachments @@ -344,12 +354,29 @@ 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 mut cancellations = Vec::new(); + let mut should_yield = false; for hook in &self.registry.pre_llm_request { - let action = hook.call(&hook_context).await; - if !matches!(action, HookPreRequestAction::Continue) { - return Ok(action.into()); + let Some(action) = hook.call_optional(&hook_context).await.map_err(|error| { + InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) + })? + else { + continue; + }; + + match action { + HookPreRequestAction::Continue => {} + HookPreRequestAction::Yield => should_yield = true, + HookPreRequestAction::Cancel(reason) => cancellations.push(reason), } } + cancellations.sort(); + if let Some(reason) = cancellations.into_iter().next() { + return Ok(PreRequestAction::Cancel(reason)); + } + if should_yield { + return Ok(PreRequestAction::Yield); + } let mut system_items: Vec = std::mem::take( &mut *pending_hook_system_items @@ -404,12 +431,35 @@ impl Interceptor for WorkerInterceptor { tool_name: info.call.name.clone(), arguments: info.call.input.clone(), }; + let mut aborts = Vec::new(); + let mut should_pause = false; + let mut denials = Vec::new(); for hook in &self.registry.pre_tool_call { - let action = hook.call(&summary).await; - if !matches!(action, HookPreToolAction::Continue) { - return Ok(action.into_worker_action(summary.call_id.clone())); + let Some(action) = hook.call_optional(&summary).await.map_err(|error| { + InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) + })? + else { + continue; + }; + + match action { + HookPreToolAction::Continue => {} + HookPreToolAction::Pause => should_pause = true, + HookPreToolAction::Deny(reason) => denials.push(reason), + HookPreToolAction::Abort(reason) => aborts.push(reason), } } + aborts.sort(); + if let Some(reason) = aborts.into_iter().next() { + return Ok(HookPreToolAction::Abort(reason).into_worker_action(summary.call_id.clone())); + } + if should_pause { + return Ok(PreToolAction::Pause); + } + denials.sort(); + if let Some(reason) = denials.into_iter().next() { + return Ok(HookPreToolAction::Deny(reason).into_worker_action(summary.call_id.clone())); + } self.tool_calls_this_turn.fetch_add(1, Ordering::Relaxed); Ok(PreToolAction::Continue) } @@ -429,12 +479,23 @@ impl Interceptor for WorkerInterceptor { attachments: Vec::new(), }, }; + let mut aborts = Vec::new(); for hook in &self.registry.post_tool_call { - let action = hook.call(&summary).await; - if !matches!(action, HookPostToolAction::Continue) { - return Ok(action.into()); + let Some(action) = hook.call_optional(&summary).await.map_err(|error| { + InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) + })? + else { + continue; + }; + + if let HookPostToolAction::Abort(reason) = action { + aborts.push(reason); } } + aborts.sort(); + if let Some(reason) = aborts.into_iter().next() { + return Ok(PostToolAction::Abort(reason)); + } Ok(PostToolAction::Continue) } @@ -455,12 +516,21 @@ impl Interceptor for WorkerInterceptor { tool_calls_count: self.tool_calls_this_turn.load(Ordering::Relaxed), final_text_preview, }; + let mut should_pause = false; for hook in &self.registry.on_turn_end { - let action = hook.call(&info).await; - if !matches!(action, HookTurnEndAction::Finish) { - return Ok(action.into()); + let Some(action) = hook.call_optional(&info).await.map_err(|error| { + InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) + })? + else { + continue; + }; + if matches!(action, HookTurnEndAction::Pause) { + should_pause = true; } } + if should_pause { + return Ok(TurnEndAction::Pause); + } Ok(TurnEndAction::Finish) } } @@ -541,9 +611,12 @@ mod tests { #[async_trait] impl Hook for CountingHook { - async fn call(&self, _info: &PreRequestContext) -> HookPreRequestAction { + async fn call( + &self, + _info: &PreRequestContext, + ) -> Result { self.0.fetch_add(1, Ordering::Relaxed); - HookPreRequestAction::Continue + Ok(HookPreRequestAction::Continue) } } @@ -582,12 +655,15 @@ mod tests { #[async_trait] impl Hook for AppendingPreRequestHook { - async fn call(&self, input: &PreRequestContext) -> HookPreRequestAction { + async fn call( + &self, + input: &PreRequestContext, + ) -> Result { if let Some(system_items) = input.system_items() { self.saw_handle.store(true, Ordering::Relaxed); system_items.append_task_reminder("hook reminder"); } - HookPreRequestAction::Continue + Ok(HookPreRequestAction::Continue) } } @@ -982,33 +1058,42 @@ mod tests { #[async_trait] impl Hook for AbortingHook { - async fn call(&self, _info: &PreRequestContext) -> HookPreRequestAction { + async fn call( + &self, + _info: &PreRequestContext, + ) -> Result { self.0.store(true, Ordering::Relaxed); - HookPreRequestAction::Cancel("nope".into()) + Ok(HookPreRequestAction::Cancel("nope".into())) } } #[tokio::test] - async fn public_pre_tool_hook_deny_becomes_synthetic_error_and_short_circuits() { + async fn public_pre_tool_hook_denials_compose_without_short_circuiting() { struct DenyToolHook(Arc); struct CountingToolHook(Arc); #[async_trait] impl Hook for DenyToolHook { - async fn call(&self, input: &ToolCallSummary) -> HookPreToolAction { + async fn call( + &self, + input: &ToolCallSummary, + ) -> Result { self.0.fetch_add(1, Ordering::Relaxed); assert_eq!(input.call_id, "call-id"); assert_eq!(input.tool_name, "TaskList"); assert_eq!(input.arguments, serde_json::json!({"scope": "all"})); - HookPreToolAction::Deny("blocked by public hook".into()) + Ok(HookPreToolAction::Deny("blocked by public hook".into())) } } #[async_trait] impl Hook for CountingToolHook { - async fn call(&self, _input: &ToolCallSummary) -> HookPreToolAction { + async fn call( + &self, + _input: &ToolCallSummary, + ) -> Result { self.0.fetch_add(1, Ordering::Relaxed); - HookPreToolAction::Continue + Ok(HookPreToolAction::Continue) } } @@ -1041,7 +1126,7 @@ mod tests { other => panic!("expected synthetic denial, got {other:?}"), } assert_eq!(first_count.load(Ordering::Relaxed), 1); - assert_eq!(second_count.load(Ordering::Relaxed), 0); + assert_eq!(second_count.load(Ordering::Relaxed), 1); } #[tokio::test] @@ -1050,14 +1135,17 @@ mod tests { #[async_trait] impl Hook for AbortAfterToolHook { - async fn call(&self, input: &ToolResultSummary) -> HookPostToolAction { + async fn call( + &self, + input: &ToolResultSummary, + ) -> Result { self.0.fetch_add(1, Ordering::Relaxed); assert_eq!(input.call_id, "call-id"); assert_eq!(input.tool_name, "TaskList"); assert!(!input.is_error); assert_eq!(input.output.summary, "ok"); assert_eq!(input.output.content.as_deref(), Some("full")); - HookPostToolAction::Abort("post tool abort".into()) + Ok(HookPostToolAction::Abort("post tool abort".into())) } } @@ -1105,12 +1193,15 @@ mod tests { #[async_trait] impl Hook for PauseTurnEndHook { - async fn call(&self, input: &TurnEndInfo) -> HookTurnEndAction { + async fn call( + &self, + input: &TurnEndInfo, + ) -> Result { self.0.fetch_add(1, Ordering::Relaxed); assert_eq!(input.turn_index, 0); assert_eq!(input.tool_calls_count, 0); assert_eq!(input.final_text_preview, "done"); - HookTurnEndAction::Pause + Ok(HookTurnEndAction::Pause) } } @@ -1468,6 +1559,6 @@ mod tests { assert!(matches!(action, PreRequestAction::Cancel(_))); assert!(first_called.load(Ordering::Relaxed)); - assert_eq!(second_count.load(Ordering::Relaxed), 0); + assert_eq!(second_count.load(Ordering::Relaxed), 1); } } diff --git a/crates/worker/src/permission.rs b/crates/worker/src/permission.rs index 22b5c831..21a9c7b5 100644 --- a/crates/worker/src/permission.rs +++ b/crates/worker/src/permission.rs @@ -45,14 +45,17 @@ impl Worker { #[async_trait] impl Hook for PermissionHook { - async fn call(&self, input: &ToolCallSummary) -> HookPreToolAction { - match self.action_for(input) { + async fn call( + &self, + input: &ToolCallSummary, + ) -> Result { + Ok(match self.action_for(input) { ToolPermissionAction::Allow => HookPreToolAction::Continue, ToolPermissionAction::Deny => HookPreToolAction::Deny(permission_denied_message(input)), ToolPermissionAction::Ask => { HookPreToolAction::Deny(permission_ask_unsupported_message(input)) } - } + }) } } @@ -174,7 +177,7 @@ mod tests { )) .await; match denied { - HookPreToolAction::Deny(message) => { + Ok(HookPreToolAction::Deny(message)) => { assert!(message.contains("permission denied")); assert!(message.contains("Bash")); } @@ -192,7 +195,7 @@ mod tests { )) .await; match asked { - HookPreToolAction::Deny(message) => { + Ok(HookPreToolAction::Deny(message)) => { assert!(message.contains("permission ask unsupported")); assert!(message.contains("denied fail-closed")); } diff --git a/crates/worker/src/shutdown_after_idle.rs b/crates/worker/src/shutdown_after_idle.rs index b718cea1..9929ba57 100644 --- a/crates/worker/src/shutdown_after_idle.rs +++ b/crates/worker/src/shutdown_after_idle.rs @@ -70,9 +70,12 @@ impl TicketIntakeReadyShutdownHook { #[async_trait] impl Hook for TicketIntakeReadyShutdownHook { - async fn call(&self, info: &ToolResultSummary) -> HookPostToolAction { + async fn call( + &self, + info: &ToolResultSummary, + ) -> Result { self.observe_tool_result(info); - HookPostToolAction::Continue + Ok(HookPostToolAction::Continue) } } diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index db3536ca..582a6a81 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -40,6 +40,7 @@ use manifest::{ use crate::compact::state::CompactState; use crate::compact::usage_tracker::UsageTracker; +use crate::feature::background::FeatureBackgroundTaskRegistry; use crate::feature::builtin::memory::WorkspaceMemoryBackendError; use crate::feature::builtin::{ MemoryExtractFeature, MemoryExtractState, SessionExploreFeature, SessionExploreState, @@ -50,7 +51,10 @@ use crate::feature::{ FeatureRegistryInstallReport, dedupe_instruction_contributions, }; use crate::hook::{ - Hook, HookRegistryBuilder, OnPromptSubmit, OnTurnEnd, PostToolCall, PreLlmRequest, PreToolCall, + BeforeSessionRewriteAction, BeforeSessionRewriteContext, Hook, HookInvocationContext, + HookRegistry, HookRegistryBuilder, OnPromptSubmit, OnTurnEnd, PostToolCall, PreLlmRequest, + PreToolCall, RunCommittedContext, RunCommittedExit, RunExitContext, SessionRewriteKind, + WorkerStoppingContext, }; use crate::in_flight::InFlightEvents; use crate::internal_worker::{ @@ -63,6 +67,15 @@ 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"; +fn hook_run_exit(exit: &EngineRunExit) -> RunCommittedExit { + match exit { + EngineRunExit::Finished => RunCommittedExit::Finished, + EngineRunExit::Paused => RunCommittedExit::Paused, + EngineRunExit::Yielded => RunCommittedExit::Yielded, + EngineRunExit::Interrupted(_) => RunCommittedExit::Interrupted, + } +} + fn worker_orchestration_instruction() -> FeatureInstructionDeclaration { FeatureInstructionDeclaration::new( FeatureInstructionId::builtin(WORKER_ORCHESTRATION_INSTRUCTION_ID), @@ -1104,6 +1117,10 @@ pub struct Worker { /// continue to use `scope`; SubWorkerSpawn validates requested child scope here. delegation_scope: DelegationScope, hook_builder: HookRegistryBuilder, + /// Frozen callback set shared by Engine interception and Worker lifecycle boundaries. + hook_registry: Option>, + /// Executable background tasks registered by successfully installed features. + feature_background_tasks: FeatureBackgroundTaskRegistry, interceptor_installed: bool, /// Shared compaction state (present when threshold is configured). compact_state: Option>, @@ -1258,6 +1275,21 @@ pub struct Worker { } impl Worker { + pub async fn stop_feature_runtime(&mut self, reason: impl Into) { + if let Some(hooks) = self.hook_registry.clone() { + let context = WorkerStoppingContext { + 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"); + } + } + if let Err(error) = self.feature_background_tasks.shutdown().await { + tracing::warn!(error = %error, "feature background task shutdown failed"); + } + } + pub async fn wait_for_memory_jobs(&mut self) { if let Some(handle) = self.memory_task.take() && let Err(e) = handle.await @@ -1295,6 +1327,8 @@ impl Worker scope: self.scope.clone(), delegation_scope: self.delegation_scope.clone(), hook_builder: HookRegistryBuilder::new(), + hook_registry: None, + feature_background_tasks: FeatureBackgroundTaskRegistry::default(), interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -1488,6 +1522,8 @@ impl Worker { scope, delegation_scope, hook_builder: HookRegistryBuilder::new(), + hook_registry: None, + feature_background_tasks: FeatureBackgroundTaskRegistry::default(), interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -1715,12 +1751,15 @@ impl Worker { /// This deliberately does not scan `.yoi/skills` locally: when a Workspace /// HTTP client is available, catalog/detail/activation authority belongs to /// the Workspace backend API. - pub fn activate_skill(&mut self, name: &str) -> Result + pub async fn activate_skill( + &mut self, + name: &str, + ) -> Result where St: Clone + 'static, { let activation = self.workspace_client().activate_skill(name)?; - self.ensure_segment_head()?; + self.ensure_segment_head().await?; let body = format!( "Agent Skill `{}` activated from {}.\n\n{}", activation.name, activation.provenance.id, activation.body @@ -1821,6 +1860,21 @@ impl Worker { self.engine.as_ref().expect("worker taken during run") } + fn hook_invocation_context(&self, run_id: Option) -> HookInvocationContext { + HookInvocationContext { + workspace_id: self + .workspace_context + .workspace_id() + .map(|id| id.as_str().to_string()), + worker_id: self.manifest.worker.name.clone(), + session_id: self.session.session_id().to_string(), + session_revision: self.session.revision(), + run_id, + turn_index: Some(self.engine().turn_count()), + call_id: None, + } + } + /// Mutable access to the underlying Engine. /// /// Use this to register tools, hooks, or subscribers before calling @@ -1845,6 +1899,10 @@ impl Worker { ) -> FeatureRegistryInstallReport { let worker = self.engine.as_mut().expect("worker taken during run"); let report = registry.install_into_engine(worker, &mut self.hook_builder); + if report.has_errors() { + return report; + } + self.feature_background_tasks = report.background_tasks.clone(); for instruction in report.installed_instruction_contributions() { self.register_feature_instruction(instruction); } @@ -1942,11 +2000,14 @@ impl Worker { } /// Truncate the current segment to just before a previously listed user input. - pub fn rewind_to( + pub async fn rewind_to( &mut self, target: RewindTargetId, expected_head_entries: usize, ) -> Result { + self.prepare_session_rewrite(SessionRewriteKind::Rewind) + .await + .map_err(|error| RewindError::Invalid(error.to_string()))?; let loc = self.segment_state.location(); if target.segment_id != loc.segment_id { return Err(RewindError::Invalid( @@ -2337,6 +2398,7 @@ impl Worker { if !self.interceptor_installed { let builder = std::mem::take(&mut self.hook_builder); let registry = Arc::new(builder.build()); + self.hook_registry = Some(registry.clone()); let (post_run_threshold, request_threshold, retained) = self .manifest @@ -2554,7 +2616,7 @@ impl Worker { self.ensure_interceptor_installed(); self.ensure_system_prompt_materialized().await?; self.cleanup_finished_memory_task(); - self.ensure_segment_head()?; + self.ensure_segment_head().await?; if self.should_pre_run_compact() { self.join_memory_task().await; } @@ -3265,11 +3327,11 @@ impl Worker { /// `ensure_system_prompt_materialized` has just rendered. Subsequent /// calls fall through to entry-count comparison, which auto-forks /// when another writer has appended behind our back. - fn ensure_segment_head(&mut self) -> Result<(), WorkerError> { - let w = self.engine.as_ref().unwrap(); + async fn ensure_segment_head(&mut self) -> Result<(), WorkerError> { let loc = self.segment_state.location(); let entries_written = self.segment_state.entries_written(); if entries_written == 0 { + let w = self.engine.as_ref().unwrap(); let initial = LogEntry::AnnotatedSegmentStart { ts: segment_log::now_millis(), session_id: loc.session_id, @@ -3305,6 +3367,9 @@ impl Worker { // state up to that turn). The new SegmentStart replaces the mirror // and is broadcast through the sink so existing subscribers reset // their view. + self.prepare_session_rewrite(SessionRewriteKind::Fork) + .await?; + let w = self.engine.as_ref().unwrap(); let fork_segment_id = session_store::new_segment_id(); let entry = LogEntry::AnnotatedSegmentStart { ts: segment_log::now_millis(), @@ -3371,10 +3436,40 @@ impl Worker { where St: Clone + 'static, { + let run_id = uuid::Uuid::now_v7().to_string(); + let hook_exit = hook_run_exit(&result); + if let Some(hooks) = self.hook_registry.clone() { + let context = RunExitContext { + invocation: self.hook_invocation_context(Some(run_id.clone())), + 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"); + } + } if matches!(&result, EngineRunExit::Interrupted(_)) { self.terminalize_orphan_tool_calls()?; } self.persist_turn(history_before, &result).await?; + let committed_invocation = self.hook_invocation_context(Some(run_id)); + if let Some(hooks) = self.hook_registry.clone() { + 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(), + }; + if let Err(error) = hooks.on_run_committed(&context).await { + tracing::warn!(error = %error, "run-committed hook requires attention"); + } + } + if let Err(error) = self + .feature_background_tasks + .start_run_committed(committed_invocation) + { + tracing::warn!(error = %error, "run-committed background task start failed"); + } if matches!(result, EngineRunExit::Yielded) { self.last_run_interrupted = true; @@ -3552,6 +3647,35 @@ impl Worker { /// The controller only calls this while Idle. Paused turns keep their /// interrupted Engine state intact and are intentionally rejected before /// this method is reached. + async fn prepare_session_rewrite( + &mut self, + kind: SessionRewriteKind, + ) -> Result<(), WorkerError> { + self.feature_background_tasks + .before_session_rewrite() + .await + .map_err(|error| WorkerError::FeatureLifecycle(error.to_string()))?; + if let Some(hooks) = self.hook_registry.clone() { + 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(), + }; + match hooks + .before_session_rewrite(&context) + .await + .map_err(|error| WorkerError::FeatureLifecycle(error.to_string()))? + { + BeforeSessionRewriteAction::Continue => {} + BeforeSessionRewriteAction::Deny(reason) => { + return Err(WorkerError::FeatureLifecycle(reason)); + } + } + } + Ok(()) + } + pub async fn manual_compact(&mut self) -> Result { if self.manifest.compaction.is_none() { let message = @@ -3568,7 +3692,7 @@ impl Worker { self.ensure_interceptor_installed(); self.cleanup_finished_memory_task(); - self.ensure_segment_head()?; + self.ensure_segment_head().await?; let state = self.compact_state.clone(); if state.as_ref().is_some_and(|s| s.is_disabled()) { @@ -3767,6 +3891,8 @@ impl Worker { /// Runs one parent-owned observable compaction service and returns the new /// Segment ID. Lifecycle revisions are committed before they are broadcast. pub async fn compact(&mut self, retained_tokens: u64) -> Result { + self.prepare_session_rewrite(SessionRewriteKind::Compact) + .await?; let mut lifecycle = CompactionLifecycle { schema_version: 2, compaction_id: uuid::Uuid::now_v7().to_string(), @@ -5302,6 +5428,8 @@ where scope, delegation_scope: common.delegation_scope, hook_builder: HookRegistryBuilder::new(), + hook_registry: None, + feature_background_tasks: FeatureBackgroundTaskRegistry::default(), interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -5386,6 +5514,8 @@ where scope, delegation_scope: common.delegation_scope, hook_builder: HookRegistryBuilder::new(), + hook_registry: None, + feature_background_tasks: FeatureBackgroundTaskRegistry::default(), interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -5505,6 +5635,8 @@ where scope, delegation_scope: common.delegation_scope, hook_builder: HookRegistryBuilder::new(), + hook_registry: None, + feature_background_tasks: FeatureBackgroundTaskRegistry::default(), interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -5879,6 +6011,8 @@ where scope, delegation_scope: common.delegation_scope, hook_builder: HookRegistryBuilder::new(), + hook_registry: None, + feature_background_tasks: FeatureBackgroundTaskRegistry::default(), interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -6553,6 +6687,9 @@ pub enum WorkerError { #[error("invalid durable Worker state: {0}")] InvalidState(String), + #[error("feature lifecycle rejected operation: {0}")] + FeatureLifecycle(String), + #[error("Flow input rejected: {0}")] FlowInput(String), @@ -7776,15 +7913,17 @@ mod build_summary_prompt_tests { async fn call( &self, _input: &crate::hook::ToolCallSummary, - ) -> crate::hook::HookPreToolAction { - if self - .should_pause - .swap(false, std::sync::atomic::Ordering::SeqCst) - { - crate::hook::HookPreToolAction::Pause - } else { - crate::hook::HookPreToolAction::Continue - } + ) -> Result { + Ok( + if self + .should_pause + .swap(false, std::sync::atomic::Ordering::SeqCst) + { + crate::hook::HookPreToolAction::Pause + } else { + crate::hook::HookPreToolAction::Continue + }, + ) } } @@ -7837,7 +7976,7 @@ mod build_summary_prompt_tests { ) .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); worker.last_run_interrupted = true; worker.engine_mut().set_active_run_turn_count(Some(3)); @@ -7880,7 +8019,7 @@ mod build_summary_prompt_tests { ) .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); worker.engine_mut().set_turn_count(7); worker.last_run_interrupted = true; worker.engine_mut().set_active_run_turn_count(Some(3)); @@ -7900,7 +8039,7 @@ mod build_summary_prompt_tests { ) .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); let fork_segment_id = worker.segment_id(); assert_ne!(fork_segment_id, source_segment_id); @@ -7944,7 +8083,7 @@ mod build_summary_prompt_tests { ) .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); let report = worker .install_runtime_flow_transition_feature() .expect("scoped Workspace Flow feature"); @@ -7978,7 +8117,7 @@ mod build_summary_prompt_tests { ) .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); let disabled = worker.prepare_flow_input(vec![Segment::Flow { selector: "builtin:coder-review".to_string(), }]); @@ -8254,7 +8393,7 @@ mod build_summary_prompt_tests { ) .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); std::fs::write( temp.path() .join(worker.session_id().to_string()) @@ -8303,7 +8442,7 @@ mod build_summary_prompt_tests { ) .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); (dir, worker) } @@ -8463,7 +8602,7 @@ mod build_summary_prompt_tests { let expected_truncate_entries = targets[0].truncate_entries; let target = targets[0].id.clone(); - let applied = worker.rewind_to(target, head_entries).unwrap(); + let applied = worker.rewind_to(target, head_entries).await.unwrap(); assert_eq!(preview_segments(&applied.input), "second message"); assert_eq!( @@ -8520,6 +8659,7 @@ mod build_summary_prompt_tests { let applied = worker .rewind_to(targets[0].id.clone(), head_entries) + .await .unwrap(); assert_eq!(applied.summary.truncated_to_entries, 5); assert!(matches!( @@ -8562,7 +8702,7 @@ mod build_summary_prompt_tests { payload: serde_json::json!({"value": true}), }, ); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); let fork_location = worker.segment_state.location(); assert_ne!(fork_location.segment_id, source_location.segment_id); let fork_entries = worker @@ -8593,6 +8733,7 @@ mod build_summary_prompt_tests { let err = worker .rewind_to(targets[0].id.clone(), head_entries) + .await .unwrap_err() .to_string(); @@ -8673,7 +8814,7 @@ mod build_summary_prompt_tests { .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); worker.wire_history_persistence(); worker.set_history_for_test(vec![ Item::tool_call("call-known", "Read", "{}"), @@ -8790,7 +8931,7 @@ mod build_summary_prompt_tests { .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); worker.wire_history_persistence(); worker.set_history_for_test(vec![Item::tool_call("call-1", "Read", "{}")]); @@ -8870,7 +9011,7 @@ mod build_summary_prompt_tests { .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); worker.wire_history_persistence(); let dangling_call = Item::tool_call("call-1", "SideEffect", "{}"); worker @@ -9186,8 +9327,8 @@ mod build_summary_prompt_tests { std::fs::create_dir_all(&cwd).unwrap(); let scope = Scope::writable(&cwd).unwrap(); let authority = WorkerFilesystemAuthority::local(cwd.clone(), cwd.clone()); - let mut worker = tokio::runtime::Runtime::new() - .unwrap() + let runtime = tokio::runtime::Runtime::new().unwrap(); + let mut worker = runtime .block_on(Worker::new( manifest, Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), @@ -9204,7 +9345,9 @@ mod build_summary_prompt_tests { )) .unwrap(); - let activation = worker.activate_skill("triage-errors").unwrap(); + let activation = runtime + .block_on(worker.activate_skill("triage-errors")) + .unwrap(); assert_eq!(activation.name, "triage-errors"); server.join().unwrap(); @@ -9274,7 +9417,7 @@ mod build_summary_prompt_tests { ) .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); worker.wire_history_persistence(); let evidence = Item::user_message( "The cancellation regression must leave this evidence available for retry.",