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