feat: add executable feature lifecycle hooks

This commit is contained in:
2026-09-04 14:50:44 +09:00
parent 9bd08a3a5b
commit af06eecfd0
12 changed files with 1969 additions and 235 deletions
+9 -3
View File
@@ -1590,7 +1590,9 @@ async fn controller_loop<C, St>(
&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<C, St>(
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<C, St>(
async fn apply_rewind<C, St>(
worker: &mut Worker<C, St>,
working_event_tx: &broadcast::Sender<Event>,
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);
+372 -95
View File
@@ -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<String>, description: impl Into<String>) -> Self {
pub fn worker_managed(name: impl Into<String>, description: impl Into<String>) -> 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<String>,
policy: HookFailurePolicy,
hook: impl Hook<OnPromptSubmit> + '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<String>,
policy: HookFailurePolicy,
hook: impl Hook<PreLlmRequest> + '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<String>,
hook: impl Hook<PreLlmRequest> + '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<String>,
policy: HookFailurePolicy,
hook: impl Hook<PreToolCall> + '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<String>,
hook: impl Hook<PreToolCall> + '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<String>,
policy: HookFailurePolicy,
hook: impl Hook<PostToolCall> + '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<String>,
hook: impl Hook<PostToolCall> + '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<String>,
policy: HookFailurePolicy,
hook: impl Hook<OnTurnEnd> + '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<String>,
hook: impl Hook<OnTurnEnd> + '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<String>,
policy: HookFailurePolicy,
hook: impl Hook<RunExit> + '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<String>,
policy: HookFailurePolicy,
hook: impl Hook<RunCommitted> + '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<String>,
policy: HookFailurePolicy,
hook: impl Hook<BeforeSessionRewrite> + '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<String>,
policy: HookFailurePolicy,
hook: impl Hook<WorkerStopping> + '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<ToolDefinition>,
installed_tool_names: &'a mut HashMap<String, FeatureId>,
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<FeatureInstallReport>,
pub services: FeatureServiceRegistry,
pub background_tasks: FeatureBackgroundTaskRegistry,
pub plan_error: Option<FeaturePlanError>,
}
@@ -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::<Vec<_>>();
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<crate::hook::HookPreToolAction, crate::hook::HookError> {
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());
+705
View File
@@ -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<AtomicBool>,
notify: Arc<Notify>,
}
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<dyn FeatureBackgroundTask>,
}
#[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::<std::collections::BTreeSet<_>>();
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<BTreeMap<u64, RunningTask>>,
diagnostics: Mutex<Vec<BackgroundTaskDiagnostic>>,
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<RegistryInner>,
}
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::<Vec<_>>(),
)
.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<BackgroundTaskStart, HookError> {
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(&registration.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::<Vec<_>>();
for (feature_id, task_name) in tasks {
let _ = self.start(&feature_id, &task_name, invocation.clone())?;
}
Ok(())
}
pub fn diagnostics(&self) -> Vec<BackgroundTaskDiagnostic> {
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::<Vec<_>>();
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<dyn FeatureBackgroundTask>,
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<AtomicUsize>,
completed: Arc<Notify>,
}
#[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);
}
}
+13 -7
View File
@@ -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<PreLlmRequest> for TaskReminderPreRequestHook {
async fn call(&self, input: &PreRequestContext) -> HookPreRequestAction {
async fn call(
&self,
input: &PreRequestContext,
) -> Result<HookPreRequestAction, crate::hook::HookError> {
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<PreToolCall> for TaskReminderToolUsageHook {
async fn call(&self, input: &ToolCallSummary) -> HookPreToolAction {
async fn call(
&self,
input: &ToolCallSummary,
) -> Result<HookPreToolAction, crate::hook::HookError> {
if is_task_management_tool(&input.tool_name) {
self.state.reminder_state.note_task_management();
}
HookPreToolAction::Continue
Ok(HookPreToolAction::Continue)
}
}
+1 -2
View File
@@ -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 {
+1 -1
View File
@@ -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"));
}
+527 -31
View File
@@ -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<String>) -> 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<E: HookEventKind>: Send + Sync {
async fn call(&self, input: &E::Input) -> E::Output;
async fn call(&self, input: &E::Input) -> Result<E::Output, HookError>;
}
// =============================================================================
// Hook Registry
// =============================================================================
/// Stable provenance attached to every Worker lifecycle callback.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct HookInvocationContext {
pub workspace_id: Option<String>,
pub worker_id: String,
pub session_id: String,
pub session_revision: u64,
pub run_id: Option<String>,
pub turn_index: Option<usize>,
pub call_id: Option<String>,
}
#[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<HistoryEntry<SessionHistoryMetadata>>,
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<HistoryEntry<SessionHistoryMetadata>>,
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<E: HookEventKind> {
owner: String,
policy: HookFailurePolicy,
hook: Box<dyn Hook<E>>,
}
impl<E: HookEventKind> RegisteredHook<E> {
pub(crate) async fn call(&self, input: &E::Input) -> Result<E::Output, HookExecutionError> {
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<Option<E::Output>, 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<Box<dyn Hook<OnPromptSubmit>>>,
pre_llm_request: Vec<Box<dyn Hook<PreLlmRequest>>>,
pre_tool_call: Vec<Box<dyn Hook<PreToolCall>>>,
post_tool_call: Vec<Box<dyn Hook<PostToolCall>>>,
on_turn_end: Vec<Box<dyn Hook<OnTurnEnd>>>,
on_prompt_submit: Vec<RegisteredHook<OnPromptSubmit>>,
pre_llm_request: Vec<RegisteredHook<PreLlmRequest>>,
pre_tool_call: Vec<RegisteredHook<PreToolCall>>,
post_tool_call: Vec<RegisteredHook<PostToolCall>>,
on_turn_end: Vec<RegisteredHook<OnTurnEnd>>,
run_exit: Vec<RegisteredHook<RunExit>>,
run_committed: Vec<RegisteredHook<RunCommitted>>,
before_session_rewrite: Vec<RegisteredHook<BeforeSessionRewrite>>,
worker_stopping: Vec<RegisteredHook<WorkerStopping>>,
}
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<String>,
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<OnPromptSubmit> + '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<PreLlmRequest> + '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<PreToolCall> + 'static) {
self.pre_tool_call.push(Box::new(hook));
}
pub fn add_post_tool_call(&mut self, hook: impl Hook<PostToolCall> + 'static) {
self.post_tool_call.push(Box::new(hook));
}
pub fn add_on_turn_end(&mut self, hook: impl Hook<OnTurnEnd> + '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<Box<dyn Hook<OnPromptSubmit>>>,
pub(crate) pre_llm_request: Vec<Box<dyn Hook<PreLlmRequest>>>,
pub(crate) pre_tool_call: Vec<Box<dyn Hook<PreToolCall>>>,
pub(crate) post_tool_call: Vec<Box<dyn Hook<PostToolCall>>>,
pub(crate) on_turn_end: Vec<Box<dyn Hook<OnTurnEnd>>>,
pub(crate) on_prompt_submit: Vec<RegisteredHook<OnPromptSubmit>>,
pub(crate) pre_llm_request: Vec<RegisteredHook<PreLlmRequest>>,
pub(crate) pre_tool_call: Vec<RegisteredHook<PreToolCall>>,
pub(crate) post_tool_call: Vec<RegisteredHook<PostToolCall>>,
pub(crate) on_turn_end: Vec<RegisteredHook<OnTurnEnd>>,
run_exit: Vec<RegisteredHook<RunExit>>,
run_committed: Vec<RegisteredHook<RunCommitted>>,
before_session_rewrite: Vec<RegisteredHook<BeforeSessionRewrite>>,
worker_stopping: Vec<RegisteredHook<WorkerStopping>>,
diagnostics: std::sync::Mutex<Vec<HookExecutionError>>,
}
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<HookExecutionError> {
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<BeforeSessionRewriteAction, HookExecutionError> {
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<BeforeSessionRewrite> for RewriteHook {
async fn call(
&self,
_input: &BeforeSessionRewriteContext,
) -> Result<BeforeSessionRewriteAction, HookError> {
Ok(self.action.clone())
}
}
struct FailingRewriteHook;
#[async_trait]
impl Hook<BeforeSessionRewrite> for FailingRewriteHook {
async fn call(
&self,
_input: &BeforeSessionRewriteContext,
) -> Result<BeforeSessionRewriteAction, HookError> {
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()));
}
}
+29 -24
View File
@@ -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)
+123 -32
View File
@@ -246,12 +246,22 @@ impl Interceptor<SessionHistoryMetadata> 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<SystemItem> = std::mem::take(
&mut *self
.pending_attachments
@@ -344,12 +354,29 @@ impl Interceptor<SessionHistoryMetadata> 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<SystemItem> = std::mem::take(
&mut *pending_hook_system_items
@@ -404,12 +431,35 @@ impl Interceptor<SessionHistoryMetadata> 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<SessionHistoryMetadata> 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<SessionHistoryMetadata> 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<PreLlmRequest> for CountingHook {
async fn call(&self, _info: &PreRequestContext) -> HookPreRequestAction {
async fn call(
&self,
_info: &PreRequestContext,
) -> Result<HookPreRequestAction, crate::hook::HookError> {
self.0.fetch_add(1, Ordering::Relaxed);
HookPreRequestAction::Continue
Ok(HookPreRequestAction::Continue)
}
}
@@ -582,12 +655,15 @@ mod tests {
#[async_trait]
impl Hook<PreLlmRequest> for AppendingPreRequestHook {
async fn call(&self, input: &PreRequestContext) -> HookPreRequestAction {
async fn call(
&self,
input: &PreRequestContext,
) -> Result<HookPreRequestAction, crate::hook::HookError> {
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<PreLlmRequest> for AbortingHook {
async fn call(&self, _info: &PreRequestContext) -> HookPreRequestAction {
async fn call(
&self,
_info: &PreRequestContext,
) -> Result<HookPreRequestAction, crate::hook::HookError> {
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<AtomicUsize>);
struct CountingToolHook(Arc<AtomicUsize>);
#[async_trait]
impl Hook<PreToolCall> for DenyToolHook {
async fn call(&self, input: &ToolCallSummary) -> HookPreToolAction {
async fn call(
&self,
input: &ToolCallSummary,
) -> Result<HookPreToolAction, crate::hook::HookError> {
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<PreToolCall> for CountingToolHook {
async fn call(&self, _input: &ToolCallSummary) -> HookPreToolAction {
async fn call(
&self,
_input: &ToolCallSummary,
) -> Result<HookPreToolAction, crate::hook::HookError> {
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<PostToolCall> for AbortAfterToolHook {
async fn call(&self, input: &ToolResultSummary) -> HookPostToolAction {
async fn call(
&self,
input: &ToolResultSummary,
) -> Result<HookPostToolAction, crate::hook::HookError> {
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<OnTurnEnd> for PauseTurnEndHook {
async fn call(&self, input: &TurnEndInfo) -> HookTurnEndAction {
async fn call(
&self,
input: &TurnEndInfo,
) -> Result<HookTurnEndAction, crate::hook::HookError> {
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);
}
}
+8 -5
View File
@@ -45,14 +45,17 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
#[async_trait]
impl Hook<PreToolCall> for PermissionHook {
async fn call(&self, input: &ToolCallSummary) -> HookPreToolAction {
match self.action_for(input) {
async fn call(
&self,
input: &ToolCallSummary,
) -> Result<HookPreToolAction, crate::hook::HookError> {
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"));
}
+5 -2
View File
@@ -70,9 +70,12 @@ impl TicketIntakeReadyShutdownHook {
#[async_trait]
impl Hook<PostToolCall> for TicketIntakeReadyShutdownHook {
async fn call(&self, info: &ToolResultSummary) -> HookPostToolAction {
async fn call(
&self,
info: &ToolResultSummary,
) -> Result<HookPostToolAction, crate::hook::HookError> {
self.observe_tool_result(info);
HookPostToolAction::Continue
Ok(HookPostToolAction::Continue)
}
}
+176 -33
View File
@@ -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<C: LlmClient, St: Store> {
/// 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<Arc<HookRegistry>>,
/// 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<Arc<CompactState>>,
@@ -1258,6 +1275,21 @@ pub struct Worker<C: LlmClient, St: Store> {
}
impl<C: LlmClient + 'static, St: Store + 'static> Worker<C, St> {
pub async fn stop_feature_runtime(&mut self, reason: impl Into<String>) {
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<C: LlmClient + Clone + 'static, St: Store + Clone + 'static> Worker<C, St>
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// 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<SkillActivationResponse, WorkerError>
pub async fn activate_skill(
&mut self,
name: &str,
) -> Result<SkillActivationResponse, WorkerError>
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
self.engine.as_ref().expect("worker taken during run")
}
fn hook_invocation_context(&self, run_id: Option<String>) -> 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<C: LlmClient + 'static, St: Store> Worker<C, St> {
) -> 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<C: LlmClient + 'static, St: Store> Worker<C, St> {
}
/// 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<RewindAppliedState, RewindError> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// `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<C: LlmClient + 'static, St: Store> Worker<C, St> {
// 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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// 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<ManualCompactResult, WorkerError> {
if self.manifest.compaction.is_none() {
let message =
@@ -3568,7 +3692,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// 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<SegmentId, WorkerError> {
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<crate::hook::HookPreToolAction, crate::hook::HookError> {
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.",