feat: add executable feature lifecycle hooks
This commit is contained in:
@@ -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
@@ -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());
|
||||
|
||||
@@ -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(®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::<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);
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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()));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"));
|
||||
}
|
||||
|
||||
@@ -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
@@ -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.",
|
||||
|
||||
Reference in New Issue
Block a user