From 68b1aa64e9e47f396373680d29ed899fb0a1c88a Mon Sep 17 00:00:00 2001 From: Hare Date: Thu, 3 Sep 2026 14:46:07 +0900 Subject: [PATCH 01/22] feat: type async interceptor failures --- crates/agen/examples/engine_cli.rs | 6 +- crates/agen/src/engine.rs | 77 ++++++- crates/agen/src/interceptor.rs | 149 ++++++++++++-- crates/agen/src/lib.rs | 4 +- crates/agen/tests/engine_state_test.rs | 202 ++++++++++++++++++- crates/agen/tests/parallel_execution_test.rs | 43 ++-- crates/session-store/tests/session_test.rs | 6 +- crates/worker/src/compact/worker.rs | 41 ++-- crates/worker/src/ipc/interceptor.rs | 96 ++++----- crates/worker/src/worker.rs | 4 + 10 files changed, 505 insertions(+), 123 deletions(-) diff --git a/crates/agen/examples/engine_cli.rs b/crates/agen/examples/engine_cli.rs index 0595cf95..2c50d354 100644 --- a/crates/agen/examples/engine_cli.rs +++ b/crates/agen/examples/engine_cli.rs @@ -40,7 +40,7 @@ use tracing_subscriber::EnvFilter; use agen::{ Engine, EngineRunExit, RunInterruptionReason, - interceptor::{Interceptor, PostToolAction, ToolResultInfo}, + interceptor::{Interceptor, InterceptorResult, PostToolAction, ToolResultInfo}, llm_client::{ LlmClient, capability::{CacheStrategy, ModelCapability, StructuredOutput, ToolCallingSupport}, @@ -280,7 +280,7 @@ impl ToolResultPrinterPolicy { #[async_trait] impl Interceptor for ToolResultPrinterPolicy { - async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction { + async fn post_tool_call(&self, info: &mut ToolResultInfo) -> InterceptorResult { let name = self .call_names .lock() @@ -294,7 +294,7 @@ impl Interceptor for ToolResultPrinterPolicy { println!(" Result ({}): ✅ {}", name, info.result.summary); } - PostToolAction::Continue + Ok(PostToolAction::Continue) } } diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index a3e79483..ae3279b2 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -15,8 +15,8 @@ use crate::{ }, handler::{ErrorKind, StatusKind, ToolUseBlockStart, UsageKind}, interceptor::{ - DefaultInterceptor, Interceptor, PostToolAction, PreRequestAction, PreToolAction, - PromptAction, ToolCallInfo, ToolResultInfo, TurnEndAction, + DefaultInterceptor, Interceptor, InterceptorFailure, InterceptorPoint, PostToolAction, + PreRequestAction, PreToolAction, PromptAction, ToolCallInfo, ToolResultInfo, TurnEndAction, }, llm_client::{ ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream, @@ -58,6 +58,9 @@ pub enum EngineError { /// A durable-history observer rejected an item before it entered history. #[error("History append failed: {0}")] HistoryAppend(String), + /// A trusted host interceptor callback failed. + #[error(transparent)] + Interceptor(#[from] InterceptorFailure), /// Tool terminalization lost its execution-attempt compare-and-set fence. #[error("Tool execution attempt fence failed: {0}")] ToolAttemptFence(String), @@ -153,6 +156,8 @@ pub enum EngineRunExit { /// A typed reason why an engine run could not finish normally. #[derive(Debug)] pub enum RunInterruptionReason { + /// A trusted host interceptor callback failed at a typed lifecycle point. + Interceptor(InterceptorFailure), LimitReached, ContextWindowExceeded, Cancelled, @@ -173,6 +178,9 @@ impl From> for EngineRunExit { } Err(EngineError::Cancelled) => Self::Interrupted(RunInterruptionReason::Cancelled), Err(EngineError::PauseRequested) => Self::Paused, + Err(EngineError::Interceptor(failure)) => { + Self::Interrupted(RunInterruptionReason::Interceptor(failure)) + } Err(error) => Self::Interrupted(RunInterruptionReason::Unexpected(error)), } } @@ -1092,7 +1100,9 @@ impl Engine { EngineError::Cancelled => "Cancelled".to_string(), _ => err.to_string(), }; - self.interceptor.on_abort(&reason).await; + if let Err(error) = self.interceptor.on_abort(&reason).await { + return Err(InterceptorFailure::new(InterceptorPoint::Abort, error).into()); + } Err(err) } } @@ -1175,7 +1185,17 @@ impl Engine { context, }; - match self.interceptor.pre_tool_call(&mut info).await { + let pre_tool_action = + self.interceptor + .pre_tool_call(&mut info) + .await + .map_err(|error| { + EngineError::from(InterceptorFailure::new( + InterceptorPoint::PreToolCall, + error, + )) + })?; + match pre_tool_action { PreToolAction::Continue => {} PreToolAction::Skip => { continue; @@ -1476,7 +1496,17 @@ impl Engine { context: context.clone(), }; - match self.interceptor.post_tool_call(&mut info).await { + let post_tool_action = + self.interceptor + .post_tool_call(&mut info) + .await + .map_err(|error| { + EngineError::from(InterceptorFailure::new( + InterceptorPoint::PostToolCall, + error, + )) + })?; + match post_tool_action { PostToolAction::Continue => {} PostToolAction::Abort(reason) => { abort_reason = Some(reason); @@ -1612,7 +1642,12 @@ impl Engine { .interceptor .pending_history_appends() .await - .map_err(EngineError::HistoryAppend)?; + .map_err(|error| { + EngineError::from(InterceptorFailure::new( + InterceptorPoint::PendingHistoryAppends, + error, + )) + })?; if !pending.is_empty() { self.append_history_items(history, pending, annotate)?; } @@ -1679,7 +1714,17 @@ impl Engine { } // Interceptor: pre_llm_request - match self.interceptor.pre_llm_request(&mut request_context).await { + let pre_request_action = self + .interceptor + .pre_llm_request(&mut request_context) + .await + .map_err(|error| { + EngineError::from(InterceptorFailure::new( + InterceptorPoint::PreLlmRequest, + error, + )) + })?; + match pre_request_action { PreRequestAction::Cancel(reason) => { info!(reason = %reason, "Aborted by interceptor"); for cb in &self.turn_end_cbs { @@ -1795,7 +1840,14 @@ impl Engine { if tool_calls.is_empty() { let turn_end_context = history.items_cloned(); - match self.interceptor.on_turn_end(&turn_end_context).await { + let turn_end_action = self + .interceptor + .on_turn_end(&turn_end_context) + .await + .map_err(|error| { + EngineError::from(InterceptorFailure::new(InterceptorPoint::TurnEnd, error)) + })?; + match turn_end_action { TurnEndAction::Finish => { return Ok(EngineResult::Finished); } @@ -2502,7 +2554,14 @@ impl Engine { // Supplying new user input abandons any paused/yielded logical run. self.active_run_turn_count = None; let mut user_item = Item::user_message(user_input); - let extras = match self.interceptor.on_prompt_submit(&mut user_item).await { + let prompt_action = match self.interceptor.on_prompt_submit(&mut user_item).await { + Ok(action) => action, + Err(error) => { + let error = InterceptorFailure::new(InterceptorPoint::PromptSubmit, error).into(); + return self.finalize_interruption(Err(error)).await; + } + }; + let extras = match prompt_action { PromptAction::Cancel(reason) => { return self .finalize_interruption(Err(EngineError::Aborted(reason))) diff --git a/crates/agen/src/interceptor.rs b/crates/agen/src/interceptor.rs index 5c5a307f..ecb1fd8f 100644 --- a/crates/agen/src/interceptor.rs +++ b/crates/agen/src/interceptor.rs @@ -11,6 +11,101 @@ use async_trait::async_trait; use crate::Item; use crate::tool::{Tool, ToolCall, ToolExecutionContext, ToolMeta, ToolResult}; +// ============================================================================= +// Failure Types +// ============================================================================= + +/// A typed failure returned by an [`Interceptor`] implementation. +/// +/// The Engine attaches the exact [`InterceptorPoint`] at which the failure was +/// observed before exposing it through the run termination boundary. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("{message}")] +pub struct InterceptorError { + message: String, +} + +impl InterceptorError { + /// Create an interceptor failure with a caller-defined message. + pub fn new(message: impl Into) -> Self { + Self { + message: message.into(), + } + } + + /// Return the failure message supplied by the interceptor. + pub fn message(&self) -> &str { + &self.message + } +} + +impl From for InterceptorError { + fn from(message: String) -> Self { + Self::new(message) + } +} + +impl From<&str> for InterceptorError { + fn from(message: &str) -> Self { + Self::new(message) + } +} + +/// The Engine lifecycle point at which an interceptor failed. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InterceptorPoint { + PromptSubmit, + PendingHistoryAppends, + PreLlmRequest, + PreToolCall, + PostToolCall, + TurnEnd, + Abort, +} + +impl std::fmt::Display for InterceptorPoint { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let name = match self { + Self::PromptSubmit => "prompt_submit", + Self::PendingHistoryAppends => "pending_history_appends", + Self::PreLlmRequest => "pre_llm_request", + Self::PreToolCall => "pre_tool_call", + Self::PostToolCall => "post_tool_call", + Self::TurnEnd => "turn_end", + Self::Abort => "abort", + }; + formatter.write_str(name) + } +} + +/// An interceptor failure bound to the exact Engine lifecycle point that ran it. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("{point} interceptor failed: {error}")] +pub struct InterceptorFailure { + point: InterceptorPoint, + #[source] + error: InterceptorError, +} + +impl InterceptorFailure { + pub(crate) fn new(point: InterceptorPoint, error: InterceptorError) -> Self { + Self { point, error } + } + + /// The lifecycle point that returned the failure. + pub fn point(&self) -> InterceptorPoint { + self.point + } + + /// The typed error returned by the interceptor. + pub fn error(&self) -> &InterceptorError { + &self.error + } +} + +/// Result returned by asynchronous interceptor lifecycle methods. +pub type InterceptorResult = Result; + // ============================================================================= // Action Enums // ============================================================================= @@ -130,14 +225,19 @@ pub struct ToolResultInfo { /// Intercepts the Engine execution loop at key decision points. /// -/// All methods have default implementations that let the Engine -/// proceed without intervention. Callers provide richer implementations for -/// approval flows, permission checks, etc. +/// Every lifecycle method is asynchronous and returns [`InterceptorResult`], +/// keeping implementation failure separate from the method's control-flow +/// action. The Engine reports a failure as a typed run interruption annotated +/// with the exact [`InterceptorPoint`] that failed. +/// +/// All methods have default implementations that let the Engine proceed +/// without intervention. Callers provide richer implementations for approval +/// flows, permission checks, and other trusted host adaptation. #[async_trait] pub trait Interceptor: Send + Sync { - /// Called after receiving user input, before adding to history. - async fn on_prompt_submit(&self, _item: &mut Item) -> PromptAction { - PromptAction::Continue + /// Called after receiving user input, before adding it to Engine history. + async fn on_prompt_submit(&self, _item: &mut Item) -> InterceptorResult { + Ok(PromptAction::Continue) } /// Items that should be **committed to `engine.history`** just @@ -158,7 +258,7 @@ pub trait Interceptor: Send + Sync { /// reproducible per-request transformations (pruning, content /// trimming, cache anchors) that depend only on the existing /// history. - async fn pending_history_appends(&self) -> Result, String> { + async fn pending_history_appends(&self) -> InterceptorResult> { Ok(Vec::new()) } @@ -170,27 +270,38 @@ pub trait Interceptor: Send + Sync { /// If an interceptor derives a human/model-visible nudge from the current /// request context, return [`PreRequestAction::ContinueWith`] so the Engine /// commits it to history before the request is sent. - async fn pre_llm_request(&self, _context: &mut Vec) -> PreRequestAction { - PreRequestAction::Continue + async fn pre_llm_request( + &self, + _context: &mut Vec, + ) -> InterceptorResult { + Ok(PreRequestAction::Continue) } /// Called before each tool is executed. - async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> PreToolAction { - PreToolAction::Continue + async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> InterceptorResult { + Ok(PreToolAction::Continue) } - /// Called after each tool completes. - async fn post_tool_call(&self, _info: &mut ToolResultInfo) -> PostToolAction { - PostToolAction::Continue + /// Called after each tool reaches one terminal result. + async fn post_tool_call( + &self, + _info: &mut ToolResultInfo, + ) -> InterceptorResult { + Ok(PostToolAction::Continue) } - /// Called when a turn ends with no tool calls. - async fn on_turn_end(&self, _history: &[Item]) -> TurnEndAction { - TurnEndAction::Finish + /// Called at the assistant boundary when a completed response has no tool calls. + /// + /// This is not the logical run termination observer. A host that needs that + /// boundary must inspect the returned [`crate::EngineRunExit`]. + async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult { + Ok(TurnEndAction::Finish) } - /// Called when execution is interrupted (abort or cancel). - async fn on_abort(&self, _reason: &str) {} + /// Called once when execution is interrupted (abort, cancellation, or failure). + async fn on_abort(&self, _reason: &str) -> InterceptorResult<()> { + Ok(()) + } } /// Default interceptor: no intervention. Engine proceeds through the loop diff --git a/crates/agen/src/lib.rs b/crates/agen/src/lib.rs index 96d4255c..52a7a246 100644 --- a/crates/agen/src/lib.rs +++ b/crates/agen/src/lib.rs @@ -26,7 +26,9 @@ pub use engine::{ }; pub use handler::ToolUseBlockStart; pub use history::{History, HistoryEntry}; -pub use interceptor::Interceptor; +pub use interceptor::{ + Interceptor, InterceptorError, InterceptorFailure, InterceptorPoint, InterceptorResult, +}; pub use message::{ContentPart, Item, Message, Role}; pub use tool::{ ToolCall, ToolExecutionContext, ToolExecutionHandle, ToolExecutionPolicy, diff --git a/crates/agen/tests/engine_state_test.rs b/crates/agen/tests/engine_state_test.rs index 527b8e94..c4d25428 100644 --- a/crates/agen/tests/engine_state_test.rs +++ b/crates/agen/tests/engine_state_test.rs @@ -10,7 +10,8 @@ use std::sync::{Arc, Mutex}; use agen::Item; use agen::interceptor::{ - Interceptor, PreRequestAction, PreToolAction, ToolCallInfo, TurnEndAction, + Interceptor, InterceptorError, InterceptorPoint, InterceptorResult, PostToolAction, + PreRequestAction, PreToolAction, PromptAction, ToolCallInfo, ToolResultInfo, TurnEndAction, }; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; @@ -613,12 +614,15 @@ struct YieldOnce { #[async_trait] impl Interceptor for YieldOnce { - async fn pre_llm_request(&self, _context: &mut Vec) -> PreRequestAction { - if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { + async fn pre_llm_request( + &self, + _context: &mut Vec, + ) -> InterceptorResult { + Ok(if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { PreRequestAction::Yield } else { PreRequestAction::Continue - } + }) } } @@ -628,12 +632,12 @@ struct PauseToolOnce { #[async_trait] impl Interceptor for PauseToolOnce { - async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> PreToolAction { - if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { + async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> InterceptorResult { + Ok(if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { PreToolAction::Pause } else { PreToolAction::Continue - } + }) } } @@ -643,13 +647,193 @@ struct ContinueTurnOnce { #[async_trait] impl Interceptor for ContinueTurnOnce { - async fn on_turn_end(&self, _history: &[Item]) -> TurnEndAction { - if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { + async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult { + Ok(if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { TurnEndAction::ContinueWithMessages(vec![Item::system_message("continue")]) } else { TurnEndAction::Finish + }) + } +} + +#[derive(Debug, Clone)] +struct FailingLifecycleInterceptor { + failure: InterceptorPoint, + calls: Arc>>, +} + +impl FailingLifecycleInterceptor { + fn new(failure: InterceptorPoint) -> Self { + Self { + failure, + calls: Arc::new(Mutex::new(Vec::new())), } } + + fn record(&self, point: InterceptorPoint, action: T) -> InterceptorResult { + self.calls.lock().unwrap().push(point); + if self.failure == point { + Err(InterceptorError::new(format!("{point} rejected"))) + } else { + Ok(action) + } + } + + fn calls(&self) -> Vec { + self.calls.lock().unwrap().clone() + } +} + +#[async_trait] +impl Interceptor for FailingLifecycleInterceptor { + async fn on_prompt_submit(&self, _item: &mut Item) -> InterceptorResult { + tokio::task::yield_now().await; + self.record(InterceptorPoint::PromptSubmit, PromptAction::Continue) + } + + async fn pending_history_appends(&self) -> InterceptorResult> { + tokio::task::yield_now().await; + self.record(InterceptorPoint::PendingHistoryAppends, Vec::new()) + } + + async fn pre_llm_request( + &self, + _context: &mut Vec, + ) -> InterceptorResult { + tokio::task::yield_now().await; + if self.failure == InterceptorPoint::Abort { + self.calls + .lock() + .unwrap() + .push(InterceptorPoint::PreLlmRequest); + Ok(PreRequestAction::Cancel( + "trigger abort callback".to_string(), + )) + } else { + self.record(InterceptorPoint::PreLlmRequest, PreRequestAction::Continue) + } + } + + async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> InterceptorResult { + tokio::task::yield_now().await; + self.record(InterceptorPoint::PreToolCall, PreToolAction::Continue) + } + + async fn post_tool_call( + &self, + _info: &mut ToolResultInfo, + ) -> InterceptorResult { + tokio::task::yield_now().await; + self.record(InterceptorPoint::PostToolCall, PostToolAction::Continue) + } + + async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult { + tokio::task::yield_now().await; + self.record(InterceptorPoint::TurnEnd, TurnEndAction::Finish) + } + + async fn on_abort(&self, _reason: &str) -> InterceptorResult<()> { + tokio::task::yield_now().await; + self.record(InterceptorPoint::Abort, ()) + } +} + +fn expected_interceptor_calls(failure: InterceptorPoint) -> Vec { + use InterceptorPoint as Point; + + match failure { + Point::PromptSubmit => vec![Point::PromptSubmit, Point::Abort], + Point::PendingHistoryAppends => { + vec![ + Point::PromptSubmit, + Point::PendingHistoryAppends, + Point::Abort, + ] + } + Point::PreLlmRequest => vec![ + Point::PromptSubmit, + Point::PendingHistoryAppends, + Point::PreLlmRequest, + Point::Abort, + ], + Point::PreToolCall => vec![ + Point::PromptSubmit, + Point::PendingHistoryAppends, + Point::PreLlmRequest, + Point::PreToolCall, + Point::Abort, + ], + Point::PostToolCall => vec![ + Point::PromptSubmit, + Point::PendingHistoryAppends, + Point::PreLlmRequest, + Point::PreToolCall, + Point::PostToolCall, + Point::Abort, + ], + Point::TurnEnd => vec![ + Point::PromptSubmit, + Point::PendingHistoryAppends, + Point::PreLlmRequest, + Point::TurnEnd, + Point::Abort, + ], + Point::Abort => vec![ + Point::PromptSubmit, + Point::PendingHistoryAppends, + Point::PreLlmRequest, + Point::Abort, + ], + } +} + +#[tokio::test] +async fn interceptor_failures_are_typed_and_each_lifecycle_point_runs_once() { + use InterceptorPoint as Point; + + for failure_point in [ + Point::PromptSubmit, + Point::PendingHistoryAppends, + Point::PreLlmRequest, + Point::PreToolCall, + Point::PostToolCall, + Point::TurnEnd, + Point::Abort, + ] { + let interceptor = FailingLifecycleInterceptor::new(failure_point); + let needs_tool = matches!(failure_point, Point::PreToolCall | Point::PostToolCall); + let events = if needs_tool { + vec![ + Event::tool_use_start(0, "call-1", "count_tool"), + Event::tool_input_delta(0, "{}"), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] + } else { + completed_text_events() + }; + let mut engine = Engine::new(MockLlmClient::new(events)); + engine.register_tool(CountingTool::new("count_tool").definition()); + engine.set_interceptor(interceptor.clone()); + let mut history = History::new(); + let mut engine = engine.lock(&history); + + let exit = engine.run(&mut history, "test").await; + let EngineRunExit::Interrupted(RunInterruptionReason::Interceptor(failure)) = exit else { + panic!("expected typed interceptor interruption at {failure_point}, got {exit:?}"); + }; + assert_eq!(failure.point(), failure_point); + assert_eq!( + failure.error().message(), + format!("{failure_point} rejected") + ); + assert_eq!( + interceptor.calls(), + expected_interceptor_calls(failure_point) + ); + } } #[tokio::test] diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index 88d2eda4..ead143b4 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -6,7 +6,9 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; -use agen::interceptor::{Interceptor, PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo}; +use agen::interceptor::{ + Interceptor, InterceptorResult, PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo, +}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::tool::{ Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult, @@ -905,24 +907,27 @@ async fn test_tool_execution_context_for_skipped_and_synthetic_paths() { #[async_trait] impl Interceptor for ContextPolicy { - async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction { + async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult { self.pre_contexts.lock().unwrap().push(info.context.clone()); - match info.call.name.as_str() { + Ok(match info.call.name.as_str() { "skip_tool" => PreToolAction::Skip, "synthetic_tool" => PreToolAction::SyntheticResult(ToolResult::from_output( &info.call.id, ToolOutput::from("synthetic result".to_string()), )), _ => PreToolAction::Continue, - } + }) } - async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction { + async fn post_tool_call( + &self, + info: &mut ToolResultInfo, + ) -> InterceptorResult { self.post_contexts .lock() .unwrap() .push(info.context.clone()); - PostToolAction::Continue + Ok(PostToolAction::Continue) } } @@ -994,12 +999,12 @@ async fn test_before_tool_call_skip() { #[async_trait] impl Interceptor for BlockingPolicy { - async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction { - if info.call.name == "blocked_tool" { + async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult { + Ok(if info.call.name == "blocked_tool" { PreToolAction::Skip } else { PreToolAction::Continue - } + }) } } @@ -1081,10 +1086,13 @@ async fn test_post_tool_call_modification() { #[async_trait] impl Interceptor for ModifyingPolicy { - async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction { + async fn post_tool_call( + &self, + info: &mut ToolResultInfo, + ) -> InterceptorResult { info.result.summary = format!("[Modified] {}", info.result.summary); *self.modified_content.lock().unwrap() = Some(info.result.summary.clone()); - PostToolAction::Continue + Ok(PostToolAction::Continue) } } @@ -1143,11 +1151,11 @@ async fn test_before_tool_call_synthetic_result_committed() { #[async_trait] impl Interceptor for SyntheticPolicy { - async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction { - PreToolAction::SyntheticResult(ToolResult::error( + async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult { + Ok(PreToolAction::SyntheticResult(ToolResult::error( info.call.id.clone(), "permission denied", - )) + ))) } } @@ -1184,8 +1192,11 @@ async fn post_tool_abort_commits_confirmed_result_before_stopping_run() { struct AbortAfterResult; #[async_trait] impl Interceptor for AbortAfterResult { - async fn post_tool_call(&self, _info: &mut ToolResultInfo) -> PostToolAction { - PostToolAction::Abort("policy stopped the run".to_string()) + async fn post_tool_call( + &self, + _info: &mut ToolResultInfo, + ) -> InterceptorResult { + Ok(PostToolAction::Abort("policy stopped the run".to_string())) } } engine.set_interceptor(AbortAfterResult); diff --git a/crates/session-store/tests/session_test.rs b/crates/session-store/tests/session_test.rs index 6f7ec3ee..2f0ffbeb 100644 --- a/crates/session-store/tests/session_test.rs +++ b/crates/session-store/tests/session_test.rs @@ -3,7 +3,7 @@ mod common; use std::ops::{Deref, DerefMut}; use std::sync::Arc; -use agen::interceptor::{Interceptor, TurnEndAction}; +use agen::interceptor::{Interceptor, InterceptorResult, TurnEndAction}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::types::{Item, RequestConfig}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; @@ -100,8 +100,8 @@ struct PausePolicy; #[async_trait] impl Interceptor for PausePolicy { - async fn on_turn_end(&self, _history: &[Item]) -> TurnEndAction { - TurnEndAction::Pause + async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult { + Ok(TurnEndAction::Pause) } } diff --git a/crates/worker/src/compact/worker.rs b/crates/worker/src/compact/worker.rs index 2f30407e..cd218ac0 100644 --- a/crates/worker/src/compact/worker.rs +++ b/crates/worker/src/compact/worker.rs @@ -22,7 +22,9 @@ use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use agen::Item; -use agen::interceptor::{Interceptor, PreRequestAction, PreToolAction, ToolCallInfo}; +use agen::interceptor::{ + Interceptor, InterceptorResult, PreRequestAction, PreToolAction, ToolCallInfo, +}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput, ToolResult}; use async_trait::async_trait; use serde::Deserialize; @@ -398,14 +400,17 @@ impl CompactWorkerInterceptor { #[async_trait] impl Interceptor for CompactWorkerInterceptor { - async fn pre_llm_request(&self, context: &mut Vec) -> PreRequestAction { + async fn pre_llm_request( + &self, + context: &mut Vec, + ) -> InterceptorResult { let records = self.usage_tracker.records(); let estimate = agen::token_counter::total_tokens(context, &records); if estimate.tokens > self.max_input_tokens { - return PreRequestAction::Cancel(format!( + return Ok(PreRequestAction::Cancel(format!( "compact worker input occupancy exceeded {} tokens", self.max_input_tokens - )); + ))); } let remaining = self.max_input_tokens.saturating_sub(estimate.tokens); @@ -413,25 +418,25 @@ impl Interceptor for CompactWorkerInterceptor { .store(remaining, Ordering::Release); if let Some(item) = self.maybe_emit_warning(remaining) { self.usage_tracker.note_request(context.len() + 1); - return PreRequestAction::ContinueWith(vec![item]); + return Ok(PreRequestAction::ContinueWith(vec![item])); } self.usage_tracker.note_request(context.len()); - PreRequestAction::Continue + Ok(PreRequestAction::Continue) } - async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction { + async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult { if self.final_reserve_tokens == 0 || info.call.name == "write_summary" { - return PreToolAction::Continue; + return Ok(PreToolAction::Continue); } let remaining = self.last_remaining_tokens.load(Ordering::Acquire); if remaining > self.final_reserve_tokens { - return PreToolAction::Continue; + return Ok(PreToolAction::Continue); } - PreToolAction::SyntheticResult(ToolResult::error( + Ok(PreToolAction::SyntheticResult(ToolResult::error( info.call.id.clone(), "compact worker final reserve reached; do not perform more exploratory tool reads. Call `write_summary` now.", - )) + ))) } } @@ -467,13 +472,13 @@ mod tests { let mut context = vec![Item::user_message("hello")]; assert!(matches!( - interceptor.pre_llm_request(&mut context).await, + interceptor.pre_llm_request(&mut context).await.unwrap(), PreRequestAction::Continue )); tracker.record_usage(&make_usage(100)); assert!(matches!( - interceptor.pre_llm_request(&mut context).await, + interceptor.pre_llm_request(&mut context).await.unwrap(), PreRequestAction::Continue )); tracker.record_usage(&make_usage(100)); @@ -481,7 +486,7 @@ mod tests { // Two 100-token requests would exceed a cumulative 150-token cap, but // current occupancy is still the latest 100-token measurement. assert!(matches!( - interceptor.pre_llm_request(&mut context).await, + interceptor.pre_llm_request(&mut context).await.unwrap(), PreRequestAction::Continue )); } @@ -503,13 +508,13 @@ mod tests { let mut context = vec![Item::user_message("hello")]; assert!(matches!( - interceptor.pre_llm_request(&mut context).await, + interceptor.pre_llm_request(&mut context).await.unwrap(), PreRequestAction::Continue )); tracker.record_usage(&make_usage(100)); assert!(matches!( - interceptor.pre_llm_request(&mut context).await, + interceptor.pre_llm_request(&mut context).await.unwrap(), PreRequestAction::ContinueWith(items) if items.len() == 1 && items[0].as_text().unwrap_or_default().contains("write_summary") )); @@ -523,13 +528,13 @@ mod tests { let mut context = vec![Item::user_message("hello")]; assert!(matches!( - interceptor.pre_llm_request(&mut context).await, + interceptor.pre_llm_request(&mut context).await.unwrap(), PreRequestAction::Continue )); tracker.record_usage(&make_usage(100)); assert!(matches!( - interceptor.pre_llm_request(&mut context).await, + interceptor.pre_llm_request(&mut context).await.unwrap(), PreRequestAction::Cancel(message) if message.contains("occupancy") )); } diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index f82444aa..a1468443 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -15,8 +15,8 @@ use std::sync::{Arc, Mutex}; use agen::Item; use agen::UsageRecord; use agen::interceptor::{ - Interceptor, PostToolAction, PreRequestAction, PreToolAction, PromptAction, ToolCallInfo, - ToolResultInfo, TurnEndAction, + Interceptor, InterceptorResult, PostToolAction, PreRequestAction, PreToolAction, PromptAction, + ToolCallInfo, ToolResultInfo, TurnEndAction, }; use agen::tool::ToolOutput; use arc_swap::ArcSwap; @@ -232,7 +232,7 @@ impl WorkerInterceptor { #[async_trait] impl Interceptor for WorkerInterceptor { - async fn on_prompt_submit(&self, item: &mut Item) -> PromptAction { + async fn on_prompt_submit(&self, item: &mut Item) -> InterceptorResult { let turn_index = self.next_turn_index.fetch_add(1, Ordering::Relaxed); self.tool_calls_this_turn.store(0, Ordering::Relaxed); @@ -243,7 +243,7 @@ impl Interceptor for WorkerInterceptor { for hook in &self.registry.on_prompt_submit { let action = hook.call(&info).await; if !matches!(action, HookPromptAction::Continue) { - return action.into(); + return Ok(action.into()); } } let mut extras: Vec = std::mem::take( @@ -252,7 +252,7 @@ impl Interceptor for WorkerInterceptor { .lock() .expect("pending_attachments poisoned"), ); - if extras.is_empty() { + Ok(if extras.is_empty() { PromptAction::Continue } else { // Commit the typed system items first, then hand the @@ -266,10 +266,10 @@ impl Interceptor for WorkerInterceptor { Ok(()) => PromptAction::ContinueWith(items), Err(error) => PromptAction::Cancel(format!("session persistence failed: {error}")), } - } + }) } - async fn pending_history_appends(&self) -> Result, String> { + async fn pending_history_appends(&self) -> InterceptorResult> { let drained = self.pending_notifies.drain(); if drained.is_empty() { return Ok(Vec::new()); @@ -295,7 +295,7 @@ impl Interceptor for WorkerInterceptor { Ok(system_item) => system_item, Err(error) => { self.pending_notifies.requeue_front(drained); - return Err(format!("failed to render notify_wrapper: {error}")); + return Err(format!("failed to render notify_wrapper: {error}").into()); } }; items.push(system_item.to_history_item()); @@ -303,15 +303,18 @@ impl Interceptor for WorkerInterceptor { } if let Err(error) = self.commit_system_items(&system_items) { self.pending_notifies.requeue_front(drained); - return Err(format!("session persistence failed: {error}")); + return Err(format!("session persistence failed: {error}").into()); } Ok(items) } - async fn pre_llm_request(&self, context: &mut Vec) -> PreRequestAction { + async fn pre_llm_request( + &self, + context: &mut Vec, + ) -> InterceptorResult { let initial_tokens = self.estimated_tokens(context); if self.request_threshold_exceeded(initial_tokens, context) { - return PreRequestAction::Yield; + return Ok(PreRequestAction::Yield); } let info = PreRequestInfo { item_count: context.len(), @@ -328,7 +331,7 @@ impl Interceptor for WorkerInterceptor { for hook in &self.registry.pre_llm_request { let action = hook.call(&hook_context).await; if !matches!(action, HookPreRequestAction::Continue) { - return action.into(); + return Ok(action.into()); } } @@ -353,28 +356,30 @@ impl Interceptor for WorkerInterceptor { if self.request_threshold_exceeded(current_tokens, effective_context.as_ref()) { if let Err(error) = self.commit_system_items(&system_items) { - return PreRequestAction::Cancel(format!("session persistence failed: {error}")); + return Ok(PreRequestAction::Cancel(format!( + "session persistence failed: {error}" + ))); } - return if appended_items.is_empty() { + return Ok(if appended_items.is_empty() { PreRequestAction::Yield } else { PreRequestAction::YieldWith(appended_items) - }; + }); } if let Some(usage_tracker) = self.usage_tracker.as_ref() { usage_tracker.note_request(effective_context.len()); } if system_items.is_empty() { - return PreRequestAction::Continue; + return Ok(PreRequestAction::Continue); } - match self.commit_system_items(&system_items) { + Ok(match self.commit_system_items(&system_items) { Ok(()) => PreRequestAction::ContinueWith(appended_items), Err(error) => PreRequestAction::Cancel(format!("session persistence failed: {error}")), - } + }) } - async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction { + async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult { let summary = ToolCallSummary { call_id: info.call.id.clone(), tool_name: info.call.name.clone(), @@ -383,14 +388,14 @@ impl Interceptor for WorkerInterceptor { for hook in &self.registry.pre_tool_call { let action = hook.call(&summary).await; if !matches!(action, HookPreToolAction::Continue) { - return action.into_worker_action(summary.call_id.clone()); + return Ok(action.into_worker_action(summary.call_id.clone())); } } self.tool_calls_this_turn.fetch_add(1, Ordering::Relaxed); - PreToolAction::Continue + Ok(PreToolAction::Continue) } - async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction { + async fn post_tool_call(&self, info: &mut ToolResultInfo) -> InterceptorResult { let summary = ToolResultSummary { call_id: info.result.tool_use_id.clone(), tool_name: info.call.name.clone(), @@ -405,13 +410,13 @@ impl Interceptor for WorkerInterceptor { for hook in &self.registry.post_tool_call { let action = hook.call(&summary).await; if !matches!(action, HookPostToolAction::Continue) { - return action.into(); + return Ok(action.into()); } } - PostToolAction::Continue + Ok(PostToolAction::Continue) } - async fn on_turn_end(&self, history: &[Item]) -> TurnEndAction { + async fn on_turn_end(&self, history: &[Item]) -> InterceptorResult { let final_text_preview = history .iter() .rev() @@ -427,19 +432,20 @@ impl Interceptor for WorkerInterceptor { for hook in &self.registry.on_turn_end { let action = hook.call(&info).await; if !matches!(action, HookTurnEndAction::Finish) { - return action.into(); + return Ok(action.into()); } } - TurnEndAction::Finish + Ok(TurnEndAction::Finish) } - async fn on_abort(&self, reason: &str) { + async fn on_abort(&self, reason: &str) -> InterceptorResult<()> { let info = AbortInfo { reason: reason.to_string(), }; for hook in &self.registry.on_abort { hook.call(&info).await; } + Ok(()) } } @@ -623,7 +629,7 @@ mod tests { None, ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await; + let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); assert!(matches!(action, PreRequestAction::Yield)); // Hook must not run when an internal mechanism short-circuits first. @@ -655,7 +661,7 @@ mod tests { })), ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await; + let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); match action { PreRequestAction::YieldWith(items) => assert_eq!(items.len(), 1), @@ -692,7 +698,7 @@ mod tests { ) .with_usage_tracker(usage_tracker); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await; + let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); assert!(matches!(action, PreRequestAction::Yield)); } @@ -716,7 +722,7 @@ mod tests { None, ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await; + let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -757,7 +763,7 @@ mod tests { None, ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await; + let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -784,7 +790,7 @@ mod tests { None, ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await; + let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -805,7 +811,7 @@ mod tests { None, ); let mut ctx: Vec = Vec::new(); - let action = interceptor.pre_llm_request(&mut ctx).await; + let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -834,7 +840,7 @@ mod tests { ); let mut ctx: Vec = Vec::new(); - let action = interceptor.pre_llm_request(&mut ctx).await; + let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); assert!(saw_handle.load(Ordering::Relaxed)); let PreRequestAction::ContinueWith(items) = action else { @@ -881,7 +887,7 @@ mod tests { ); let mut ctx: Vec = Vec::new(); - let action = interceptor.pre_llm_request(&mut ctx).await; + let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); assert!(!saw_handle.load(Ordering::Relaxed)); assert!(matches!(action, PreRequestAction::Continue)); @@ -938,7 +944,7 @@ mod tests { ); let mut info = task_tool_call_info("TaskList", serde_json::json!({"scope": "all"})); - let action = interceptor.pre_tool_call(&mut info).await; + let action = interceptor.pre_tool_call(&mut info).await.unwrap(); match action { PreToolAction::SyntheticResult(result) => { @@ -1000,7 +1006,7 @@ mod tests { context: info.context, }; - let action = interceptor.post_tool_call(&mut result_info).await; + let action = interceptor.post_tool_call(&mut result_info).await.unwrap(); assert_eq!(action, PostToolAction::Abort("post tool abort".to_string())); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -1036,7 +1042,7 @@ mod tests { ); let history = vec![Item::user_message("hi"), Item::assistant_message("done")]; - let action = interceptor.on_turn_end(&history).await; + let action = interceptor.on_turn_end(&history).await.unwrap(); assert!(matches!(action, TurnEndAction::Pause)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -1073,7 +1079,7 @@ mod tests { let ctx_items = vec![Item::user_message("hi")]; for _ in 0..23 { let mut ctx = ctx_items.clone(); - let action = interceptor.pre_llm_request(&mut ctx).await; + let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); assert!(matches!(action, PreRequestAction::Continue)); usage_tracker.record_usage(&agen::event::UsageEvent { input_tokens: Some(10), @@ -1085,7 +1091,7 @@ mod tests { } let mut ctx = ctx_items.clone(); - let action = interceptor.pre_llm_request(&mut ctx).await; + let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); let appended_len = match action { PreRequestAction::ContinueWith(items) => items.len(), other => panic!("expected reminder append, got {other:?}"), @@ -1210,7 +1216,7 @@ mod tests { let error = interceptor.pending_history_appends().await.unwrap_err(); - assert!(error.contains("failed to render notify_wrapper")); + assert!(error.message().contains("failed to render notify_wrapper")); let requeued = buffer.drain(); assert_eq!(requeued.len(), 1); } @@ -1269,7 +1275,7 @@ mod tests { None, ); let mut ctx: Vec = vec![Item::user_message("hi")]; - let action = interceptor.pre_llm_request(&mut ctx).await; + let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(ctx.len(), 1, "pre_llm_request must not append notifies"); @@ -1299,7 +1305,7 @@ mod tests { None, ); let mut ctx: Vec = Vec::new(); - let action = interceptor.pre_llm_request(&mut ctx).await; + let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); assert!(matches!(action, PreRequestAction::Cancel(_))); assert!(first_called.load(Ordering::Relaxed)); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index d8b84640..f98a9b5a 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -3733,6 +3733,7 @@ impl Worker { | EngineRunExit::Yielded | EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) | EngineRunExit::Interrupted(RunInterruptionReason::ContextWindowExceeded) + | EngineRunExit::Interrupted(RunInterruptionReason::Interceptor(_)) | EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(_)) ); let active_run_turn_count = self.engine.as_ref().unwrap().active_run_turn_count(); @@ -6129,12 +6130,14 @@ fn run_interruption_reason_error_code(reason: &RunInterruptionReason) -> ErrorCo RunInterruptionReason::Unexpected(EngineError::Tool(_)) => ErrorCode::ToolError, RunInterruptionReason::LimitReached | RunInterruptionReason::Cancelled + | RunInterruptionReason::Interceptor(_) | RunInterruptionReason::Unexpected( EngineError::Aborted(_) | EngineError::Cancelled | EngineError::PauseRequested | EngineError::ConfigWarnings(_) | EngineError::HistoryAppend(_) + | EngineError::Interceptor(_) | EngineError::ToolAttemptFence(_), ) => ErrorCode::Internal, } @@ -6145,6 +6148,7 @@ fn run_interruption_reason_message(reason: &RunInterruptionReason) -> String { RunInterruptionReason::LimitReached => "engine turn limit reached".to_string(), RunInterruptionReason::ContextWindowExceeded => "model context window reached".to_string(), RunInterruptionReason::Cancelled => "engine run cancelled".to_string(), + RunInterruptionReason::Interceptor(failure) => failure.to_string(), RunInterruptionReason::Unexpected(error) => format!("unexpected engine failure: {error}"), } } From 0245980ea5720b1c414337fb35e62218431931cc Mon Sep 17 00:00:00 2001 From: Hare Date: Thu, 3 Sep 2026 15:33:15 +0900 Subject: [PATCH 02/22] refactor: finalize interceptor lifecycle boundaries --- crates/agen/examples/engine_cli.rs | 2 +- crates/agen/src/engine.rs | 210 ++++++------- crates/agen/src/interceptor.rs | 73 +++-- crates/agen/src/lib.rs | 3 +- crates/agen/tests/engine_state_test.rs | 313 ++++++++++++++++--- crates/agen/tests/parallel_execution_test.rs | 62 ++-- crates/session-store/tests/session_test.rs | 12 +- crates/worker/src/compact/worker.rs | 55 +++- crates/worker/src/hook.rs | 22 +- crates/worker/src/ipc/interceptor.rs | 112 +++++-- crates/worker/src/worker.rs | 12 +- 11 files changed, 587 insertions(+), 289 deletions(-) diff --git a/crates/agen/examples/engine_cli.rs b/crates/agen/examples/engine_cli.rs index 2c50d354..10eb28ef 100644 --- a/crates/agen/examples/engine_cli.rs +++ b/crates/agen/examples/engine_cli.rs @@ -280,7 +280,7 @@ impl ToolResultPrinterPolicy { #[async_trait] impl Interceptor for ToolResultPrinterPolicy { - async fn post_tool_call(&self, info: &mut ToolResultInfo) -> InterceptorResult { + async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult { let name = self .call_names .lock() diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index ae3279b2..276ee3a4 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -15,8 +15,10 @@ use crate::{ }, handler::{ErrorKind, StatusKind, ToolUseBlockStart, UsageKind}, interceptor::{ - DefaultInterceptor, Interceptor, InterceptorFailure, InterceptorPoint, PostToolAction, - PreRequestAction, PreToolAction, PromptAction, ToolCallInfo, ToolResultInfo, TurnEndAction, + AssistantTurnEndContext, DefaultInterceptor, Interceptor, InterceptorFailure, + InterceptorPoint, PostToolAction, PreLlmRequestContext, PreRequestAction, PreToolAction, + PromptAction, PromptSubmitContext, RunExitContext, ToolCallInfo, ToolResultInfo, + TurnEndAction, }, llm_client::{ ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream, @@ -156,8 +158,6 @@ pub enum EngineRunExit { /// A typed reason why an engine run could not finish normally. #[derive(Debug)] pub enum RunInterruptionReason { - /// A trusted host interceptor callback failed at a typed lifecycle point. - Interceptor(InterceptorFailure), LimitReached, ContextWindowExceeded, Cancelled, @@ -178,9 +178,6 @@ impl From> for EngineRunExit { } Err(EngineError::Cancelled) => Self::Interrupted(RunInterruptionReason::Cancelled), Err(EngineError::PauseRequested) => Self::Paused, - Err(EngineError::Interceptor(failure)) => { - Self::Interrupted(RunInterruptionReason::Interceptor(failure)) - } Err(error) => Self::Interrupted(RunInterruptionReason::Unexpected(error)), } } @@ -438,11 +435,8 @@ impl Engine { self.active_run_turn_count.get_or_insert(0); } - fn finish_logical_run(&mut self, result: &Result) { - if !matches!( - result, - Ok(EngineResult::Paused | EngineResult::Yielded) | Err(EngineError::PauseRequested) - ) { + fn finish_logical_run(&mut self, exit: &EngineRunExit) { + if !matches!(exit, EngineRunExit::Paused | EngineRunExit::Yielded) { self.active_run_turn_count = None; } } @@ -1086,26 +1080,23 @@ impl Engine { request } - /// Hooks: on_prompt_submit - /// - async fn finalize_interruption( + async fn finalize_run_exit( &mut self, - result: Result, - ) -> Result { - match result { - Ok(value) => Ok(value), - Err(err) => { - let reason = match &err { - EngineError::Aborted(reason) => reason.clone(), - EngineError::Cancelled => "Cancelled".to_string(), - _ => err.to_string(), - }; - if let Err(error) = self.interceptor.on_abort(&reason).await { - return Err(InterceptorFailure::new(InterceptorPoint::Abort, error).into()); - } - Err(err) - } - } + result: Result, + ) -> EngineRunExit { + let exit = EngineRunExit::from(result); + let exit = match self + .interceptor + .on_run_exit(RunExitContext { exit: &exit }) + .await + { + Ok(()) => exit, + Err(error) => EngineRunExit::Interrupted(RunInterruptionReason::Unexpected( + InterceptorFailure::new(InterceptorPoint::RunExit, error).into(), + )), + }; + self.finish_logical_run(&exit); + exit } /// Check for pending tool calls (for resuming from Pause) @@ -1486,41 +1477,13 @@ impl Engine { } let call_info = call_info_map.get(&tool_result.tool_use_id); - let mut abort_reason = None; - if let Some((tool_call, meta, tool, context)) = call_info { - let mut info = ToolResultInfo { - call: tool_call.clone(), - result: tool_result, - meta: meta.clone(), - tool: tool.clone(), - context: context.clone(), - }; - - let post_tool_action = - self.interceptor - .post_tool_call(&mut info) - .await - .map_err(|error| { - EngineError::from(InterceptorFailure::new( - InterceptorPoint::PostToolCall, - error, - )) - })?; - match post_tool_action { - PostToolAction::Continue => {} - PostToolAction::Abort(reason) => { - abort_reason = Some(reason); - } - } - tool_result = info.result; - } if tool_result.is_error && tool_result.disposition.is_success() { tool_result.disposition = ToolResultDisposition::Error; } tool_result.is_error = !tool_result.disposition.is_success(); - // Cap content only after post_tool_call so interceptors still observe - // the full payload and any content they inject is bounded too. + // Bound the terminal payload before committing it so the post-tool + // interceptor observes exactly the model-visible durable result. if let (Some(limits), Some((tool_call, _, _, _)), Some(content)) = ( self.tool_output_limits.as_ref(), call_info, @@ -1573,9 +1536,30 @@ impl Engine { "Tool execution terminalized" ); self.emit_tool_result(&tool_result); - if let Some(reason) = abort_reason { - return Err(EngineError::Aborted(reason)); + + if let Some((tool_call, meta, tool, context)) = call_info { + let info = ToolResultInfo { + call: tool_call.clone(), + result: tool_result, + meta: meta.clone(), + tool: tool.clone(), + context: context.clone(), + }; + let post_tool_action = + self.interceptor + .post_tool_call(&info) + .await + .map_err(|error| { + EngineError::from(InterceptorFailure::new( + InterceptorPoint::PostToolCall, + error, + )) + })?; + if let PostToolAction::Abort(reason) = post_tool_action { + return Err(EngineError::Aborted(reason)); + } } + Ok(true) } @@ -1716,7 +1700,9 @@ impl Engine { // Interceptor: pre_llm_request let pre_request_action = self .interceptor - .pre_llm_request(&mut request_context) + .pre_llm_request(PreLlmRequestContext { + items: &mut request_context, + }) .await .map_err(|error| { EngineError::from(InterceptorFailure::new( @@ -1836,28 +1822,37 @@ impl Engine { let tool_calls = self.tool_call_collector.take_collected(); let assistant_items = self.build_assistant_items(&reasoning_items, &text_blocks, &tool_calls); + let committed_assistant_items = assistant_items.clone(); self.append_history_items(history, assistant_items, annotate)?; - if tool_calls.is_empty() { - let turn_end_context = history.items_cloned(); - let turn_end_action = self - .interceptor - .on_turn_end(&turn_end_context) - .await - .map_err(|error| { - EngineError::from(InterceptorFailure::new(InterceptorPoint::TurnEnd, error)) - })?; - match turn_end_action { - TurnEndAction::Finish => { - return Ok(EngineResult::Finished); - } - TurnEndAction::ContinueWithMessages(additional) => { - self.append_history_items(history, additional, annotate)?; + let assistant_turn_history = history.items_cloned(); + let assistant_turn_action = self + .interceptor + .on_assistant_turn_end(AssistantTurnEndContext { + assistant_items: &committed_assistant_items, + history: &assistant_turn_history, + tool_calls: &tool_calls, + }) + .await + .map_err(|error| { + EngineError::from(InterceptorFailure::new( + InterceptorPoint::AssistantTurnEnd, + error, + )) + })?; + match assistant_turn_action { + TurnEndAction::Finish if tool_calls.is_empty() => { + return Ok(EngineResult::Finished); + } + TurnEndAction::Finish => {} + TurnEndAction::ContinueWithMessages(additional) => { + self.append_history_items(history, additional, annotate)?; + if tool_calls.is_empty() { continue; } - TurnEndAction::Pause => { - return Ok(EngineResult::Paused); - } + } + TurnEndAction::Pause => { + return Ok(EngineResult::Paused); } } @@ -2540,9 +2535,10 @@ impl Engine { user_input: impl Into, annotate: &mut impl FnMut(&Item) -> Result, ) -> EngineRunExit { - self.run_result_with_annotation(history, user_input.into(), annotate) - .await - .into() + let result = self + .run_result_with_annotation(history, user_input.into(), annotate) + .await; + self.finalize_run_exit(result).await } async fn run_result_with_annotation( @@ -2554,19 +2550,20 @@ impl Engine { // Supplying new user input abandons any paused/yielded logical run. self.active_run_turn_count = None; let mut user_item = Item::user_message(user_input); - let prompt_action = match self.interceptor.on_prompt_submit(&mut user_item).await { - Ok(action) => action, - Err(error) => { - let error = InterceptorFailure::new(InterceptorPoint::PromptSubmit, error).into(); - return self.finalize_interruption(Err(error)).await; - } - }; + let prompt_action = self + .interceptor + .on_prompt_submit(PromptSubmitContext { + item: &mut user_item, + }) + .await + .map_err(|error| { + EngineError::from(InterceptorFailure::new( + InterceptorPoint::PromptSubmit, + error, + )) + })?; let extras = match prompt_action { - PromptAction::Cancel(reason) => { - return self - .finalize_interruption(Err(EngineError::Aborted(reason))) - .await; - } + PromptAction::Cancel(reason) => return Err(EngineError::Aborted(reason)), PromptAction::Continue => Vec::new(), PromptAction::ContinueWith(items) => items, }; @@ -2575,13 +2572,10 @@ impl Engine { self.append_history_items(history, extras, annotate)?; } self.start_logical_run(); - let result = match self.run_turn_loop(history, annotate).await { + match self.run_turn_loop(history, annotate).await { Err(EngineError::PauseRequested) => Ok(EngineResult::Paused), other => other, - }; - let result = self.finalize_interruption(result).await; - self.finish_logical_run(&result); - result + } } /// Resume execution (from Paused state). @@ -2590,9 +2584,8 @@ impl Engine { history: &mut History, annotate: &mut impl FnMut(&Item) -> Result, ) -> EngineRunExit { - self.resume_result_with_annotation(history, annotate) - .await - .into() + let result = self.resume_result_with_annotation(history, annotate).await; + self.finalize_run_exit(result).await } async fn resume_result_with_annotation( @@ -2601,13 +2594,10 @@ impl Engine { annotate: &mut impl FnMut(&Item) -> Result, ) -> Result { self.ensure_logical_run(); - let result = match self.run_turn_loop(history, annotate).await { + match self.run_turn_loop(history, annotate).await { Err(EngineError::PauseRequested) => Ok(EngineResult::Paused), other => other, - }; - let result = self.finalize_interruption(result).await; - self.finish_logical_run(&result); - result + } } /// Get the prefix length at lock time diff --git a/crates/agen/src/interceptor.rs b/crates/agen/src/interceptor.rs index ecb1fd8f..b85e0d9e 100644 --- a/crates/agen/src/interceptor.rs +++ b/crates/agen/src/interceptor.rs @@ -9,6 +9,7 @@ use std::sync::Arc; use async_trait::async_trait; use crate::Item; +use crate::engine::EngineRunExit; use crate::tool::{Tool, ToolCall, ToolExecutionContext, ToolMeta, ToolResult}; // ============================================================================= @@ -59,8 +60,8 @@ pub enum InterceptorPoint { PreLlmRequest, PreToolCall, PostToolCall, - TurnEnd, - Abort, + AssistantTurnEnd, + RunExit, } impl std::fmt::Display for InterceptorPoint { @@ -71,8 +72,8 @@ impl std::fmt::Display for InterceptorPoint { Self::PreLlmRequest => "pre_llm_request", Self::PreToolCall => "pre_tool_call", Self::PostToolCall => "post_tool_call", - Self::TurnEnd => "turn_end", - Self::Abort => "abort", + Self::AssistantTurnEnd => "assistant_turn_end", + Self::RunExit => "run_exit", }; formatter.write_str(name) } @@ -106,6 +107,35 @@ impl InterceptorFailure { /// Result returned by asynchronous interceptor lifecycle methods. pub type InterceptorResult = Result; +// ============================================================================= +// Lifecycle Contexts +// ============================================================================= + +/// Mutable prompt input presented before it is committed to Engine history. +pub struct PromptSubmitContext<'a> { + pub item: &'a mut Item, +} + +/// Mutable provider-visible item projection presented before an LLM request. +pub struct PreLlmRequestContext<'a> { + pub items: &'a mut Vec, +} + +/// A terminalized and committed assistant response at the next-phase boundary. +pub struct AssistantTurnEndContext<'a> { + /// The exact assistant items committed for this response. + pub assistant_items: &'a [Item], + /// The committed Engine history after the assistant items were appended. + pub history: &'a [Item], + /// Terminal tool calls collected from the response, if any. + pub tool_calls: &'a [ToolCall], +} + +/// The one terminal outcome produced by a public Engine run or resume call. +pub struct RunExitContext<'a> { + pub exit: &'a EngineRunExit, +} + // ============================================================================= // Action Enums // ============================================================================= @@ -181,9 +211,9 @@ pub enum PostToolAction { /// Action at the end of a turn (when LLM produces no tool calls). #[derive(Debug, Clone)] pub enum TurnEndAction { - /// Turn is finished, return to caller. + /// Accept the Engine's natural next phase: execute tools, or finish when none exist. Finish, - /// Continue with additional messages injected into history. + /// Commit additional messages, then continue through the natural next phase. ContinueWithMessages(Vec), /// Pause execution (can be resumed later). Pause, @@ -209,7 +239,7 @@ pub struct ToolCallInfo { pub struct ToolResultInfo { /// Original tool call. pub call: ToolCall, - /// Tool execution result (modifiable). + /// Committed terminal tool execution result. pub result: ToolResult, /// Tool meta information. pub meta: ToolMeta, @@ -236,7 +266,10 @@ pub struct ToolResultInfo { #[async_trait] pub trait Interceptor: Send + Sync { /// Called after receiving user input, before adding it to Engine history. - async fn on_prompt_submit(&self, _item: &mut Item) -> InterceptorResult { + async fn on_prompt_submit( + &self, + _context: PromptSubmitContext<'_>, + ) -> InterceptorResult { Ok(PromptAction::Continue) } @@ -272,7 +305,7 @@ pub trait Interceptor: Send + Sync { /// commits it to history before the request is sent. async fn pre_llm_request( &self, - _context: &mut Vec, + _context: PreLlmRequestContext<'_>, ) -> InterceptorResult { Ok(PreRequestAction::Continue) } @@ -282,24 +315,22 @@ pub trait Interceptor: Send + Sync { Ok(PreToolAction::Continue) } - /// Called after each tool reaches one terminal result. - async fn post_tool_call( - &self, - _info: &mut ToolResultInfo, - ) -> InterceptorResult { + /// Called after each tool reaches one terminal result and that result is committed. + async fn post_tool_call(&self, _info: &ToolResultInfo) -> InterceptorResult { Ok(PostToolAction::Continue) } - /// Called at the assistant boundary when a completed response has no tool calls. - /// - /// This is not the logical run termination observer. A host that needs that - /// boundary must inspect the returned [`crate::EngineRunExit`]. - async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult { + /// Called after every terminal assistant response is committed and before + /// the Engine decides whether to execute tools, continue, or finish. + async fn on_assistant_turn_end( + &self, + _context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { Ok(TurnEndAction::Finish) } - /// Called once when execution is interrupted (abort, cancellation, or failure). - async fn on_abort(&self, _reason: &str) -> InterceptorResult<()> { + /// Called once for the terminal outcome of each public run or resume call. + async fn on_run_exit(&self, _context: RunExitContext<'_>) -> InterceptorResult<()> { Ok(()) } } diff --git a/crates/agen/src/lib.rs b/crates/agen/src/lib.rs index 52a7a246..27bed76e 100644 --- a/crates/agen/src/lib.rs +++ b/crates/agen/src/lib.rs @@ -27,7 +27,8 @@ pub use engine::{ pub use handler::ToolUseBlockStart; pub use history::{History, HistoryEntry}; pub use interceptor::{ - Interceptor, InterceptorError, InterceptorFailure, InterceptorPoint, InterceptorResult, + AssistantTurnEndContext, Interceptor, InterceptorError, InterceptorFailure, InterceptorPoint, + InterceptorResult, PreLlmRequestContext, PromptSubmitContext, RunExitContext, }; pub use message::{ContentPart, Item, Message, Role}; pub use tool::{ diff --git a/crates/agen/tests/engine_state_test.rs b/crates/agen/tests/engine_state_test.rs index c4d25428..ab810e9e 100644 --- a/crates/agen/tests/engine_state_test.rs +++ b/crates/agen/tests/engine_state_test.rs @@ -10,10 +10,14 @@ use std::sync::{Arc, Mutex}; use agen::Item; use agen::interceptor::{ - Interceptor, InterceptorError, InterceptorPoint, InterceptorResult, PostToolAction, - PreRequestAction, PreToolAction, PromptAction, ToolCallInfo, ToolResultInfo, TurnEndAction, + AssistantTurnEndContext, Interceptor, InterceptorError, InterceptorPoint, InterceptorResult, + PostToolAction, PreLlmRequestContext, PreRequestAction, PreToolAction, PromptAction, + PromptSubmitContext, RunExitContext, ToolCallInfo, ToolResultInfo, TurnEndAction, +}; +use agen::llm_client::{ + ClientError, LlmClient, Request, ResponseStream, + event::{Event, ResponseStatus, StatusEvent}, }; -use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::{Engine, EngineError, EngineRunExit, History, RunInterruptionReason}; use async_trait::async_trait; @@ -616,7 +620,7 @@ struct YieldOnce { impl Interceptor for YieldOnce { async fn pre_llm_request( &self, - _context: &mut Vec, + _context: PreLlmRequestContext<'_>, ) -> InterceptorResult { Ok(if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { PreRequestAction::Yield @@ -647,7 +651,10 @@ struct ContinueTurnOnce { #[async_trait] impl Interceptor for ContinueTurnOnce { - async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult { + async fn on_assistant_turn_end( + &self, + _context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { Ok(if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { TurnEndAction::ContinueWithMessages(vec![Item::system_message("continue")]) } else { @@ -686,7 +693,10 @@ impl FailingLifecycleInterceptor { #[async_trait] impl Interceptor for FailingLifecycleInterceptor { - async fn on_prompt_submit(&self, _item: &mut Item) -> InterceptorResult { + async fn on_prompt_submit( + &self, + _context: PromptSubmitContext<'_>, + ) -> InterceptorResult { tokio::task::yield_now().await; self.record(InterceptorPoint::PromptSubmit, PromptAction::Continue) } @@ -698,20 +708,10 @@ impl Interceptor for FailingLifecycleInterceptor { async fn pre_llm_request( &self, - _context: &mut Vec, + _context: PreLlmRequestContext<'_>, ) -> InterceptorResult { tokio::task::yield_now().await; - if self.failure == InterceptorPoint::Abort { - self.calls - .lock() - .unwrap() - .push(InterceptorPoint::PreLlmRequest); - Ok(PreRequestAction::Cancel( - "trigger abort callback".to_string(), - )) - } else { - self.record(InterceptorPoint::PreLlmRequest, PreRequestAction::Continue) - } + self.record(InterceptorPoint::PreLlmRequest, PreRequestAction::Continue) } async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> InterceptorResult { @@ -719,76 +719,83 @@ impl Interceptor for FailingLifecycleInterceptor { self.record(InterceptorPoint::PreToolCall, PreToolAction::Continue) } - async fn post_tool_call( - &self, - _info: &mut ToolResultInfo, - ) -> InterceptorResult { + async fn post_tool_call(&self, _info: &ToolResultInfo) -> InterceptorResult { tokio::task::yield_now().await; self.record(InterceptorPoint::PostToolCall, PostToolAction::Continue) } - async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult { + async fn on_assistant_turn_end( + &self, + context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { tokio::task::yield_now().await; - self.record(InterceptorPoint::TurnEnd, TurnEndAction::Finish) + assert!(context.history.ends_with(context.assistant_items)); + if !context.tool_calls.is_empty() { + assert_eq!( + context + .assistant_items + .iter() + .filter(|item| matches!(item, Item::ToolCall { .. })) + .count(), + context.tool_calls.len() + ); + } + self.record(InterceptorPoint::AssistantTurnEnd, TurnEndAction::Finish) } - async fn on_abort(&self, _reason: &str) -> InterceptorResult<()> { + async fn on_run_exit(&self, _context: RunExitContext<'_>) -> InterceptorResult<()> { tokio::task::yield_now().await; - self.record(InterceptorPoint::Abort, ()) + self.record(InterceptorPoint::RunExit, ()) } } fn expected_interceptor_calls(failure: InterceptorPoint) -> Vec { use InterceptorPoint as Point; - match failure { - Point::PromptSubmit => vec![Point::PromptSubmit, Point::Abort], + let mut calls = match failure { + Point::PromptSubmit => vec![Point::PromptSubmit], Point::PendingHistoryAppends => { - vec![ - Point::PromptSubmit, - Point::PendingHistoryAppends, - Point::Abort, - ] + vec![Point::PromptSubmit, Point::PendingHistoryAppends] } Point::PreLlmRequest => vec![ Point::PromptSubmit, Point::PendingHistoryAppends, Point::PreLlmRequest, - Point::Abort, ], Point::PreToolCall => vec![ Point::PromptSubmit, Point::PendingHistoryAppends, Point::PreLlmRequest, + Point::AssistantTurnEnd, Point::PreToolCall, - Point::Abort, ], Point::PostToolCall => vec![ Point::PromptSubmit, Point::PendingHistoryAppends, Point::PreLlmRequest, + Point::AssistantTurnEnd, Point::PreToolCall, Point::PostToolCall, - Point::Abort, ], - Point::TurnEnd => vec![ + Point::AssistantTurnEnd => vec![ Point::PromptSubmit, Point::PendingHistoryAppends, Point::PreLlmRequest, - Point::TurnEnd, - Point::Abort, + Point::AssistantTurnEnd, ], - Point::Abort => vec![ + Point::RunExit => vec![ Point::PromptSubmit, Point::PendingHistoryAppends, Point::PreLlmRequest, - Point::Abort, + Point::AssistantTurnEnd, ], - } + }; + calls.push(Point::RunExit); + calls } #[tokio::test] -async fn interceptor_failures_are_typed_and_each_lifecycle_point_runs_once() { +async fn interceptor_failures_are_typed_unexpected_and_each_lifecycle_point_runs_once() { use InterceptorPoint as Point; for failure_point in [ @@ -797,8 +804,8 @@ async fn interceptor_failures_are_typed_and_each_lifecycle_point_runs_once() { Point::PreLlmRequest, Point::PreToolCall, Point::PostToolCall, - Point::TurnEnd, - Point::Abort, + Point::AssistantTurnEnd, + Point::RunExit, ] { let interceptor = FailingLifecycleInterceptor::new(failure_point); let needs_tool = matches!(failure_point, Point::PreToolCall | Point::PostToolCall); @@ -821,7 +828,10 @@ async fn interceptor_failures_are_typed_and_each_lifecycle_point_runs_once() { let mut engine = engine.lock(&history); let exit = engine.run(&mut history, "test").await; - let EngineRunExit::Interrupted(RunInterruptionReason::Interceptor(failure)) = exit else { + let EngineRunExit::Interrupted(RunInterruptionReason::Unexpected( + EngineError::Interceptor(failure), + )) = exit + else { panic!("expected typed interceptor interruption at {failure_point}, got {exit:?}"); }; assert_eq!(failure.point(), failure_point); @@ -833,9 +843,220 @@ async fn interceptor_failures_are_typed_and_each_lifecycle_point_runs_once() { interceptor.calls(), expected_interceptor_calls(failure_point) ); + if failure_point == Point::PostToolCall { + assert!( + history + .items() + .any(|item| matches!(item, Item::ToolResult { .. })), + "post-tool failure must not precede terminal output commit" + ); + } } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TerminalMode { + Finish, + PauseOnce, + Yield, +} + +#[derive(Debug, Clone)] +struct RecordingTerminalInterceptor { + mode: TerminalMode, + assistant_turns: Arc, + exits: Arc>>, +} + +impl RecordingTerminalInterceptor { + fn new(mode: TerminalMode) -> Self { + Self { + mode, + assistant_turns: Arc::new(AtomicUsize::new(0)), + exits: Arc::new(Mutex::new(Vec::new())), + } + } + + fn exits(&self) -> Vec<&'static str> { + self.exits.lock().unwrap().clone() + } +} + +#[async_trait] +impl Interceptor for RecordingTerminalInterceptor { + async fn pre_llm_request( + &self, + _context: PreLlmRequestContext<'_>, + ) -> InterceptorResult { + Ok(if self.mode == TerminalMode::Yield { + PreRequestAction::Yield + } else { + PreRequestAction::Continue + }) + } + + async fn on_assistant_turn_end( + &self, + context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { + assert!(!context.assistant_items.is_empty()); + assert!( + context.history.ends_with(context.assistant_items), + "assistant-turn callback must observe committed terminal items" + ); + let turn = self.assistant_turns.fetch_add(1, Ordering::SeqCst); + Ok(if self.mode == TerminalMode::PauseOnce && turn == 0 { + TurnEndAction::Pause + } else { + TurnEndAction::Finish + }) + } + + async fn on_run_exit(&self, context: RunExitContext<'_>) -> InterceptorResult<()> { + let kind = match context.exit { + EngineRunExit::Finished => "finished", + EngineRunExit::Paused => "paused", + EngineRunExit::Yielded => "yielded", + EngineRunExit::Interrupted(RunInterruptionReason::LimitReached) => "limit", + EngineRunExit::Interrupted(RunInterruptionReason::ContextWindowExceeded) => "context", + EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) => "cancelled", + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(_)) => "unexpected", + }; + self.exits.lock().unwrap().push(kind); + Ok(()) + } +} + +#[derive(Clone)] +struct ContextWindowClient; + +#[async_trait] +impl LlmClient for ContextWindowClient { + async fn stream(&self, _request: Request) -> Result { + Err(ClientError::ContextWindowExceeded) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +#[tokio::test] +async fn terminal_observer_runs_once_for_every_exit_and_interruption_kind() { + let finished = RecordingTerminalInterceptor::new(TerminalMode::Finish); + let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); + engine.set_interceptor(finished.clone()); + let mut history = History::new(); + assert!(matches!( + engine.lock(&history).run(&mut history, "finish").await, + EngineRunExit::Finished + )); + assert_eq!(finished.exits(), ["finished"]); + + let yielded = RecordingTerminalInterceptor::new(TerminalMode::Yield); + let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); + engine.set_interceptor(yielded.clone()); + let mut history = History::new(); + assert!(matches!( + engine.lock(&history).run(&mut history, "yield").await, + EngineRunExit::Yielded + )); + assert_eq!(yielded.exits(), ["yielded"]); + + let limited = RecordingTerminalInterceptor::new(TerminalMode::Finish); + let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); + engine.set_max_turns(Some(0)); + engine.set_interceptor(limited.clone()); + let mut history = History::new(); + assert!(matches!( + engine.lock(&history).run(&mut history, "limit").await, + EngineRunExit::Interrupted(RunInterruptionReason::LimitReached) + )); + assert_eq!(limited.exits(), ["limit"]); + + let cancelled = RecordingTerminalInterceptor::new(TerminalMode::Finish); + let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); + engine.set_interceptor(cancelled.clone()); + engine.cancel(); + let mut history = History::new(); + assert!(matches!( + engine.lock(&history).run(&mut history, "cancel").await, + EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) + )); + assert_eq!(cancelled.exits(), ["cancelled"]); + + let context = RecordingTerminalInterceptor::new(TerminalMode::Finish); + let mut engine = Engine::new(ContextWindowClient); + engine.set_interceptor(context.clone()); + let mut history = History::new(); + assert!(matches!( + engine.lock(&history).run(&mut history, "context").await, + EngineRunExit::Interrupted(RunInterruptionReason::ContextWindowExceeded) + )); + assert_eq!(context.exits(), ["context"]); + + let unexpected = FailingLifecycleInterceptor::new(InterceptorPoint::PromptSubmit); + let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); + engine.set_interceptor(unexpected.clone()); + let mut history = History::new(); + assert!(matches!( + engine.lock(&history).run(&mut history, "fail").await, + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(EngineError::Interceptor( + _ + ))) + )); + assert_eq!( + unexpected + .calls() + .iter() + .filter(|point| **point == InterceptorPoint::RunExit) + .count(), + 1 + ); +} + +#[tokio::test] +async fn terminal_observer_does_not_duplicate_on_resume() { + let interceptor = RecordingTerminalInterceptor::new(TerminalMode::PauseOnce); + let first_response = vec![ + Event::tool_use_start(0, "call-1", "count_tool"), + Event::tool_input_delta(0, "{}"), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]; + let client = MockLlmClient::with_responses(vec![first_response, completed_text_events()]); + let tool = CountingTool::new("count_tool"); + let mut engine = Engine::new(client); + engine.register_tool(tool.definition()); + engine.set_interceptor(interceptor.clone()); + let mut history = History::new(); + let mut engine = engine.lock(&history); + + assert!(matches!( + engine.run(&mut history, "pause").await, + EngineRunExit::Paused + )); + assert_eq!(interceptor.exits(), ["paused"]); + assert_eq!( + tool.call_count(), + 0, + "pause must retain the pending tool phase" + ); + + assert!(matches!( + engine.resume(&mut history).await, + EngineRunExit::Finished + )); + assert_eq!(interceptor.exits(), ["paused", "finished"]); + assert_eq!( + tool.call_count(), + 1, + "resume must execute the retained tool once" + ); +} + #[tokio::test] async fn max_turns_is_scoped_to_each_fresh_run() { let mut history: History = History::new(); diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index ead143b4..f63b20a1 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -919,10 +919,7 @@ async fn test_tool_execution_context_for_skipped_and_synthetic_paths() { }) } - async fn post_tool_call( - &self, - info: &mut ToolResultInfo, - ) -> InterceptorResult { + async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult { self.post_contexts .lock() .unwrap() @@ -1026,9 +1023,9 @@ async fn test_before_tool_call_skip() { ); } -/// Hook: post_tool_call - verify that results can be modified +/// Hook: post_tool_call - verify that the committed terminal result is observed. #[tokio::test] -async fn test_post_tool_call_modification() { +async fn test_post_tool_call_observes_committed_result() { // Prepare responses for multiple requests let client = MockLlmClient::with_responses(vec![ // First request: tool call @@ -1079,43 +1076,39 @@ async fn test_post_tool_call_modification() { engine.register_tool(simple_tool_definition()); - // Policy to modify results - struct ModifyingPolicy { - modified_content: Arc>>, + // Policy to observe the committed terminal result. + struct ObservingPolicy { + observed_content: Arc>>, } #[async_trait] - impl Interceptor for ModifyingPolicy { - async fn post_tool_call( - &self, - info: &mut ToolResultInfo, - ) -> InterceptorResult { - info.result.summary = format!("[Modified] {}", info.result.summary); - *self.modified_content.lock().unwrap() = Some(info.result.summary.clone()); + impl Interceptor for ObservingPolicy { + async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult { + *self.observed_content.lock().unwrap() = Some(info.result.summary.clone()); Ok(PostToolAction::Continue) } } - let modified_content = Arc::new(std::sync::Mutex::new(None)); - engine.set_interceptor(ModifyingPolicy { - modified_content: modified_content.clone(), + let observed_content = Arc::new(std::sync::Mutex::new(None)); + engine.set_interceptor(ObservingPolicy { + observed_content: observed_content.clone(), }); // Mutable::run consumes self, returns (Locked, EngineResult) - let result = engine.run(&mut history, "Test modification").await; + let result = engine.run(&mut history, "Test observation").await; assert!( matches!(result.result, agen::EngineRunExit::Finished), "Engine should complete" ); - // Verify hook was called and content was modified - let content = modified_content.lock().unwrap().clone(); - assert!(content.is_some(), "Hook should have been called"); - assert!( - content.unwrap().contains("[Modified]"), - "Result should be modified" - ); + // Verify the interceptor observed the exact committed result. + let observed = observed_content.lock().unwrap().clone(); + assert_eq!(observed.as_deref(), Some("Original Result")); + assert!(history.items().any(|item| matches!( + item, + Item::ToolResult { summary, .. } if summary == "Original Result" + ))); } /// Hook: pre_tool_call synthetic result - skipped tool gets an error result in history. @@ -1189,19 +1182,24 @@ async fn post_tool_abort_commits_confirmed_result_before_stopping_run() { let tool = SlowTool::new("confirmed", 1); engine.register_tool(tool.definition()); - struct AbortAfterResult; + let observed = Arc::new(Mutex::new(Vec::<&'static str>::new())); + struct AbortAfterResult { + lifecycle: Arc>>, + } #[async_trait] impl Interceptor for AbortAfterResult { async fn post_tool_call( &self, - _info: &mut ToolResultInfo, + _info: &ToolResultInfo, ) -> InterceptorResult { + self.lifecycle.lock().unwrap().push("post_tool_call"); Ok(PostToolAction::Abort("policy stopped the run".to_string())) } } - engine.set_interceptor(AbortAfterResult); + engine.set_interceptor(AbortAfterResult { + lifecycle: observed.clone(), + }); - let observed = Arc::new(Mutex::new(Vec::<&'static str>::new())); let published = observed.clone(); engine.on_tool_result(move |_| published.lock().unwrap().push("published")); let committed = observed.clone(); @@ -1221,7 +1219,7 @@ async fn post_tool_abort_commits_confirmed_result_before_stopping_run() { assert_eq!(tool.call_count(), 1); assert_eq!( observed.lock().unwrap().as_slice(), - ["committed", "published", "run-returned"] + ["committed", "published", "post_tool_call", "run-returned"] ); assert!(matches!( output.result, diff --git a/crates/session-store/tests/session_test.rs b/crates/session-store/tests/session_test.rs index 2f0ffbeb..dced4ca9 100644 --- a/crates/session-store/tests/session_test.rs +++ b/crates/session-store/tests/session_test.rs @@ -3,7 +3,7 @@ mod common; use std::ops::{Deref, DerefMut}; use std::sync::Arc; -use agen::interceptor::{Interceptor, InterceptorResult, TurnEndAction}; +use agen::interceptor::{AssistantTurnEndContext, Interceptor, InterceptorResult, TurnEndAction}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::types::{Item, RequestConfig}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; @@ -100,7 +100,10 @@ struct PausePolicy; #[async_trait] impl Interceptor for PausePolicy { - async fn on_turn_end(&self, _history: &[Item]) -> InterceptorResult { + async fn on_assistant_turn_end( + &self, + _context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { Ok(TurnEndAction::Pause) } } @@ -350,7 +353,8 @@ async fn session_run_with_tool_call() { async fn session_resume_after_pause() { let (_dir, store) = make_store(); - // First run: tool call with pause policy → Paused + // First terminal assistant response requests a tool; the assistant-turn + // interceptor pauses before the Engine enters the tool phase. let client = MockLlmClient::with_responses(tool_call_events()); let mut worker = TestWorker::new(Engine::new(client)); worker.register_tool(weather_tool_definition()); @@ -386,7 +390,7 @@ async fn session_resume_after_pause() { // Restore state and verify let state = session_store::restore(&store, sid, segid).unwrap(); assert!(state.last_run_interrupted); - assert_eq!(state.active_run_turn_count, Some(2)); + assert_eq!(state.active_run_turn_count, Some(1)); } #[tokio::test] diff --git a/crates/worker/src/compact/worker.rs b/crates/worker/src/compact/worker.rs index cd218ac0..158c68ff 100644 --- a/crates/worker/src/compact/worker.rs +++ b/crates/worker/src/compact/worker.rs @@ -23,7 +23,8 @@ use std::sync::{Arc, Mutex}; use agen::Item; use agen::interceptor::{ - Interceptor, InterceptorResult, PreRequestAction, PreToolAction, ToolCallInfo, + Interceptor, InterceptorResult, PreLlmRequestContext, PreRequestAction, PreToolAction, + ToolCallInfo, }; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput, ToolResult}; use async_trait::async_trait; @@ -402,8 +403,9 @@ impl CompactWorkerInterceptor { impl Interceptor for CompactWorkerInterceptor { async fn pre_llm_request( &self, - context: &mut Vec, + context: PreLlmRequestContext<'_>, ) -> InterceptorResult { + let context = context.items; let records = self.usage_tracker.records(); let estimate = agen::token_counter::total_tokens(context, &records); if estimate.tokens > self.max_input_tokens { @@ -472,13 +474,23 @@ mod tests { let mut context = vec![Item::user_message("hello")]; assert!(matches!( - interceptor.pre_llm_request(&mut context).await.unwrap(), + interceptor + .pre_llm_request(PreLlmRequestContext { + items: &mut context, + }) + .await + .unwrap(), PreRequestAction::Continue )); tracker.record_usage(&make_usage(100)); assert!(matches!( - interceptor.pre_llm_request(&mut context).await.unwrap(), + interceptor + .pre_llm_request(PreLlmRequestContext { + items: &mut context, + }) + .await + .unwrap(), PreRequestAction::Continue )); tracker.record_usage(&make_usage(100)); @@ -486,7 +498,12 @@ mod tests { // Two 100-token requests would exceed a cumulative 150-token cap, but // current occupancy is still the latest 100-token measurement. assert!(matches!( - interceptor.pre_llm_request(&mut context).await.unwrap(), + interceptor + .pre_llm_request(PreLlmRequestContext { + items: &mut context, + }) + .await + .unwrap(), PreRequestAction::Continue )); } @@ -508,13 +525,23 @@ mod tests { let mut context = vec![Item::user_message("hello")]; assert!(matches!( - interceptor.pre_llm_request(&mut context).await.unwrap(), + interceptor + .pre_llm_request(PreLlmRequestContext { + items: &mut context, + }) + .await + .unwrap(), PreRequestAction::Continue )); tracker.record_usage(&make_usage(100)); assert!(matches!( - interceptor.pre_llm_request(&mut context).await.unwrap(), + interceptor + .pre_llm_request(PreLlmRequestContext { + items: &mut context, + }) + .await + .unwrap(), PreRequestAction::ContinueWith(items) if items.len() == 1 && items[0].as_text().unwrap_or_default().contains("write_summary") )); @@ -528,13 +555,23 @@ mod tests { let mut context = vec![Item::user_message("hello")]; assert!(matches!( - interceptor.pre_llm_request(&mut context).await.unwrap(), + interceptor + .pre_llm_request(PreLlmRequestContext { + items: &mut context, + }) + .await + .unwrap(), PreRequestAction::Continue )); tracker.record_usage(&make_usage(100)); assert!(matches!( - interceptor.pre_llm_request(&mut context).await.unwrap(), + interceptor + .pre_llm_request(PreLlmRequestContext { + items: &mut context, + }) + .await + .unwrap(), PreRequestAction::Cancel(message) if message.contains("occupancy") )); } diff --git a/crates/worker/src/hook.rs b/crates/worker/src/hook.rs index d34ec5cc..4318f9c3 100644 --- a/crates/worker/src/hook.rs +++ b/crates/worker/src/hook.rs @@ -285,12 +285,6 @@ pub struct TurnEndInfo { pub final_text_preview: String, } -/// Information passed to `OnAbort` hooks. -pub struct AbortInfo { - /// Reason supplied by the aborter. - pub reason: String, -} - // ============================================================================= // Hook Event Kinds // ============================================================================= @@ -315,10 +309,8 @@ pub struct PreLlmRequest; pub struct PreToolCall; /// After each tool completes; observational except it may abort the run. pub struct PostToolCall; -/// When a turn ends with no tool calls; observational except it may pause. +/// After every terminal assistant response is committed; observational except it may pause. pub struct OnTurnEnd; -/// When execution is interrupted; observational only. -pub struct OnAbort; impl HookEventKind for OnPromptSubmit { type Input = PromptSubmitInfo; @@ -345,11 +337,6 @@ impl HookEventKind for OnTurnEnd { type Output = HookTurnEndAction; } -impl HookEventKind for OnAbort { - type Input = AbortInfo; - type Output = (); -} - // ============================================================================= // Hook Trait // ============================================================================= @@ -380,7 +367,6 @@ pub struct HookRegistryBuilder { pre_tool_call: Vec>>, post_tool_call: Vec>>, on_turn_end: Vec>>, - on_abort: Vec>>, } impl HookRegistryBuilder { @@ -408,10 +394,6 @@ impl HookRegistryBuilder { self.on_turn_end.push(Box::new(hook)); } - pub fn add_on_abort(&mut self, hook: impl Hook + 'static) { - self.on_abort.push(Box::new(hook)); - } - /// Freeze the builder into an immutable registry. pub fn build(self) -> HookRegistry { HookRegistry { @@ -420,7 +402,6 @@ impl HookRegistryBuilder { pre_tool_call: self.pre_tool_call, post_tool_call: self.post_tool_call, on_turn_end: self.on_turn_end, - on_abort: self.on_abort, } } } @@ -432,7 +413,6 @@ pub struct HookRegistry { pub(crate) pre_tool_call: Vec>>, pub(crate) post_tool_call: Vec>>, pub(crate) on_turn_end: Vec>>, - pub(crate) on_abort: Vec>>, } #[cfg(test)] diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index a1468443..8d380ea7 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -15,8 +15,9 @@ use std::sync::{Arc, Mutex}; use agen::Item; use agen::UsageRecord; use agen::interceptor::{ - Interceptor, InterceptorResult, PostToolAction, PreRequestAction, PreToolAction, PromptAction, - ToolCallInfo, ToolResultInfo, TurnEndAction, + AssistantTurnEndContext, Interceptor, InterceptorResult, PostToolAction, PreLlmRequestContext, + PreRequestAction, PreToolAction, PromptAction, PromptSubmitContext, ToolCallInfo, + ToolResultInfo, TurnEndAction, }; use agen::tool::ToolOutput; use arc_swap::ArcSwap; @@ -28,9 +29,9 @@ use crate::compact::usage_tracker::UsageTracker; use session_store::SystemItem; use crate::hook::{ - AbortInfo, HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookPromptAction, - HookRegistry, HookTurnEndAction, PreRequestContext, PreRequestInfo, PromptSubmitInfo, - SystemItemAppendHandle, ToolCallSummary, ToolResultSummary, TurnEndInfo, + HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookPromptAction, HookRegistry, + HookTurnEndAction, PreRequestContext, PreRequestInfo, PromptSubmitInfo, SystemItemAppendHandle, + ToolCallSummary, ToolResultSummary, TurnEndInfo, }; use crate::ipc::notify_buffer::{NotifyBuffer, build_system_item_with_provenance}; use crate::prompt::catalog::PromptCatalog; @@ -232,7 +233,11 @@ impl WorkerInterceptor { #[async_trait] impl Interceptor for WorkerInterceptor { - async fn on_prompt_submit(&self, item: &mut Item) -> InterceptorResult { + async fn on_prompt_submit( + &self, + context: PromptSubmitContext<'_>, + ) -> InterceptorResult { + let item = context.item; let turn_index = self.next_turn_index.fetch_add(1, Ordering::Relaxed); self.tool_calls_this_turn.store(0, Ordering::Relaxed); @@ -310,8 +315,9 @@ impl Interceptor for WorkerInterceptor { async fn pre_llm_request( &self, - context: &mut Vec, + context: PreLlmRequestContext<'_>, ) -> InterceptorResult { + let context = context.items; let initial_tokens = self.estimated_tokens(context); if self.request_threshold_exceeded(initial_tokens, context) { return Ok(PreRequestAction::Yield); @@ -395,7 +401,7 @@ impl Interceptor for WorkerInterceptor { Ok(PreToolAction::Continue) } - async fn post_tool_call(&self, info: &mut ToolResultInfo) -> InterceptorResult { + async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult { let summary = ToolResultSummary { call_id: info.result.tool_use_id.clone(), tool_name: info.call.name.clone(), @@ -416,7 +422,11 @@ impl Interceptor for WorkerInterceptor { Ok(PostToolAction::Continue) } - async fn on_turn_end(&self, history: &[Item]) -> InterceptorResult { + async fn on_assistant_turn_end( + &self, + context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { + let history = context.history; let final_text_preview = history .iter() .rev() @@ -437,16 +447,6 @@ impl Interceptor for WorkerInterceptor { } Ok(TurnEndAction::Finish) } - - async fn on_abort(&self, reason: &str) -> InterceptorResult<()> { - let info = AbortInfo { - reason: reason.to_string(), - }; - for hook in &self.registry.on_abort { - hook.call(&info).await; - } - Ok(()) - } } struct ContextShape { @@ -629,7 +629,10 @@ mod tests { None, ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Yield)); // Hook must not run when an internal mechanism short-circuits first. @@ -661,7 +664,10 @@ mod tests { })), ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); match action { PreRequestAction::YieldWith(items) => assert_eq!(items.len(), 1), @@ -698,7 +704,10 @@ mod tests { ) .with_usage_tracker(usage_tracker); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Yield)); } @@ -722,7 +731,10 @@ mod tests { None, ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -763,7 +775,10 @@ mod tests { None, ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -790,7 +805,10 @@ mod tests { None, ); let mut ctx = ctx_items; - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -811,7 +829,10 @@ mod tests { None, ); let mut ctx: Vec = Vec::new(); - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -840,7 +861,10 @@ mod tests { ); let mut ctx: Vec = Vec::new(); - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(saw_handle.load(Ordering::Relaxed)); let PreRequestAction::ContinueWith(items) = action else { @@ -887,7 +911,10 @@ mod tests { ); let mut ctx: Vec = Vec::new(); - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(!saw_handle.load(Ordering::Relaxed)); assert!(matches!(action, PreRequestAction::Continue)); @@ -1042,7 +1069,14 @@ mod tests { ); let history = vec![Item::user_message("hi"), Item::assistant_message("done")]; - let action = interceptor.on_turn_end(&history).await.unwrap(); + let action = interceptor + .on_assistant_turn_end(AssistantTurnEndContext { + assistant_items: &[], + history: &history, + tool_calls: &[], + }) + .await + .unwrap(); assert!(matches!(action, TurnEndAction::Pause)); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -1079,7 +1113,10 @@ mod tests { let ctx_items = vec![Item::user_message("hi")]; for _ in 0..23 { let mut ctx = ctx_items.clone(); - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Continue)); usage_tracker.record_usage(&agen::event::UsageEvent { input_tokens: Some(10), @@ -1091,7 +1128,10 @@ mod tests { } let mut ctx = ctx_items.clone(); - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); let appended_len = match action { PreRequestAction::ContinueWith(items) => items.len(), other => panic!("expected reminder append, got {other:?}"), @@ -1275,7 +1315,10 @@ mod tests { None, ); let mut ctx: Vec = vec![Item::user_message("hi")]; - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Continue)); assert_eq!(ctx.len(), 1, "pre_llm_request must not append notifies"); @@ -1305,7 +1348,10 @@ mod tests { None, ); let mut ctx: Vec = Vec::new(); - let action = interceptor.pre_llm_request(&mut ctx).await.unwrap(); + let action = interceptor + .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .await + .unwrap(); assert!(matches!(action, PreRequestAction::Cancel(_))); assert!(first_called.load(Ordering::Relaxed)); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index f98a9b5a..332ab2e5 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -50,8 +50,7 @@ use crate::feature::{ FeatureRegistryInstallReport, dedupe_instruction_contributions, }; use crate::hook::{ - Hook, HookRegistryBuilder, OnAbort, OnPromptSubmit, OnTurnEnd, PostToolCall, PreLlmRequest, - PreToolCall, + Hook, HookRegistryBuilder, OnPromptSubmit, OnTurnEnd, PostToolCall, PreLlmRequest, PreToolCall, }; use crate::in_flight::InFlightEvents; use crate::internal_worker::{ @@ -2328,12 +2327,6 @@ impl Worker { self.hook_builder.add_on_turn_end(hook); } - /// Register a hook that runs when execution is aborted. - pub fn add_on_abort_hook(&mut self, hook: impl Hook + 'static) { - self.assert_hooks_open(); - self.hook_builder.add_on_abort(hook); - } - /// Install the hook-based interceptor on the Engine if not already done. /// /// When either compaction threshold (`threshold` or @@ -3733,7 +3726,6 @@ impl Worker { | EngineRunExit::Yielded | EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) | EngineRunExit::Interrupted(RunInterruptionReason::ContextWindowExceeded) - | EngineRunExit::Interrupted(RunInterruptionReason::Interceptor(_)) | EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(_)) ); let active_run_turn_count = self.engine.as_ref().unwrap().active_run_turn_count(); @@ -6130,7 +6122,6 @@ fn run_interruption_reason_error_code(reason: &RunInterruptionReason) -> ErrorCo RunInterruptionReason::Unexpected(EngineError::Tool(_)) => ErrorCode::ToolError, RunInterruptionReason::LimitReached | RunInterruptionReason::Cancelled - | RunInterruptionReason::Interceptor(_) | RunInterruptionReason::Unexpected( EngineError::Aborted(_) | EngineError::Cancelled @@ -6148,7 +6139,6 @@ fn run_interruption_reason_message(reason: &RunInterruptionReason) -> String { RunInterruptionReason::LimitReached => "engine turn limit reached".to_string(), RunInterruptionReason::ContextWindowExceeded => "model context window reached".to_string(), RunInterruptionReason::Cancelled => "engine run cancelled".to_string(), - RunInterruptionReason::Interceptor(failure) => failure.to_string(), RunInterruptionReason::Unexpected(error) => format!("unexpected engine failure: {error}"), } } From eac4a0c071a22082f7180371fb7fb1621a00ff2d Mon Sep 17 00:00:00 2001 From: Hare Date: Thu, 3 Sep 2026 16:46:02 +0900 Subject: [PATCH 03/22] fix: harden interceptor lifecycle contracts --- crates/agen/examples/engine_cli.rs | 5 +- crates/agen/src/engine.rs | 204 ++++++++++++++--- crates/agen/src/interceptor.rs | 218 +++++++++++++------ crates/agen/src/lib.rs | 7 +- crates/agen/tests/annotated_history_test.rs | 126 +++++++++++ crates/agen/tests/engine_state_test.rs | 164 +++++++++++--- crates/agen/tests/parallel_execution_test.rs | 117 +++++++++- crates/worker/src/compact/worker.rs | 37 +++- crates/worker/src/feature.rs | 2 +- crates/worker/src/ipc/interceptor.rs | 185 +++++++++++++--- crates/worker/src/worker.rs | 2 +- 11 files changed, 877 insertions(+), 190 deletions(-) diff --git a/crates/agen/examples/engine_cli.rs b/crates/agen/examples/engine_cli.rs index 10eb28ef..ee90ca7e 100644 --- a/crates/agen/examples/engine_cli.rs +++ b/crates/agen/examples/engine_cli.rs @@ -280,7 +280,10 @@ impl ToolResultPrinterPolicy { #[async_trait] impl Interceptor for ToolResultPrinterPolicy { - async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult { + async fn post_tool_call( + &self, + info: &ToolResultInfo<'_, ()>, + ) -> InterceptorResult { let name = self .call_names .lock() diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 276ee3a4..644cf161 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -15,10 +15,12 @@ use crate::{ }, handler::{ErrorKind, StatusKind, ToolUseBlockStart, UsageKind}, interceptor::{ - AssistantTurnEndContext, DefaultInterceptor, Interceptor, InterceptorFailure, - InterceptorPoint, PostToolAction, PreLlmRequestContext, PreRequestAction, PreToolAction, - PromptAction, PromptSubmitContext, RunExitContext, ToolCallInfo, ToolResultInfo, - TurnEndAction, + AssistantTurnEndContext, DefaultInterceptor, Interceptor, InterceptorCallId, + InterceptorCounter, InterceptorCounters, InterceptorError, InterceptorErrorCategory, + InterceptorFailure, InterceptorInvocation, InterceptorPhase, InterceptorRunId, + InterceptorTurnId, PendingHistoryAppendsContext, PostToolAction, PreLlmRequestContext, + PreRequestAction, PreToolAction, PromptAction, PromptSubmitContext, RunExitContext, + ToolCallInfo, ToolResultInfo, TurnEndAction, }, llm_client::{ ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream, @@ -186,7 +188,7 @@ impl From> for EngineRunExit { /// Result of [`Engine::run`] or [`Engine::resume`]. /// /// Contains the `Locked` Engine (ready for subsequent runs) and the outcome. -pub struct EngineRunOutput { +pub struct EngineRunOutput { /// The Engine, now in Locked state. pub engine: Engine, /// Outcome of the turn. @@ -310,7 +312,7 @@ enum StreamCompletion { Interrupted { reason: String }, } -pub struct Engine { +pub struct Engine { /// LLM client client: C, /// Retry policy for opening an LLM response stream. @@ -327,7 +329,7 @@ pub struct Engine { /// Tool server handle tool_server: ToolServerHandle, /// Interceptor for control-flow decisions - interceptor: Box, + interceptor: Box>, /// System prompt system_prompt: Option, /// History length at lock time (only meaningful in Locked state) @@ -346,6 +348,11 @@ pub struct Engine { /// `max_turns` is enforced against this run-scoped count rather than the /// cumulative `turn_count` above. active_run_turn_count: Option, + /// Identity retained across pause/yield and resume. + active_run_id: Option, + next_run_id: u64, + interceptor_invocation_count: usize, + last_run_exit_observer_failure: Option, /// LlmCall count (per-Engine running counter, monotonic). Unlike /// `turn_count` this never collapses retries. llm_call_count: usize, @@ -426,18 +433,57 @@ pub struct Engine { _state: PhantomData<(S, A)>, } -impl Engine { +impl Engine { fn start_logical_run(&mut self) { self.active_run_turn_count = Some(0); + self.active_run_id = Some(InterceptorRunId(self.next_run_id)); + self.next_run_id = self.next_run_id.wrapping_add(1).max(1); + self.interceptor_invocation_count = 0; + self.last_run_exit_observer_failure = None; } fn ensure_logical_run(&mut self) { self.active_run_turn_count.get_or_insert(0); + if self.active_run_id.is_none() { + self.active_run_id = Some(InterceptorRunId(self.next_run_id)); + self.next_run_id = self.next_run_id.wrapping_add(1).max(1); + self.interceptor_invocation_count = 0; + } + } + + fn interceptor_invocation( + &mut self, + phase: InterceptorPhase, + turn_id: Option, + call_id: Option, + tool_call: usize, + ) -> InterceptorInvocation { + let invocation = self.interceptor_invocation_count; + self.interceptor_invocation_count = self.interceptor_invocation_count.saturating_add(1); + InterceptorInvocation { + run_id: self + .active_run_id + .expect("logical run identity must exist before interception"), + turn_id: turn_id.map(|value| InterceptorTurnId(value as u64)), + call_id, + phase, + counters: InterceptorCounters { + invocation: InterceptorCounter::from_usize(invocation), + engine_turn: InterceptorCounter::from_usize(self.turn_count), + run_turn: InterceptorCounter::from_usize( + self.active_run_turn_count.unwrap_or_default(), + ), + llm_call: InterceptorCounter::from_usize(self.llm_call_count), + tool_batch: InterceptorCounter::from_usize(self.tool_execution_batch_count), + tool_call: InterceptorCounter::from_usize(tool_call), + }, + } } fn finish_logical_run(&mut self, exit: &EngineRunExit) { if !matches!(exit, EngineRunExit::Paused | EngineRunExit::Yielded) { self.active_run_turn_count = None; + self.active_run_id = None; } } @@ -743,7 +789,7 @@ impl Engine { /// The interceptor governs approval, skip, pause, and abort decisions /// at key points in the execution loop. If not set, the default /// interceptor is used (all Continue / Finish). - pub fn set_interceptor(&mut self, interceptor: impl Interceptor + 'static) { + pub fn set_interceptor(&mut self, interceptor: impl Interceptor + 'static) { self.interceptor = Box::new(interceptor); } @@ -844,6 +890,10 @@ impl Engine { /// /// `Some` is retained only while Pause or Yield permits a later /// [`resume`](Self::resume). Terminal outcomes return this to `None`. + pub fn last_run_exit_observer_failure(&self) -> Option<&InterceptorFailure> { + self.last_run_exit_observer_failure.as_ref() + } + pub fn active_run_turn_count(&self) -> Option { self.active_run_turn_count } @@ -855,6 +905,13 @@ impl Engine { /// [`resume`](Self::resume) starts a fresh budget. pub fn set_active_run_turn_count(&mut self, turn_count: Option) { self.active_run_turn_count = turn_count; + if turn_count.is_none() { + self.active_run_id = None; + } else if self.active_run_id.is_none() { + self.active_run_id = Some(InterceptorRunId(self.next_run_id)); + self.next_run_id = self.next_run_id.wrapping_add(1).max(1); + self.interceptor_invocation_count = 0; + } } /// Get the current LlmCall count (per-Engine running counter, never @@ -1082,19 +1139,24 @@ impl Engine { async fn finalize_run_exit( &mut self, + history: &History, result: Result, ) -> EngineRunExit { let exit = EngineRunExit::from(result); - let exit = match self + let invocation = self.interceptor_invocation(InterceptorPhase::RunExit, None, None, 0); + self.last_run_exit_observer_failure = None; + if let Err(error) = self .interceptor - .on_run_exit(RunExitContext { exit: &exit }) + .on_run_exit(RunExitContext { + invocation, + exit: &exit, + history: history.entries(), + }) .await { - Ok(()) => exit, - Err(error) => EngineRunExit::Interrupted(RunInterruptionReason::Unexpected( - InterceptorFailure::new(InterceptorPoint::RunExit, error).into(), - )), - }; + self.last_run_exit_observer_failure = + Some(InterceptorFailure::new(InterceptorPhase::RunExit, error)); + } self.finish_logical_run(&exit); exit } @@ -1167,9 +1229,18 @@ impl Engine { // Phase 1: Apply pre_tool_call interceptor (determine skip/abort/synthetic result) let mut approved_calls = Vec::new(); for (call_index, mut tool_call) in tool_calls.into_iter().enumerate() { + let expected_tool_use_id = tool_call.id.clone(); let context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index); if let Some((meta, tool)) = self.tool_server.get_tool(&tool_call.name) { + let invocation = self.interceptor_invocation( + InterceptorPhase::PreToolCall, + Some(self.turn_count.saturating_sub(1)), + Some(InterceptorCallId::Tool(expected_tool_use_id.clone())), + call_index, + ); let mut info = ToolCallInfo { + invocation, + history: history.entries(), call: tool_call.clone(), meta, tool, @@ -1182,16 +1253,36 @@ impl Engine { .await .map_err(|error| { EngineError::from(InterceptorFailure::new( - InterceptorPoint::PreToolCall, + InterceptorPhase::PreToolCall, error, )) })?; + if info.call.id != expected_tool_use_id { + return Err(InterceptorFailure::new( + InterceptorPhase::PreToolCall, + InterceptorError::new( + InterceptorErrorCategory::ContractViolation, + "pre-tool interceptor changed immutable tool call identity", + ), + ) + .into()); + } match pre_tool_action { PreToolAction::Continue => {} PreToolAction::Skip => { continue; } PreToolAction::SyntheticResult(result) => { + if result.tool_use_id != expected_tool_use_id { + return Err(InterceptorFailure::new( + InterceptorPhase::PreToolCall, + InterceptorError::new( + InterceptorErrorCategory::ContractViolation, + "synthetic tool result changed immutable tool call identity", + ), + ) + .into()); + } let tool_call = info.call; let mut context = info.context; context.call_id = tool_call.id.clone(); @@ -1538,7 +1629,15 @@ impl Engine { self.emit_tool_result(&tool_result); if let Some((tool_call, meta, tool, context)) = call_info { + let invocation = self.interceptor_invocation( + InterceptorPhase::PostToolCall, + Some(self.turn_count.saturating_sub(1)), + Some(InterceptorCallId::Tool(tool_call.id.clone())), + context.call_index, + ); let info = ToolResultInfo { + invocation, + history: history.entries(), call: tool_call.clone(), result: tool_result, meta: meta.clone(), @@ -1551,7 +1650,7 @@ impl Engine { .await .map_err(|error| { EngineError::from(InterceptorFailure::new( - InterceptorPoint::PostToolCall, + InterceptorPhase::PostToolCall, error, )) })?; @@ -1622,13 +1721,22 @@ impl Engine { // These are committed *before* the per-request clone so they // participate in the LLM request below and get persisted by // the caller that owns durable history. + let pending_invocation = self.interceptor_invocation( + InterceptorPhase::PendingHistoryAppends, + Some(current_turn), + None, + 0, + ); let pending = self .interceptor - .pending_history_appends() + .pending_history_appends(PendingHistoryAppendsContext { + invocation: pending_invocation, + history: history.entries(), + }) .await .map_err(|error| { EngineError::from(InterceptorFailure::new( - InterceptorPoint::PendingHistoryAppends, + InterceptorPhase::PendingHistoryAppends, error, )) })?; @@ -1698,15 +1806,23 @@ impl Engine { } // Interceptor: pre_llm_request + let request_invocation = self.interceptor_invocation( + InterceptorPhase::PreLlmRequest, + Some(current_turn), + Some(InterceptorCallId::Llm(self.llm_call_count as u64)), + 0, + ); let pre_request_action = self .interceptor .pre_llm_request(PreLlmRequestContext { + invocation: request_invocation, items: &mut request_context, + history: history.entries(), }) .await .map_err(|error| { EngineError::from(InterceptorFailure::new( - InterceptorPoint::PreLlmRequest, + InterceptorPhase::PreLlmRequest, error, )) })?; @@ -1822,21 +1938,29 @@ impl Engine { let tool_calls = self.tool_call_collector.take_collected(); let assistant_items = self.build_assistant_items(&reasoning_items, &text_blocks, &tool_calls); - let committed_assistant_items = assistant_items.clone(); + let assistant_start = history.len(); self.append_history_items(history, assistant_items, annotate)?; - let assistant_turn_history = history.items_cloned(); + let assistant_invocation = self.interceptor_invocation( + InterceptorPhase::AssistantTurnEnd, + Some(current_turn), + Some(InterceptorCallId::Llm( + self.llm_call_count.saturating_sub(1) as u64, + )), + 0, + ); let assistant_turn_action = self .interceptor .on_assistant_turn_end(AssistantTurnEndContext { - assistant_items: &committed_assistant_items, - history: &assistant_turn_history, + invocation: assistant_invocation, + assistant_entries: &history.entries()[assistant_start..], + history: history.entries(), tool_calls: &tool_calls, }) .await .map_err(|error| { EngineError::from(InterceptorFailure::new( - InterceptorPoint::AssistantTurnEnd, + InterceptorPhase::AssistantTurnEnd, error, )) })?; @@ -2145,7 +2269,7 @@ impl Engine { } } -impl Engine { +impl Engine { /// Create a new annotated Engine (in Mutable state). pub fn new_annotated(client: C) -> Self { let text_block_collector = TextBlockCollector::new(); @@ -2173,6 +2297,10 @@ impl Engine { locked_prefix_len: 0, turn_count: 0, active_run_turn_count: None, + active_run_id: None, + next_run_id: 1, + interceptor_invocation_count: 0, + last_run_exit_observer_failure: None, llm_call_count: 0, tool_execution_batch_count: 0, max_turns: None, @@ -2448,6 +2576,10 @@ impl Engine { locked_prefix_len, turn_count: self.turn_count, active_run_turn_count: self.active_run_turn_count, + active_run_id: self.active_run_id, + next_run_id: self.next_run_id, + interceptor_invocation_count: self.interceptor_invocation_count, + last_run_exit_observer_failure: self.last_run_exit_observer_failure, llm_call_count: self.llm_call_count, tool_execution_batch_count: self.tool_execution_batch_count, max_turns: self.max_turns, @@ -2524,7 +2656,7 @@ impl Engine { } } -impl Engine { +impl Engine { /// Execute a turn /// /// Adds a new user message to history and sends a request to the LLM. @@ -2538,7 +2670,7 @@ impl Engine { let result = self .run_result_with_annotation(history, user_input.into(), annotate) .await; - self.finalize_run_exit(result).await + self.finalize_run_exit(history, result).await } async fn run_result_with_annotation( @@ -2549,16 +2681,21 @@ impl Engine { ) -> Result { // Supplying new user input abandons any paused/yielded logical run. self.active_run_turn_count = None; + self.active_run_id = None; + self.start_logical_run(); let mut user_item = Item::user_message(user_input); + let invocation = self.interceptor_invocation(InterceptorPhase::PromptSubmit, None, None, 0); let prompt_action = self .interceptor .on_prompt_submit(PromptSubmitContext { + invocation, item: &mut user_item, + history: history.entries(), }) .await .map_err(|error| { EngineError::from(InterceptorFailure::new( - InterceptorPoint::PromptSubmit, + InterceptorPhase::PromptSubmit, error, )) })?; @@ -2571,7 +2708,6 @@ impl Engine { if !extras.is_empty() { self.append_history_items(history, extras, annotate)?; } - self.start_logical_run(); match self.run_turn_loop(history, annotate).await { Err(EngineError::PauseRequested) => Ok(EngineResult::Paused), other => other, @@ -2585,7 +2721,7 @@ impl Engine { annotate: &mut impl FnMut(&Item) -> Result, ) -> EngineRunExit { let result = self.resume_result_with_annotation(history, annotate).await; - self.finalize_run_exit(result).await + self.finalize_run_exit(history, result).await } async fn resume_result_with_annotation( @@ -2623,6 +2759,10 @@ impl Engine { locked_prefix_len: 0, turn_count: self.turn_count, active_run_turn_count: self.active_run_turn_count, + active_run_id: self.active_run_id, + next_run_id: self.next_run_id, + interceptor_invocation_count: self.interceptor_invocation_count, + last_run_exit_observer_failure: self.last_run_exit_observer_failure, llm_call_count: self.llm_call_count, tool_execution_batch_count: self.tool_execution_batch_count, max_turns: self.max_turns, diff --git a/crates/agen/src/interceptor.rs b/crates/agen/src/interceptor.rs index b85e0d9e..28b2afa6 100644 --- a/crates/agen/src/interceptor.rs +++ b/crates/agen/src/interceptor.rs @@ -10,51 +10,73 @@ use async_trait::async_trait; use crate::Item; use crate::engine::EngineRunExit; +use crate::history::HistoryEntry; use crate::tool::{Tool, ToolCall, ToolExecutionContext, ToolMeta, ToolResult}; // ============================================================================= -// Failure Types +// Typed lifecycle metadata and failures // ============================================================================= -/// A typed failure returned by an [`Interceptor`] implementation. -/// -/// The Engine attaches the exact [`InterceptorPoint`] at which the failure was -/// observed before exposing it through the run termination boundary. +/// Maximum UTF-8 byte length retained for interceptor diagnostics. +pub const MAX_INTERCEPTOR_DIAGNOSTIC_BYTES: usize = 1024; + +/// Stable category for the source of an interceptor failure. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InterceptorErrorCategory { + Policy, + Dependency, + ContractViolation, + Internal, +} + +impl std::fmt::Display for InterceptorErrorCategory { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(match self { + Self::Policy => "policy", + Self::Dependency => "dependency", + Self::ContractViolation => "contract_violation", + Self::Internal => "internal", + }) + } +} + +/// A typed, bounded failure returned by an [`Interceptor`] implementation. #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] -#[error("{message}")] +#[error("{category}: {diagnostic}")] pub struct InterceptorError { - message: String, + category: InterceptorErrorCategory, + diagnostic: String, } impl InterceptorError { - /// Create an interceptor failure with a caller-defined message. - pub fn new(message: impl Into) -> Self { + pub fn new(category: InterceptorErrorCategory, diagnostic: impl Into) -> Self { + let mut diagnostic = diagnostic.into(); + if diagnostic.len() > MAX_INTERCEPTOR_DIAGNOSTIC_BYTES { + let mut end = MAX_INTERCEPTOR_DIAGNOSTIC_BYTES; + while !diagnostic.is_char_boundary(end) { + end -= 1; + } + diagnostic.truncate(end); + } Self { - message: message.into(), + category, + diagnostic, } } - /// Return the failure message supplied by the interceptor. - pub fn message(&self) -> &str { - &self.message + pub fn category(&self) -> InterceptorErrorCategory { + self.category + } + + pub fn diagnostic(&self) -> &str { + &self.diagnostic } } -impl From for InterceptorError { - fn from(message: String) -> Self { - Self::new(message) - } -} - -impl From<&str> for InterceptorError { - fn from(message: &str) -> Self { - Self::new(message) - } -} - -/// The Engine lifecycle point at which an interceptor failed. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum InterceptorPoint { +/// The lifecycle phase at which an interceptor callback executes. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum InterceptorPhase { + #[default] PromptSubmit, PendingHistoryAppends, PreLlmRequest, @@ -64,9 +86,9 @@ pub enum InterceptorPoint { RunExit, } -impl std::fmt::Display for InterceptorPoint { +impl std::fmt::Display for InterceptorPhase { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let name = match self { + formatter.write_str(match self { Self::PromptSubmit => "prompt_submit", Self::PendingHistoryAppends => "pending_history_appends", Self::PreLlmRequest => "pre_llm_request", @@ -74,66 +96,113 @@ impl std::fmt::Display for InterceptorPoint { Self::PostToolCall => "post_tool_call", Self::AssistantTurnEnd => "assistant_turn_end", Self::RunExit => "run_exit", - }; - formatter.write_str(name) + }) } } -/// An interceptor failure bound to the exact Engine lifecycle point that ran it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +pub struct InterceptorRunId(pub u64); + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct InterceptorTurnId(pub u64); + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum InterceptorCallId { + Llm(u64), + Tool(String), +} + +/// Saturating public counter used by interceptor contexts. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)] +pub struct InterceptorCounter(u32); + +impl InterceptorCounter { + pub fn from_usize(value: usize) -> Self { + Self(u32::try_from(value).unwrap_or(u32::MAX)) + } + + pub fn get(self) -> u32 { + self.0 + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct InterceptorCounters { + pub invocation: InterceptorCounter, + pub engine_turn: InterceptorCounter, + pub run_turn: InterceptorCounter, + pub llm_call: InterceptorCounter, + pub tool_batch: InterceptorCounter, + pub tool_call: InterceptorCounter, +} + +/// Identity, phase, and bounded counters common to every lifecycle callback. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct InterceptorInvocation { + pub run_id: InterceptorRunId, + pub turn_id: Option, + pub call_id: Option, + pub phase: InterceptorPhase, + pub counters: InterceptorCounters, +} + +/// An interceptor failure bound to the exact Engine lifecycle phase that ran it. #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] -#[error("{point} interceptor failed: {error}")] +#[error("{phase} interceptor failed: {error}")] pub struct InterceptorFailure { - point: InterceptorPoint, + phase: InterceptorPhase, #[source] error: InterceptorError, } impl InterceptorFailure { - pub(crate) fn new(point: InterceptorPoint, error: InterceptorError) -> Self { - Self { point, error } + pub(crate) fn new(phase: InterceptorPhase, error: InterceptorError) -> Self { + Self { phase, error } } - /// The lifecycle point that returned the failure. - pub fn point(&self) -> InterceptorPoint { - self.point + pub fn phase(&self) -> InterceptorPhase { + self.phase } - /// The typed error returned by the interceptor. pub fn error(&self) -> &InterceptorError { &self.error } } -/// Result returned by asynchronous interceptor lifecycle methods. pub type InterceptorResult = Result; // ============================================================================= // Lifecycle Contexts // ============================================================================= -/// Mutable prompt input presented before it is committed to Engine history. -pub struct PromptSubmitContext<'a> { +pub struct PromptSubmitContext<'a, A = ()> { + pub invocation: InterceptorInvocation, pub item: &'a mut Item, + pub history: &'a [HistoryEntry], } -/// Mutable provider-visible item projection presented before an LLM request. -pub struct PreLlmRequestContext<'a> { +pub struct PendingHistoryAppendsContext<'a, A = ()> { + pub invocation: InterceptorInvocation, + pub history: &'a [HistoryEntry], +} + +pub struct PreLlmRequestContext<'a, A = ()> { + pub invocation: InterceptorInvocation, pub items: &'a mut Vec, + pub history: &'a [HistoryEntry], } -/// A terminalized and committed assistant response at the next-phase boundary. -pub struct AssistantTurnEndContext<'a> { - /// The exact assistant items committed for this response. - pub assistant_items: &'a [Item], - /// The committed Engine history after the assistant items were appended. - pub history: &'a [Item], - /// Terminal tool calls collected from the response, if any. +pub struct AssistantTurnEndContext<'a, A = ()> { + pub invocation: InterceptorInvocation, + pub assistant_entries: &'a [HistoryEntry], + pub history: &'a [HistoryEntry], pub tool_calls: &'a [ToolCall], } -/// The one terminal outcome produced by a public Engine run or resume call. -pub struct RunExitContext<'a> { +pub struct RunExitContext<'a, A = ()> { + pub invocation: InterceptorInvocation, pub exit: &'a EngineRunExit, + pub history: &'a [HistoryEntry], } // ============================================================================= @@ -224,8 +293,9 @@ pub enum TurnEndAction { // ============================================================================= /// Context for pre-tool-call decisions. -pub struct ToolCallInfo { - /// Tool call information (modifiable). +pub struct ToolCallInfo<'a, A = ()> { + pub invocation: InterceptorInvocation, + pub history: &'a [HistoryEntry], pub call: ToolCall, /// Tool meta information. pub meta: ToolMeta, @@ -236,8 +306,9 @@ pub struct ToolCallInfo { } /// Context for post-tool-call decisions. -pub struct ToolResultInfo { - /// Original tool call. +pub struct ToolResultInfo<'a, A = ()> { + pub invocation: InterceptorInvocation, + pub history: &'a [HistoryEntry], pub call: ToolCall, /// Committed terminal tool execution result. pub result: ToolResult, @@ -258,17 +329,17 @@ pub struct ToolResultInfo { /// Every lifecycle method is asynchronous and returns [`InterceptorResult`], /// keeping implementation failure separate from the method's control-flow /// action. The Engine reports a failure as a typed run interruption annotated -/// with the exact [`InterceptorPoint`] that failed. +/// with the exact [`InterceptorPhase`] that failed. /// /// All methods have default implementations that let the Engine proceed /// without intervention. Callers provide richer implementations for approval /// flows, permission checks, and other trusted host adaptation. #[async_trait] -pub trait Interceptor: Send + Sync { +pub trait Interceptor: Send + Sync { /// Called after receiving user input, before adding it to Engine history. async fn on_prompt_submit( &self, - _context: PromptSubmitContext<'_>, + _context: PromptSubmitContext<'_, A>, ) -> InterceptorResult { Ok(PromptAction::Continue) } @@ -291,7 +362,10 @@ pub trait Interceptor: Send + Sync { /// reproducible per-request transformations (pruning, content /// trimming, cache anchors) that depend only on the existing /// history. - async fn pending_history_appends(&self) -> InterceptorResult> { + async fn pending_history_appends( + &self, + _context: PendingHistoryAppendsContext<'_, A>, + ) -> InterceptorResult> { Ok(Vec::new()) } @@ -305,18 +379,24 @@ pub trait Interceptor: Send + Sync { /// commits it to history before the request is sent. async fn pre_llm_request( &self, - _context: PreLlmRequestContext<'_>, + _context: PreLlmRequestContext<'_, A>, ) -> InterceptorResult { Ok(PreRequestAction::Continue) } /// Called before each tool is executed. - async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> InterceptorResult { + async fn pre_tool_call( + &self, + _info: &mut ToolCallInfo<'_, A>, + ) -> InterceptorResult { Ok(PreToolAction::Continue) } /// Called after each tool reaches one terminal result and that result is committed. - async fn post_tool_call(&self, _info: &ToolResultInfo) -> InterceptorResult { + async fn post_tool_call( + &self, + _info: &ToolResultInfo<'_, A>, + ) -> InterceptorResult { Ok(PostToolAction::Continue) } @@ -324,13 +404,13 @@ pub trait Interceptor: Send + Sync { /// the Engine decides whether to execute tools, continue, or finish. async fn on_assistant_turn_end( &self, - _context: AssistantTurnEndContext<'_>, + _context: AssistantTurnEndContext<'_, A>, ) -> InterceptorResult { Ok(TurnEndAction::Finish) } /// Called once for the terminal outcome of each public run or resume call. - async fn on_run_exit(&self, _context: RunExitContext<'_>) -> InterceptorResult<()> { + async fn on_run_exit(&self, _context: RunExitContext<'_, A>) -> InterceptorResult<()> { Ok(()) } } @@ -340,4 +420,4 @@ pub trait Interceptor: Send + Sync { pub(crate) struct DefaultInterceptor; #[async_trait] -impl Interceptor for DefaultInterceptor {} +impl Interceptor for DefaultInterceptor {} diff --git a/crates/agen/src/lib.rs b/crates/agen/src/lib.rs index 27bed76e..89bdaee9 100644 --- a/crates/agen/src/lib.rs +++ b/crates/agen/src/lib.rs @@ -27,8 +27,11 @@ pub use engine::{ pub use handler::ToolUseBlockStart; pub use history::{History, HistoryEntry}; pub use interceptor::{ - AssistantTurnEndContext, Interceptor, InterceptorError, InterceptorFailure, InterceptorPoint, - InterceptorResult, PreLlmRequestContext, PromptSubmitContext, RunExitContext, + AssistantTurnEndContext, Interceptor, InterceptorCallId, InterceptorCounter, + InterceptorCounters, InterceptorError, InterceptorErrorCategory, InterceptorFailure, + InterceptorInvocation, InterceptorPhase, InterceptorResult, InterceptorRunId, + InterceptorTurnId, MAX_INTERCEPTOR_DIAGNOSTIC_BYTES, PendingHistoryAppendsContext, + PreLlmRequestContext, PromptSubmitContext, RunExitContext, }; pub use message::{ContentPart, Item, Message, Role}; pub use tool::{ diff --git a/crates/agen/tests/annotated_history_test.rs b/crates/agen/tests/annotated_history_test.rs index 09a5019d..cf1be3c0 100644 --- a/crates/agen/tests/annotated_history_test.rs +++ b/crates/agen/tests/annotated_history_test.rs @@ -1,8 +1,15 @@ mod common; +use agen::interceptor::{ + AssistantTurnEndContext, Interceptor, InterceptorCallId, InterceptorInvocation, + InterceptorPhase, InterceptorResult, PendingHistoryAppendsContext, PreLlmRequestContext, + PreRequestAction, PromptAction, PromptSubmitContext, RunExitContext, TurnEndAction, +}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::{Engine, EngineError, History, HistoryEntry, Item, Role}; +use async_trait::async_trait; use common::MockLlmClient; +use std::sync::{Arc, Mutex}; fn completed_text_events(text: &str) -> Vec { vec![ @@ -47,6 +54,125 @@ async fn run_preserves_item_annotations_without_projecting_them() { assert_eq!(history.items_cloned().len(), 2); } +#[derive(Clone)] +struct AnnotationObservingInterceptor { + observed: Arc)>>>, +} + +impl AnnotationObservingInterceptor { + fn record(&self, invocation: &InterceptorInvocation, history: &[HistoryEntry]) { + self.observed.lock().unwrap().push(( + invocation.clone(), + history + .iter() + .map(|entry| entry.annotation.clone()) + .collect(), + )); + } +} + +#[async_trait] +impl Interceptor for AnnotationObservingInterceptor { + async fn on_prompt_submit( + &self, + context: PromptSubmitContext<'_, String>, + ) -> InterceptorResult { + self.record(&context.invocation, context.history); + Ok(PromptAction::Continue) + } + + async fn pending_history_appends( + &self, + context: PendingHistoryAppendsContext<'_, String>, + ) -> InterceptorResult> { + self.record(&context.invocation, context.history); + Ok(Vec::new()) + } + + async fn pre_llm_request( + &self, + context: PreLlmRequestContext<'_, String>, + ) -> InterceptorResult { + self.record(&context.invocation, context.history); + Ok(PreRequestAction::Continue) + } + + async fn on_assistant_turn_end( + &self, + context: AssistantTurnEndContext<'_, String>, + ) -> InterceptorResult { + assert_eq!(context.assistant_entries.len(), 1); + assert_eq!(context.assistant_entries[0].annotation, "2:assistant"); + self.record(&context.invocation, context.history); + Ok(TurnEndAction::Finish) + } + + async fn on_run_exit(&self, context: RunExitContext<'_, String>) -> InterceptorResult<()> { + self.record(&context.invocation, context.history); + Ok(()) + } +} + +#[tokio::test] +async fn interceptor_contexts_preserve_annotations_and_typed_lifecycle_identity() { + let client = MockLlmClient::new(completed_text_events("assistant reply")); + let mut engine = Engine::<_, agen::state::Mutable, String>::new_annotated(client); + let observed = Arc::new(Mutex::new(Vec::new())); + engine.set_interceptor(AnnotationObservingInterceptor { + observed: observed.clone(), + }); + let mut history = History::::new(); + let mut next = 0usize; + let mut annotate = |item: &Item| { + next += 1; + let kind = if item.is_assistant_message() { + "assistant" + } else { + "user" + }; + Ok(format!("{next}:{kind}")) + }; + + let output = engine + .run_with_annotation(&mut history, "hello", &mut annotate) + .await; + assert!(matches!(output.result, agen::EngineRunExit::Finished)); + + let observed = observed.lock().unwrap(); + let phases: Vec<_> = observed + .iter() + .map(|(invocation, _)| invocation.phase) + .collect(); + assert_eq!( + phases, + [ + InterceptorPhase::PromptSubmit, + InterceptorPhase::PendingHistoryAppends, + InterceptorPhase::PreLlmRequest, + InterceptorPhase::AssistantTurnEnd, + InterceptorPhase::RunExit, + ] + ); + assert!( + observed + .iter() + .all(|(invocation, _)| invocation.run_id == observed[0].0.run_id) + ); + assert_eq!( + observed + .iter() + .map(|(invocation, _)| invocation.counters.invocation.get()) + .collect::>(), + [0, 1, 2, 3, 4] + ); + assert_eq!(observed[2].0.call_id, Some(InterceptorCallId::Llm(0))); + assert_eq!(observed[3].0.call_id, Some(InterceptorCallId::Llm(0))); + assert_eq!(observed[1].1, ["1:user"]); + assert_eq!(observed[2].1, ["1:user"]); + assert_eq!(observed[3].1, ["1:user", "2:assistant"]); + assert_eq!(observed[4].1, ["1:user", "2:assistant"]); +} + #[test] fn append_failure_does_not_make_item_live() { let client = MockLlmClient::new(vec![]); diff --git a/crates/agen/tests/engine_state_test.rs b/crates/agen/tests/engine_state_test.rs index ab810e9e..5060015d 100644 --- a/crates/agen/tests/engine_state_test.rs +++ b/crates/agen/tests/engine_state_test.rs @@ -10,9 +10,11 @@ use std::sync::{Arc, Mutex}; use agen::Item; use agen::interceptor::{ - AssistantTurnEndContext, Interceptor, InterceptorError, InterceptorPoint, InterceptorResult, - PostToolAction, PreLlmRequestContext, PreRequestAction, PreToolAction, PromptAction, - PromptSubmitContext, RunExitContext, ToolCallInfo, ToolResultInfo, TurnEndAction, + AssistantTurnEndContext, Interceptor, InterceptorError, InterceptorErrorCategory, + InterceptorPhase as InterceptorPoint, InterceptorResult, MAX_INTERCEPTOR_DIAGNOSTIC_BYTES, + PendingHistoryAppendsContext, PostToolAction, PreLlmRequestContext, PreRequestAction, + PreToolAction, PromptAction, PromptSubmitContext, RunExitContext, ToolCallInfo, ToolResultInfo, + TurnEndAction, }; use agen::llm_client::{ ClientError, LlmClient, Request, ResponseStream, @@ -620,7 +622,7 @@ struct YieldOnce { impl Interceptor for YieldOnce { async fn pre_llm_request( &self, - _context: PreLlmRequestContext<'_>, + _context: PreLlmRequestContext<'_, ()>, ) -> InterceptorResult { Ok(if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { PreRequestAction::Yield @@ -636,7 +638,10 @@ struct PauseToolOnce { #[async_trait] impl Interceptor for PauseToolOnce { - async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> InterceptorResult { + async fn pre_tool_call( + &self, + _info: &mut ToolCallInfo<'_, ()>, + ) -> InterceptorResult { Ok(if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { PreToolAction::Pause } else { @@ -653,7 +658,7 @@ struct ContinueTurnOnce { impl Interceptor for ContinueTurnOnce { async fn on_assistant_turn_end( &self, - _context: AssistantTurnEndContext<'_>, + _context: AssistantTurnEndContext<'_, ()>, ) -> InterceptorResult { Ok(if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { TurnEndAction::ContinueWithMessages(vec![Item::system_message("continue")]) @@ -680,7 +685,10 @@ impl FailingLifecycleInterceptor { fn record(&self, point: InterceptorPoint, action: T) -> InterceptorResult { self.calls.lock().unwrap().push(point); if self.failure == point { - Err(InterceptorError::new(format!("{point} rejected"))) + Err(InterceptorError::new( + InterceptorErrorCategory::Policy, + format!("{point} rejected"), + )) } else { Ok(action) } @@ -695,47 +703,56 @@ impl FailingLifecycleInterceptor { impl Interceptor for FailingLifecycleInterceptor { async fn on_prompt_submit( &self, - _context: PromptSubmitContext<'_>, + _context: PromptSubmitContext<'_, ()>, ) -> InterceptorResult { tokio::task::yield_now().await; self.record(InterceptorPoint::PromptSubmit, PromptAction::Continue) } - async fn pending_history_appends(&self) -> InterceptorResult> { + async fn pending_history_appends( + &self, + _context: PendingHistoryAppendsContext<'_, ()>, + ) -> InterceptorResult> { tokio::task::yield_now().await; self.record(InterceptorPoint::PendingHistoryAppends, Vec::new()) } async fn pre_llm_request( &self, - _context: PreLlmRequestContext<'_>, + _context: PreLlmRequestContext<'_, ()>, ) -> InterceptorResult { tokio::task::yield_now().await; self.record(InterceptorPoint::PreLlmRequest, PreRequestAction::Continue) } - async fn pre_tool_call(&self, _info: &mut ToolCallInfo) -> InterceptorResult { + async fn pre_tool_call( + &self, + _info: &mut ToolCallInfo<'_, ()>, + ) -> InterceptorResult { tokio::task::yield_now().await; self.record(InterceptorPoint::PreToolCall, PreToolAction::Continue) } - async fn post_tool_call(&self, _info: &ToolResultInfo) -> InterceptorResult { + async fn post_tool_call( + &self, + _info: &ToolResultInfo<'_, ()>, + ) -> InterceptorResult { tokio::task::yield_now().await; self.record(InterceptorPoint::PostToolCall, PostToolAction::Continue) } async fn on_assistant_turn_end( &self, - context: AssistantTurnEndContext<'_>, + context: AssistantTurnEndContext<'_, ()>, ) -> InterceptorResult { tokio::task::yield_now().await; - assert!(context.history.ends_with(context.assistant_items)); + assert!(context.history.ends_with(context.assistant_entries)); if !context.tool_calls.is_empty() { assert_eq!( context - .assistant_items + .assistant_entries .iter() - .filter(|item| matches!(item, Item::ToolCall { .. })) + .filter(|entry| matches!(&entry.item, Item::ToolCall { .. })) .count(), context.tool_calls.len() ); @@ -743,7 +760,7 @@ impl Interceptor for FailingLifecycleInterceptor { self.record(InterceptorPoint::AssistantTurnEnd, TurnEndAction::Finish) } - async fn on_run_exit(&self, _context: RunExitContext<'_>) -> InterceptorResult<()> { + async fn on_run_exit(&self, _context: RunExitContext<'_, ()>) -> InterceptorResult<()> { tokio::task::yield_now().await; self.record(InterceptorPoint::RunExit, ()) } @@ -795,7 +812,7 @@ fn expected_interceptor_calls(failure: InterceptorPoint) -> Vec, + ) -> InterceptorResult { + Ok(if self.pause { + TurnEndAction::Pause + } else { + TurnEndAction::Finish + }) + } + + async fn on_run_exit(&self, _context: RunExitContext<'_, ()>) -> InterceptorResult<()> { + Err(InterceptorError::new( + InterceptorErrorCategory::Dependency, + "terminal audit unavailable", + )) + } +} + +#[tokio::test] +async fn terminal_observer_failure_preserves_paused_and_interrupted_exits() { + let mut paused_engine = Engine::new(MockLlmClient::new(completed_text_events())); + paused_engine.set_interceptor(FailingRunExitObserver { pause: true }); + let mut paused_history = History::new(); + let mut paused_engine = paused_engine.lock(&paused_history); + assert!(matches!( + paused_engine.run(&mut paused_history, "pause").await, + EngineRunExit::Paused + )); + assert_eq!( + paused_engine + .last_run_exit_observer_failure() + .expect("paused observer diagnostic") + .error() + .category(), + InterceptorErrorCategory::Dependency + ); + + let mut interrupted_engine = Engine::new(MockLlmClient::new(completed_text_events())); + interrupted_engine.set_max_turns(Some(0)); + interrupted_engine.set_interceptor(FailingRunExitObserver { pause: false }); + let mut interrupted_history = History::new(); + let mut interrupted_engine = interrupted_engine.lock(&interrupted_history); + assert!(matches!( + interrupted_engine + .run(&mut interrupted_history, "limit") + .await, + EngineRunExit::Interrupted(RunInterruptionReason::LimitReached) + )); + assert_eq!( + interrupted_engine + .last_run_exit_observer_failure() + .expect("interrupted observer diagnostic") + .phase(), + InterceptorPoint::RunExit + ); +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum TerminalMode { Finish, @@ -886,7 +990,7 @@ impl RecordingTerminalInterceptor { impl Interceptor for RecordingTerminalInterceptor { async fn pre_llm_request( &self, - _context: PreLlmRequestContext<'_>, + _context: PreLlmRequestContext<'_, ()>, ) -> InterceptorResult { Ok(if self.mode == TerminalMode::Yield { PreRequestAction::Yield @@ -897,11 +1001,11 @@ impl Interceptor for RecordingTerminalInterceptor { async fn on_assistant_turn_end( &self, - context: AssistantTurnEndContext<'_>, + context: AssistantTurnEndContext<'_, ()>, ) -> InterceptorResult { - assert!(!context.assistant_items.is_empty()); + assert!(!context.assistant_entries.is_empty()); assert!( - context.history.ends_with(context.assistant_items), + context.history.ends_with(context.assistant_entries), "assistant-turn callback must observe committed terminal items" ); let turn = self.assistant_turns.fetch_add(1, Ordering::SeqCst); @@ -912,7 +1016,7 @@ impl Interceptor for RecordingTerminalInterceptor { }) } - async fn on_run_exit(&self, context: RunExitContext<'_>) -> InterceptorResult<()> { + async fn on_run_exit(&self, context: RunExitContext<'_, ()>) -> InterceptorResult<()> { let kind = match context.exit { EngineRunExit::Finished => "finished", EngineRunExit::Paused => "paused", diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index f63b20a1..c79c0a2e 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -7,14 +7,17 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use agen::interceptor::{ - Interceptor, InterceptorResult, PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo, + Interceptor, InterceptorErrorCategory, InterceptorPhase, InterceptorResult, PostToolAction, + PreToolAction, ToolCallInfo, ToolResultInfo, }; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::tool::{ Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult, ToolResultDisposition, }; -use agen::{Engine, History, Item, ToolExecutionPolicy}; +use agen::{ + Engine, EngineError, EngineRunExit, History, Item, RunInterruptionReason, ToolExecutionPolicy, +}; use async_trait::async_trait; mod common; @@ -907,7 +910,10 @@ async fn test_tool_execution_context_for_skipped_and_synthetic_paths() { #[async_trait] impl Interceptor for ContextPolicy { - async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult { + async fn pre_tool_call( + &self, + info: &mut ToolCallInfo<'_, ()>, + ) -> InterceptorResult { self.pre_contexts.lock().unwrap().push(info.context.clone()); Ok(match info.call.name.as_str() { "skip_tool" => PreToolAction::Skip, @@ -919,7 +925,10 @@ async fn test_tool_execution_context_for_skipped_and_synthetic_paths() { }) } - async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult { + async fn post_tool_call( + &self, + info: &ToolResultInfo<'_, ()>, + ) -> InterceptorResult { self.post_contexts .lock() .unwrap() @@ -996,7 +1005,10 @@ async fn test_before_tool_call_skip() { #[async_trait] impl Interceptor for BlockingPolicy { - async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult { + async fn pre_tool_call( + &self, + info: &mut ToolCallInfo<'_, ()>, + ) -> InterceptorResult { Ok(if info.call.name == "blocked_tool" { PreToolAction::Skip } else { @@ -1083,7 +1095,19 @@ async fn test_post_tool_call_observes_committed_result() { #[async_trait] impl Interceptor for ObservingPolicy { - async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult { + async fn post_tool_call( + &self, + info: &ToolResultInfo<'_, ()>, + ) -> InterceptorResult { + assert_eq!(info.invocation.phase, InterceptorPhase::PostToolCall); + assert_eq!( + info.invocation.call_id, + Some(agen::InterceptorCallId::Tool(info.call.id.clone())) + ); + assert!(matches!( + info.history.last().map(|entry| &entry.item), + Some(Item::ToolResult { call_id, .. }) if call_id == &info.call.id + )); *self.observed_content.lock().unwrap() = Some(info.result.summary.clone()); Ok(PostToolAction::Continue) } @@ -1144,7 +1168,10 @@ async fn test_before_tool_call_synthetic_result_committed() { #[async_trait] impl Interceptor for SyntheticPolicy { - async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult { + async fn pre_tool_call( + &self, + info: &mut ToolCallInfo<'_, ()>, + ) -> InterceptorResult { Ok(PreToolAction::SyntheticResult(ToolResult::error( info.call.id.clone(), "permission denied", @@ -1168,6 +1195,80 @@ async fn test_before_tool_call_synthetic_result_committed() { ))); } +#[derive(Clone, Copy)] +enum InvalidIdentityMode { + ContinuedCall, + SyntheticResult, +} + +struct InvalidIdentityPolicy(InvalidIdentityMode); + +#[async_trait] +impl Interceptor for InvalidIdentityPolicy { + async fn pre_tool_call( + &self, + info: &mut ToolCallInfo<'_, ()>, + ) -> InterceptorResult { + assert_eq!(info.invocation.phase, InterceptorPhase::PreToolCall); + assert_eq!( + info.invocation.call_id, + Some(agen::InterceptorCallId::Tool("call_1".to_string())) + ); + assert!(matches!( + info.history.last().map(|entry| &entry.item), + Some(Item::ToolCall { call_id, .. }) if call_id == "call_1" + )); + Ok(match self.0 { + InvalidIdentityMode::ContinuedCall => { + info.call.id = "different-call".to_string(); + PreToolAction::Continue + } + InvalidIdentityMode::SyntheticResult => PreToolAction::SyntheticResult( + ToolResult::error("different-call", "invalid synthetic result"), + ), + }) + } +} + +#[tokio::test] +async fn interceptor_cannot_change_tool_call_identity() { + for mode in [ + InvalidIdentityMode::ContinuedCall, + InvalidIdentityMode::SyntheticResult, + ] { + let client = MockLlmClient::new(vec![ + Event::tool_use_start(0, "call_1", "echo"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]); + let mut engine = Engine::new(client); + engine.register_tool(SlowTool::new("echo", 1).definition()); + engine.set_interceptor(InvalidIdentityPolicy(mode)); + let mut history = History::new(); + + let result = engine.run(&mut history, "identity").await; + let EngineRunExit::Interrupted(RunInterruptionReason::Unexpected( + EngineError::Interceptor(failure), + )) = result.result + else { + panic!("invalid tool identity must interrupt with a typed failure"); + }; + assert_eq!(failure.phase(), InterceptorPhase::PreToolCall); + assert_eq!( + failure.error().category(), + InterceptorErrorCategory::ContractViolation + ); + assert!( + !history + .items() + .any(|item| matches!(item, Item::ToolResult { .. })) + ); + } +} + #[tokio::test] async fn post_tool_abort_commits_confirmed_result_before_stopping_run() { let client = MockLlmClient::new(vec![ @@ -1190,7 +1291,7 @@ async fn post_tool_abort_commits_confirmed_result_before_stopping_run() { impl Interceptor for AbortAfterResult { async fn post_tool_call( &self, - _info: &ToolResultInfo, + _info: &ToolResultInfo<'_, ()>, ) -> InterceptorResult { self.lifecycle.lock().unwrap().push("post_tool_call"); Ok(PostToolAction::Abort("policy stopped the run".to_string())) diff --git a/crates/worker/src/compact/worker.rs b/crates/worker/src/compact/worker.rs index 158c68ff..5aef61df 100644 --- a/crates/worker/src/compact/worker.rs +++ b/crates/worker/src/compact/worker.rs @@ -400,10 +400,10 @@ impl CompactWorkerInterceptor { } #[async_trait] -impl Interceptor for CompactWorkerInterceptor { +impl Interceptor for CompactWorkerInterceptor { async fn pre_llm_request( &self, - context: PreLlmRequestContext<'_>, + context: PreLlmRequestContext<'_, A>, ) -> InterceptorResult { let context = context.items; let records = self.usage_tracker.records(); @@ -427,7 +427,10 @@ impl Interceptor for CompactWorkerInterceptor { Ok(PreRequestAction::Continue) } - async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult { + async fn pre_tool_call( + &self, + info: &mut ToolCallInfo<'_, A>, + ) -> InterceptorResult { if self.final_reserve_tokens == 0 || info.call.name == "write_summary" { return Ok(PreToolAction::Continue); } @@ -475,8 +478,10 @@ mod tests { assert!(matches!( interceptor - .pre_llm_request(PreLlmRequestContext { + .pre_llm_request(PreLlmRequestContext::<()> { + invocation: Default::default(), items: &mut context, + history: &[], }) .await .unwrap(), @@ -486,8 +491,10 @@ mod tests { assert!(matches!( interceptor - .pre_llm_request(PreLlmRequestContext { + .pre_llm_request(PreLlmRequestContext::<()> { + invocation: Default::default(), items: &mut context, + history: &[], }) .await .unwrap(), @@ -499,8 +506,10 @@ mod tests { // current occupancy is still the latest 100-token measurement. assert!(matches!( interceptor - .pre_llm_request(PreLlmRequestContext { + .pre_llm_request(PreLlmRequestContext::<()> { + invocation: Default::default(), items: &mut context, + history: &[], }) .await .unwrap(), @@ -526,8 +535,10 @@ mod tests { assert!(matches!( interceptor - .pre_llm_request(PreLlmRequestContext { + .pre_llm_request(PreLlmRequestContext::<()> { + invocation: Default::default(), items: &mut context, + history: &[], }) .await .unwrap(), @@ -537,8 +548,10 @@ mod tests { assert!(matches!( interceptor - .pre_llm_request(PreLlmRequestContext { + .pre_llm_request(PreLlmRequestContext::<()> { + invocation: Default::default(), items: &mut context, + history: &[], }) .await .unwrap(), @@ -556,8 +569,10 @@ mod tests { assert!(matches!( interceptor - .pre_llm_request(PreLlmRequestContext { + .pre_llm_request(PreLlmRequestContext::<()> { + invocation: Default::default(), items: &mut context, + history: &[], }) .await .unwrap(), @@ -567,8 +582,10 @@ mod tests { assert!(matches!( interceptor - .pre_llm_request(PreLlmRequestContext { + .pre_llm_request(PreLlmRequestContext::<()> { + invocation: Default::default(), items: &mut context, + history: &[], }) .await .unwrap(), diff --git a/crates/worker/src/feature.rs b/crates/worker/src/feature.rs index f0231d15..5698e84b 100644 --- a/crates/worker/src/feature.rs +++ b/crates/worker/src/feature.rs @@ -1795,7 +1795,7 @@ impl FeatureRegistryBuilder { } /// Install modules into the existing Engine tool path and hook builder. - pub(crate) fn install_into_engine( + pub(crate) fn install_into_engine( self, worker: &mut Engine, hook_builder: &mut HookRegistryBuilder, diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index 8d380ea7..2706f715 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -15,7 +15,8 @@ use std::sync::{Arc, Mutex}; use agen::Item; use agen::UsageRecord; use agen::interceptor::{ - AssistantTurnEndContext, Interceptor, InterceptorResult, PostToolAction, PreLlmRequestContext, + AssistantTurnEndContext, Interceptor, InterceptorError, InterceptorErrorCategory, + InterceptorResult, PendingHistoryAppendsContext, PostToolAction, PreLlmRequestContext, PreRequestAction, PreToolAction, PromptAction, PromptSubmitContext, ToolCallInfo, ToolResultInfo, TurnEndAction, }; @@ -232,10 +233,10 @@ impl WorkerInterceptor { } #[async_trait] -impl Interceptor for WorkerInterceptor { +impl Interceptor for WorkerInterceptor { async fn on_prompt_submit( &self, - context: PromptSubmitContext<'_>, + context: PromptSubmitContext<'_, SessionHistoryMetadata>, ) -> InterceptorResult { let item = context.item; let turn_index = self.next_turn_index.fetch_add(1, Ordering::Relaxed); @@ -274,7 +275,10 @@ impl Interceptor for WorkerInterceptor { }) } - async fn pending_history_appends(&self) -> InterceptorResult> { + async fn pending_history_appends( + &self, + _context: PendingHistoryAppendsContext<'_, SessionHistoryMetadata>, + ) -> InterceptorResult> { let drained = self.pending_notifies.drain(); if drained.is_empty() { return Ok(Vec::new()); @@ -300,7 +304,10 @@ impl Interceptor for WorkerInterceptor { Ok(system_item) => system_item, Err(error) => { self.pending_notifies.requeue_front(drained); - return Err(format!("failed to render notify_wrapper: {error}").into()); + return Err(InterceptorError::new( + InterceptorErrorCategory::Dependency, + format!("failed to render notify_wrapper: {error}"), + )); } }; items.push(system_item.to_history_item()); @@ -308,14 +315,17 @@ impl Interceptor for WorkerInterceptor { } if let Err(error) = self.commit_system_items(&system_items) { self.pending_notifies.requeue_front(drained); - return Err(format!("session persistence failed: {error}").into()); + return Err(InterceptorError::new( + InterceptorErrorCategory::Dependency, + format!("session persistence failed: {error}"), + )); } Ok(items) } async fn pre_llm_request( &self, - context: PreLlmRequestContext<'_>, + context: PreLlmRequestContext<'_, SessionHistoryMetadata>, ) -> InterceptorResult { let context = context.items; let initial_tokens = self.estimated_tokens(context); @@ -385,7 +395,10 @@ impl Interceptor for WorkerInterceptor { }) } - async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> InterceptorResult { + async fn pre_tool_call( + &self, + info: &mut ToolCallInfo<'_, SessionHistoryMetadata>, + ) -> InterceptorResult { let summary = ToolCallSummary { call_id: info.call.id.clone(), tool_name: info.call.name.clone(), @@ -401,7 +414,10 @@ impl Interceptor for WorkerInterceptor { Ok(PreToolAction::Continue) } - async fn post_tool_call(&self, info: &ToolResultInfo) -> InterceptorResult { + async fn post_tool_call( + &self, + info: &ToolResultInfo<'_, SessionHistoryMetadata>, + ) -> InterceptorResult { let summary = ToolResultSummary { call_id: info.result.tool_use_id.clone(), tool_name: info.call.name.clone(), @@ -424,14 +440,14 @@ impl Interceptor for WorkerInterceptor { async fn on_assistant_turn_end( &self, - context: AssistantTurnEndContext<'_>, + context: AssistantTurnEndContext<'_, SessionHistoryMetadata>, ) -> InterceptorResult { let history = context.history; let final_text_preview = history .iter() .rev() - .find(|i| i.is_assistant_message()) - .and_then(extract_message_text) + .find(|entry| entry.item.is_assistant_message()) + .and_then(|entry| extract_message_text(&entry.item)) .map(|t| preview(&t, FINAL_TEXT_PREVIEW_LIMIT)) .unwrap_or_default(); let info = TurnEndInfo { @@ -515,6 +531,7 @@ mod tests { Hook, HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookRegistryBuilder, HookTurnEndAction, OnTurnEnd, PostToolCall, PreLlmRequest, PreToolCall, }; + use crate::session_history::{WorkerHistoryProvenance, history_entry}; fn test_prompts() -> Arc> { Arc::new(ArcSwap::from(PromptCatalog::builtins_only().unwrap())) @@ -574,7 +591,10 @@ mod tests { } } - fn task_tool_call_info(name: &str, input: serde_json::Value) -> ToolCallInfo { + fn task_tool_call_info( + name: &str, + input: serde_json::Value, + ) -> ToolCallInfo<'static, SessionHistoryMetadata> { let def = crate::feature::builtin::task::task_tools( crate::feature::builtin::task::TaskStore::new(), ) @@ -586,6 +606,8 @@ mod tests { .expect("task tool definition"); let (meta, tool) = def(); ToolCallInfo { + invocation: Default::default(), + history: &[], call: agen::tool::ToolCall { id: "call-id".into(), name: name.into(), @@ -630,7 +652,11 @@ mod tests { ); let mut ctx = ctx_items; let action = interceptor - .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) .await .unwrap(); @@ -665,7 +691,11 @@ mod tests { ); let mut ctx = ctx_items; let action = interceptor - .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) .await .unwrap(); @@ -705,7 +735,11 @@ mod tests { .with_usage_tracker(usage_tracker); let mut ctx = ctx_items; let action = interceptor - .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) .await .unwrap(); @@ -732,7 +766,11 @@ mod tests { ); let mut ctx = ctx_items; let action = interceptor - .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) .await .unwrap(); @@ -776,7 +814,11 @@ mod tests { ); let mut ctx = ctx_items; let action = interceptor - .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) .await .unwrap(); @@ -806,7 +848,11 @@ mod tests { ); let mut ctx = ctx_items; let action = interceptor - .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) .await .unwrap(); @@ -830,7 +876,11 @@ mod tests { ); let mut ctx: Vec = Vec::new(); let action = interceptor - .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) .await .unwrap(); @@ -862,7 +912,11 @@ mod tests { let mut ctx: Vec = Vec::new(); let action = interceptor - .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) .await .unwrap(); @@ -912,7 +966,11 @@ mod tests { let mut ctx: Vec = Vec::new(); let action = interceptor - .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) .await .unwrap(); @@ -1017,7 +1075,9 @@ mod tests { None, ); let info = task_tool_call_info("TaskList", serde_json::json!({})); - let mut result_info = ToolResultInfo { + let result_info = ToolResultInfo { + invocation: Default::default(), + history: &[], call: info.call, result: agen::tool::ToolResult::from_output( "call-id", @@ -1033,7 +1093,7 @@ mod tests { context: info.context, }; - let action = interceptor.post_tool_call(&mut result_info).await.unwrap(); + let action = interceptor.post_tool_call(&result_info).await.unwrap(); assert_eq!(action, PostToolAction::Abort("post tool abort".to_string())); assert_eq!(count.load(Ordering::Relaxed), 1); @@ -1067,11 +1127,20 @@ mod tests { test_prompts(), None, ); - let history = vec![Item::user_message("hi"), Item::assistant_message("done")]; - + let history = vec![ + history_entry( + Item::user_message("hi"), + WorkerHistoryProvenance::LegacyUnknown, + ), + history_entry( + Item::assistant_message("done"), + WorkerHistoryProvenance::LegacyUnknown, + ), + ]; let action = interceptor .on_assistant_turn_end(AssistantTurnEndContext { - assistant_items: &[], + invocation: Default::default(), + assistant_entries: &history[1..], history: &history, tool_calls: &[], }) @@ -1114,7 +1183,11 @@ mod tests { for _ in 0..23 { let mut ctx = ctx_items.clone(); let action = interceptor - .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) .await .unwrap(); assert!(matches!(action, PreRequestAction::Continue)); @@ -1129,7 +1202,11 @@ mod tests { let mut ctx = ctx_items.clone(); let action = interceptor - .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) .await .unwrap(); let appended_len = match action { @@ -1204,7 +1281,13 @@ mod tests { )); buffer.push_notify("updated".to_string(), false); - let appends = interceptor.pending_history_appends().await.unwrap(); + let appends = interceptor + .pending_history_appends(PendingHistoryAppendsContext { + invocation: Default::default(), + history: &[], + }) + .await + .unwrap(); assert_eq!(appends.len(), 1); assert!(format!("{:?}", appends[0]).contains("CURRENT-PROJECTION updated")); let committed = committed.lock().unwrap(); @@ -1254,9 +1337,19 @@ mod tests { )); buffer.push_notify("must persist".to_string(), false); - let error = interceptor.pending_history_appends().await.unwrap_err(); + let error = interceptor + .pending_history_appends(PendingHistoryAppendsContext { + invocation: Default::default(), + history: &[], + }) + .await + .unwrap_err(); - assert!(error.message().contains("failed to render notify_wrapper")); + assert!( + error + .diagnostic() + .contains("failed to render notify_wrapper") + ); let requeued = buffer.drain(); assert_eq!(requeued.len(), 1); } @@ -1278,7 +1371,13 @@ mod tests { None, ); - let items = interceptor.pending_history_appends().await.unwrap(); + let items = interceptor + .pending_history_appends(PendingHistoryAppendsContext { + invocation: Default::default(), + history: &[], + }) + .await + .unwrap(); assert_eq!(items.len(), 2); let first = items[0].as_text().unwrap_or_default(); let second = items[1].as_text().unwrap_or_default(); @@ -1292,7 +1391,13 @@ mod tests { ); // Empty buffer → empty Vec (no synthesised items). - let again = interceptor.pending_history_appends().await.unwrap(); + let again = interceptor + .pending_history_appends(PendingHistoryAppendsContext { + invocation: Default::default(), + history: &[], + }) + .await + .unwrap(); assert!(again.is_empty()); } @@ -1316,7 +1421,11 @@ mod tests { ); let mut ctx: Vec = vec![Item::user_message("hi")]; let action = interceptor - .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) .await .unwrap(); @@ -1349,7 +1458,11 @@ mod tests { ); let mut ctx: Vec = Vec::new(); let action = interceptor - .pre_llm_request(PreLlmRequestContext { items: &mut ctx }) + .pre_llm_request(PreLlmRequestContext { + invocation: Default::default(), + items: &mut ctx, + history: &[], + }) .await .unwrap(); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 332ab2e5..db3536ca 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -5996,7 +5996,7 @@ where /// Note: `system_prompt` is intentionally not applied here. It is a /// minijinja template that is parsed by `Worker::from_manifest` and /// rendered once at first turn in `ensure_system_prompt_materialized`. -pub fn apply_worker_manifest( +pub fn apply_worker_manifest( worker: &mut Engine, wm: &manifest::EngineManifest, ) { From 9b48b1ff5d43f57ff38e932aa5ab269415f00c30 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 12:27:59 +0900 Subject: [PATCH 04/22] fix: terminalize parallel tool siblings on interceptor stop --- crates/agen/src/engine.rs | 63 +++++++++++---- crates/agen/tests/parallel_execution_test.rs | 80 +++++++++++++++++++- 2 files changed, 125 insertions(+), 18 deletions(-) diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 644cf161..6c23e0cd 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -1389,20 +1389,29 @@ impl Engine { let mut terminal_call_ids = HashSet::new(); let mut pause_requested = false; let mut pause_deadline = None; + let mut batch_error = None; for result in synthetic_results { - self.finalize_and_commit_tool_result( - history, - annotate, - result, - None, - &call_info_map, - &mut attempt_fence, - &mut terminal_call_ids, - ) - .await?; + if let Err(error) = self + .finalize_and_commit_tool_result( + history, + annotate, + result, + None, + &call_info_map, + &mut attempt_fence, + &mut terminal_call_ids, + ) + .await + && batch_error.is_none() + { + batch_error = Some(error); + } } let mut futures = futures; + if batch_error.is_some() && !futures.is_empty() { + let _ = self.cancel_tx.try_send(()); + } while !futures.is_empty() { tokio::select! { // If cancellation and a completed result are both ready, drain @@ -1412,7 +1421,7 @@ impl Engine { result = futures.next() => { let (attempt_id, result) = result.expect("non-empty FuturesUnordered returns a result"); - self.finalize_and_commit_tool_result( + if let Err(error) = self.finalize_and_commit_tool_result( history, annotate, result, @@ -1420,7 +1429,14 @@ impl Engine { &call_info_map, &mut attempt_fence, &mut terminal_call_ids, - ).await?; + ).await { + if batch_error.is_none() { + batch_error = Some(error); + } + if !futures.is_empty() { + let _ = self.cancel_tx.try_send(()); + } + } } pause = self.pause_rx.recv(), if !pause_requested => { if pause.is_some() { @@ -1482,7 +1498,7 @@ impl Engine { result = futures.next() => { let (attempt_id, result) = result.expect("non-empty FuturesUnordered returns a result"); - self.finalize_and_commit_tool_result( + if let Err(error) = self.finalize_and_commit_tool_result( history, annotate, result, @@ -1490,7 +1506,11 @@ impl Engine { &call_info_map, &mut attempt_fence, &mut terminal_call_ids, - ).await?; + ).await + && batch_error.is_none() + { + batch_error = Some(error); + } } _ = tokio::time::sleep_until(deadline) => break, } @@ -1504,7 +1524,7 @@ impl Engine { if let Some(handle) = execution_handles.get(call_id) { handle.force_close(); } - self.finalize_and_commit_tool_result( + if let Err(error) = self.finalize_and_commit_tool_result( history, annotate, ToolResult::outcome_unknown(call_id), @@ -1512,11 +1532,18 @@ impl Engine { &call_info_map, &mut attempt_fence, &mut terminal_call_ids, - ).await?; + ).await + && batch_error.is_none() + { + batch_error = Some(error); + } } } self.timeline.abort_current_block(); + if let Some(error) = batch_error.take() { + return Err(error); + } if pause_requested { return Ok(ToolExecutionResult::Paused); } @@ -1525,6 +1552,10 @@ impl Engine { } } + if let Some(error) = batch_error { + self.timeline.abort_current_block(); + return Err(error); + } Ok(if pause_requested { ToolExecutionResult::Paused } else { diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index c79c0a2e..3b7c9c18 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -7,8 +7,8 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use agen::interceptor::{ - Interceptor, InterceptorErrorCategory, InterceptorPhase, InterceptorResult, PostToolAction, - PreToolAction, ToolCallInfo, ToolResultInfo, + Interceptor, InterceptorError, InterceptorErrorCategory, InterceptorPhase, InterceptorResult, + PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo, }; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::tool::{ @@ -1349,3 +1349,79 @@ async fn post_tool_abort_commits_confirmed_result_before_stopping_run() { } if call_id == "call_confirmed" ))); } + +#[derive(Clone, Copy)] +enum PostToolStopMode { + Abort, + Failure, +} + +struct StopFirstParallelResult(PostToolStopMode); + +#[async_trait] +impl Interceptor for StopFirstParallelResult { + async fn post_tool_call( + &self, + info: &ToolResultInfo<'_, ()>, + ) -> InterceptorResult { + if info.call.id != "call_fast" { + return Ok(PostToolAction::Continue); + } + match self.0 { + PostToolStopMode::Abort => Ok(PostToolAction::Abort("stop parallel batch".to_string())), + PostToolStopMode::Failure => Err(InterceptorError::new( + InterceptorErrorCategory::Policy, + "reject parallel batch", + )), + } + } +} + +#[tokio::test] +async fn post_tool_stop_terminalizes_started_parallel_siblings_before_returning() { + for mode in [PostToolStopMode::Abort, PostToolStopMode::Failure] { + let client = MockLlmClient::new(vec![ + Event::tool_use_start(0, "call_fast", "fast"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::tool_use_start(1, "call_slow", "slow"), + Event::tool_input_delta(1, r#"{}"#), + Event::tool_use_stop(1), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]); + let mut engine = Engine::new(client); + engine.register_tool(SlowTool::new("fast", 1).definition()); + engine.register_tool(SlowTool::new("slow", 10_000).definition()); + engine.set_interceptor(StopFirstParallelResult(mode)); + let mut history = History::new(); + + let output = engine.run(&mut history, "parallel stop").await; + match mode { + PostToolStopMode::Abort => assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected( + EngineError::Aborted(ref reason) + )) if reason == "stop parallel batch" + )), + PostToolStopMode::Failure => assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected( + EngineError::Interceptor(ref failure) + )) if failure.phase() == InterceptorPhase::PostToolCall + )), + } + + let terminal_ids: Vec<_> = history + .iter() + .filter_map(|entry| match &entry.item { + Item::ToolResult { call_id, .. } => Some(call_id.as_str()), + _ => None, + }) + .collect(); + assert_eq!(terminal_ids.len(), 2); + assert!(terminal_ids.contains(&"call_fast")); + assert!(terminal_ids.contains(&"call_slow")); + } +} From 5d61da481bc0f3a70c99a114656e8683230c7dd6 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 13:01:37 +0900 Subject: [PATCH 05/22] fix: consume internal tool batch cancellation --- crates/agen/src/engine.rs | 10 ++++++++ crates/agen/tests/parallel_execution_test.rs | 26 +++++++++++++++----- 2 files changed, 30 insertions(+), 6 deletions(-) diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 6c23e0cd..da3cead9 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -1390,6 +1390,7 @@ impl Engine { let mut pause_requested = false; let mut pause_deadline = None; let mut batch_error = None; + let mut locally_enqueued_cancel = false; for result in synthetic_results { if let Err(error) = self .finalize_and_commit_tool_result( @@ -1411,6 +1412,7 @@ impl Engine { let mut futures = futures; if batch_error.is_some() && !futures.is_empty() { let _ = self.cancel_tx.try_send(()); + locally_enqueued_cancel = true; } while !futures.is_empty() { tokio::select! { @@ -1435,6 +1437,7 @@ impl Engine { } if !futures.is_empty() { let _ = self.cancel_tx.try_send(()); + locally_enqueued_cancel = true; } } } @@ -1453,6 +1456,7 @@ impl Engine { _ = tokio::time::sleep_until(pause_deadline.unwrap_or_else(TokioInstant::now)), if pause_deadline.is_some() => { pause_deadline = None; let _ = self.cancel_tx.try_send(()); + locally_enqueued_cancel = true; } cancel = self.cancel_rx.recv() => { if cancel.is_some() { @@ -1552,6 +1556,12 @@ impl Engine { } } + // A result-biased ready sibling can empty the batch before the local + // cancel signal is selected. Never let that current-batch signal leak + // into the next run or resume call. + if locally_enqueued_cancel { + let _ = self.cancel_rx.try_recv(); + } if let Some(error) = batch_error { self.timeline.abort_current_block(); return Err(error); diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index 3b7c9c18..5df9f386 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -1367,6 +1367,7 @@ impl Interceptor for StopFirstParallelResult { if info.call.id != "call_fast" { return Ok(PostToolAction::Continue); } + tokio::time::sleep(Duration::from_millis(5)).await; match self.0 { PostToolStopMode::Abort => Ok(PostToolAction::Abort("stop parallel batch".to_string())), PostToolStopMode::Failure => Err(InterceptorError::new( @@ -1380,20 +1381,29 @@ impl Interceptor for StopFirstParallelResult { #[tokio::test] async fn post_tool_stop_terminalizes_started_parallel_siblings_before_returning() { for mode in [PostToolStopMode::Abort, PostToolStopMode::Failure] { - let client = MockLlmClient::new(vec![ + let first_response = vec![ Event::tool_use_start(0, "call_fast", "fast"), Event::tool_input_delta(0, r#"{}"#), Event::tool_use_stop(0), - Event::tool_use_start(1, "call_slow", "slow"), + Event::tool_use_start(1, "call_ready", "ready"), Event::tool_input_delta(1, r#"{}"#), Event::tool_use_stop(1), Event::Status(StatusEvent { status: ResponseStatus::Completed, }), - ]); + ]; + let second_response = vec![ + Event::text_block_start(0), + Event::text_delta(0, "next run completed"), + Event::text_block_stop(0, None), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]; + let client = MockLlmClient::with_responses(vec![first_response, second_response]); let mut engine = Engine::new(client); - engine.register_tool(SlowTool::new("fast", 1).definition()); - engine.register_tool(SlowTool::new("slow", 10_000).definition()); + engine.register_tool(SlowTool::new("fast", 0).definition()); + engine.register_tool(SlowTool::new("ready", 1).definition()); engine.set_interceptor(StopFirstParallelResult(mode)); let mut history = History::new(); @@ -1422,6 +1432,10 @@ async fn post_tool_stop_terminalizes_started_parallel_siblings_before_returning( .collect(); assert_eq!(terminal_ids.len(), 2); assert!(terminal_ids.contains(&"call_fast")); - assert!(terminal_ids.contains(&"call_slow")); + assert!(terminal_ids.contains(&"call_ready")); + + let mut engine = output.engine; + let next = engine.run(&mut history, "next run").await; + assert!(matches!(next, EngineRunExit::Finished)); } } From 4390554477a6f20db731482ca7414ab0ba26660c Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 13:30:29 +0900 Subject: [PATCH 06/22] feat: share Skill REST DTO authority --- crates/worker/src/skill.rs | 153 +---- crates/workspace-api/Cargo.toml | 4 + .../examples/generate_skill_api_types.rs | 3 + crates/workspace-api/src/lib.rs | 601 ++++++++++++++++++ crates/workspace-server/src/skills.rs | 67 +- web/workspace/src/lib/generated/skill-api.ts | 87 +++ .../src/lib/workspace/api/http.test.ts | 17 +- web/workspace/src/lib/workspace/api/http.ts | 64 +- web/workspace/src/lib/workspace/skills/api.ts | 455 +++++++++++++ web/workspace/tests/skill-api.test.ts | 221 +++++++ 10 files changed, 1487 insertions(+), 185 deletions(-) create mode 100644 crates/workspace-api/examples/generate_skill_api_types.rs create mode 100644 web/workspace/src/lib/generated/skill-api.ts create mode 100644 web/workspace/src/lib/workspace/skills/api.ts create mode 100644 web/workspace/tests/skill-api.test.ts diff --git a/crates/worker/src/skill.rs b/crates/worker/src/skill.rs index 055cb9fe..ecb47ced 100644 --- a/crates/worker/src/skill.rs +++ b/crates/worker/src/skill.rs @@ -2,127 +2,11 @@ use serde::{Deserialize, Serialize}; use crate::worker::WorkspaceClient; -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -pub enum SkillDiagnosticSeverity { - Error, - Warning, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct SkillDiagnostic { - pub severity: SkillDiagnosticSeverity, - pub code: String, - pub message: String, - /// Path-free authority/provenance label such as `builtin:foo` or `workspace:foo`. - #[serde(skip_serializing_if = "Option::is_none")] - pub source: Option, -} - -impl SkillDiagnostic { - pub fn error( - code: impl Into, - message: impl Into, - source: Option, - ) -> Self { - Self { - severity: SkillDiagnosticSeverity::Error, - code: code.into(), - message: message.into(), - source, - } - } - - pub fn warning( - code: impl Into, - message: impl Into, - source: Option, - ) -> Self { - Self { - severity: SkillDiagnosticSeverity::Warning, - code: code.into(), - message: message.into(), - source, - } - } -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -pub enum SkillSourceKind { - Builtin, - Workspace, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct SkillProvenance { - pub kind: SkillSourceKind, - /// Stable id: `builtin:` or `workspace:`. - pub id: String, - /// Virtual config/resource path. Never an absolute host filesystem path. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub virtual_path: Option, - /// Active Workspace config revision for Workspace Skills. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub revision: Option, - /// Digest of the immutable `SKILL.md` source. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub source_digest: Option, - /// Digest of the active virtual config tree snapshot. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub tree_digest: Option, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct SkillResourceRef { - pub kind: String, - /// Skill-relative resource name/path. Never an absolute filesystem path. - pub name: String, - pub supported: bool, - #[serde(skip_serializing_if = "Option::is_none")] - pub diagnostic: Option, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct SkillCatalogEntry { - pub name: String, - pub description: String, - pub provenance: SkillProvenance, - #[serde(default)] - pub overrides: Vec, - #[serde(default)] - pub diagnostics: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct SkillCatalogResponse { - /// Authority label for diagnostics; callers must not interpret it as a path. - pub authority: String, - #[serde(default)] - pub entries: Vec, - #[serde(default)] - pub diagnostics: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct SkillDetailResponse { - pub name: String, - pub description: String, - pub provenance: SkillProvenance, - #[serde(default)] - pub overrides: Vec, - #[serde(default)] - pub diagnostics: Vec, - /// Imported Markdown content with YAML frontmatter delimiters removed. - /// This is intentionally omitted from catalog responses. - pub body: String, - #[serde(default)] - pub allowed_tools: Vec, - /// Explicitly documents that allowed-tools is parsed only as an experimental hint. - pub allowed_tools_status: String, - #[serde(default)] - pub resources: Vec, -} +pub use workspace_api::{ + SkillActivationStatus, SkillCatalogEntry, SkillCatalogResponse, SkillDetailResponse, + SkillDiagnostic, SkillDiagnosticSeverity, SkillProjectionIdentity, SkillProjectionStatus, + SkillProvenance, SkillResourceRef, SkillSourceKind, +}; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct SkillActivationResponse { @@ -144,6 +28,8 @@ pub enum SkillClientError { Request(#[from] crate::worker::WorkspaceClientError), #[error("Skill API response JSON is invalid: {0}")] Json(#[from] serde_json::Error), + #[error("Skill API response violates the shared contract: {0}")] + InvalidResponse(#[from] workspace_api::SkillApiValidationError), #[error("Skill API returned HTTP {status}: {body}")] Http { status: reqwest::StatusCode, @@ -155,11 +41,15 @@ pub enum SkillClientError { impl dyn WorkspaceClient + '_ { pub fn list_skills(&self) -> Result { - self.get_skill_json("skills") + let response: SkillCatalogResponse = self.get_skill_json("skills")?; + response.validate()?; + Ok(response) } pub fn read_skill(&self, name: &str) -> Result { - self.get_skill_json(&format!("skills/{name}")) + let response: SkillDetailResponse = self.get_skill_json(&format!("skills/{name}"))?; + response.validate()?; + Ok(response) } pub fn activate_skill(&self, name: &str) -> Result { @@ -229,11 +119,24 @@ mod tests { assert_eq!(worker_header, None); assert_eq!(authorization, None); let body = serde_json::json!({ - "authority": "workspace-backend-skills-v0", + "authority": "workspace-config-skills-v1", + "projection": { + "config_revision": 7, + "tree_digest": "tree-digest" + }, "entries": [{ "name": "triage-errors", "description": "Use when triaging errors.", - "provenance": { "kind": "workspace", "id": "workspace:triage-errors" }, + "activation_status": "active", + "projection_status": "valid", + "provenance": { + "kind": "workspace", + "id": "workspace:triage-errors", + "virtual_path": "skills/triage-errors/SKILL.md", + "revision": 7, + "source_digest": "source-digest", + "tree_digest": "tree-digest" + }, "overrides": [], "diagnostics": [] }], diff --git a/crates/workspace-api/Cargo.toml b/crates/workspace-api/Cargo.toml index a4e5f31e..468c6547 100644 --- a/crates/workspace-api/Cargo.toml +++ b/crates/workspace-api/Cargo.toml @@ -38,6 +38,10 @@ required-features = ["typescript"] name = "generate_memory_api_types" required-features = ["typescript"] +[[example]] +name = "generate_skill_api_types" +required-features = ["typescript"] + [[example]] name = "generate_auth_api_types" required-features = ["typescript"] diff --git a/crates/workspace-api/examples/generate_skill_api_types.rs b/crates/workspace-api/examples/generate_skill_api_types.rs new file mode 100644 index 00000000..4fa47ded --- /dev/null +++ b/crates/workspace-api/examples/generate_skill_api_types.rs @@ -0,0 +1,3 @@ +fn main() { + print!("{}", workspace_api::skill_api_typescript()); +} diff --git a/crates/workspace-api/src/lib.rs b/crates/workspace-api/src/lib.rs index ea076429..e9b16045 100644 --- a/crates/workspace-api/src/lib.rs +++ b/crates/workspace-api/src/lib.rs @@ -1934,6 +1934,407 @@ pub struct RepositoryAccessProjection { pub bindings: Vec, } +pub const SKILL_CATALOG_AUTHORITY: &str = "workspace-config-skills-v1"; +pub const SKILL_API_MAX_SAFE_INTEGER: u64 = 9_007_199_254_740_991; +pub const SKILL_API_MAX_CATALOG_ENTRIES: usize = 500; +pub const SKILL_API_MAX_OVERRIDES: usize = 64; +pub const SKILL_API_MAX_DIAGNOSTICS: usize = 100; +pub const SKILL_API_MAX_RESOURCES: usize = 500; +pub const SKILL_API_MAX_ALLOWED_TOOLS: usize = 100; +pub const SKILL_API_MAX_NAME_BYTES: usize = 128; +pub const SKILL_API_MAX_LABEL_BYTES: usize = 4_096; +pub const SKILL_API_MAX_BODY_BYTES: usize = 1_048_576; +pub const SKILL_API_MAX_PATH_BYTES: usize = 1_024; +pub const SKILL_API_MAX_DIGEST_BYTES: usize = 128; + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[cfg_attr(feature = "typescript", ts(rename_all = "snake_case"))] +#[serde(rename_all = "snake_case")] +pub enum SkillDiagnosticSeverity { + Error, + Warning, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(deny_unknown_fields)] +pub struct SkillDiagnostic { + pub severity: SkillDiagnosticSeverity, + pub code: String, + pub message: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "typescript", ts(optional))] + pub source: Option, +} + +impl SkillDiagnostic { + pub fn error( + code: impl Into, + message: impl Into, + source: Option, + ) -> Self { + Self { + severity: SkillDiagnosticSeverity::Error, + code: code.into(), + message: message.into(), + source, + } + } + + pub fn warning( + code: impl Into, + message: impl Into, + source: Option, + ) -> Self { + Self { + severity: SkillDiagnosticSeverity::Warning, + code: code.into(), + message: message.into(), + source, + } + } +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[cfg_attr(feature = "typescript", ts(rename_all = "snake_case"))] +#[serde(rename_all = "snake_case")] +pub enum SkillSourceKind { + Builtin, + Workspace, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(deny_unknown_fields)] +pub struct SkillProvenance { + pub kind: SkillSourceKind, + pub id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "typescript", ts(optional))] + pub virtual_path: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "typescript", ts(optional, type = "number"))] + pub revision: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "typescript", ts(optional))] + pub source_digest: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "typescript", ts(optional))] + pub tree_digest: Option, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[cfg_attr(feature = "typescript", ts(rename_all = "snake_case"))] +#[serde(rename_all = "snake_case")] +pub enum SkillActivationStatus { + Active, + Inactive, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[cfg_attr(feature = "typescript", ts(rename_all = "snake_case"))] +#[serde(rename_all = "snake_case")] +pub enum SkillProjectionStatus { + Valid, + Invalid, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(deny_unknown_fields)] +pub struct SkillProjectionIdentity { + #[cfg_attr(feature = "typescript", ts(type = "number"))] + pub config_revision: u64, + pub tree_digest: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(deny_unknown_fields)] +pub struct SkillResourceRef { + pub kind: String, + pub name: String, + pub supported: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "typescript", ts(optional))] + pub diagnostic: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(deny_unknown_fields)] +pub struct SkillCatalogEntry { + pub name: String, + pub description: String, + pub activation_status: SkillActivationStatus, + pub projection_status: SkillProjectionStatus, + pub provenance: SkillProvenance, + pub overrides: Vec, + pub diagnostics: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(deny_unknown_fields)] +pub struct SkillCatalogResponse { + pub authority: String, + pub projection: SkillProjectionIdentity, + pub entries: Vec, + pub diagnostics: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(deny_unknown_fields)] +pub struct SkillDetailResponse { + pub authority: String, + pub projection: SkillProjectionIdentity, + pub name: String, + pub description: String, + pub provenance: SkillProvenance, + pub overrides: Vec, + pub diagnostics: Vec, + pub activation_status: SkillActivationStatus, + pub projection_status: SkillProjectionStatus, + pub body: String, + pub allowed_tools: Vec, + pub allowed_tools_status: String, + pub resources: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SkillApiValidationError { + CollectionTooLarge, + StringTooLarge, + InvalidProjectionIdentity, + InvalidProvenance, + InvalidVirtualPath, + StaleProjection, +} + +impl std::fmt::Display for SkillApiValidationError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let message = match self { + Self::CollectionTooLarge => "Skill API collection exceeds its limit", + Self::StringTooLarge => "Skill API string exceeds its limit", + Self::InvalidProjectionIdentity => "Skill API projection identity is invalid", + Self::InvalidProvenance => "Skill API provenance is invalid", + Self::InvalidVirtualPath => "Skill API virtual path is invalid", + Self::StaleProjection => "Workspace Skill projection is stale", + }; + formatter.write_str(message) + } +} + +impl std::error::Error for SkillApiValidationError {} + +impl SkillProjectionIdentity { + fn validate(&self) -> Result<(), SkillApiValidationError> { + validate_safe_integer(self.config_revision)?; + validate_nonempty_string(&self.tree_digest, SKILL_API_MAX_DIGEST_BYTES) + .map_err(|_| SkillApiValidationError::InvalidProjectionIdentity) + } +} + +impl SkillProvenance { + fn validate( + &self, + projection: &SkillProjectionIdentity, + ) -> Result<(), SkillApiValidationError> { + validate_nonempty_string(&self.id, SKILL_API_MAX_LABEL_BYTES)?; + validate_optional_string(self.virtual_path.as_deref(), SKILL_API_MAX_PATH_BYTES)?; + validate_optional_string(self.source_digest.as_deref(), SKILL_API_MAX_DIGEST_BYTES)?; + validate_optional_string(self.tree_digest.as_deref(), SKILL_API_MAX_DIGEST_BYTES)?; + if let Some(revision) = self.revision { + validate_safe_integer(revision)?; + } + + let expected_prefix = match self.kind { + SkillSourceKind::Builtin => "builtin:", + SkillSourceKind::Workspace => "workspace:", + }; + if !self.id.starts_with(expected_prefix) + || self + .virtual_path + .as_deref() + .is_none_or(|path| !is_virtual_path(path)) + || self.source_digest.is_none() + { + return Err(SkillApiValidationError::InvalidProvenance); + } + + match self.kind { + SkillSourceKind::Builtin => { + if self.revision.is_some() || self.tree_digest.is_some() { + return Err(SkillApiValidationError::InvalidProvenance); + } + } + SkillSourceKind::Workspace => { + let Some(revision) = self.revision else { + return Err(SkillApiValidationError::InvalidProvenance); + }; + let Some(tree_digest) = self.tree_digest.as_deref() else { + return Err(SkillApiValidationError::InvalidProvenance); + }; + if revision != projection.config_revision || tree_digest != projection.tree_digest { + return Err(SkillApiValidationError::StaleProjection); + } + } + } + Ok(()) + } +} + +impl SkillCatalogEntry { + fn validate( + &self, + projection: &SkillProjectionIdentity, + ) -> Result<(), SkillApiValidationError> { + validate_nonempty_string(&self.name, SKILL_API_MAX_NAME_BYTES)?; + validate_string(&self.description, SKILL_API_MAX_LABEL_BYTES)?; + validate_collection(&self.overrides, SKILL_API_MAX_OVERRIDES)?; + validate_diagnostics(&self.diagnostics)?; + self.provenance.validate(projection)?; + for provenance in &self.overrides { + provenance.validate(projection)?; + } + Ok(()) + } +} + +impl SkillCatalogResponse { + pub fn validate(&self) -> Result<(), SkillApiValidationError> { + validate_nonempty_string(&self.authority, SKILL_API_MAX_LABEL_BYTES)?; + if self.authority != SKILL_CATALOG_AUTHORITY { + return Err(SkillApiValidationError::InvalidProjectionIdentity); + } + self.projection.validate()?; + validate_collection(&self.entries, SKILL_API_MAX_CATALOG_ENTRIES)?; + validate_diagnostics(&self.diagnostics)?; + for entry in &self.entries { + entry.validate(&self.projection)?; + } + Ok(()) + } +} + +impl SkillDetailResponse { + pub fn validate(&self) -> Result<(), SkillApiValidationError> { + validate_nonempty_string(&self.authority, SKILL_API_MAX_LABEL_BYTES)?; + if self.authority != SKILL_CATALOG_AUTHORITY { + return Err(SkillApiValidationError::InvalidProjectionIdentity); + } + self.projection.validate()?; + validate_nonempty_string(&self.name, SKILL_API_MAX_NAME_BYTES)?; + validate_string(&self.description, SKILL_API_MAX_LABEL_BYTES)?; + validate_string(&self.body, SKILL_API_MAX_BODY_BYTES)?; + validate_strings( + &self.allowed_tools, + SKILL_API_MAX_ALLOWED_TOOLS, + SKILL_API_MAX_LABEL_BYTES, + )?; + validate_nonempty_string(&self.allowed_tools_status, SKILL_API_MAX_LABEL_BYTES)?; + validate_resources(&self.resources)?; + validate_collection(&self.overrides, SKILL_API_MAX_OVERRIDES)?; + validate_diagnostics(&self.diagnostics)?; + self.provenance.validate(&self.projection)?; + for provenance in &self.overrides { + provenance.validate(&self.projection)?; + } + Ok(()) + } +} + +fn validate_safe_integer(value: u64) -> Result<(), SkillApiValidationError> { + if value <= SKILL_API_MAX_SAFE_INTEGER { + Ok(()) + } else { + Err(SkillApiValidationError::InvalidProjectionIdentity) + } +} + +fn validate_collection(values: &[T], limit: usize) -> Result<(), SkillApiValidationError> { + if values.len() <= limit { + Ok(()) + } else { + Err(SkillApiValidationError::CollectionTooLarge) + } +} + +fn validate_string(value: &str, limit: usize) -> Result<(), SkillApiValidationError> { + if value.len() <= limit { + Ok(()) + } else { + Err(SkillApiValidationError::StringTooLarge) + } +} + +fn validate_nonempty_string(value: &str, limit: usize) -> Result<(), SkillApiValidationError> { + validate_string(value, limit)?; + if value.is_empty() { + Err(SkillApiValidationError::StringTooLarge) + } else { + Ok(()) + } +} + +fn validate_optional_string( + value: Option<&str>, + limit: usize, +) -> Result<(), SkillApiValidationError> { + if let Some(value) = value { + validate_nonempty_string(value, limit)?; + } + Ok(()) +} + +fn validate_strings( + values: &[String], + collection_limit: usize, + string_limit: usize, +) -> Result<(), SkillApiValidationError> { + validate_collection(values, collection_limit)?; + for value in values { + validate_nonempty_string(value, string_limit)?; + } + Ok(()) +} + +fn validate_resources(resources: &[SkillResourceRef]) -> Result<(), SkillApiValidationError> { + validate_collection(resources, SKILL_API_MAX_RESOURCES)?; + for resource in resources { + validate_nonempty_string(&resource.kind, SKILL_API_MAX_LABEL_BYTES)?; + validate_nonempty_string(&resource.name, SKILL_API_MAX_PATH_BYTES)?; + if !is_virtual_path(&resource.name) { + return Err(SkillApiValidationError::InvalidVirtualPath); + } + validate_optional_string(resource.diagnostic.as_deref(), SKILL_API_MAX_LABEL_BYTES)?; + } + Ok(()) +} + +fn is_virtual_path(value: &str) -> bool { + !value.starts_with('/') + && !value.contains('\\') + && value + .split('/') + .all(|component| !component.is_empty() && component != "." && component != "..") +} + +fn validate_diagnostics(diagnostics: &[SkillDiagnostic]) -> Result<(), SkillApiValidationError> { + validate_collection(diagnostics, SKILL_API_MAX_DIAGNOSTICS)?; + for diagnostic in diagnostics { + validate_nonempty_string(&diagnostic.code, SKILL_API_MAX_LABEL_BYTES)?; + validate_nonempty_string(&diagnostic.message, SKILL_API_MAX_LABEL_BYTES)?; + validate_optional_string(diagnostic.source.as_deref(), SKILL_API_MAX_PATH_BYTES)?; + } + Ok(()) +} + #[cfg(feature = "typescript")] pub fn catalog_typescript() -> String { use ts_rs::TS; @@ -2005,6 +2406,37 @@ pub fn repository_access_api_typescript() -> String { ) } +#[cfg(feature = "typescript")] +pub fn skill_api_typescript() -> String { + use ts_rs::TS; + + let config = ts_rs::Config::default(); + let declarations = [ + SkillDiagnosticSeverity::decl(&config), + SkillDiagnostic::decl(&config), + SkillSourceKind::decl(&config), + SkillProvenance::decl(&config), + SkillActivationStatus::decl(&config), + SkillProjectionStatus::decl(&config), + SkillProjectionIdentity::decl(&config), + SkillResourceRef::decl(&config), + SkillCatalogEntry::decl(&config), + SkillCatalogResponse::decl(&config), + SkillDetailResponse::decl(&config), + ]; + let limits = format!( + "export const SKILL_API_AUTHORITY = \"{SKILL_CATALOG_AUTHORITY}\" as const;\n\nexport const SKILL_API_LIMITS = {{\n maxSafeInteger: {SKILL_API_MAX_SAFE_INTEGER},\n maxCatalogEntries: {SKILL_API_MAX_CATALOG_ENTRIES},\n maxOverrides: {SKILL_API_MAX_OVERRIDES},\n maxDiagnostics: {SKILL_API_MAX_DIAGNOSTICS},\n maxResources: {SKILL_API_MAX_RESOURCES},\n maxAllowedTools: {SKILL_API_MAX_ALLOWED_TOOLS},\n maxNameBytes: {SKILL_API_MAX_NAME_BYTES},\n maxLabelBytes: {SKILL_API_MAX_LABEL_BYTES},\n maxBodyBytes: {SKILL_API_MAX_BODY_BYTES},\n maxPathBytes: {SKILL_API_MAX_PATH_BYTES},\n maxDigestBytes: {SKILL_API_MAX_DIGEST_BYTES},\n}} as const;" + ); + format!( + "// Generated from workspace-api. Do not edit by hand.\n// Regenerate: cargo run -q -p workspace-api --features typescript --example generate_skill_api_types > web/workspace/src/lib/generated/skill-api.ts\n\n{limits}\n\n{}\n", + declarations + .into_iter() + .map(|declaration| format!("export {declaration}")) + .collect::>() + .join("\n\n") + ) +} + #[cfg(feature = "typescript")] pub fn auth_api_typescript() -> String { use ts_rs::TS; @@ -2172,6 +2604,36 @@ mod memory_typescript_tests { } } +#[cfg(all(test, feature = "typescript"))] +mod skill_typescript_tests { + #[test] + fn generated_skill_api_contract_is_current() { + let expected = super::skill_api_typescript(); + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../../web/workspace/src/lib/generated/skill-api.ts"); + let actual = std::fs::read_to_string(&path) + .unwrap_or_else(|error| panic!("failed to read {}: {error}", path.display())); + assert_eq!( + normalize(&actual), + normalize(&expected), + "regenerate Skill API TypeScript types with `cargo run -q -p workspace-api --features typescript --example generate_skill_api_types > web/workspace/src/lib/generated/skill-api.ts` and format the generated file", + ); + } + + fn normalize(value: &str) -> String { + value + .chars() + .filter_map(|character| match character { + character if character.is_whitespace() => None, + ',' => Some(';'), + character => Some(character), + }) + .collect::() + .replace("=|", "=") + .replace(";}", "}") + } +} + #[cfg(all(test, feature = "typescript"))] mod workdir_typescript_tests { #[test] @@ -2205,6 +2667,145 @@ mod workdir_typescript_tests { mod tests { use super::*; + fn skill_projection() -> SkillProjectionIdentity { + SkillProjectionIdentity { + config_revision: 42, + tree_digest: "tree-digest".to_string(), + } + } + + fn builtin_skill_provenance() -> SkillProvenance { + SkillProvenance { + kind: SkillSourceKind::Builtin, + id: "builtin:errors".to_string(), + virtual_path: Some("skills/errors/SKILL.md".to_string()), + revision: None, + source_digest: Some("builtin-source-digest".to_string()), + tree_digest: None, + } + } + + fn workspace_skill_provenance() -> SkillProvenance { + SkillProvenance { + kind: SkillSourceKind::Workspace, + id: "workspace:skills/release/SKILL.md".to_string(), + virtual_path: Some("skills/release/SKILL.md".to_string()), + revision: Some(42), + source_digest: Some("workspace-source-digest".to_string()), + tree_digest: Some("tree-digest".to_string()), + } + } + + #[test] + fn skill_catalog_round_trips_builtin_workspace_and_invalid_projection_entries() { + let response = SkillCatalogResponse { + authority: "workspace-config-skills-v1".to_string(), + projection: skill_projection(), + entries: vec![ + SkillCatalogEntry { + name: "errors".to_string(), + description: "Builtin guidance".to_string(), + activation_status: SkillActivationStatus::Active, + projection_status: SkillProjectionStatus::Valid, + provenance: builtin_skill_provenance(), + overrides: vec![], + diagnostics: vec![], + }, + SkillCatalogEntry { + name: "release".to_string(), + description: "Workspace guidance".to_string(), + activation_status: SkillActivationStatus::Inactive, + projection_status: SkillProjectionStatus::Invalid, + provenance: workspace_skill_provenance(), + overrides: vec![builtin_skill_provenance()], + diagnostics: vec![SkillDiagnostic { + severity: SkillDiagnosticSeverity::Error, + code: "invalid_projection".to_string(), + message: "invalid projected Skill".to_string(), + source: Some("skills/release/SKILL.md".to_string()), + }], + }, + ], + diagnostics: vec![], + }; + + response.validate().expect("fixture should be valid"); + let json = serde_json::to_string(&response).expect("serialize Skill catalog"); + let decoded: SkillCatalogResponse = + serde_json::from_str(&json).expect("deserialize Skill catalog"); + assert_eq!(decoded, response); + assert!(!json.contains("\"revision\":null")); + assert!(!json.contains("\"tree_digest\":null")); + } + + #[test] + fn skill_detail_round_trips_shared_response() { + let response = SkillDetailResponse { + authority: "workspace-config-skills-v1".to_string(), + projection: skill_projection(), + name: "release".to_string(), + description: "Workspace guidance".to_string(), + body: "# Release\n".to_string(), + allowed_tools: vec!["Bash".to_string()], + allowed_tools_status: "experimental_hint_only".to_string(), + resources: vec![], + activation_status: SkillActivationStatus::Active, + projection_status: SkillProjectionStatus::Valid, + provenance: workspace_skill_provenance(), + overrides: vec![], + diagnostics: vec![], + }; + + response.validate().expect("fixture should be valid"); + let decoded: SkillDetailResponse = serde_json::from_value( + serde_json::to_value(&response).expect("serialize Skill detail"), + ) + .expect("deserialize Skill detail"); + assert_eq!(decoded, response); + } + + #[test] + fn skill_projection_validation_detects_stale_workspace_revision() { + let mut provenance = workspace_skill_provenance(); + provenance.revision = Some(41); + let response = SkillCatalogResponse { + authority: "workspace-config-skills-v1".to_string(), + projection: skill_projection(), + entries: vec![SkillCatalogEntry { + name: "release".to_string(), + description: String::new(), + activation_status: SkillActivationStatus::Active, + projection_status: SkillProjectionStatus::Valid, + provenance, + overrides: vec![], + diagnostics: vec![], + }], + diagnostics: vec![], + }; + + assert_eq!( + response.validate(), + Err(SkillApiValidationError::StaleProjection) + ); + } + + #[test] + fn skill_dto_rejects_unknown_fields_and_unknown_provenance_kind() { + let unknown_field = serde_json::json!({ + "authority": "workspace-config-skills-v1", + "projection": {"config_revision": 42, "tree_digest": "tree-digest"}, + "entries": [], + "diagnostics": [], + "body": "must not be accepted" + }); + assert!(serde_json::from_value::(unknown_field).is_err()); + + let mut provenance = + serde_json::to_value(workspace_skill_provenance()).expect("serialize provenance"); + provenance["kind"] = serde_json::Value::String("newer_source_kind".to_string()); + assert!(serde_json::from_value::(provenance).is_err()); + } + #[test] fn memory_evidence_origins_round_trip_as_typed_provenance() { let kinds = [ diff --git a/crates/workspace-server/src/skills.rs b/crates/workspace-server/src/skills.rs index 03bcc793..91703108 100644 --- a/crates/workspace-server/src/skills.rs +++ b/crates/workspace-server/src/skills.rs @@ -4,9 +4,11 @@ use config_source::{ ConfigSchemaContribution, MarkdownDocumentProjection, VirtualPath, project_markdown_document, }; use serde::Deserialize; -use worker::skill::{ - SkillActivationResponse, SkillCatalogEntry, SkillCatalogResponse, SkillDetailResponse, - SkillDiagnostic, SkillDiagnosticSeverity, SkillProvenance, SkillResourceRef, SkillSourceKind, +use worker::skill::SkillActivationResponse; +use workspace_api::{ + SKILL_CATALOG_AUTHORITY, SkillActivationStatus, SkillCatalogEntry, SkillCatalogResponse, + SkillDetailResponse, SkillDiagnostic, SkillDiagnosticSeverity, SkillProjectionIdentity, + SkillProjectionStatus, SkillProvenance, SkillResourceRef, SkillSourceKind, }; use crate::config_source::{ @@ -19,7 +21,6 @@ const BUILTIN_SKILL_VIRTUAL_PATH: &str = "builtin/skills/agent-skills/SKILL.md"; const SKILL_SCHEMA_PROVIDER_ID: &str = "builtin:skills"; const SKILL_SCHEMA_NAMESPACE: &str = "skills"; const SKILL_SCHEMA_VERSION: &str = "1"; -const SKILL_CATALOG_AUTHORITY: &str = "workspace-config-skills-v1"; /// Skill documents are values imported from `SKILL.md`. Known Agent Skills /// frontmatter is typed while extension keys remain concrete values. @@ -101,32 +102,54 @@ pub fn catalog(state: &WorkspaceConfigState) -> Result Result { catalog(state) } +fn projection_identity(state: &WorkspaceConfigState) -> SkillProjectionIdentity { + SkillProjectionIdentity { + config_revision: state.snapshot.revision, + tree_digest: state.snapshot.digest.clone(), + } +} + pub fn detail(state: &WorkspaceConfigState, name: &str) -> Result { let skill = merged_skills(state)? .remove(name) .ok_or_else(|| SkillError::NotFound(name.to_string()))?; - Ok(SkillDetailResponse { + let activation_status = skill.activation_status(); + let projection_status = skill.projection_status(); + let response = SkillDetailResponse { + authority: SKILL_CATALOG_AUTHORITY.to_string(), + projection: projection_identity(state), name: skill.name, description: skill.description, provenance: skill.provenance, overrides: skill.overrides, diagnostics: skill.diagnostics, + activation_status, + projection_status, body: skill.body, allowed_tools: skill.allowed_tools, allowed_tools_status: "experimental_hint_only".to_string(), resources: skill.resources, - }) + }; + response + .validate() + .map_err(|error| SkillError::InvalidProjection(error.to_string()))?; + Ok(response) } pub fn activation( @@ -415,10 +438,28 @@ impl ParsedSkill { .any(|diagnostic| diagnostic.severity == SkillDiagnosticSeverity::Error) } + fn activation_status(&self) -> SkillActivationStatus { + if self.has_errors() { + SkillActivationStatus::Inactive + } else { + SkillActivationStatus::Active + } + } + + fn projection_status(&self) -> SkillProjectionStatus { + if self.has_errors() { + SkillProjectionStatus::Invalid + } else { + SkillProjectionStatus::Valid + } + } + fn catalog_entry(&self) -> SkillCatalogEntry { SkillCatalogEntry { name: self.name.clone(), description: self.description.clone(), + activation_status: self.activation_status(), + projection_status: self.projection_status(), provenance: self.provenance.clone(), overrides: self.overrides.clone(), diagnostics: self.diagnostics.clone(), @@ -513,12 +554,20 @@ mod tests { .unwrap(); assert_eq!(item.provenance.kind, SkillSourceKind::Workspace); assert_eq!(item.provenance.revision, Some(9)); + assert_eq!(catalog.projection.config_revision, 9); + assert_eq!(catalog.projection.tree_digest, state.snapshot.digest); + assert_eq!(item.activation_status, SkillActivationStatus::Active); + assert_eq!(item.projection_status, SkillProjectionStatus::Valid); assert!( item.diagnostics .iter() .all(|diagnostic| diagnostic.severity != SkillDiagnosticSeverity::Error) ); let detail = detail(&state, "debug-rust").unwrap(); + assert_eq!(detail.authority, SKILL_CATALOG_AUTHORITY); + assert_eq!(detail.projection.config_revision, 9); + assert_eq!(detail.activation_status, SkillActivationStatus::Active); + assert_eq!(detail.projection_status, SkillProjectionStatus::Valid); assert_eq!(detail.body, "# Debug Rust\n"); assert_eq!(detail.allowed_tools, vec!["Read", "Grep"]); assert_eq!( @@ -579,6 +628,8 @@ mod tests { .into_iter() .find(|item| item.name == "debug-rust") .unwrap(); + assert_eq!(item.activation_status, SkillActivationStatus::Inactive); + assert_eq!(item.projection_status, SkillProjectionStatus::Invalid); assert!( item.diagnostics .iter() diff --git a/web/workspace/src/lib/generated/skill-api.ts b/web/workspace/src/lib/generated/skill-api.ts new file mode 100644 index 00000000..5847d06e --- /dev/null +++ b/web/workspace/src/lib/generated/skill-api.ts @@ -0,0 +1,87 @@ +// Generated from workspace-api. Do not edit by hand. +// Regenerate: cargo run -q -p workspace-api --features typescript --example generate_skill_api_types > web/workspace/src/lib/generated/skill-api.ts + +export const SKILL_API_AUTHORITY = "workspace-config-skills-v1" as const; + +export const SKILL_API_LIMITS = { + maxSafeInteger: 9007199254740991, + maxCatalogEntries: 500, + maxOverrides: 64, + maxDiagnostics: 100, + maxResources: 500, + maxAllowedTools: 100, + maxNameBytes: 128, + maxLabelBytes: 4096, + maxBodyBytes: 1048576, + maxPathBytes: 1024, + maxDigestBytes: 128, +} as const; + +export type SkillDiagnosticSeverity = "error" | "warning"; + +export type SkillDiagnostic = { + severity: SkillDiagnosticSeverity; + code: string; + message: string; + source?: string; +}; + +export type SkillSourceKind = "builtin" | "workspace"; + +export type SkillProvenance = { + kind: SkillSourceKind; + id: string; + virtual_path?: string; + revision?: number; + source_digest?: string; + tree_digest?: string; +}; + +export type SkillActivationStatus = "active" | "inactive"; + +export type SkillProjectionStatus = "valid" | "invalid"; + +export type SkillProjectionIdentity = { + config_revision: number; + tree_digest: string; +}; + +export type SkillResourceRef = { + kind: string; + name: string; + supported: boolean; + diagnostic?: string; +}; + +export type SkillCatalogEntry = { + name: string; + description: string; + activation_status: SkillActivationStatus; + projection_status: SkillProjectionStatus; + provenance: SkillProvenance; + overrides: Array; + diagnostics: Array; +}; + +export type SkillCatalogResponse = { + authority: string; + projection: SkillProjectionIdentity; + entries: Array; + diagnostics: Array; +}; + +export type SkillDetailResponse = { + authority: string; + projection: SkillProjectionIdentity; + name: string; + description: string; + provenance: SkillProvenance; + overrides: Array; + diagnostics: Array; + activation_status: SkillActivationStatus; + projection_status: SkillProjectionStatus; + body: string; + allowed_tools: Array; + allowed_tools_status: string; + resources: Array; +}; diff --git a/web/workspace/src/lib/workspace/api/http.test.ts b/web/workspace/src/lib/workspace/api/http.test.ts index eefa309f..f7a8419f 100644 --- a/web/workspace/src/lib/workspace/api/http.test.ts +++ b/web/workspace/src/lib/workspace/api/http.test.ts @@ -89,11 +89,24 @@ Deno.test("loadWorkspaceSkillCatalog fetches lightweight catalog", async () => { return Promise.resolve( new Response( JSON.stringify({ - authority: "workspace-backend-skills-v0", + authority: "workspace-config-skills-v1", + projection: { + config_revision: 7, + tree_digest: "tree-digest", + }, entries: [{ name: "triage-errors", description: "Use when triaging errors.", - provenance: { kind: "workspace", id: "workspace:triage-errors" }, + activation_status: "active", + projection_status: "valid", + provenance: { + kind: "workspace", + id: "workspace:triage-errors", + virtual_path: "skills/triage-errors/SKILL.md", + revision: 7, + source_digest: "source-digest", + tree_digest: "tree-digest", + }, overrides: [], diagnostics: [], }], diff --git a/web/workspace/src/lib/workspace/api/http.ts b/web/workspace/src/lib/workspace/api/http.ts index bd5c51b8..c86beb48 100644 --- a/web/workspace/src/lib/workspace/api/http.ts +++ b/web/workspace/src/lib/workspace/api/http.ts @@ -1,58 +1,18 @@ +import type { + SkillCatalogResponse, + SkillDetailResponse, +} from "$lib/generated/skill-api.ts"; +import { + parseSkillCatalogResponse, + parseSkillDetailResponse, +} from "$lib/workspace/skills/api.ts"; + export type ApiResult = { data: T | null; error: string | null; }; -export type SkillDiagnosticSeverity = "error" | "warning"; - -export type SkillDiagnostic = { - severity: SkillDiagnosticSeverity; - code: string; - message: string; - source?: string; -}; - -export type SkillProvenance = { - kind: "builtin" | "workspace"; - id: string; - virtual_path?: string; - revision?: number; - source_digest?: string; - tree_digest?: string; -}; - -export type SkillCatalogEntry = { - name: string; - description: string; - provenance: SkillProvenance; - overrides: SkillProvenance[]; - diagnostics: SkillDiagnostic[]; -}; - -export type SkillCatalogResponse = { - authority: string; - entries: SkillCatalogEntry[]; - diagnostics: SkillDiagnostic[]; -}; - -export type SkillResourceRef = { - kind: string; - name: string; - supported: boolean; - diagnostic?: string; -}; - -export type SkillDetailResponse = { - name: string; - description: string; - provenance: SkillProvenance; - overrides: SkillProvenance[]; - diagnostics: SkillDiagnostic[]; - body: string; - allowed_tools: string[]; - allowed_tools_status: string; - resources: SkillResourceRef[]; -}; +export type { SkillCatalogResponse, SkillDetailResponse }; function normalizePath(path: string): string { if (!path || path === "/") return ""; @@ -95,6 +55,8 @@ export async function loadWorkspaceSkillCatalog( return loadJson( fetchFn, workspaceSkillCatalogPath(workspaceId), + undefined, + parseSkillCatalogResponse, ); } @@ -106,6 +68,8 @@ export async function loadWorkspaceSkillDetail( return loadJson( fetchFn, workspaceSkillDetailPath(workspaceId, name), + undefined, + parseSkillDetailResponse, ); } diff --git a/web/workspace/src/lib/workspace/skills/api.ts b/web/workspace/src/lib/workspace/skills/api.ts new file mode 100644 index 00000000..4f5012c3 --- /dev/null +++ b/web/workspace/src/lib/workspace/skills/api.ts @@ -0,0 +1,455 @@ +import { + SKILL_API_AUTHORITY, + SKILL_API_LIMITS, + type SkillActivationStatus, + type SkillCatalogEntry, + type SkillCatalogResponse, + type SkillDetailResponse, + type SkillDiagnostic, + type SkillDiagnosticSeverity, + type SkillProjectionIdentity, + type SkillProjectionStatus, + type SkillProvenance, + type SkillResourceRef, + type SkillSourceKind, +} from "$lib/generated/skill-api.ts"; + +export class SkillApiContractError extends Error { + constructor(message: string) { + super(message); + this.name = "SkillApiContractError"; + } +} + +export function parseSkillCatalogResponse( + value: unknown, +): SkillCatalogResponse { + const record = strictObject(value, [ + "authority", + "projection", + "entries", + "diagnostics", + ], "Skill catalog response"); + const authority = boundedString( + record.authority, + "Skill catalog authority", + SKILL_API_LIMITS.maxLabelBytes, + false, + ); + if (authority !== SKILL_API_AUTHORITY) { + throw contractError("unsupported Skill catalog authority"); + } + const projection = parseProjection(record.projection); + return { + authority, + projection, + entries: boundedArray( + record.entries, + "Skill catalog entries", + SKILL_API_LIMITS.maxCatalogEntries, + ).map((entry) => parseCatalogEntry(entry, projection)), + diagnostics: parseDiagnostics(record.diagnostics), + }; +} + +export function parseSkillDetailResponse(value: unknown): SkillDetailResponse { + const record = strictObject(value, [ + "authority", + "projection", + "name", + "description", + "provenance", + "overrides", + "diagnostics", + "activation_status", + "projection_status", + "body", + "allowed_tools", + "allowed_tools_status", + "resources", + ], "Skill detail response"); + const authority = boundedString( + record.authority, + "Skill detail authority", + SKILL_API_LIMITS.maxLabelBytes, + false, + ); + if (authority !== SKILL_API_AUTHORITY) { + throw contractError("unsupported Skill detail authority"); + } + const projection = parseProjection(record.projection); + return { + authority, + projection, + name: boundedString( + record.name, + "Skill name", + SKILL_API_LIMITS.maxNameBytes, + false, + ), + description: boundedString( + record.description, + "Skill description", + SKILL_API_LIMITS.maxLabelBytes, + true, + ), + provenance: parseProvenance(record.provenance, projection), + overrides: parseProvenances(record.overrides, projection), + diagnostics: parseDiagnostics(record.diagnostics), + activation_status: activationStatus(record.activation_status), + projection_status: projectionStatus(record.projection_status), + body: boundedString( + record.body, + "Skill body", + SKILL_API_LIMITS.maxBodyBytes, + true, + ), + allowed_tools: boundedArray( + record.allowed_tools, + "Skill allowed tools", + SKILL_API_LIMITS.maxAllowedTools, + ).map((tool) => + boundedString( + tool, + "Skill allowed tool", + SKILL_API_LIMITS.maxLabelBytes, + false, + ) + ), + allowed_tools_status: boundedString( + record.allowed_tools_status, + "Skill allowed-tools status", + SKILL_API_LIMITS.maxLabelBytes, + false, + ), + resources: boundedArray( + record.resources, + "Skill resources", + SKILL_API_LIMITS.maxResources, + ).map(parseResource), + }; +} + +function parseCatalogEntry( + value: unknown, + projection: SkillProjectionIdentity, +): SkillCatalogEntry { + const record = strictObject(value, [ + "name", + "description", + "activation_status", + "projection_status", + "provenance", + "overrides", + "diagnostics", + ], "Skill catalog entry"); + return { + name: boundedString( + record.name, + "Skill name", + SKILL_API_LIMITS.maxNameBytes, + false, + ), + description: boundedString( + record.description, + "Skill description", + SKILL_API_LIMITS.maxLabelBytes, + true, + ), + activation_status: activationStatus(record.activation_status), + projection_status: projectionStatus(record.projection_status), + provenance: parseProvenance(record.provenance, projection), + overrides: parseProvenances(record.overrides, projection), + diagnostics: parseDiagnostics(record.diagnostics), + }; +} + +function parseProjection(value: unknown): SkillProjectionIdentity { + const record = strictObject( + value, + ["config_revision", "tree_digest"], + "Skill projection identity", + ); + return { + config_revision: safeInteger( + record.config_revision, + "Skill config revision", + ), + tree_digest: boundedString( + record.tree_digest, + "Skill tree digest", + SKILL_API_LIMITS.maxDigestBytes, + false, + ), + }; +} + +function parseProvenances( + value: unknown, + projection: SkillProjectionIdentity, +): SkillProvenance[] { + return boundedArray( + value, + "Skill overrides", + SKILL_API_LIMITS.maxOverrides, + ).map((provenance) => parseProvenance(provenance, projection)); +} + +function parseProvenance( + value: unknown, + projection: SkillProjectionIdentity, +): SkillProvenance { + const record = strictObject( + value, + [ + "kind", + "id", + "virtual_path", + "revision", + "source_digest", + "tree_digest", + ], + "Skill provenance", + [ + "virtual_path", + "revision", + "source_digest", + "tree_digest", + ], + ); + const kind = sourceKind(record.kind); + const id = boundedString( + record.id, + "Skill provenance id", + SKILL_API_LIMITS.maxLabelBytes, + false, + ); + const virtualPath = optionalBoundedString( + record.virtual_path, + "Skill virtual path", + SKILL_API_LIMITS.maxPathBytes, + ); + const sourceDigest = optionalBoundedString( + record.source_digest, + "Skill source digest", + SKILL_API_LIMITS.maxDigestBytes, + ); + const treeDigest = optionalBoundedString( + record.tree_digest, + "Skill provenance tree digest", + SKILL_API_LIMITS.maxDigestBytes, + ); + const revision = record.revision === undefined + ? undefined + : safeInteger(record.revision, "Skill provenance revision"); + + if ( + !id.startsWith(`${kind}:`) || virtualPath === undefined || + sourceDigest === undefined || !isVirtualPath(virtualPath) + ) { + throw contractError("invalid Skill provenance"); + } + if (kind === "builtin") { + if (revision !== undefined || treeDigest !== undefined) { + throw contractError("invalid built-in Skill provenance"); + } + } else { + if (revision === undefined || treeDigest === undefined) { + throw contractError("incomplete Workspace Skill provenance"); + } + if ( + revision !== projection.config_revision || + treeDigest !== projection.tree_digest + ) { + throw contractError("stale Workspace Skill projection"); + } + } + + return { + kind, + id, + virtual_path: virtualPath, + revision, + source_digest: sourceDigest, + tree_digest: treeDigest, + }; +} + +function parseDiagnostics(value: unknown): SkillDiagnostic[] { + return boundedArray( + value, + "Skill diagnostics", + SKILL_API_LIMITS.maxDiagnostics, + ).map((diagnostic) => { + const record = strictObject( + diagnostic, + [ + "severity", + "code", + "message", + "source", + ], + "Skill diagnostic", + ["source"], + ); + return { + severity: diagnosticSeverity(record.severity), + code: boundedString( + record.code, + "Skill diagnostic code", + SKILL_API_LIMITS.maxLabelBytes, + false, + ), + message: boundedString( + record.message, + "Skill diagnostic message", + SKILL_API_LIMITS.maxLabelBytes, + false, + ), + source: optionalBoundedString( + record.source, + "Skill diagnostic source", + SKILL_API_LIMITS.maxPathBytes, + ), + }; + }); +} + +function parseResource(value: unknown): SkillResourceRef { + const record = strictObject( + value, + [ + "kind", + "name", + "supported", + "diagnostic", + ], + "Skill resource", + ["diagnostic"], + ); + if (typeof record.supported !== "boolean") { + throw contractError("Skill resource supported must be a boolean"); + } + const name = boundedString( + record.name, + "Skill resource name", + SKILL_API_LIMITS.maxPathBytes, + false, + ); + if (!isVirtualPath(name)) { + throw contractError("invalid Skill resource virtual path"); + } + return { + kind: boundedString( + record.kind, + "Skill resource kind", + SKILL_API_LIMITS.maxLabelBytes, + false, + ), + name, + supported: record.supported, + diagnostic: optionalBoundedString( + record.diagnostic, + "Skill resource diagnostic", + SKILL_API_LIMITS.maxLabelBytes, + ), + }; +} + +function sourceKind(value: unknown): SkillSourceKind { + if (value === "builtin" || value === "workspace") return value; + throw contractError("unsupported Skill provenance kind"); +} + +function diagnosticSeverity(value: unknown): SkillDiagnosticSeverity { + if (value === "error" || value === "warning") return value; + throw contractError("unsupported Skill diagnostic severity"); +} + +function activationStatus(value: unknown): SkillActivationStatus { + if (value === "active" || value === "inactive") return value; + throw contractError("unsupported Skill activation status"); +} + +function projectionStatus(value: unknown): SkillProjectionStatus { + if (value === "valid" || value === "invalid") return value; + throw contractError("unsupported Skill projection status"); +} + +function safeInteger(value: unknown, label: string): number { + if ( + typeof value !== "number" || !Number.isSafeInteger(value) || value < 0 || + value > SKILL_API_LIMITS.maxSafeInteger + ) { + throw contractError(`${label} must be a non-negative safe integer`); + } + return value; +} + +function boundedArray( + value: unknown, + label: string, + limit: number, +): unknown[] { + if (!Array.isArray(value) || value.length > limit) { + throw contractError(`${label} must be a bounded array`); + } + return value; +} + +function optionalBoundedString( + value: unknown, + label: string, + limit: number, +): string | undefined { + return value === undefined + ? undefined + : boundedString(value, label, limit, false); +} + +function boundedString( + value: unknown, + label: string, + limit: number, + allowEmpty: boolean, +): string { + if ( + typeof value !== "string" || (!allowEmpty && value.length === 0) || + new TextEncoder().encode(value).length > limit + ) { + throw contractError(`${label} must be a bounded string`); + } + return value; +} + +function isVirtualPath(value: string): boolean { + return !value.startsWith("/") && !value.includes("\\") && + value.split("/").every((part) => + part !== "" && part !== "." && part !== ".." + ); +} + +function strictObject( + value: unknown, + allowedKeys: readonly string[], + label: string, + optionalKeys: readonly string[] = [], +): Record { + if (typeof value !== "object" || value === null || Array.isArray(value)) { + throw contractError(`${label} must be an object`); + } + const record = value as Record; + const allowed = new Set(allowedKeys); + if (Object.keys(record).some((key) => !allowed.has(key))) { + throw contractError(`${label} contains unknown fields`); + } + const optional = new Set(optionalKeys); + if (allowedKeys.some((key) => !optional.has(key) && !(key in record))) { + throw contractError(`${label} is missing required fields`); + } + return record; +} + +function contractError(message: string): SkillApiContractError { + return new SkillApiContractError(message.slice(0, 256)); +} diff --git a/web/workspace/tests/skill-api.test.ts b/web/workspace/tests/skill-api.test.ts new file mode 100644 index 00000000..e8e6a610 --- /dev/null +++ b/web/workspace/tests/skill-api.test.ts @@ -0,0 +1,221 @@ +import { + parseSkillCatalogResponse, + parseSkillDetailResponse, + SkillApiContractError, +} from "../src/lib/workspace/skills/api.ts"; +import { SKILL_API_LIMITS } from "../src/lib/generated/skill-api.ts"; + +declare const Deno: { + test(name: string, fn: () => Promise | void): void; +}; + +function assert(condition: unknown, message: string): asserts condition { + if (!condition) throw new Error(message); +} + +function assertEquals(actual: T, expected: T): void { + if (JSON.stringify(actual) !== JSON.stringify(expected)) { + throw new Error( + `expected ${JSON.stringify(expected)}, got ${JSON.stringify(actual)}`, + ); + } +} + +function assertContractError(value: () => unknown, expected: string): void { + try { + value(); + } catch (error) { + assert( + error instanceof SkillApiContractError, + "expected SkillApiContractError", + ); + assert( + error.message.includes(expected), + `expected bounded diagnostic containing ${expected}, got ${error.message}`, + ); + assert(error.message.length <= 256, "diagnostic must remain bounded"); + return; + } + throw new Error("expected parser to reject malformed Skill response"); +} + +function builtinProvenance() { + return { + kind: "builtin", + id: "builtin:errors", + virtual_path: "skills/errors/SKILL.md", + source_digest: "builtin-source-digest", + }; +} + +function workspaceProvenance() { + return { + kind: "workspace", + id: "workspace:release", + virtual_path: "skills/release/SKILL.md", + revision: 42, + source_digest: "workspace-source-digest", + tree_digest: "tree-digest", + }; +} + +function catalogFixture(): Record { + return { + authority: "workspace-config-skills-v1", + projection: { config_revision: 42, tree_digest: "tree-digest" }, + entries: [{ + name: "errors", + description: "Builtin guidance", + activation_status: "active", + projection_status: "valid", + provenance: builtinProvenance(), + overrides: [], + diagnostics: [], + }, { + name: "release", + description: "Workspace guidance", + activation_status: "inactive", + projection_status: "invalid", + provenance: workspaceProvenance(), + overrides: [builtinProvenance()], + diagnostics: [{ + severity: "error", + code: "invalid_projection", + message: "invalid projected Skill", + source: "workspace:release", + }], + }], + diagnostics: [], + }; +} + +function detailFixture(): Record { + return { + authority: "workspace-config-skills-v1", + projection: { config_revision: 42, tree_digest: "tree-digest" }, + name: "release", + description: "Workspace guidance", + provenance: workspaceProvenance(), + overrides: [], + diagnostics: [], + activation_status: "active", + projection_status: "valid", + body: "# Release\n", + allowed_tools: ["Bash"], + allowed_tools_status: "experimental_hint_only", + resources: [{ + kind: "reference", + name: "skills/release/references/checklist.md", + supported: true, + }], + }; +} + +Deno.test("Skill catalog parser accepts generated builtin, Workspace, and invalid projection shapes", () => { + const parsed = parseSkillCatalogResponse(catalogFixture()); + assertEquals(parsed.entries.length, 2); + assertEquals(parsed.entries[0].provenance.kind, "builtin"); + assertEquals(parsed.entries[1].activation_status, "inactive"); + assertEquals(parsed.entries[1].projection_status, "invalid"); + assertEquals(parsed.projection.config_revision, 42); +}); + +Deno.test("Skill detail parser preserves shared generated DTO fields", () => { + const parsed = parseSkillDetailResponse(detailFixture()); + assertEquals(parsed.name, "release"); + assertEquals(parsed.allowed_tools, ["Bash"]); + assertEquals(parsed.resources[0].supported, true); +}); + +Deno.test("Skill parser rejects stale Workspace projection revision and digest", () => { + const staleRevision = catalogFixture(); + (staleRevision.projection as Record).config_revision = 43; + assertContractError( + () => parseSkillCatalogResponse(staleRevision), + "stale Workspace Skill projection", + ); + + const staleDigest = catalogFixture(); + (staleDigest.projection as Record).tree_digest = "new-tree"; + assertContractError( + () => parseSkillCatalogResponse(staleDigest), + "stale Workspace Skill projection", + ); +}); + +Deno.test("Skill parser fails closed on unknown fields and newer enum values", () => { + const unknownField = catalogFixture(); + unknownField.unexpected = true; + assertContractError( + () => parseSkillCatalogResponse(unknownField), + "unknown fields", + ); + + const newerProvenance = catalogFixture(); + const entries = newerProvenance.entries as Record[]; + (entries[0].provenance as Record).kind = "remote_catalog"; + assertContractError( + () => parseSkillCatalogResponse(newerProvenance), + "unsupported Skill provenance kind", + ); + + const newerStatus = catalogFixture(); + const newerEntries = newerStatus.entries as Record[]; + newerEntries[0].projection_status = "stale"; + assertContractError( + () => parseSkillCatalogResponse(newerStatus), + "unsupported Skill projection status", + ); +}); + +Deno.test("Skill parser rejects unsafe revisions and oversized collections or strings", () => { + const unsafeRevision = catalogFixture(); + (unsafeRevision.projection as Record).config_revision = + Number.MAX_SAFE_INTEGER + 1; + assertContractError( + () => parseSkillCatalogResponse(unsafeRevision), + "safe integer", + ); + + const oversizedCatalog = catalogFixture(); + const firstEntry = (oversizedCatalog.entries as unknown[])[0]; + oversizedCatalog.entries = Array.from( + { length: SKILL_API_LIMITS.maxCatalogEntries + 1 }, + () => firstEntry, + ); + assertContractError( + () => parseSkillCatalogResponse(oversizedCatalog), + "bounded array", + ); + + const oversizedDetail = detailFixture(); + oversizedDetail.body = "x".repeat(SKILL_API_LIMITS.maxBodyBytes + 1); + assertContractError( + () => parseSkillDetailResponse(oversizedDetail), + "bounded string", + ); +}); + +Deno.test("Skill parser diagnostics never include rejected Skill body content", () => { + const secret = "SENSITIVE-SKILL-BODY-CONTENT"; + const malformed = detailFixture(); + malformed.body = secret; + malformed.provenance = { + ...workspaceProvenance(), + kind: "newer_source_kind", + }; + try { + parseSkillDetailResponse(malformed); + throw new Error("expected malformed provenance to fail"); + } catch (error) { + assert( + error instanceof SkillApiContractError, + "expected SkillApiContractError", + ); + assert( + !error.message.includes(secret), + "diagnostic leaked Skill body content", + ); + assert(error.message.length <= 256, "diagnostic must remain bounded"); + } +}); From 3d66247e1129369d4f054a0e33a167e644c563c4 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 13:36:36 +0900 Subject: [PATCH 07/22] test: include Skill API contract suite --- web/workspace/deno.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/web/workspace/deno.json b/web/workspace/deno.json index 8714b347..a8b34b18 100644 --- a/web/workspace/deno.json +++ b/web/workspace/deno.json @@ -6,7 +6,7 @@ "dev": "deno run -A npm:vite@7.2.7 dev", "dev:backend": "cd ../.. && cargo run -p yoi-workspace-server --bin yoi-server -- serve --listen 127.0.0.1:8787", "check": "deno run -A npm:@sveltejs/kit@2.49.4 sync && deno run -A npm:svelte-check@4.3.4 --tsconfig ./tsconfig.json", - "test": "deno test --allow-read=src,test,tests --allow-env=LOG,VSCODE_TEXTMATE_DEBUG,NODE_ENV tests/workspace-model.test.ts tests/workspace-catalog.test.ts tests/profile-api.test.ts src/lib/workspace/auth/model.test.ts tests/auth-api.test.ts src/lib/workspace/api/http.test.ts src/lib/workspace/api/workers.test.ts src/lib/workspace/header/breadcrumb-model.test.ts src/lib/workspace/console/chat-submit.test.ts test/composer-history.test.ts tests/composer-paste.test.ts src/lib/workspace/console/composer-command.test.ts src/lib/workspace/console/composer-draft.test.ts src/lib/workspace/console/composer-completion.test.ts src/lib/workspace/console/markdown.test.ts test/console/ansi.test.ts src/lib/workspace/console/model.test.ts src/lib/workspace/companion/api.test.ts tests/workdir-api.test.ts src/lib/workspace/console/tasks.test.ts test/ticket-detail-route-reuse.test.ts test/repositories/ui.test.ts src/lib/workspace/console/worker-console.ui.test.ts src/lib/workspace/settings/model.test.ts src/lib/workspace/sidebar/override-stack.test.ts src/lib/workspace/sidebar/workers.test.ts src/lib/workspace/sidebar/workspace-switcher.test.ts src/lib/workspace/sidebar/worker-subscription.test.ts src/lib/workspace/sidebar/worker-launch.test.ts test/sidebar/worker-actions.test.ts src/lib/workspace/tickets/merge-request-resources.test.ts src/lib/workspace/tickets/ticket-panel.test.ts test/merge-request-status.test.ts test/config-source/decodal-grammar.test.ts test/config-source/editor-state.test.ts test/config-source/fixed-schema-wrapper.test.ts test/config-source/toolchain.test.ts test/config-source/wasm-parity.test.ts test/repository-access/api.test.ts test/repository-access/loader.test.ts test/repository-access/ui.test.ts", + "test": "deno test --allow-read=src,test,tests --allow-env=LOG,VSCODE_TEXTMATE_DEBUG,NODE_ENV tests/workspace-model.test.ts tests/workspace-catalog.test.ts tests/profile-api.test.ts tests/skill-api.test.ts src/lib/workspace/auth/model.test.ts tests/auth-api.test.ts src/lib/workspace/api/http.test.ts src/lib/workspace/api/workers.test.ts src/lib/workspace/header/breadcrumb-model.test.ts src/lib/workspace/console/chat-submit.test.ts test/composer-history.test.ts tests/composer-paste.test.ts src/lib/workspace/console/composer-command.test.ts src/lib/workspace/console/composer-draft.test.ts src/lib/workspace/console/composer-completion.test.ts src/lib/workspace/console/markdown.test.ts test/console/ansi.test.ts src/lib/workspace/console/model.test.ts src/lib/workspace/companion/api.test.ts tests/workdir-api.test.ts src/lib/workspace/console/tasks.test.ts test/ticket-detail-route-reuse.test.ts test/repositories/ui.test.ts src/lib/workspace/console/worker-console.ui.test.ts src/lib/workspace/settings/model.test.ts src/lib/workspace/sidebar/override-stack.test.ts src/lib/workspace/sidebar/workers.test.ts src/lib/workspace/sidebar/workspace-switcher.test.ts src/lib/workspace/sidebar/worker-subscription.test.ts src/lib/workspace/sidebar/worker-launch.test.ts test/sidebar/worker-actions.test.ts src/lib/workspace/tickets/merge-request-resources.test.ts src/lib/workspace/tickets/ticket-panel.test.ts test/merge-request-status.test.ts test/config-source/decodal-grammar.test.ts test/config-source/editor-state.test.ts test/config-source/fixed-schema-wrapper.test.ts test/config-source/toolchain.test.ts test/config-source/wasm-parity.test.ts test/repository-access/api.test.ts test/repository-access/loader.test.ts test/repository-access/ui.test.ts", "build": "deno run -A npm:vite@7.2.7 build", "preview": "deno run -A npm:vite@7.2.7 preview" }, From 9bd08a3a5bdbd5052090533f53a407dff611e640 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 13:55:13 +0900 Subject: [PATCH 08/22] fix: bound Skill API response handling --- crates/workspace-api/src/lib.rs | 3 +- web/workspace/src/lib/generated/skill-api.ts | 1 + .../src/lib/workspace/api/http.test.ts | 43 ++++++++++ web/workspace/src/lib/workspace/api/http.ts | 81 ++++++++++++++++++- 4 files changed, 126 insertions(+), 2 deletions(-) diff --git a/crates/workspace-api/src/lib.rs b/crates/workspace-api/src/lib.rs index e9b16045..b77db77a 100644 --- a/crates/workspace-api/src/lib.rs +++ b/crates/workspace-api/src/lib.rs @@ -1946,6 +1946,7 @@ pub const SKILL_API_MAX_LABEL_BYTES: usize = 4_096; pub const SKILL_API_MAX_BODY_BYTES: usize = 1_048_576; pub const SKILL_API_MAX_PATH_BYTES: usize = 1_024; pub const SKILL_API_MAX_DIGEST_BYTES: usize = 128; +pub const SKILL_API_MAX_RESPONSE_BYTES: usize = 2_097_152; #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] #[cfg_attr(feature = "typescript", derive(ts_rs::TS))] @@ -2425,7 +2426,7 @@ pub fn skill_api_typescript() -> String { SkillDetailResponse::decl(&config), ]; let limits = format!( - "export const SKILL_API_AUTHORITY = \"{SKILL_CATALOG_AUTHORITY}\" as const;\n\nexport const SKILL_API_LIMITS = {{\n maxSafeInteger: {SKILL_API_MAX_SAFE_INTEGER},\n maxCatalogEntries: {SKILL_API_MAX_CATALOG_ENTRIES},\n maxOverrides: {SKILL_API_MAX_OVERRIDES},\n maxDiagnostics: {SKILL_API_MAX_DIAGNOSTICS},\n maxResources: {SKILL_API_MAX_RESOURCES},\n maxAllowedTools: {SKILL_API_MAX_ALLOWED_TOOLS},\n maxNameBytes: {SKILL_API_MAX_NAME_BYTES},\n maxLabelBytes: {SKILL_API_MAX_LABEL_BYTES},\n maxBodyBytes: {SKILL_API_MAX_BODY_BYTES},\n maxPathBytes: {SKILL_API_MAX_PATH_BYTES},\n maxDigestBytes: {SKILL_API_MAX_DIGEST_BYTES},\n}} as const;" + "export const SKILL_API_AUTHORITY = \"{SKILL_CATALOG_AUTHORITY}\" as const;\n\nexport const SKILL_API_LIMITS = {{\n maxSafeInteger: {SKILL_API_MAX_SAFE_INTEGER},\n maxCatalogEntries: {SKILL_API_MAX_CATALOG_ENTRIES},\n maxOverrides: {SKILL_API_MAX_OVERRIDES},\n maxDiagnostics: {SKILL_API_MAX_DIAGNOSTICS},\n maxResources: {SKILL_API_MAX_RESOURCES},\n maxAllowedTools: {SKILL_API_MAX_ALLOWED_TOOLS},\n maxNameBytes: {SKILL_API_MAX_NAME_BYTES},\n maxLabelBytes: {SKILL_API_MAX_LABEL_BYTES},\n maxBodyBytes: {SKILL_API_MAX_BODY_BYTES},\n maxPathBytes: {SKILL_API_MAX_PATH_BYTES},\n maxDigestBytes: {SKILL_API_MAX_DIGEST_BYTES},\n maxResponseBytes: {SKILL_API_MAX_RESPONSE_BYTES},\n}} as const;" ); format!( "// Generated from workspace-api. Do not edit by hand.\n// Regenerate: cargo run -q -p workspace-api --features typescript --example generate_skill_api_types > web/workspace/src/lib/generated/skill-api.ts\n\n{limits}\n\n{}\n", diff --git a/web/workspace/src/lib/generated/skill-api.ts b/web/workspace/src/lib/generated/skill-api.ts index 5847d06e..b7812b03 100644 --- a/web/workspace/src/lib/generated/skill-api.ts +++ b/web/workspace/src/lib/generated/skill-api.ts @@ -15,6 +15,7 @@ export const SKILL_API_LIMITS = { maxBodyBytes: 1048576, maxPathBytes: 1024, maxDigestBytes: 128, + maxResponseBytes: 2097152, } as const; export type SkillDiagnosticSeverity = "error" | "warning"; diff --git a/web/workspace/src/lib/workspace/api/http.test.ts b/web/workspace/src/lib/workspace/api/http.test.ts index f7a8419f..89b1a5e0 100644 --- a/web/workspace/src/lib/workspace/api/http.test.ts +++ b/web/workspace/src/lib/workspace/api/http.test.ts @@ -1,11 +1,13 @@ import { loadWorkspaceSkillCatalog, + loadWorkspaceSkillDetail, workspaceApiPath, workspaceRoute, workspaceSkillActivationPath, workspaceSkillCatalogPath, workspaceSkillDetailPath, } from "./http.ts"; +import { SKILL_API_LIMITS } from "$lib/generated/skill-api.ts"; declare const Deno: { test(name: string, fn: () => Promise | void): void; @@ -123,3 +125,44 @@ Deno.test("loadWorkspaceSkillCatalog fetches lightweight catalog", async () => { assertEquals(result.data?.entries[0].name, "triage-errors"); assertEquals(JSON.stringify(result.data).includes("SKILL.md body"), false); }); + +Deno.test("Skill loaders redact and bound non-success response diagnostics", async () => { + const secret = "SENSITIVE-SKILL-BODY-CONTENT".repeat(300); + const result = await loadWorkspaceSkillCatalog( + (() => + Promise.resolve(new Response(secret, { status: 500 }))) as typeof fetch, + "ws-1", + ); + + assertEquals(result.data, null); + assertEquals(result.error, "Skill API request failed with HTTP 500"); + assert( + !result.error?.includes(secret.slice(0, 64)), + "Skill API diagnostic must not expose response body content", + ); + assert( + (result.error?.length ?? 0) <= 256, + "Skill API diagnostic must remain bounded", + ); +}); + +Deno.test("Skill loaders stop reading success responses above the wire byte limit", async () => { + const oversized = `{"body":"${ + "x".repeat(SKILL_API_LIMITS.maxResponseBytes + 1) + }"}`; + const result = await loadWorkspaceSkillDetail( + (() => + Promise.resolve( + new Response(oversized, { status: 200 }), + )) as typeof fetch, + "ws-1", + "release", + ); + + assertEquals(result.data, null); + assertEquals(result.error, "Skill API response exceeds its byte limit"); + assert( + !result.error?.includes(oversized.slice(0, 64)), + "Skill API diagnostic must not expose oversized response content", + ); +}); diff --git a/web/workspace/src/lib/workspace/api/http.ts b/web/workspace/src/lib/workspace/api/http.ts index c86beb48..2f207fac 100644 --- a/web/workspace/src/lib/workspace/api/http.ts +++ b/web/workspace/src/lib/workspace/api/http.ts @@ -1,3 +1,4 @@ +import { SKILL_API_LIMITS } from "$lib/generated/skill-api.ts"; import type { SkillCatalogResponse, SkillDetailResponse, @@ -5,6 +6,7 @@ import type { import { parseSkillCatalogResponse, parseSkillDetailResponse, + SkillApiContractError, } from "$lib/workspace/skills/api.ts"; export type ApiResult = { @@ -14,6 +16,18 @@ export type ApiResult = { export type { SkillCatalogResponse, SkillDetailResponse }; +type JsonLoadPolicy = { + diagnosticLabel: string; + maxResponseBytes: number; +}; + +const SKILL_API_LOAD_POLICY: JsonLoadPolicy = { + diagnosticLabel: "Skill API", + maxResponseBytes: SKILL_API_LIMITS.maxResponseBytes, +}; + +class ResponseByteLimitError extends Error {} + function normalizePath(path: string): string { if (!path || path === "/") return ""; return path.startsWith("/") ? path : `/${path}`; @@ -57,6 +71,7 @@ export async function loadWorkspaceSkillCatalog( workspaceSkillCatalogPath(workspaceId), undefined, parseSkillCatalogResponse, + SKILL_API_LOAD_POLICY, ); } @@ -70,6 +85,7 @@ export async function loadWorkspaceSkillDetail( workspaceSkillDetailPath(workspaceId, name), undefined, parseSkillDetailResponse, + SKILL_API_LOAD_POLICY, ); } @@ -78,19 +94,38 @@ export async function loadJson( path: string, init?: RequestInit, parse: (value: unknown) => T = (value) => value as T, + policy?: JsonLoadPolicy, ): Promise> { try { const response = await fetchFn(path, init); if (!response.ok) { + if (policy) { + await response.body?.cancel(); + return { + data: null, + error: + `${policy.diagnosticLabel} request failed with HTTP ${response.status}`, + }; + } const text = await response.text(); return { data: null, error: text || `${path} request failed (${response.status})`, }; } - const payload: unknown = await response.json(); + const payload: unknown = policy + ? await readBoundedJson(response, policy.maxResponseBytes) + : await response.json(); return { data: parse(payload), error: null }; } catch (error) { + if (policy) { + const diagnostic = error instanceof SkillApiContractError + ? error.message + : error instanceof ResponseByteLimitError + ? `${policy.diagnosticLabel} response exceeds its byte limit` + : `${policy.diagnosticLabel} response is invalid`; + return { data: null, error: diagnostic.slice(0, 256) }; + } return { data: null, error: error instanceof Error ? error.message : `${path} request failed`, @@ -98,6 +133,50 @@ export async function loadJson( } } +async function readBoundedJson( + response: Response, + maxBytes: number, +): Promise { + const contentLength = response.headers.get("content-length"); + if (contentLength !== null) { + const parsedLength = Number(contentLength); + if (Number.isFinite(parsedLength) && parsedLength > maxBytes) { + await response.body?.cancel(); + throw new ResponseByteLimitError(); + } + } + + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("response body is unavailable"); + } + const chunks: Uint8Array[] = []; + let totalBytes = 0; + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + totalBytes += value.byteLength; + if (totalBytes > maxBytes) { + await reader.cancel(); + throw new ResponseByteLimitError(); + } + chunks.push(value); + } + } finally { + reader.releaseLock(); + } + + const bytes = new Uint8Array(totalBytes); + let offset = 0; + for (const chunk of chunks) { + bytes.set(chunk, offset); + offset += chunk.byteLength; + } + const text = new TextDecoder("utf-8", { fatal: true }).decode(bytes); + return JSON.parse(text) as unknown; +} + async function requireJson(response: Response, path: string): Promise { if (!response.ok) { const text = await response.text(); From af06eecfd04df503591f8e6b0f9645853021e496 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 14:50:18 +0900 Subject: [PATCH 09/22] feat: add executable feature lifecycle hooks --- crates/worker/src/controller.rs | 12 +- crates/worker/src/feature.rs | 467 +++++++++--- crates/worker/src/feature/background.rs | 705 ++++++++++++++++++ crates/worker/src/feature/builtin/task/mod.rs | 20 +- crates/worker/src/feature/mcp.rs | 3 +- crates/worker/src/feature/plugin.rs | 2 +- crates/worker/src/hook.rs | 558 +++++++++++++- crates/worker/src/internal_worker.rs | 53 +- crates/worker/src/ipc/interceptor.rs | 155 +++- crates/worker/src/permission.rs | 13 +- crates/worker/src/shutdown_after_idle.rs | 7 +- crates/worker/src/worker.rs | 209 +++++- 12 files changed, 1969 insertions(+), 235 deletions(-) create mode 100644 crates/worker/src/feature/background.rs diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 5d7ae971..5e848f2c 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -1590,7 +1590,9 @@ async fn controller_loop( &working_event_tx, target, expected_head_entries, - ) { + ) + .await + { worker.clear_in_flight_events(); shared_state.set_status(WorkerStatus::Idle); let _ = working_event_tx.send(Event::Status { @@ -1711,6 +1713,10 @@ async fn controller_loop( tracing::warn!(%error, "Worker runtime socket cleanup failed"); } + // Feature callbacks and tasks share the Worker scope. Stop them before + // Memory/Workdir teardown so they cannot observe a partially closed Worker. + worker.stop_feature_runtime("controller shutdown").await; + // Background memory jobs own extract/consolidate workers after a // turn completes. Join them before closing the Workdir session so no // Worker-owned task can outlive its operation attachment. @@ -1993,7 +1999,7 @@ where } } -fn apply_rewind( +async fn apply_rewind( worker: &mut Worker, working_event_tx: &broadcast::Sender, target: RewindTargetId, @@ -2003,7 +2009,7 @@ where C: LlmClient + 'static, St: Store, { - match worker.rewind_to(target, expected_head_entries) { + match worker.rewind_to(target, expected_head_entries).await { Ok(applied) => { let session = session_store::public_snapshot::project_current_session_snapshot(&applied.entries); diff --git a/crates/worker/src/feature.rs b/crates/worker/src/feature.rs index 5698e84b..c13bb603 100644 --- a/crates/worker/src/feature.rs +++ b/crates/worker/src/feature.rs @@ -23,7 +23,14 @@ use agen::tool::ToolDefinition; use serde::{Deserialize, Serialize}; use thiserror::Error; -use crate::hook::{Hook, HookRegistryBuilder, OnTurnEnd, PostToolCall, PreLlmRequest, PreToolCall}; +use crate::hook::{ + BeforeSessionRewrite, Hook, HookFailurePolicy, HookRegistryBuilder, OnPromptSubmit, OnTurnEnd, + PostToolCall, PreLlmRequest, PreToolCall, RunCommitted, RunExit, WorkerStopping, +}; +use background::{ + BackgroundTaskSpec, FeatureBackgroundTask, FeatureBackgroundTaskRegistry, + FeatureBackgroundTaskRegistryBuilder, +}; /// Stable source-qualified identifier for a feature module. #[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)] @@ -253,10 +260,15 @@ pub enum FeatureRuntimeKind { #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum FeatureHookPoint { - PreRequest, + PromptSubmit, + PreLlmRequest, PreToolCall, - ToolResult, - TurnEnd, + PostToolCall, + AssistantTurnEnd, + RunExit, + RunCommitted, + BeforeSessionRewrite, + WorkerStopping, } /// Serializable declaration of a tool contribution. The executable factory is @@ -379,16 +391,17 @@ impl FeatureInstructionContribution { } } -/// Background task lifecycle phase represented by this registry slice. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +/// Background tasks are always Worker-managed and execute inside the owning +/// feature scope. Report-only and detached host-managed declarations are not +/// accepted by the current contract. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum BackgroundTaskLifecycle { - DescriptorOnly, - HostManaged, + WorkerManaged, } -/// Declaration for a feature-provided background task. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +/// Declaration for a feature-provided executable background task. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] pub struct BackgroundTaskDeclaration { pub name: String, pub description: String, @@ -396,11 +409,11 @@ pub struct BackgroundTaskDeclaration { } impl BackgroundTaskDeclaration { - pub fn descriptor_only(name: impl Into, description: impl Into) -> Self { + pub fn worker_managed(name: impl Into, description: impl Into) -> Self { Self { name: name.into(), description: description.into(), - lifecycle: BackgroundTaskLifecycle::DescriptorOnly, + lifecycle: BackgroundTaskLifecycle::WorkerManaged, } } } @@ -772,6 +785,15 @@ impl FeatureInstallReport { } } + fn clear_installed_contributions(&mut self) { + self.installed = false; + self.installed_tools.clear(); + self.installed_hooks.clear(); + self.installed_instructions.clear(); + self.declared_background_tasks.clear(); + self.provided_services.clear(); + } + fn mark_skipped( &mut self, kind: FeatureContributionKind, @@ -1042,15 +1064,68 @@ impl HookContributionRegistrar<'_> { )) } + fn record(&mut self, declaration: HookDeclaration) { + if !self.report.installed_hooks.contains(&declaration) { + self.report.installed_hooks.push(declaration); + } + } + + pub fn add_prompt_submit( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + hook: impl Hook + 'static, + ) -> Result<(), FeatureInstallError> { + let declaration = HookDeclaration::new(name, FeatureHookPoint::PromptSubmit); + self.require_declared(&declaration)?; + self.hook_builder.add_named_on_prompt_submit( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ); + self.record(declaration); + Ok(()) + } + + pub fn add_pre_llm_request( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + hook: impl Hook + 'static, + ) -> Result<(), FeatureInstallError> { + let declaration = HookDeclaration::new(name, FeatureHookPoint::PreLlmRequest); + self.require_declared(&declaration)?; + self.hook_builder.add_named_pre_llm_request( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ); + self.record(declaration); + Ok(()) + } + pub fn add_pre_request( &mut self, name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - let declaration = HookDeclaration::new(name, FeatureHookPoint::PreRequest); + self.add_pre_llm_request(name, HookFailurePolicy::FailClosed, hook) + } + + pub fn add_pre_tool_call_with_policy( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + hook: impl Hook + 'static, + ) -> Result<(), FeatureInstallError> { + let declaration = HookDeclaration::new(name, FeatureHookPoint::PreToolCall); self.require_declared(&declaration)?; - self.hook_builder.add_pre_llm_request(hook); - self.report.installed_hooks.push(declaration); + self.hook_builder.add_named_pre_tool_call( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ); + self.record(declaration); Ok(()) } @@ -1059,10 +1134,23 @@ impl HookContributionRegistrar<'_> { name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - let declaration = HookDeclaration::new(name, FeatureHookPoint::PreToolCall); + self.add_pre_tool_call_with_policy(name, HookFailurePolicy::FailClosed, hook) + } + + pub fn add_post_tool_call( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + hook: impl Hook + 'static, + ) -> Result<(), FeatureInstallError> { + let declaration = HookDeclaration::new(name, FeatureHookPoint::PostToolCall); self.require_declared(&declaration)?; - self.hook_builder.add_pre_tool_call(hook); - self.report.installed_hooks.push(declaration); + self.hook_builder.add_named_post_tool_call( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ); + self.record(declaration); Ok(()) } @@ -1071,10 +1159,23 @@ impl HookContributionRegistrar<'_> { name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - let declaration = HookDeclaration::new(name, FeatureHookPoint::ToolResult); + self.add_post_tool_call(name, HookFailurePolicy::FailClosed, hook) + } + + pub fn add_assistant_turn_end( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + hook: impl Hook + 'static, + ) -> Result<(), FeatureInstallError> { + let declaration = HookDeclaration::new(name, FeatureHookPoint::AssistantTurnEnd); self.require_declared(&declaration)?; - self.hook_builder.add_post_tool_call(hook); - self.report.installed_hooks.push(declaration); + self.hook_builder.add_named_on_turn_end( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ); + self.record(declaration); Ok(()) } @@ -1083,10 +1184,74 @@ impl HookContributionRegistrar<'_> { name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - let declaration = HookDeclaration::new(name, FeatureHookPoint::TurnEnd); + self.add_assistant_turn_end(name, HookFailurePolicy::FailClosed, hook) + } + + pub fn add_run_exit( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + hook: impl Hook + 'static, + ) -> Result<(), FeatureInstallError> { + let declaration = HookDeclaration::new(name, FeatureHookPoint::RunExit); self.require_declared(&declaration)?; - self.hook_builder.add_on_turn_end(hook); - self.report.installed_hooks.push(declaration); + self.hook_builder.add_named_run_exit( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ); + self.record(declaration); + Ok(()) + } + + pub fn add_run_committed( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + hook: impl Hook + 'static, + ) -> Result<(), FeatureInstallError> { + let declaration = HookDeclaration::new(name, FeatureHookPoint::RunCommitted); + self.require_declared(&declaration)?; + self.hook_builder.add_named_run_committed( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ); + self.record(declaration); + Ok(()) + } + + pub fn add_before_session_rewrite( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + hook: impl Hook + 'static, + ) -> Result<(), FeatureInstallError> { + let declaration = HookDeclaration::new(name, FeatureHookPoint::BeforeSessionRewrite); + self.require_declared(&declaration)?; + self.hook_builder.add_named_before_session_rewrite( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ); + self.record(declaration); + Ok(()) + } + + pub fn add_worker_stopping( + &mut self, + name: impl Into, + policy: HookFailurePolicy, + hook: impl Hook + 'static, + ) -> Result<(), FeatureInstallError> { + let declaration = HookDeclaration::new(name, FeatureHookPoint::WorkerStopping); + self.require_declared(&declaration)?; + self.hook_builder.add_named_worker_stopping( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ); + self.record(declaration); Ok(()) } } @@ -1124,33 +1289,40 @@ impl FeatureInstructionRegistrar<'_> { } } -/// Background task registrar for descriptor/report-only contributions. +/// Registrar for executable, Worker-managed background task contributions. pub struct BackgroundTaskRegistrar<'a> { feature_id: &'a FeatureId, declarations: &'a FeatureContributionDeclarations, + registry: &'a mut FeatureBackgroundTaskRegistryBuilder, report: &'a mut FeatureInstallReport, } impl BackgroundTaskRegistrar<'_> { - pub fn declare( + pub fn register( &mut self, - declaration: BackgroundTaskDeclaration, + spec: BackgroundTaskSpec, + task: impl FeatureBackgroundTask + 'static, ) -> Result<(), FeatureInstallError> { - if !self.declarations.contains_background_task(&declaration) { + if !self + .declarations + .contains_background_task(&spec.declaration) + { return Err(reject_undeclared_contribution( self.feature_id, self.report, FeatureContributionKind::BackgroundTask, - declaration.name, + spec.declaration.name, )); } + self.registry + .register(self.feature_id.clone(), spec.clone(), task)?; if !self .report .declared_background_tasks .iter() - .any(|task| task.name == declaration.name) + .any(|task| task.name == spec.declaration.name) { - self.report.declared_background_tasks.push(declaration); + self.report.declared_background_tasks.push(spec.declaration); } Ok(()) } @@ -1330,15 +1502,17 @@ impl ProtocolProviderRegistrar<'_> { } } - for task in background_tasks { - if !self - .report - .declared_background_tasks - .iter() - .any(|declared| declared.name == task.name) - { - self.report.declared_background_tasks.push(task); - } + if let Some(task) = background_tasks.first() { + let reason = format!( + "protocol provider background task `{}` has no executable Worker-managed handler", + task.name + ); + self.report.mark_skipped( + FeatureContributionKind::BackgroundTask, + task.name.clone(), + reason.clone(), + ); + return Err(FeatureInstallError::InvalidDescriptor(reason)); } Ok(()) @@ -1352,6 +1526,7 @@ pub struct FeatureInstallContext<'a> { pending_tools: &'a mut Vec, installed_tool_names: &'a mut HashMap, hook_builder: &'a mut HookRegistryBuilder, + background_task_builder: &'a mut FeatureBackgroundTaskRegistryBuilder, service_registry: &'a mut FeatureServiceRegistry, report: &'a mut FeatureInstallReport, } @@ -1392,6 +1567,7 @@ impl FeatureInstallContext<'_> { BackgroundTaskRegistrar { feature_id: self.feature_id, declarations: self.declarations, + registry: self.background_task_builder, report: self.report, } } @@ -1440,6 +1616,7 @@ impl FeatureInstallContext<'_> { pub struct FeatureRegistryInstallReport { pub reports: Vec, pub services: FeatureServiceRegistry, + pub background_tasks: FeatureBackgroundTaskRegistry, pub plan_error: Option, } @@ -1861,12 +2038,16 @@ impl FeatureRegistryBuilder { return FeatureRegistryInstallReport { reports, services: FeatureServiceRegistry::default(), + background_tasks: FeatureBackgroundTaskRegistry::default(), plan_error: Some(error), }; } }; let mut service_registry = FeatureServiceRegistry::default(); + let mut background_task_builder = FeatureBackgroundTaskRegistryBuilder::default(); let mut reports = Vec::with_capacity(plan.ordered_indices.len()); + let install_hook_checkpoint = hook_builder.checkpoint(); + let install_tool_checkpoint = pending_tools.len(); let mut modules = self.modules.into_iter().map(Some).collect::>(); let ordered_modules = plan .ordered_indices @@ -1884,6 +2065,11 @@ impl FeatureRegistryBuilder { for (module, descriptor) in ordered_modules { let declarations = FeatureContributionDeclarations::from_descriptor(&descriptor); let mut report = FeatureInstallReport::new(&descriptor); + let hook_checkpoint = hook_builder.checkpoint(); + let background_checkpoint = background_task_builder.checkpoint(); + let service_checkpoint = service_registry.clone(); + let tool_checkpoint = pending_tools.len(); + let installed_tool_checkpoint = installed_tool_names.clone(); let mut required_service_failed = false; for requirement in descriptor.requires_services.iter().cloned() { @@ -1920,10 +2106,6 @@ impl FeatureRegistryBuilder { continue; } - for background_task in descriptor.background_tasks.iter().cloned() { - report.declared_background_tasks.push(background_task); - } - let install_result = { let mut context = FeatureInstallContext { feature_id: &descriptor.id, @@ -1931,6 +2113,7 @@ impl FeatureRegistryBuilder { pending_tools, installed_tool_names: &mut installed_tool_names, hook_builder, + background_task_builder: &mut background_task_builder, service_registry: &mut service_registry, report: &mut report, }; @@ -1940,18 +2123,81 @@ impl FeatureRegistryBuilder { match install_result { Ok(()) => report.installed = true, Err(error) => { + hook_builder.rollback_to(hook_checkpoint); + background_task_builder.rollback_to(&background_checkpoint); + service_registry = service_checkpoint.clone(); + pending_tools.truncate(tool_checkpoint); + installed_tool_names = installed_tool_checkpoint.clone(); + report.clear_installed_contributions(); report .diagnostics .push(FeatureDiagnostic::error(error.to_string())); } } + if report.installed { + for hook in &descriptor.hooks { + if !report.installed_hooks.contains(hook) { + report.diagnostics.push(FeatureDiagnostic::error(format!( + "feature `{}` declared hook `{}` at {:?} but did not register it", + descriptor.id, hook.name, hook.point + ))); + } + } + for task in &descriptor.background_tasks { + if !report.declared_background_tasks.contains(task) { + report.diagnostics.push(FeatureDiagnostic::error(format!( + "feature `{}` declared background task `{}` but did not register an executable handler", + descriptor.id, task.name + ))); + } + } + if report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.severity == FeatureDiagnosticSeverity::Error) + { + hook_builder.rollback_to(hook_checkpoint); + background_task_builder.rollback_to(&background_checkpoint); + service_registry = service_checkpoint.clone(); + pending_tools.truncate(tool_checkpoint); + installed_tool_names = installed_tool_checkpoint.clone(); + report.clear_installed_contributions(); + report.clear_installed_contributions(); + } + } reports.push(report); } - FeatureRegistryInstallReport { - reports, - services: service_registry, - plan_error: None, + let failed = reports.iter().any(|report| { + report + .diagnostics + .iter() + .any(|diagnostic| diagnostic.severity == FeatureDiagnosticSeverity::Error) + }); + if failed { + hook_builder.rollback_to(install_hook_checkpoint); + pending_tools.truncate(install_tool_checkpoint); + for report in &mut reports { + if report.installed { + report.clear_installed_contributions(); + report.diagnostics.push(FeatureDiagnostic::warning( + "feature scope rolled back because another contribution failed", + )); + } + } + FeatureRegistryInstallReport { + reports, + services: FeatureServiceRegistry::default(), + background_tasks: FeatureBackgroundTaskRegistry::default(), + plan_error: None, + } + } else { + FeatureRegistryInstallReport { + reports, + services: service_registry, + background_tasks: background_task_builder.build(), + plan_error: None, + } } } } @@ -1996,6 +2242,7 @@ pub enum FeatureInstallError { Install(String), } +pub mod background; pub mod builtin; pub mod mcp; pub mod plugin; @@ -2398,13 +2645,9 @@ mod tests { } #[test] - fn descriptor_contributions_are_recorded() { + fn executable_contributions_are_recorded() { let descriptor = FeatureDescriptor::builtin("dummy", "Dummy") - .with_tool(ToolDeclaration::new("Dummy", "dummy tool")) - .with_background_task(BackgroundTaskDeclaration::descriptor_only( - "daily", - "descriptor-only background task", - )); + .with_tool(ToolDeclaration::new("Dummy", "dummy tool")); let mut hook_builder = HookRegistryBuilder::default(); let mut pending_tools = Vec::new(); let report = FeatureRegistryBuilder::new() @@ -2420,7 +2663,7 @@ mod tests { let feature_report = &report.reports[0]; assert!(feature_report.installed); assert_eq!(feature_report.installed_tools, vec!["Dummy"]); - assert_eq!(feature_report.declared_background_tasks[0].name, "daily"); + assert!(feature_report.declared_background_tasks.is_empty()); } #[test] @@ -2480,8 +2723,9 @@ mod tests { }) .install_into_pending(&mut pending_tools, &mut hook_builder); - assert_eq!(pending_tools.len(), 1); - assert!(report.reports[0].installed); + assert!(pending_tools.is_empty()); + assert!(!report.reports[0].installed); + assert!(report.reports[0].installed_tools.is_empty()); assert!(!report.reports[1].installed); assert!( report.reports[1] @@ -2558,7 +2802,7 @@ mod tests { "1.0.0", "startup-discovered service", )) - .with_background_task(BackgroundTaskDeclaration::descriptor_only( + .with_background_task(BackgroundTaskDeclaration::worker_managed( "provider-poller", "provider lifecycle poller", )) @@ -2568,7 +2812,7 @@ mod tests { } #[test] - fn protocol_provider_registers_startup_discovered_contributions_through_worker_path() { + fn protocol_provider_report_only_background_task_is_rejected_atomically() { let provider = ProtocolProviderDeclaration::new( ProviderId::builtin("dynamic-provider"), "test-protocol", @@ -2599,30 +2843,18 @@ mod tests { .collect(); let feature_report = &report.reports[0]; - assert!(feature_report.installed); - assert_eq!(feature_report.installed_tools, vec!["DynamicTool"]); - assert_eq!(tool_names, vec!["DynamicTool"]); + assert!(!feature_report.installed); + assert!(feature_report.installed_tools.is_empty()); + assert!(tool_names.is_empty()); assert_eq!(calls.load(Ordering::SeqCst), 1); - assert_eq!(feature_report.provided_services.len(), 1); - assert_eq!( - feature_report.provided_services[0].id, - ServiceId::builtin("dynamic-service") - ); - assert_eq!( - feature_report.declared_background_tasks[0].name, - "provider-poller" - ); + assert!(feature_report.provided_services.is_empty()); + assert!(feature_report.declared_background_tasks.is_empty()); assert_eq!(feature_report.protocol_providers.len(), 1); - assert_eq!( - feature_report.protocol_providers[0].state, - ProtocolProviderLifecycleState::Ready - ); - assert!( - feature_report - .diagnostics - .iter() - .any(|diagnostic| diagnostic.message.contains("startup discovery completed")) - ); + assert!(feature_report.diagnostics.iter().any(|diagnostic| { + diagnostic + .message + .contains("has no executable Worker-managed handler") + })); } #[test] @@ -2779,8 +3011,8 @@ mod tests { async fn call( &self, _input: &crate::hook::ToolCallSummary, - ) -> crate::hook::HookPreToolAction { - crate::hook::HookPreToolAction::Continue + ) -> Result { + Ok(crate::hook::HookPreToolAction::Continue) } } @@ -2804,6 +3036,19 @@ mod tests { } } + struct NoopBackgroundTask; + + #[async_trait] + impl FeatureBackgroundTask for NoopBackgroundTask { + async fn run( + &self, + _context: background::BackgroundTaskContext, + _cancellation: background::BackgroundTaskCancellation, + ) -> Result<(), crate::hook::HookError> { + Ok(()) + } + } + struct BackgroundFeature { descriptor: FeatureDescriptor, task_name: &'static str, @@ -2818,12 +3063,22 @@ mod tests { &self, context: &mut FeatureInstallContext<'_>, ) -> Result<(), FeatureInstallError> { - context - .background_tasks() - .declare(BackgroundTaskDeclaration::descriptor_only( - self.task_name, - "runtime background task", - )) + let declaration = self + .descriptor + .background_tasks + .iter() + .find(|task| task.name == self.task_name) + .cloned() + .unwrap_or_else(|| { + BackgroundTaskDeclaration::worker_managed( + self.task_name, + "undeclared background task", + ) + }); + context.background_tasks().register( + BackgroundTaskSpec::single_flight(declaration, std::time::Duration::from_secs(1)), + NoopBackgroundTask, + ) } } @@ -2985,25 +3240,44 @@ mod tests { ); } - #[test] - fn background_task_declaration_is_descriptor_contribution() { + #[tokio::test] + async fn executable_background_task_is_registered_in_worker_scope() { let descriptor = FeatureDescriptor::builtin("background", "Background") - .with_background_task(BackgroundTaskDeclaration::descriptor_only( + .with_background_task(BackgroundTaskDeclaration::worker_managed( "declared-task", "descriptor contribution", )); let mut hook_builder = HookRegistryBuilder::default(); let mut pending_tools = Vec::new(); let report = FeatureRegistryBuilder::new() - .with_module(ServiceFeature { descriptor }) + .with_module(BackgroundFeature { + descriptor, + task_name: "declared-task", + }) .install_into_pending(&mut pending_tools, &mut hook_builder); - assert!(report.reports[0].installed); assert_eq!( report.reports[0].declared_background_tasks[0].name, "declared-task" ); assert!(report.reports[0].skipped.is_empty()); + assert!(matches!( + report + .background_tasks + .start( + &FeatureId::builtin("background"), + "declared-task", + crate::hook::HookInvocationContext::default(), + ) + .unwrap(), + background::BackgroundTaskStart::Started { .. } + )); + report.background_tasks.shutdown().await.unwrap(); + assert!(matches!( + report.background_tasks.diagnostics()[0].outcome, + background::BackgroundTaskOutcome::Completed + | background::BackgroundTaskOutcome::Cancelled + )); } #[test] @@ -3118,7 +3392,10 @@ mod tests { assert_eq!(descriptor.runtime, FeatureRuntimeKind::Builtin); assert_eq!( hook_points, - vec![FeatureHookPoint::PreRequest, FeatureHookPoint::PreToolCall] + vec![ + FeatureHookPoint::PreLlmRequest, + FeatureHookPoint::PreToolCall + ] ); assert!(descriptor.background_tasks.is_empty()); assert!(descriptor.provides_services.is_empty()); diff --git a/crates/worker/src/feature/background.rs b/crates/worker/src/feature/background.rs new file mode 100644 index 00000000..76254e5e --- /dev/null +++ b/crates/worker/src/feature/background.rs @@ -0,0 +1,705 @@ +//! Executable, scope-owned background tasks contributed by Worker features. +//! +//! Tasks never receive provider handles, credentials, or raw Workdir paths from +//! this registry. Callers pass only stable Worker/session provenance. The task +//! implementation obtains any additional authority through the services its +//! feature was explicitly granted at install time. + +use std::collections::BTreeMap; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use tokio::sync::Notify; +use tokio::task::JoinHandle; + +use super::{BackgroundTaskDeclaration, FeatureId, FeatureInstallError}; +use crate::hook::{HookError, HookErrorCategory, HookInvocationContext}; + +const MAX_TASK_CONCURRENCY: u16 = 64; +const MAX_TASK_ATTEMPTS: u16 = 16; +const MAX_TASK_TIMEOUT_MS: u64 = 24 * 60 * 60 * 1_000; +const MAX_RETAINED_DIAGNOSTICS: usize = 128; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BackgroundTaskRewritePolicy { + CancelAndWait, + Wait, + Block, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BackgroundTaskShutdownPolicy { + CancelAndWait, + Wait, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BackgroundTaskRetryPolicy { + Never, + Bounded { max_attempts: u16, delay_ms: u64 }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BackgroundTaskTrigger { + Manual, + RunCommitted, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct BackgroundTaskSpec { + pub declaration: BackgroundTaskDeclaration, + pub trigger: BackgroundTaskTrigger, + pub max_concurrency: u16, + pub timeout_ms: u64, + pub retry: BackgroundTaskRetryPolicy, + pub rewrite: BackgroundTaskRewritePolicy, + pub shutdown: BackgroundTaskShutdownPolicy, +} + +impl BackgroundTaskSpec { + pub fn single_flight(declaration: BackgroundTaskDeclaration, timeout: Duration) -> Self { + Self { + declaration, + trigger: BackgroundTaskTrigger::Manual, + max_concurrency: 1, + timeout_ms: u64::try_from(timeout.as_millis()).unwrap_or(u64::MAX), + retry: BackgroundTaskRetryPolicy::Never, + rewrite: BackgroundTaskRewritePolicy::CancelAndWait, + shutdown: BackgroundTaskShutdownPolicy::CancelAndWait, + } + } + + fn validate(&self) -> Result<(), FeatureInstallError> { + if self.max_concurrency == 0 || self.max_concurrency > MAX_TASK_CONCURRENCY { + return Err(FeatureInstallError::InvalidDescriptor(format!( + "background task `{}` max_concurrency must be within 1..={MAX_TASK_CONCURRENCY}", + self.declaration.name + ))); + } + if self.timeout_ms == 0 || self.timeout_ms > MAX_TASK_TIMEOUT_MS { + return Err(FeatureInstallError::InvalidDescriptor(format!( + "background task `{}` timeout_ms must be within 1..={MAX_TASK_TIMEOUT_MS}", + self.declaration.name + ))); + } + if let BackgroundTaskRetryPolicy::Bounded { max_attempts, .. } = self.retry + && (max_attempts == 0 || max_attempts > MAX_TASK_ATTEMPTS) + { + return Err(FeatureInstallError::InvalidDescriptor(format!( + "background task `{}` max_attempts must be within 1..={MAX_TASK_ATTEMPTS}", + self.declaration.name + ))); + } + Ok(()) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct BackgroundTaskContext { + pub invocation: HookInvocationContext, + pub feature_id: FeatureId, + pub task_name: String, + pub execution_id: u64, + pub attempt: u16, +} + +#[derive(Clone, Default)] +pub struct BackgroundTaskCancellation { + cancelled: Arc, + notify: Arc, +} + +impl BackgroundTaskCancellation { + pub fn is_cancelled(&self) -> bool { + self.cancelled.load(Ordering::Acquire) + } + + pub async fn cancelled(&self) { + if self.is_cancelled() { + return; + } + self.notify.notified().await; + } + + fn cancel(&self) { + if !self.cancelled.swap(true, Ordering::AcqRel) { + self.notify.notify_one(); + } + } +} + +#[async_trait] +pub trait FeatureBackgroundTask: Send + Sync { + async fn run( + &self, + context: BackgroundTaskContext, + cancellation: BackgroundTaskCancellation, + ) -> Result<(), HookError>; +} + +struct Registration { + feature_id: FeatureId, + spec: BackgroundTaskSpec, + task: Arc, +} + +#[derive(Default)] +pub struct FeatureBackgroundTaskRegistryBuilder { + registrations: BTreeMap<(FeatureId, String), Registration>, +} + +impl FeatureBackgroundTaskRegistryBuilder { + pub fn register( + &mut self, + feature_id: FeatureId, + spec: BackgroundTaskSpec, + task: impl FeatureBackgroundTask + 'static, + ) -> Result<(), FeatureInstallError> { + spec.validate()?; + let key = (feature_id.clone(), spec.declaration.name.clone()); + if self.registrations.contains_key(&key) { + return Err(FeatureInstallError::InvalidDescriptor(format!( + "feature `{feature_id}` registered background task `{}` more than once", + spec.declaration.name + ))); + } + self.registrations.insert( + key, + Registration { + feature_id, + spec, + task: Arc::new(task), + }, + ); + Ok(()) + } + + pub(crate) fn checkpoint(&self) -> Vec<(FeatureId, String)> { + self.registrations.keys().cloned().collect() + } + + pub(crate) fn rollback_to(&mut self, checkpoint: &[(FeatureId, String)]) { + let retained = checkpoint + .iter() + .cloned() + .collect::>(); + self.registrations.retain(|key, _| retained.contains(key)); + } + + pub fn build(self) -> FeatureBackgroundTaskRegistry { + FeatureBackgroundTaskRegistry { + inner: Arc::new(RegistryInner { + registrations: self.registrations, + running: Mutex::new(BTreeMap::new()), + diagnostics: Mutex::new(Vec::new()), + next_execution_id: AtomicU64::new(1), + accepting: AtomicBool::new(true), + }), + } + } +} + +struct RunningTask { + feature_id: FeatureId, + task_name: String, + cancellation: BackgroundTaskCancellation, + handle: JoinHandle<()>, +} + +struct RegistryInner { + registrations: BTreeMap<(FeatureId, String), Registration>, + running: Mutex>, + diagnostics: Mutex>, + next_execution_id: AtomicU64, + accepting: AtomicBool, +} + +impl Drop for RegistryInner { + fn drop(&mut self) { + if let Ok(running) = self.running.get_mut() { + for task in running.values() { + task.cancellation.cancel(); + task.handle.abort(); + } + running.clear(); + } + } +} + +#[derive(Clone)] +pub struct FeatureBackgroundTaskRegistry { + inner: Arc, +} + +impl Default for FeatureBackgroundTaskRegistry { + fn default() -> Self { + FeatureBackgroundTaskRegistryBuilder::default().build() + } +} + +impl std::fmt::Debug for FeatureBackgroundTaskRegistry { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("FeatureBackgroundTaskRegistry") + .field( + "registrations", + &self.inner.registrations.keys().collect::>(), + ) + .finish_non_exhaustive() + } +} + +impl PartialEq for FeatureBackgroundTaskRegistry { + fn eq(&self, other: &Self) -> bool { + self.inner + .registrations + .keys() + .eq(other.inner.registrations.keys()) + } +} + +impl Eq for FeatureBackgroundTaskRegistry {} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum BackgroundTaskStart { + Started { execution_id: u64 }, + AtCapacity, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum BackgroundTaskOutcome { + Completed, + Cancelled, + TimedOut, + Failed(HookError), + JoinFailed, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct BackgroundTaskDiagnostic { + pub execution_id: u64, + pub feature_id: FeatureId, + pub task_name: String, + pub attempts: u16, + pub outcome: BackgroundTaskOutcome, +} + +impl FeatureBackgroundTaskRegistry { + /// Starts one execution immediately. The Worker scope intentionally owns no + /// hidden queue: reaching `max_concurrency` returns `AtCapacity`, so callers + /// must retain durable work in their domain authority and retry explicitly. + pub fn start( + &self, + feature_id: &FeatureId, + task_name: &str, + invocation: HookInvocationContext, + ) -> Result { + if !self.inner.accepting.load(Ordering::Acquire) { + return Err(HookError::new( + HookErrorCategory::ScopeDisposed, + "background task scope is stopping", + )); + } + let key = (feature_id.clone(), task_name.to_string()); + let registration = self.inner.registrations.get(&key).ok_or_else(|| { + HookError::new( + HookErrorCategory::InvalidInput, + format!("unknown background task `{feature_id}/{task_name}`"), + ) + })?; + + let mut running = self + .inner + .running + .lock() + .expect("background tasks poisoned"); + running.retain(|_, running| !running.handle.is_finished()); + let active = running + .values() + .filter(|running| running.feature_id == *feature_id && running.task_name == task_name) + .count(); + if active >= usize::from(registration.spec.max_concurrency) { + return Ok(BackgroundTaskStart::AtCapacity); + } + + let execution_id = self.inner.next_execution_id.fetch_add(1, Ordering::Relaxed); + let cancellation = BackgroundTaskCancellation::default(); + let task_cancellation = cancellation.clone(); + let task = Arc::clone(®istration.task); + let spec = registration.spec.clone(); + let feature_id = registration.feature_id.clone(); + let task_name = task_name.to_string(); + let weak_inner = Arc::downgrade(&self.inner); + let task_feature_id = feature_id.clone(); + let task_task_name = task_name.clone(); + let handle = tokio::spawn(async move { + let (attempts, outcome) = execute_task( + task, + spec, + invocation, + task_feature_id.clone(), + task_task_name.clone(), + execution_id, + task_cancellation, + ) + .await; + if let Some(inner) = weak_inner.upgrade() { + let mut diagnostics = inner.diagnostics.lock().expect("diagnostics poisoned"); + diagnostics.push(BackgroundTaskDiagnostic { + execution_id, + feature_id: task_feature_id, + task_name: task_task_name, + attempts, + outcome, + }); + if diagnostics.len() > MAX_RETAINED_DIAGNOSTICS { + let remove = diagnostics.len() - MAX_RETAINED_DIAGNOSTICS; + diagnostics.drain(..remove); + } + } + }); + running.insert( + execution_id, + RunningTask { + feature_id, + task_name, + cancellation, + handle, + }, + ); + Ok(BackgroundTaskStart::Started { execution_id }) + } + + /// Starts every task explicitly bound to the committed-run boundary in + /// deterministic `(FeatureId, task name)` order. Capacity is an expected + /// single-flight outcome and leaves the already-running execution intact. + pub fn start_run_committed(&self, invocation: HookInvocationContext) -> Result<(), HookError> { + let tasks = self + .inner + .registrations + .iter() + .filter(|(_, registration)| { + registration.spec.trigger == BackgroundTaskTrigger::RunCommitted + }) + .map(|((feature_id, task_name), _)| (feature_id.clone(), task_name.clone())) + .collect::>(); + for (feature_id, task_name) in tasks { + let _ = self.start(&feature_id, &task_name, invocation.clone())?; + } + Ok(()) + } + + pub fn diagnostics(&self) -> Vec { + self.inner + .diagnostics + .lock() + .expect("diagnostics poisoned") + .clone() + } + + pub async fn before_session_rewrite(&self) -> Result<(), HookError> { + self.settle(false).await + } + + pub async fn shutdown(&self) -> Result<(), HookError> { + self.inner.accepting.store(false, Ordering::Release); + self.settle(true).await + } + + async fn settle(&self, shutdown: bool) -> Result<(), HookError> { + let mut waiting = Vec::new(); + { + let mut running = self + .inner + .running + .lock() + .expect("background tasks poisoned"); + if !shutdown { + for task in running.values() { + let registration = self + .inner + .registrations + .get(&(task.feature_id.clone(), task.task_name.clone())) + .expect("running background task must retain registration"); + if registration.spec.rewrite == BackgroundTaskRewritePolicy::Block { + return Err(HookError::new( + HookErrorCategory::Dependency, + "session rewrite blocked by a running feature background task", + )); + } + } + } + let ids = running.keys().copied().collect::>(); + for id in ids { + let Some(task) = running.remove(&id) else { + continue; + }; + let registration = self + .inner + .registrations + .get(&(task.feature_id.clone(), task.task_name.clone())) + .expect("running background task must retain registration"); + let cancel = if shutdown { + registration.spec.shutdown == BackgroundTaskShutdownPolicy::CancelAndWait + } else { + match registration.spec.rewrite { + BackgroundTaskRewritePolicy::CancelAndWait => true, + BackgroundTaskRewritePolicy::Wait => false, + BackgroundTaskRewritePolicy::Block => unreachable!( + "blocking rewrite policies are rejected before task handles are drained" + ), + } + }; + if cancel { + task.cancellation.cancel(); + } + waiting.push(( + id, + task.feature_id.clone(), + task.task_name.clone(), + task.handle, + )); + } + } + + let mut join_failed = false; + for (execution_id, feature_id, task_name, handle) in waiting { + if handle.await.is_err() { + join_failed = true; + let mut diagnostics = self.inner.diagnostics.lock().expect("diagnostics poisoned"); + diagnostics.push(BackgroundTaskDiagnostic { + execution_id, + feature_id, + task_name, + attempts: 0, + outcome: BackgroundTaskOutcome::JoinFailed, + }); + if diagnostics.len() > MAX_RETAINED_DIAGNOSTICS { + let remove = diagnostics.len() - MAX_RETAINED_DIAGNOSTICS; + diagnostics.drain(..remove); + } + } + } + if join_failed { + return Err(HookError::new( + HookErrorCategory::Internal, + "feature background task join failed", + )); + } + Ok(()) + } +} + +async fn execute_task( + task: Arc, + spec: BackgroundTaskSpec, + invocation: HookInvocationContext, + feature_id: FeatureId, + task_name: String, + execution_id: u64, + cancellation: BackgroundTaskCancellation, +) -> (u16, BackgroundTaskOutcome) { + let (max_attempts, delay_ms) = match spec.retry { + BackgroundTaskRetryPolicy::Never => (1, 0), + BackgroundTaskRetryPolicy::Bounded { + max_attempts, + delay_ms, + } => (max_attempts, delay_ms), + }; + for attempt in 1..=max_attempts { + if cancellation.is_cancelled() { + return (attempt, BackgroundTaskOutcome::Cancelled); + } + let context = BackgroundTaskContext { + invocation: invocation.clone(), + feature_id: feature_id.clone(), + task_name: task_name.clone(), + execution_id, + attempt, + }; + let result = tokio::time::timeout( + Duration::from_millis(spec.timeout_ms), + task.run(context, cancellation.clone()), + ) + .await; + match result { + Ok(Ok(())) => return (attempt, BackgroundTaskOutcome::Completed), + Ok(Err(_error)) if cancellation.is_cancelled() => { + return (attempt, BackgroundTaskOutcome::Cancelled); + } + Ok(Err(error)) if attempt == max_attempts => { + return (attempt, BackgroundTaskOutcome::Failed(error)); + } + Ok(Err(_)) => {} + Err(_) => return (attempt, BackgroundTaskOutcome::TimedOut), + } + if delay_ms > 0 { + tokio::select! { + () = tokio::time::sleep(Duration::from_millis(delay_ms)) => {} + () = cancellation.cancelled() => { + return (attempt, BackgroundTaskOutcome::Cancelled); + } + } + } + } + ( + max_attempts, + BackgroundTaskOutcome::Failed(HookError::new( + HookErrorCategory::Internal, + "background task exhausted retry policy", + )), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + fn invocation() -> HookInvocationContext { + HookInvocationContext { + workspace_id: Some("workspace".into()), + worker_id: "worker".into(), + session_id: "session".into(), + session_revision: 3, + run_id: Some("run".into()), + turn_index: Some(2), + call_id: None, + } + } + + struct WaitForCancellation; + + #[async_trait] + impl FeatureBackgroundTask for WaitForCancellation { + async fn run( + &self, + _context: BackgroundTaskContext, + cancellation: BackgroundTaskCancellation, + ) -> Result<(), HookError> { + cancellation.cancelled().await; + Err(HookError::new(HookErrorCategory::Cancelled, "cancelled")) + } + } + + #[tokio::test] + async fn single_flight_rejects_overlap_and_shutdown_joins_the_task() { + let feature = FeatureId::builtin("background-test"); + let declaration = BackgroundTaskDeclaration::worker_managed("extract", "extract"); + let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); + builder + .register( + feature.clone(), + BackgroundTaskSpec::single_flight(declaration, Duration::from_secs(1)), + WaitForCancellation, + ) + .unwrap(); + let registry = builder.build(); + + assert!(matches!( + registry.start(&feature, "extract", invocation()).unwrap(), + BackgroundTaskStart::Started { .. } + )); + assert_eq!( + registry.start(&feature, "extract", invocation()).unwrap(), + BackgroundTaskStart::AtCapacity + ); + + registry.shutdown().await.unwrap(); + assert_eq!(registry.diagnostics().len(), 1); + assert_eq!( + registry.diagnostics()[0].outcome, + BackgroundTaskOutcome::Cancelled + ); + assert!(matches!( + registry.start(&feature, "extract", invocation()), + Err(HookError { + category: HookErrorCategory::ScopeDisposed, + .. + }) + )); + } + + struct FailTwice { + calls: Arc, + completed: Arc, + } + + #[async_trait] + impl FeatureBackgroundTask for FailTwice { + async fn run( + &self, + _context: BackgroundTaskContext, + _cancellation: BackgroundTaskCancellation, + ) -> Result<(), HookError> { + let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1; + if call < 3 { + return Err(HookError::new(HookErrorCategory::Dependency, "retry")); + } + self.completed.notify_one(); + Ok(()) + } + } + + #[tokio::test] + async fn bounded_retry_records_attempts_without_task_output() { + let feature = FeatureId::builtin("retry-test"); + let declaration = BackgroundTaskDeclaration::worker_managed("retry", "retry"); + let calls = Arc::new(AtomicUsize::new(0)); + let completed = Arc::new(Notify::new()); + let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); + let mut spec = BackgroundTaskSpec::single_flight(declaration, Duration::from_secs(1)); + spec.trigger = BackgroundTaskTrigger::RunCommitted; + spec.retry = BackgroundTaskRetryPolicy::Bounded { + max_attempts: 3, + delay_ms: 0, + }; + builder + .register( + feature.clone(), + spec, + FailTwice { + calls: Arc::clone(&calls), + completed: Arc::clone(&completed), + }, + ) + .unwrap(); + let registry = builder.build(); + registry.start_run_committed(invocation()).unwrap(); + completed.notified().await; + registry.shutdown().await.unwrap(); + + assert_eq!(calls.load(Ordering::SeqCst), 3); + assert_eq!(registry.diagnostics()[0].attempts, 3); + assert_eq!( + registry.diagnostics()[0].outcome, + BackgroundTaskOutcome::Completed + ); + } + + #[tokio::test] + async fn block_policy_fences_rewrite_without_detaching_the_task() { + let feature = FeatureId::builtin("rewrite-test"); + let declaration = BackgroundTaskDeclaration::worker_managed("rewrite", "rewrite"); + let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); + let mut spec = BackgroundTaskSpec::single_flight(declaration, Duration::from_secs(1)); + spec.rewrite = BackgroundTaskRewritePolicy::Block; + builder + .register(feature.clone(), spec, WaitForCancellation) + .unwrap(); + let registry = builder.build(); + registry.start(&feature, "rewrite", invocation()).unwrap(); + + let error = registry.before_session_rewrite().await.unwrap_err(); + assert_eq!(error.category, HookErrorCategory::Dependency); + registry.shutdown().await.unwrap(); + assert_eq!(registry.diagnostics().len(), 1); + } +} diff --git a/crates/worker/src/feature/builtin/task/mod.rs b/crates/worker/src/feature/builtin/task/mod.rs index 9df849bf..9ab38be3 100644 --- a/crates/worker/src/feature/builtin/task/mod.rs +++ b/crates/worker/src/feature/builtin/task/mod.rs @@ -125,7 +125,7 @@ impl FeatureModule for TaskFeature { )) .with_hook(HookDeclaration::new( "task-reminder-pre-request", - FeatureHookPoint::PreRequest, + FeatureHookPoint::PreLlmRequest, )) .with_hook(HookDeclaration::new( "task-reminder-tool-usage", @@ -209,24 +209,27 @@ struct TaskReminderPreRequestHook { #[async_trait] impl Hook for TaskReminderPreRequestHook { - async fn call(&self, input: &PreRequestContext) -> HookPreRequestAction { + async fn call( + &self, + input: &PreRequestContext, + ) -> Result { let tasks = self.state.task_store.list(); if tasks.is_empty() { - return HookPreRequestAction::Continue; + return Ok(HookPreRequestAction::Continue); } let (since_task_management, since_reminder) = self.state.reminder_state.note_request(); if since_task_management < TASK_REMINDER_REQUEST_THRESHOLD || since_reminder < TASK_REMINDER_COOLDOWN_REQUESTS { - return HookPreRequestAction::Continue; + return Ok(HookPreRequestAction::Continue); } if let Some(system_items) = input.system_items() { self.state.reminder_state.note_reminder(); system_items.append_task_reminder(render_task_reminder_body(&tasks)); } - HookPreRequestAction::Continue + Ok(HookPreRequestAction::Continue) } } @@ -236,11 +239,14 @@ struct TaskReminderToolUsageHook { #[async_trait] impl Hook for TaskReminderToolUsageHook { - async fn call(&self, input: &ToolCallSummary) -> HookPreToolAction { + async fn call( + &self, + input: &ToolCallSummary, + ) -> Result { if is_task_management_tool(&input.tool_name) { self.state.reminder_state.note_task_management(); } - HookPreToolAction::Continue + Ok(HookPreToolAction::Continue) } } diff --git a/crates/worker/src/feature/mcp.rs b/crates/worker/src/feature/mcp.rs index cfdd8f27..ca13ae13 100644 --- a/crates/worker/src/feature/mcp.rs +++ b/crates/worker/src/feature/mcp.rs @@ -2115,8 +2115,7 @@ read exit_notification || true meta.name }) .collect(); - assert!(!names.iter().any(|name| name == "Mcp_demo_search_files")); - assert!(names.iter().any(|name| name == "Mcp_demo_unique")); + assert!(names.is_empty()); } fn shell_tool_server(response: &str) -> McpStdioServerSpec { diff --git a/crates/worker/src/feature/plugin.rs b/crates/worker/src/feature/plugin.rs index 143d8e9f..39dbfecf 100644 --- a/crates/worker/src/feature/plugin.rs +++ b/crates/worker/src/feature/plugin.rs @@ -7621,7 +7621,7 @@ mod tests { ))) .install_into_pending(&mut pending, &mut hooks); - assert_eq!(pending.len(), 1); + assert!(pending.is_empty()); assert_eq!(skipped_count(&report), 1); assert!(has_diagnostic(&report, "duplicate tool contribution")); } diff --git a/crates/worker/src/hook.rs b/crates/worker/src/hook.rs index 4318f9c3..b326697a 100644 --- a/crates/worker/src/hook.rs +++ b/crates/worker/src/hook.rs @@ -18,13 +18,78 @@ use std::ops::Deref; use std::sync::{Arc, Mutex}; +use agen::HistoryEntry; use agen::interceptor::{ PostToolAction, PreRequestAction, PreToolAction, PromptAction, TurnEndAction, }; use agen::tool::{ToolOutput, ToolResult}; use async_trait::async_trait; +use serde::{Deserialize, Serialize}; use serde_json::Value; use session_store::{SystemItem, SystemReminder}; +use thiserror::Error; + +use crate::session_history::SessionHistoryMetadata; + +const HOOK_DIAGNOSTIC_MAX_BYTES: usize = 1_024; + +/// Failure category exposed by the safe Worker hook boundary. +/// +/// Categories are intentionally closed and payload-free so extensions cannot +/// smuggle provider, credential, or history data into diagnostics. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum HookErrorCategory { + InvalidInput, + Dependency, + Timeout, + Cancelled, + Trap, + ScopeDisposed, + Internal, +} + +/// Bounded hook callback failure. Raw tool arguments, output, prompts, and +/// credentials must never be placed in `diagnostic`. +#[derive(Clone, Debug, PartialEq, Eq, Error)] +#[error("{category:?}: {diagnostic}")] +pub struct HookError { + pub category: HookErrorCategory, + pub diagnostic: String, +} + +impl HookError { + pub fn new(category: HookErrorCategory, diagnostic: impl Into) -> Self { + Self { + category, + diagnostic: bounded_utf8(diagnostic.into(), HOOK_DIAGNOSTIC_MAX_BYTES), + } + } +} + +/// Failure behavior declared when a hook is registered. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum HookFailurePolicy { + /// Gate the current operation when the hook cannot decide safely. + FailClosed, + /// Keep the already-authorized operation moving and emit a diagnostic. + FailOpenWithDiagnostic, + /// Keep committed state intact and mark the failure for operator attention. + AttentionRequired, +} + +fn bounded_utf8(mut value: String, max_bytes: usize) -> String { + if value.len() <= max_bytes { + return value; + } + let mut end = max_bytes; + while end > 0 && !value.is_char_boundary(end) { + end -= 1; + } + value.truncate(end); + value +} /// Hook-facing prompt-submit action. /// @@ -349,24 +414,175 @@ impl HookEventKind for OnTurnEnd { /// short-circuit on the first non-continue action. #[async_trait] pub trait Hook: Send + Sync { - async fn call(&self, input: &E::Input) -> E::Output; + async fn call(&self, input: &E::Input) -> Result; } // ============================================================================= // Hook Registry // ============================================================================= +/// Stable provenance attached to every Worker lifecycle callback. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct HookInvocationContext { + pub workspace_id: Option, + pub worker_id: String, + pub session_id: String, + pub session_revision: u64, + pub run_id: Option, + pub turn_index: Option, + pub call_id: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RunCommittedExit { + Finished, + Paused, + Yielded, + Interrupted, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RunExitContext { + pub invocation: HookInvocationContext, + pub exit: RunCommittedExit, + pub history_len: usize, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct RunCommittedContext { + pub invocation: HookInvocationContext, + pub exit: RunCommittedExit, + pub committed_history: Vec>, + pub committed_history_len: usize, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SessionRewriteKind { + Rewind, + Compact, + Fork, + Restore, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct BeforeSessionRewriteContext { + pub invocation: HookInvocationContext, + pub kind: SessionRewriteKind, + pub current_history: Vec>, + pub current_history_len: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum BeforeSessionRewriteAction { + Continue, + Deny(String), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct WorkerStoppingContext { + pub invocation: HookInvocationContext, + pub reason: String, +} + +pub struct RunExit; +pub struct RunCommitted; +pub struct BeforeSessionRewrite; +pub struct WorkerStopping; + +impl HookEventKind for RunExit { + type Input = RunExitContext; + type Output = (); +} + +impl HookEventKind for RunCommitted { + type Input = RunCommittedContext; + type Output = (); +} + +impl HookEventKind for BeforeSessionRewrite { + type Input = BeforeSessionRewriteContext; + type Output = BeforeSessionRewriteAction; +} + +impl HookEventKind for WorkerStopping { + type Input = WorkerStoppingContext; + type Output = (); +} + +pub(crate) struct RegisteredHook { + owner: String, + policy: HookFailurePolicy, + hook: Box>, +} + +impl RegisteredHook { + pub(crate) async fn call(&self, input: &E::Input) -> Result { + self.hook + .call(input) + .await + .map_err(|source| HookExecutionError { + owner: self.owner.clone(), + policy: self.policy, + source, + }) + } + + pub(crate) async fn call_optional( + &self, + input: &E::Input, + ) -> Result, HookExecutionError> { + match self.call(input).await { + Ok(output) => Ok(Some(output)), + Err(error) if error.policy == HookFailurePolicy::FailOpenWithDiagnostic => { + tracing::warn!(owner = %error.owner, error = %error.source, "inline hook failed open"); + Ok(None) + } + Err(error) => Err(error), + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Error)] +#[error("hook `{owner}` failed under {policy:?}: {source}")] +pub struct HookExecutionError { + pub owner: String, + pub policy: HookFailurePolicy, + pub source: HookError, +} + /// Builder for constructing a frozen `HookRegistry`. -/// -/// Hooks are added during setup, then `build()` produces an immutable -/// registry that can be shared via `Arc`. #[derive(Default)] pub struct HookRegistryBuilder { - on_prompt_submit: Vec>>, - pre_llm_request: Vec>>, - pre_tool_call: Vec>>, - post_tool_call: Vec>>, - on_turn_end: Vec>>, + on_prompt_submit: Vec>, + pre_llm_request: Vec>, + pre_tool_call: Vec>, + post_tool_call: Vec>, + on_turn_end: Vec>, + run_exit: Vec>, + run_committed: Vec>, + before_session_rewrite: Vec>, + worker_stopping: Vec>, +} + +macro_rules! add_hook_methods { + ($default:ident, $named:ident, $field:ident, $event:ty) => { + pub fn $default(&mut self, hook: impl Hook<$event> + 'static) { + self.$named("worker.host", HookFailurePolicy::FailClosed, hook); + } + + pub fn $named( + &mut self, + owner: impl Into, + policy: HookFailurePolicy, + hook: impl Hook<$event> + 'static, + ) { + self.$field.push(RegisteredHook { + owner: owner.into(), + policy, + hook: Box::new(hook), + }); + } + }; } impl HookRegistryBuilder { @@ -374,27 +590,82 @@ impl HookRegistryBuilder { Self::default() } - pub fn add_on_prompt_submit(&mut self, hook: impl Hook + 'static) { - self.on_prompt_submit.push(Box::new(hook)); + add_hook_methods!( + add_on_prompt_submit, + add_named_on_prompt_submit, + on_prompt_submit, + OnPromptSubmit + ); + add_hook_methods!( + add_pre_llm_request, + add_named_pre_llm_request, + pre_llm_request, + PreLlmRequest + ); + add_hook_methods!( + add_pre_tool_call, + add_named_pre_tool_call, + pre_tool_call, + PreToolCall + ); + add_hook_methods!( + add_post_tool_call, + add_named_post_tool_call, + post_tool_call, + PostToolCall + ); + add_hook_methods!( + add_on_turn_end, + add_named_on_turn_end, + on_turn_end, + OnTurnEnd + ); + add_hook_methods!(add_run_exit, add_named_run_exit, run_exit, RunExit); + add_hook_methods!( + add_run_committed, + add_named_run_committed, + run_committed, + RunCommitted + ); + add_hook_methods!( + add_before_session_rewrite, + add_named_before_session_rewrite, + before_session_rewrite, + BeforeSessionRewrite + ); + add_hook_methods!( + add_worker_stopping, + add_named_worker_stopping, + worker_stopping, + WorkerStopping + ); + + pub(crate) fn checkpoint(&self) -> [usize; 9] { + [ + self.on_prompt_submit.len(), + self.pre_llm_request.len(), + self.pre_tool_call.len(), + self.post_tool_call.len(), + self.on_turn_end.len(), + self.run_exit.len(), + self.run_committed.len(), + self.before_session_rewrite.len(), + self.worker_stopping.len(), + ] } - pub fn add_pre_llm_request(&mut self, hook: impl Hook + 'static) { - self.pre_llm_request.push(Box::new(hook)); + pub(crate) fn rollback_to(&mut self, checkpoint: [usize; 9]) { + self.on_prompt_submit.truncate(checkpoint[0]); + self.pre_llm_request.truncate(checkpoint[1]); + self.pre_tool_call.truncate(checkpoint[2]); + self.post_tool_call.truncate(checkpoint[3]); + self.on_turn_end.truncate(checkpoint[4]); + self.run_exit.truncate(checkpoint[5]); + self.run_committed.truncate(checkpoint[6]); + self.before_session_rewrite.truncate(checkpoint[7]); + self.worker_stopping.truncate(checkpoint[8]); } - pub fn add_pre_tool_call(&mut self, hook: impl Hook + 'static) { - self.pre_tool_call.push(Box::new(hook)); - } - - pub fn add_post_tool_call(&mut self, hook: impl Hook + 'static) { - self.post_tool_call.push(Box::new(hook)); - } - - pub fn add_on_turn_end(&mut self, hook: impl Hook + 'static) { - self.on_turn_end.push(Box::new(hook)); - } - - /// Freeze the builder into an immutable registry. pub fn build(self) -> HookRegistry { HookRegistry { on_prompt_submit: self.on_prompt_submit, @@ -402,17 +673,129 @@ impl HookRegistryBuilder { pre_tool_call: self.pre_tool_call, post_tool_call: self.post_tool_call, on_turn_end: self.on_turn_end, + run_exit: self.run_exit, + run_committed: self.run_committed, + before_session_rewrite: self.before_session_rewrite, + worker_stopping: self.worker_stopping, + diagnostics: std::sync::Mutex::new(Vec::new()), } } } /// Frozen registry of hooks. Constructed via [`HookRegistryBuilder::build()`]. pub struct HookRegistry { - pub(crate) on_prompt_submit: Vec>>, - pub(crate) pre_llm_request: Vec>>, - pub(crate) pre_tool_call: Vec>>, - pub(crate) post_tool_call: Vec>>, - pub(crate) on_turn_end: Vec>>, + pub(crate) on_prompt_submit: Vec>, + pub(crate) pre_llm_request: Vec>, + pub(crate) pre_tool_call: Vec>, + pub(crate) post_tool_call: Vec>, + pub(crate) on_turn_end: Vec>, + run_exit: Vec>, + run_committed: Vec>, + before_session_rewrite: Vec>, + worker_stopping: Vec>, + diagnostics: std::sync::Mutex>, +} + +impl HookRegistry { + fn record_diagnostic(&self, error: HookExecutionError) { + let mut diagnostics = self.diagnostics.lock().expect("hook diagnostics poisoned"); + diagnostics.push(error); + if diagnostics.len() > 128 { + let remove = diagnostics.len() - 128; + diagnostics.drain(..remove); + } + } + + pub fn diagnostics(&self) -> Vec { + self.diagnostics + .lock() + .expect("hook diagnostics poisoned") + .clone() + } + pub async fn on_run_exit(&self, context: &RunExitContext) -> Result<(), HookExecutionError> { + for registration in &self.run_exit { + if let Err(error) = registration.call(context).await { + self.record_diagnostic(error.clone()); + match error.policy { + HookFailurePolicy::FailOpenWithDiagnostic => { + tracing::warn!(owner = %error.owner, error = %error.source, "run-exit hook failed open"); + } + HookFailurePolicy::FailClosed | HookFailurePolicy::AttentionRequired => { + return Err(error); + } + } + } + } + Ok(()) + } + + pub async fn on_run_committed( + &self, + context: &RunCommittedContext, + ) -> Result<(), HookExecutionError> { + for registration in &self.run_committed { + if let Err(error) = registration.call(context).await { + self.record_diagnostic(error.clone()); + match error.policy { + HookFailurePolicy::FailOpenWithDiagnostic => { + tracing::warn!(owner = %error.owner, error = %error.source, "run-committed hook failed open"); + } + HookFailurePolicy::FailClosed | HookFailurePolicy::AttentionRequired => { + return Err(error); + } + } + } + } + Ok(()) + } + + pub async fn before_session_rewrite( + &self, + context: &BeforeSessionRewriteContext, + ) -> Result { + let mut denials = Vec::new(); + for registration in &self.before_session_rewrite { + match registration.call(context).await { + Ok(BeforeSessionRewriteAction::Continue) => {} + Ok(BeforeSessionRewriteAction::Deny(reason)) => { + denials.push((registration.owner.clone(), reason)); + } + Err(error) if error.policy == HookFailurePolicy::FailOpenWithDiagnostic => { + self.record_diagnostic(error.clone()); + tracing::warn!(owner = %error.owner, error = %error.source, "session-rewrite hook failed open"); + } + Err(error) => { + self.record_diagnostic(error.clone()); + return Err(error); + } + } + } + denials.sort_by(|left, right| left.0.cmp(&right.0)); + Ok(denials + .into_iter() + .next() + .map(|(_, reason)| BeforeSessionRewriteAction::Deny(reason)) + .unwrap_or(BeforeSessionRewriteAction::Continue)) + } + + pub async fn on_worker_stopping( + &self, + context: &WorkerStoppingContext, + ) -> Result<(), HookExecutionError> { + for registration in &self.worker_stopping { + if let Err(error) = registration.call(context).await { + self.record_diagnostic(error.clone()); + match error.policy { + HookFailurePolicy::FailOpenWithDiagnostic + | HookFailurePolicy::AttentionRequired => { + tracing::warn!(owner = %error.owner, error = %error.source, "worker-stopping hook requires attention"); + } + HookFailurePolicy::FailClosed => return Err(error), + } + } + } + Ok(()) + } } #[cfg(test)] @@ -474,4 +857,117 @@ mod tests { let pause_action = HookPreToolAction::Pause.into_worker_action("call_4".into()); assert!(matches!(pause_action, PreToolAction::Pause)); } + + struct RewriteHook { + action: BeforeSessionRewriteAction, + } + + #[async_trait] + impl Hook for RewriteHook { + async fn call( + &self, + _input: &BeforeSessionRewriteContext, + ) -> Result { + Ok(self.action.clone()) + } + } + + struct FailingRewriteHook; + + #[async_trait] + impl Hook for FailingRewriteHook { + async fn call( + &self, + _input: &BeforeSessionRewriteContext, + ) -> Result { + Err(HookError::new( + HookErrorCategory::Dependency, + "provider unavailable", + )) + } + } + + fn rewrite_context() -> BeforeSessionRewriteContext { + BeforeSessionRewriteContext { + invocation: HookInvocationContext { + workspace_id: Some("workspace".into()), + worker_id: "worker".into(), + session_id: "session".into(), + session_revision: 4, + run_id: None, + turn_index: None, + call_id: None, + }, + kind: SessionRewriteKind::Compact, + current_history: Vec::new(), + current_history_len: 8, + } + } + + #[tokio::test] + async fn rewrite_denials_are_resolved_by_owner_not_registration_order() { + let mut builder = HookRegistryBuilder::new(); + builder.add_named_before_session_rewrite( + "z-feature", + HookFailurePolicy::FailClosed, + RewriteHook { + action: BeforeSessionRewriteAction::Deny("z denied".into()), + }, + ); + builder.add_named_before_session_rewrite( + "a-feature", + HookFailurePolicy::FailClosed, + RewriteHook { + action: BeforeSessionRewriteAction::Deny("a denied".into()), + }, + ); + + assert_eq!( + builder + .build() + .before_session_rewrite(&rewrite_context()) + .await + .unwrap(), + BeforeSessionRewriteAction::Deny("a denied".into()) + ); + } + + #[tokio::test] + async fn hook_failure_policy_is_applied_at_the_registry_boundary() { + let mut fail_open = HookRegistryBuilder::new(); + fail_open.add_named_before_session_rewrite( + "feature", + HookFailurePolicy::FailOpenWithDiagnostic, + FailingRewriteHook, + ); + let fail_open = fail_open.build(); + assert_eq!( + fail_open + .before_session_rewrite(&rewrite_context()) + .await + .unwrap(), + BeforeSessionRewriteAction::Continue + ); + assert_eq!(fail_open.diagnostics().len(), 1); + + let mut fail_closed = HookRegistryBuilder::new(); + fail_closed.add_named_before_session_rewrite( + "feature", + HookFailurePolicy::FailClosed, + FailingRewriteHook, + ); + let error = fail_closed + .build() + .before_session_rewrite(&rewrite_context()) + .await + .unwrap_err(); + assert_eq!(error.source.category, HookErrorCategory::Dependency); + } + + #[test] + fn hook_diagnostics_are_utf8_bounded() { + let error = HookError::new(HookErrorCategory::Internal, "界".repeat(1_000)); + assert!(error.diagnostic.len() <= HOOK_DIAGNOSTIC_MAX_BYTES); + assert!(error.diagnostic.is_char_boundary(error.diagnostic.len())); + } } diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index 6dc23c51..b69ac42c 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -210,7 +210,7 @@ where let segment_id = worker.segment_id(); on_cancel_sender(worker.engine_mut().cancel_sender()); - match worker.run_text(&input).await { + let outcome = match worker.run_text(&input).await { Ok(lifecycle @ WorkerRunResult::Finished) | Ok(lifecycle @ WorkerRunResult::Paused) | Ok(lifecycle @ WorkerRunResult::RolledBack) => Ok(InternalWorkerResult { @@ -239,7 +239,11 @@ where identity, history_entries: store.entries_count(session_id, segment_id), }), - } + }; + worker + .stop_feature_runtime("internal Worker terminal outcome") + .await; + outcome } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -782,7 +786,8 @@ pub(crate) async fn prepare_internal_worker_session( }; tokio::spawn(async move { - while let Some(command) = command_rx.recv().await { + let mut stop_done = None; + 'actor: while let Some(command) = command_rx.recv().await { match command { InternalWorkerSessionCommand::Run(input) => { actor_in_flight.clear(); @@ -825,21 +830,16 @@ pub(crate) async fn prepare_internal_worker_session( Some(InternalWorkerSessionCommand::Stop(done)) => { let _ = cancel_sender.send(()).await; let _ = (&mut run).await; - actor_in_flight.clear(); - status.store(InternalWorkerSessionStatus::Stopped.encode(), std::sync::atomic::Ordering::Release); - let _ = event_tx.send(Event::Status { status: WorkerStatus::Stopped }); - let _ = event_tx.send(Event::Shutdown); - state_changed.notify_waiters(); - let _ = done.send(()); - return; + stop_done = Some(done); + break 'actor; } Some(InternalWorkerSessionCommand::Run(_)) => { // `send` reserves Running atomically, so a second Run cannot be enqueued. } None => { let _ = cancel_sender.send(()).await; - actor_in_flight.clear(); - return; + let _ = (&mut run).await; + break 'actor; } } } @@ -847,22 +847,27 @@ pub(crate) async fn prepare_internal_worker_session( } } InternalWorkerSessionCommand::Stop(done) => { - actor_in_flight.clear(); - status.store( - InternalWorkerSessionStatus::Stopped.encode(), - std::sync::atomic::Ordering::Release, - ); - let _ = event_tx.send(Event::Status { - status: WorkerStatus::Stopped, - }); - let _ = event_tx.send(Event::Shutdown); - state_changed.notify_waiters(); - let _ = done.send(()); - return; + stop_done = Some(done); + break; } } } + worker + .stop_feature_runtime("internal Worker session stopped") + .await; actor_in_flight.clear(); + status.store( + InternalWorkerSessionStatus::Stopped.encode(), + std::sync::atomic::Ordering::Release, + ); + let _ = event_tx.send(Event::Status { + status: WorkerStatus::Stopped, + }); + let _ = event_tx.send(Event::Shutdown); + state_changed.notify_waiters(); + if let Some(done) = stop_done { + let _ = done.send(()); + } }); Ok(handle) diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index 2706f715..0e5c5b40 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -246,12 +246,22 @@ impl Interceptor for WorkerInterceptor { input_text: extract_message_text(item).unwrap_or_default(), turn_index, }; + let mut cancellations = Vec::new(); for hook in &self.registry.on_prompt_submit { - let action = hook.call(&info).await; - if !matches!(action, HookPromptAction::Continue) { - return Ok(action.into()); + let Some(action) = hook.call_optional(&info).await.map_err(|error| { + InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) + })? + else { + continue; + }; + if let HookPromptAction::Cancel(reason) = action { + cancellations.push(reason); } } + cancellations.sort(); + if let Some(reason) = cancellations.into_iter().next() { + return Ok(PromptAction::Cancel(reason)); + } let mut extras: Vec = std::mem::take( &mut *self .pending_attachments @@ -344,12 +354,29 @@ impl Interceptor for WorkerInterceptor { .as_ref() .map(|_| SystemItemAppendHandle::new(Arc::clone(&pending_hook_system_items))); let hook_context = PreRequestContext::new(info, system_item_sink); + let mut cancellations = Vec::new(); + let mut should_yield = false; for hook in &self.registry.pre_llm_request { - let action = hook.call(&hook_context).await; - if !matches!(action, HookPreRequestAction::Continue) { - return Ok(action.into()); + let Some(action) = hook.call_optional(&hook_context).await.map_err(|error| { + InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) + })? + else { + continue; + }; + + match action { + HookPreRequestAction::Continue => {} + HookPreRequestAction::Yield => should_yield = true, + HookPreRequestAction::Cancel(reason) => cancellations.push(reason), } } + cancellations.sort(); + if let Some(reason) = cancellations.into_iter().next() { + return Ok(PreRequestAction::Cancel(reason)); + } + if should_yield { + return Ok(PreRequestAction::Yield); + } let mut system_items: Vec = std::mem::take( &mut *pending_hook_system_items @@ -404,12 +431,35 @@ impl Interceptor for WorkerInterceptor { tool_name: info.call.name.clone(), arguments: info.call.input.clone(), }; + let mut aborts = Vec::new(); + let mut should_pause = false; + let mut denials = Vec::new(); for hook in &self.registry.pre_tool_call { - let action = hook.call(&summary).await; - if !matches!(action, HookPreToolAction::Continue) { - return Ok(action.into_worker_action(summary.call_id.clone())); + let Some(action) = hook.call_optional(&summary).await.map_err(|error| { + InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) + })? + else { + continue; + }; + + match action { + HookPreToolAction::Continue => {} + HookPreToolAction::Pause => should_pause = true, + HookPreToolAction::Deny(reason) => denials.push(reason), + HookPreToolAction::Abort(reason) => aborts.push(reason), } } + aborts.sort(); + if let Some(reason) = aborts.into_iter().next() { + return Ok(HookPreToolAction::Abort(reason).into_worker_action(summary.call_id.clone())); + } + if should_pause { + return Ok(PreToolAction::Pause); + } + denials.sort(); + if let Some(reason) = denials.into_iter().next() { + return Ok(HookPreToolAction::Deny(reason).into_worker_action(summary.call_id.clone())); + } self.tool_calls_this_turn.fetch_add(1, Ordering::Relaxed); Ok(PreToolAction::Continue) } @@ -429,12 +479,23 @@ impl Interceptor for WorkerInterceptor { attachments: Vec::new(), }, }; + let mut aborts = Vec::new(); for hook in &self.registry.post_tool_call { - let action = hook.call(&summary).await; - if !matches!(action, HookPostToolAction::Continue) { - return Ok(action.into()); + let Some(action) = hook.call_optional(&summary).await.map_err(|error| { + InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) + })? + else { + continue; + }; + + if let HookPostToolAction::Abort(reason) = action { + aborts.push(reason); } } + aborts.sort(); + if let Some(reason) = aborts.into_iter().next() { + return Ok(PostToolAction::Abort(reason)); + } Ok(PostToolAction::Continue) } @@ -455,12 +516,21 @@ impl Interceptor for WorkerInterceptor { tool_calls_count: self.tool_calls_this_turn.load(Ordering::Relaxed), final_text_preview, }; + let mut should_pause = false; for hook in &self.registry.on_turn_end { - let action = hook.call(&info).await; - if !matches!(action, HookTurnEndAction::Finish) { - return Ok(action.into()); + let Some(action) = hook.call_optional(&info).await.map_err(|error| { + InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) + })? + else { + continue; + }; + if matches!(action, HookTurnEndAction::Pause) { + should_pause = true; } } + if should_pause { + return Ok(TurnEndAction::Pause); + } Ok(TurnEndAction::Finish) } } @@ -541,9 +611,12 @@ mod tests { #[async_trait] impl Hook for CountingHook { - async fn call(&self, _info: &PreRequestContext) -> HookPreRequestAction { + async fn call( + &self, + _info: &PreRequestContext, + ) -> Result { self.0.fetch_add(1, Ordering::Relaxed); - HookPreRequestAction::Continue + Ok(HookPreRequestAction::Continue) } } @@ -582,12 +655,15 @@ mod tests { #[async_trait] impl Hook for AppendingPreRequestHook { - async fn call(&self, input: &PreRequestContext) -> HookPreRequestAction { + async fn call( + &self, + input: &PreRequestContext, + ) -> Result { if let Some(system_items) = input.system_items() { self.saw_handle.store(true, Ordering::Relaxed); system_items.append_task_reminder("hook reminder"); } - HookPreRequestAction::Continue + Ok(HookPreRequestAction::Continue) } } @@ -982,33 +1058,42 @@ mod tests { #[async_trait] impl Hook for AbortingHook { - async fn call(&self, _info: &PreRequestContext) -> HookPreRequestAction { + async fn call( + &self, + _info: &PreRequestContext, + ) -> Result { self.0.store(true, Ordering::Relaxed); - HookPreRequestAction::Cancel("nope".into()) + Ok(HookPreRequestAction::Cancel("nope".into())) } } #[tokio::test] - async fn public_pre_tool_hook_deny_becomes_synthetic_error_and_short_circuits() { + async fn public_pre_tool_hook_denials_compose_without_short_circuiting() { struct DenyToolHook(Arc); struct CountingToolHook(Arc); #[async_trait] impl Hook for DenyToolHook { - async fn call(&self, input: &ToolCallSummary) -> HookPreToolAction { + async fn call( + &self, + input: &ToolCallSummary, + ) -> Result { self.0.fetch_add(1, Ordering::Relaxed); assert_eq!(input.call_id, "call-id"); assert_eq!(input.tool_name, "TaskList"); assert_eq!(input.arguments, serde_json::json!({"scope": "all"})); - HookPreToolAction::Deny("blocked by public hook".into()) + Ok(HookPreToolAction::Deny("blocked by public hook".into())) } } #[async_trait] impl Hook for CountingToolHook { - async fn call(&self, _input: &ToolCallSummary) -> HookPreToolAction { + async fn call( + &self, + _input: &ToolCallSummary, + ) -> Result { self.0.fetch_add(1, Ordering::Relaxed); - HookPreToolAction::Continue + Ok(HookPreToolAction::Continue) } } @@ -1041,7 +1126,7 @@ mod tests { other => panic!("expected synthetic denial, got {other:?}"), } assert_eq!(first_count.load(Ordering::Relaxed), 1); - assert_eq!(second_count.load(Ordering::Relaxed), 0); + assert_eq!(second_count.load(Ordering::Relaxed), 1); } #[tokio::test] @@ -1050,14 +1135,17 @@ mod tests { #[async_trait] impl Hook for AbortAfterToolHook { - async fn call(&self, input: &ToolResultSummary) -> HookPostToolAction { + async fn call( + &self, + input: &ToolResultSummary, + ) -> Result { self.0.fetch_add(1, Ordering::Relaxed); assert_eq!(input.call_id, "call-id"); assert_eq!(input.tool_name, "TaskList"); assert!(!input.is_error); assert_eq!(input.output.summary, "ok"); assert_eq!(input.output.content.as_deref(), Some("full")); - HookPostToolAction::Abort("post tool abort".into()) + Ok(HookPostToolAction::Abort("post tool abort".into())) } } @@ -1105,12 +1193,15 @@ mod tests { #[async_trait] impl Hook for PauseTurnEndHook { - async fn call(&self, input: &TurnEndInfo) -> HookTurnEndAction { + async fn call( + &self, + input: &TurnEndInfo, + ) -> Result { self.0.fetch_add(1, Ordering::Relaxed); assert_eq!(input.turn_index, 0); assert_eq!(input.tool_calls_count, 0); assert_eq!(input.final_text_preview, "done"); - HookTurnEndAction::Pause + Ok(HookTurnEndAction::Pause) } } @@ -1468,6 +1559,6 @@ mod tests { assert!(matches!(action, PreRequestAction::Cancel(_))); assert!(first_called.load(Ordering::Relaxed)); - assert_eq!(second_count.load(Ordering::Relaxed), 0); + assert_eq!(second_count.load(Ordering::Relaxed), 1); } } diff --git a/crates/worker/src/permission.rs b/crates/worker/src/permission.rs index 22b5c831..21a9c7b5 100644 --- a/crates/worker/src/permission.rs +++ b/crates/worker/src/permission.rs @@ -45,14 +45,17 @@ impl Worker { #[async_trait] impl Hook for PermissionHook { - async fn call(&self, input: &ToolCallSummary) -> HookPreToolAction { - match self.action_for(input) { + async fn call( + &self, + input: &ToolCallSummary, + ) -> Result { + Ok(match self.action_for(input) { ToolPermissionAction::Allow => HookPreToolAction::Continue, ToolPermissionAction::Deny => HookPreToolAction::Deny(permission_denied_message(input)), ToolPermissionAction::Ask => { HookPreToolAction::Deny(permission_ask_unsupported_message(input)) } - } + }) } } @@ -174,7 +177,7 @@ mod tests { )) .await; match denied { - HookPreToolAction::Deny(message) => { + Ok(HookPreToolAction::Deny(message)) => { assert!(message.contains("permission denied")); assert!(message.contains("Bash")); } @@ -192,7 +195,7 @@ mod tests { )) .await; match asked { - HookPreToolAction::Deny(message) => { + Ok(HookPreToolAction::Deny(message)) => { assert!(message.contains("permission ask unsupported")); assert!(message.contains("denied fail-closed")); } diff --git a/crates/worker/src/shutdown_after_idle.rs b/crates/worker/src/shutdown_after_idle.rs index b718cea1..9929ba57 100644 --- a/crates/worker/src/shutdown_after_idle.rs +++ b/crates/worker/src/shutdown_after_idle.rs @@ -70,9 +70,12 @@ impl TicketIntakeReadyShutdownHook { #[async_trait] impl Hook for TicketIntakeReadyShutdownHook { - async fn call(&self, info: &ToolResultSummary) -> HookPostToolAction { + async fn call( + &self, + info: &ToolResultSummary, + ) -> Result { self.observe_tool_result(info); - HookPostToolAction::Continue + Ok(HookPostToolAction::Continue) } } diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index db3536ca..582a6a81 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -40,6 +40,7 @@ use manifest::{ use crate::compact::state::CompactState; use crate::compact::usage_tracker::UsageTracker; +use crate::feature::background::FeatureBackgroundTaskRegistry; use crate::feature::builtin::memory::WorkspaceMemoryBackendError; use crate::feature::builtin::{ MemoryExtractFeature, MemoryExtractState, SessionExploreFeature, SessionExploreState, @@ -50,7 +51,10 @@ use crate::feature::{ FeatureRegistryInstallReport, dedupe_instruction_contributions, }; use crate::hook::{ - Hook, HookRegistryBuilder, OnPromptSubmit, OnTurnEnd, PostToolCall, PreLlmRequest, PreToolCall, + BeforeSessionRewriteAction, BeforeSessionRewriteContext, Hook, HookInvocationContext, + HookRegistry, HookRegistryBuilder, OnPromptSubmit, OnTurnEnd, PostToolCall, PreLlmRequest, + PreToolCall, RunCommittedContext, RunCommittedExit, RunExitContext, SessionRewriteKind, + WorkerStoppingContext, }; use crate::in_flight::InFlightEvents; use crate::internal_worker::{ @@ -63,6 +67,15 @@ const LARGE_PASTE_INLINE_MAX_BYTES: usize = 32 * 1024; const WORKER_ORCHESTRATION_INSTRUCTION_ID: &str = "worker.orchestration"; const WORKER_ORCHESTRATION_PROMPT_REF: &str = "common.worker_orchestration"; +fn hook_run_exit(exit: &EngineRunExit) -> RunCommittedExit { + match exit { + EngineRunExit::Finished => RunCommittedExit::Finished, + EngineRunExit::Paused => RunCommittedExit::Paused, + EngineRunExit::Yielded => RunCommittedExit::Yielded, + EngineRunExit::Interrupted(_) => RunCommittedExit::Interrupted, + } +} + fn worker_orchestration_instruction() -> FeatureInstructionDeclaration { FeatureInstructionDeclaration::new( FeatureInstructionId::builtin(WORKER_ORCHESTRATION_INSTRUCTION_ID), @@ -1104,6 +1117,10 @@ pub struct Worker { /// continue to use `scope`; SubWorkerSpawn validates requested child scope here. delegation_scope: DelegationScope, hook_builder: HookRegistryBuilder, + /// Frozen callback set shared by Engine interception and Worker lifecycle boundaries. + hook_registry: Option>, + /// Executable background tasks registered by successfully installed features. + feature_background_tasks: FeatureBackgroundTaskRegistry, interceptor_installed: bool, /// Shared compaction state (present when threshold is configured). compact_state: Option>, @@ -1258,6 +1275,21 @@ pub struct Worker { } impl Worker { + pub async fn stop_feature_runtime(&mut self, reason: impl Into) { + if let Some(hooks) = self.hook_registry.clone() { + let context = WorkerStoppingContext { + invocation: self.hook_invocation_context(None), + reason: reason.into(), + }; + if let Err(error) = hooks.on_worker_stopping(&context).await { + tracing::warn!(error = %error, "worker-stopping hook requires attention"); + } + } + if let Err(error) = self.feature_background_tasks.shutdown().await { + tracing::warn!(error = %error, "feature background task shutdown failed"); + } + } + pub async fn wait_for_memory_jobs(&mut self) { if let Some(handle) = self.memory_task.take() && let Err(e) = handle.await @@ -1295,6 +1327,8 @@ impl Worker scope: self.scope.clone(), delegation_scope: self.delegation_scope.clone(), hook_builder: HookRegistryBuilder::new(), + hook_registry: None, + feature_background_tasks: FeatureBackgroundTaskRegistry::default(), interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -1488,6 +1522,8 @@ impl Worker { scope, delegation_scope, hook_builder: HookRegistryBuilder::new(), + hook_registry: None, + feature_background_tasks: FeatureBackgroundTaskRegistry::default(), interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -1715,12 +1751,15 @@ impl Worker { /// This deliberately does not scan `.yoi/skills` locally: when a Workspace /// HTTP client is available, catalog/detail/activation authority belongs to /// the Workspace backend API. - pub fn activate_skill(&mut self, name: &str) -> Result + pub async fn activate_skill( + &mut self, + name: &str, + ) -> Result where St: Clone + 'static, { let activation = self.workspace_client().activate_skill(name)?; - self.ensure_segment_head()?; + self.ensure_segment_head().await?; let body = format!( "Agent Skill `{}` activated from {}.\n\n{}", activation.name, activation.provenance.id, activation.body @@ -1821,6 +1860,21 @@ impl Worker { self.engine.as_ref().expect("worker taken during run") } + fn hook_invocation_context(&self, run_id: Option) -> HookInvocationContext { + HookInvocationContext { + workspace_id: self + .workspace_context + .workspace_id() + .map(|id| id.as_str().to_string()), + worker_id: self.manifest.worker.name.clone(), + session_id: self.session.session_id().to_string(), + session_revision: self.session.revision(), + run_id, + turn_index: Some(self.engine().turn_count()), + call_id: None, + } + } + /// Mutable access to the underlying Engine. /// /// Use this to register tools, hooks, or subscribers before calling @@ -1845,6 +1899,10 @@ impl Worker { ) -> FeatureRegistryInstallReport { let worker = self.engine.as_mut().expect("worker taken during run"); let report = registry.install_into_engine(worker, &mut self.hook_builder); + if report.has_errors() { + return report; + } + self.feature_background_tasks = report.background_tasks.clone(); for instruction in report.installed_instruction_contributions() { self.register_feature_instruction(instruction); } @@ -1942,11 +2000,14 @@ impl Worker { } /// Truncate the current segment to just before a previously listed user input. - pub fn rewind_to( + pub async fn rewind_to( &mut self, target: RewindTargetId, expected_head_entries: usize, ) -> Result { + self.prepare_session_rewrite(SessionRewriteKind::Rewind) + .await + .map_err(|error| RewindError::Invalid(error.to_string()))?; let loc = self.segment_state.location(); if target.segment_id != loc.segment_id { return Err(RewindError::Invalid( @@ -2337,6 +2398,7 @@ impl Worker { if !self.interceptor_installed { let builder = std::mem::take(&mut self.hook_builder); let registry = Arc::new(builder.build()); + self.hook_registry = Some(registry.clone()); let (post_run_threshold, request_threshold, retained) = self .manifest @@ -2554,7 +2616,7 @@ impl Worker { self.ensure_interceptor_installed(); self.ensure_system_prompt_materialized().await?; self.cleanup_finished_memory_task(); - self.ensure_segment_head()?; + self.ensure_segment_head().await?; if self.should_pre_run_compact() { self.join_memory_task().await; } @@ -3265,11 +3327,11 @@ impl Worker { /// `ensure_system_prompt_materialized` has just rendered. Subsequent /// calls fall through to entry-count comparison, which auto-forks /// when another writer has appended behind our back. - fn ensure_segment_head(&mut self) -> Result<(), WorkerError> { - let w = self.engine.as_ref().unwrap(); + async fn ensure_segment_head(&mut self) -> Result<(), WorkerError> { let loc = self.segment_state.location(); let entries_written = self.segment_state.entries_written(); if entries_written == 0 { + let w = self.engine.as_ref().unwrap(); let initial = LogEntry::AnnotatedSegmentStart { ts: segment_log::now_millis(), session_id: loc.session_id, @@ -3305,6 +3367,9 @@ impl Worker { // state up to that turn). The new SegmentStart replaces the mirror // and is broadcast through the sink so existing subscribers reset // their view. + self.prepare_session_rewrite(SessionRewriteKind::Fork) + .await?; + let w = self.engine.as_ref().unwrap(); let fork_segment_id = session_store::new_segment_id(); let entry = LogEntry::AnnotatedSegmentStart { ts: segment_log::now_millis(), @@ -3371,10 +3436,40 @@ impl Worker { where St: Clone + 'static, { + let run_id = uuid::Uuid::now_v7().to_string(); + let hook_exit = hook_run_exit(&result); + if let Some(hooks) = self.hook_registry.clone() { + let context = RunExitContext { + invocation: self.hook_invocation_context(Some(run_id.clone())), + exit: hook_exit, + history_len: self.session.history().len(), + }; + if let Err(error) = hooks.on_run_exit(&context).await { + tracing::warn!(error = %error, "run-exit hook failed; preserving terminal commit"); + } + } if matches!(&result, EngineRunExit::Interrupted(_)) { self.terminalize_orphan_tool_calls()?; } self.persist_turn(history_before, &result).await?; + let committed_invocation = self.hook_invocation_context(Some(run_id)); + if let Some(hooks) = self.hook_registry.clone() { + let context = RunCommittedContext { + invocation: committed_invocation.clone(), + exit: hook_exit, + committed_history: self.session.history().entries().to_vec(), + committed_history_len: self.session.history().len(), + }; + if let Err(error) = hooks.on_run_committed(&context).await { + tracing::warn!(error = %error, "run-committed hook requires attention"); + } + } + if let Err(error) = self + .feature_background_tasks + .start_run_committed(committed_invocation) + { + tracing::warn!(error = %error, "run-committed background task start failed"); + } if matches!(result, EngineRunExit::Yielded) { self.last_run_interrupted = true; @@ -3552,6 +3647,35 @@ impl Worker { /// The controller only calls this while Idle. Paused turns keep their /// interrupted Engine state intact and are intentionally rejected before /// this method is reached. + async fn prepare_session_rewrite( + &mut self, + kind: SessionRewriteKind, + ) -> Result<(), WorkerError> { + self.feature_background_tasks + .before_session_rewrite() + .await + .map_err(|error| WorkerError::FeatureLifecycle(error.to_string()))?; + if let Some(hooks) = self.hook_registry.clone() { + let context = BeforeSessionRewriteContext { + invocation: self.hook_invocation_context(None), + kind, + current_history: self.session.history().entries().to_vec(), + current_history_len: self.session.history().len(), + }; + match hooks + .before_session_rewrite(&context) + .await + .map_err(|error| WorkerError::FeatureLifecycle(error.to_string()))? + { + BeforeSessionRewriteAction::Continue => {} + BeforeSessionRewriteAction::Deny(reason) => { + return Err(WorkerError::FeatureLifecycle(reason)); + } + } + } + Ok(()) + } + pub async fn manual_compact(&mut self) -> Result { if self.manifest.compaction.is_none() { let message = @@ -3568,7 +3692,7 @@ impl Worker { self.ensure_interceptor_installed(); self.cleanup_finished_memory_task(); - self.ensure_segment_head()?; + self.ensure_segment_head().await?; let state = self.compact_state.clone(); if state.as_ref().is_some_and(|s| s.is_disabled()) { @@ -3767,6 +3891,8 @@ impl Worker { /// Runs one parent-owned observable compaction service and returns the new /// Segment ID. Lifecycle revisions are committed before they are broadcast. pub async fn compact(&mut self, retained_tokens: u64) -> Result { + self.prepare_session_rewrite(SessionRewriteKind::Compact) + .await?; let mut lifecycle = CompactionLifecycle { schema_version: 2, compaction_id: uuid::Uuid::now_v7().to_string(), @@ -5302,6 +5428,8 @@ where scope, delegation_scope: common.delegation_scope, hook_builder: HookRegistryBuilder::new(), + hook_registry: None, + feature_background_tasks: FeatureBackgroundTaskRegistry::default(), interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -5386,6 +5514,8 @@ where scope, delegation_scope: common.delegation_scope, hook_builder: HookRegistryBuilder::new(), + hook_registry: None, + feature_background_tasks: FeatureBackgroundTaskRegistry::default(), interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -5505,6 +5635,8 @@ where scope, delegation_scope: common.delegation_scope, hook_builder: HookRegistryBuilder::new(), + hook_registry: None, + feature_background_tasks: FeatureBackgroundTaskRegistry::default(), interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -5879,6 +6011,8 @@ where scope, delegation_scope: common.delegation_scope, hook_builder: HookRegistryBuilder::new(), + hook_registry: None, + feature_background_tasks: FeatureBackgroundTaskRegistry::default(), interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -6553,6 +6687,9 @@ pub enum WorkerError { #[error("invalid durable Worker state: {0}")] InvalidState(String), + #[error("feature lifecycle rejected operation: {0}")] + FeatureLifecycle(String), + #[error("Flow input rejected: {0}")] FlowInput(String), @@ -7776,15 +7913,17 @@ mod build_summary_prompt_tests { async fn call( &self, _input: &crate::hook::ToolCallSummary, - ) -> crate::hook::HookPreToolAction { - if self - .should_pause - .swap(false, std::sync::atomic::Ordering::SeqCst) - { - crate::hook::HookPreToolAction::Pause - } else { - crate::hook::HookPreToolAction::Continue - } + ) -> Result { + Ok( + if self + .should_pause + .swap(false, std::sync::atomic::Ordering::SeqCst) + { + crate::hook::HookPreToolAction::Pause + } else { + crate::hook::HookPreToolAction::Continue + }, + ) } } @@ -7837,7 +7976,7 @@ mod build_summary_prompt_tests { ) .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); worker.last_run_interrupted = true; worker.engine_mut().set_active_run_turn_count(Some(3)); @@ -7880,7 +8019,7 @@ mod build_summary_prompt_tests { ) .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); worker.engine_mut().set_turn_count(7); worker.last_run_interrupted = true; worker.engine_mut().set_active_run_turn_count(Some(3)); @@ -7900,7 +8039,7 @@ mod build_summary_prompt_tests { ) .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); let fork_segment_id = worker.segment_id(); assert_ne!(fork_segment_id, source_segment_id); @@ -7944,7 +8083,7 @@ mod build_summary_prompt_tests { ) .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); let report = worker .install_runtime_flow_transition_feature() .expect("scoped Workspace Flow feature"); @@ -7978,7 +8117,7 @@ mod build_summary_prompt_tests { ) .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); let disabled = worker.prepare_flow_input(vec![Segment::Flow { selector: "builtin:coder-review".to_string(), }]); @@ -8254,7 +8393,7 @@ mod build_summary_prompt_tests { ) .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); std::fs::write( temp.path() .join(worker.session_id().to_string()) @@ -8303,7 +8442,7 @@ mod build_summary_prompt_tests { ) .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); (dir, worker) } @@ -8463,7 +8602,7 @@ mod build_summary_prompt_tests { let expected_truncate_entries = targets[0].truncate_entries; let target = targets[0].id.clone(); - let applied = worker.rewind_to(target, head_entries).unwrap(); + let applied = worker.rewind_to(target, head_entries).await.unwrap(); assert_eq!(preview_segments(&applied.input), "second message"); assert_eq!( @@ -8520,6 +8659,7 @@ mod build_summary_prompt_tests { let applied = worker .rewind_to(targets[0].id.clone(), head_entries) + .await .unwrap(); assert_eq!(applied.summary.truncated_to_entries, 5); assert!(matches!( @@ -8562,7 +8702,7 @@ mod build_summary_prompt_tests { payload: serde_json::json!({"value": true}), }, ); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); let fork_location = worker.segment_state.location(); assert_ne!(fork_location.segment_id, source_location.segment_id); let fork_entries = worker @@ -8593,6 +8733,7 @@ mod build_summary_prompt_tests { let err = worker .rewind_to(targets[0].id.clone(), head_entries) + .await .unwrap_err() .to_string(); @@ -8673,7 +8814,7 @@ mod build_summary_prompt_tests { .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); worker.wire_history_persistence(); worker.set_history_for_test(vec![ Item::tool_call("call-known", "Read", "{}"), @@ -8790,7 +8931,7 @@ mod build_summary_prompt_tests { .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); worker.wire_history_persistence(); worker.set_history_for_test(vec![Item::tool_call("call-1", "Read", "{}")]); @@ -8870,7 +9011,7 @@ mod build_summary_prompt_tests { .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); worker.wire_history_persistence(); let dangling_call = Item::tool_call("call-1", "SideEffect", "{}"); worker @@ -9186,8 +9327,8 @@ mod build_summary_prompt_tests { std::fs::create_dir_all(&cwd).unwrap(); let scope = Scope::writable(&cwd).unwrap(); let authority = WorkerFilesystemAuthority::local(cwd.clone(), cwd.clone()); - let mut worker = tokio::runtime::Runtime::new() - .unwrap() + let runtime = tokio::runtime::Runtime::new().unwrap(); + let mut worker = runtime .block_on(Worker::new( manifest, Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), @@ -9204,7 +9345,9 @@ mod build_summary_prompt_tests { )) .unwrap(); - let activation = worker.activate_skill("triage-errors").unwrap(); + let activation = runtime + .block_on(worker.activate_skill("triage-errors")) + .unwrap(); assert_eq!(activation.name, "triage-errors"); server.join().unwrap(); @@ -9274,7 +9417,7 @@ mod build_summary_prompt_tests { ) .await .unwrap(); - worker.ensure_segment_head().unwrap(); + worker.ensure_segment_head().await.unwrap(); worker.wire_history_persistence(); let evidence = Item::user_message( "The cancellation regression must leave this evidence available for retry.", From eecb116709b3849bc8b4aa4f910f43d2142b9169 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 16:04:10 +0900 Subject: [PATCH 10/22] fix: bound feature lifecycle execution --- crates/worker/src/feature.rs | 190 +++++++++------------- crates/worker/src/feature/background.rs | 187 ++++++++++++++++++--- crates/worker/src/hook.rs | 205 +++++++++++++++++++----- crates/worker/src/ipc/interceptor.rs | 61 ++++--- crates/worker/src/lib.rs | 1 + crates/worker/src/session_capture.rs | 8 +- crates/worker/src/worker.rs | 88 +++++++--- 7 files changed, 518 insertions(+), 222 deletions(-) diff --git a/crates/worker/src/feature.rs b/crates/worker/src/feature.rs index c13bb603..e91d44a6 100644 --- a/crates/worker/src/feature.rs +++ b/crates/worker/src/feature.rs @@ -24,8 +24,8 @@ use serde::{Deserialize, Serialize}; use thiserror::Error; use crate::hook::{ - BeforeSessionRewrite, Hook, HookFailurePolicy, HookRegistryBuilder, OnPromptSubmit, OnTurnEnd, - PostToolCall, PreLlmRequest, PreToolCall, RunCommitted, RunExit, WorkerStopping, + BeforeSessionRewrite, Hook, HookExecutionPolicy, HookRegistryBuilder, OnPromptSubmit, + OnTurnEnd, PostToolCall, PreLlmRequest, PreToolCall, RunCommitted, RunExit, WorkerStopping, }; use background::{ BackgroundTaskSpec, FeatureBackgroundTask, FeatureBackgroundTaskRegistry, @@ -903,46 +903,6 @@ fn reject_undeclared_contribution( error } -/// Model-visible durable notification sink skeleton. The first slice exposes -/// the boundary without implementing a new event channel. -pub struct FeatureNotificationSink<'a> { - report: &'a mut FeatureInstallReport, -} - -impl FeatureNotificationSink<'_> { - pub fn notify_model(&mut self, message: impl Into) -> Result<(), FeatureInstallError> { - let message = message.into(); - self.report.diagnostics.push(FeatureDiagnostic::warning(format!( - "model notification requested during feature installation but no durable Notify host is attached: {message}" - ))); - self.report.mark_skipped( - FeatureContributionKind::Notification, - "notify_model", - "durable Notify/SystemItem host is not connected during feature installation", - ); - Ok(()) - } -} - -/// Transient human-facing alert sink skeleton. -pub struct FeatureAlertSink<'a> { - report: &'a mut FeatureInstallReport, -} - -impl FeatureAlertSink<'_> { - pub fn alert(&mut self, message: impl Into) { - let message = message.into(); - self.report - .diagnostics - .push(FeatureDiagnostic::info(format!("feature alert: {message}"))); - self.report.mark_skipped( - FeatureContributionKind::Alert, - "alert", - "transient alert host is not connected during feature installation", - ); - } -} - /// Diagnostic sink available to feature installers. pub struct FeatureDiagnosticSink<'a> { report: &'a mut FeatureInstallReport, @@ -1073,16 +1033,18 @@ impl HookContributionRegistrar<'_> { pub fn add_prompt_submit( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::PromptSubmit); self.require_declared(&declaration)?; - self.hook_builder.add_named_on_prompt_submit( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_on_prompt_submit( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1090,16 +1052,18 @@ impl HookContributionRegistrar<'_> { pub fn add_pre_llm_request( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::PreLlmRequest); self.require_declared(&declaration)?; - self.hook_builder.add_named_pre_llm_request( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_pre_llm_request( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1109,22 +1073,24 @@ impl HookContributionRegistrar<'_> { name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - self.add_pre_llm_request(name, HookFailurePolicy::FailClosed, hook) + self.add_pre_llm_request(name, HookExecutionPolicy::fail_closed(), hook) } pub fn add_pre_tool_call_with_policy( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::PreToolCall); self.require_declared(&declaration)?; - self.hook_builder.add_named_pre_tool_call( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_pre_tool_call( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1134,22 +1100,24 @@ impl HookContributionRegistrar<'_> { name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - self.add_pre_tool_call_with_policy(name, HookFailurePolicy::FailClosed, hook) + self.add_pre_tool_call_with_policy(name, HookExecutionPolicy::fail_closed(), hook) } pub fn add_post_tool_call( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::PostToolCall); self.require_declared(&declaration)?; - self.hook_builder.add_named_post_tool_call( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_post_tool_call( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1159,22 +1127,24 @@ impl HookContributionRegistrar<'_> { name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - self.add_post_tool_call(name, HookFailurePolicy::FailClosed, hook) + self.add_post_tool_call(name, HookExecutionPolicy::fail_closed(), hook) } pub fn add_assistant_turn_end( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::AssistantTurnEnd); self.require_declared(&declaration)?; - self.hook_builder.add_named_on_turn_end( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_on_turn_end( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1184,22 +1154,24 @@ impl HookContributionRegistrar<'_> { name: impl Into, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { - self.add_assistant_turn_end(name, HookFailurePolicy::FailClosed, hook) + self.add_assistant_turn_end(name, HookExecutionPolicy::fail_closed(), hook) } pub fn add_run_exit( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::RunExit); self.require_declared(&declaration)?; - self.hook_builder.add_named_run_exit( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_run_exit( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1207,16 +1179,18 @@ impl HookContributionRegistrar<'_> { pub fn add_run_committed( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::RunCommitted); self.require_declared(&declaration)?; - self.hook_builder.add_named_run_committed( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_run_committed( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1224,16 +1198,18 @@ impl HookContributionRegistrar<'_> { pub fn add_before_session_rewrite( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::BeforeSessionRewrite); self.require_declared(&declaration)?; - self.hook_builder.add_named_before_session_rewrite( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_before_session_rewrite( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1241,16 +1217,18 @@ impl HookContributionRegistrar<'_> { pub fn add_worker_stopping( &mut self, name: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook + 'static, ) -> Result<(), FeatureInstallError> { let declaration = HookDeclaration::new(name, FeatureHookPoint::WorkerStopping); self.require_declared(&declaration)?; - self.hook_builder.add_named_worker_stopping( - format!("{}:{}", self.feature_id, declaration.name), - policy, - hook, - ); + self.hook_builder + .add_named_worker_stopping( + format!("{}:{}", self.feature_id, declaration.name), + policy, + hook, + ) + .map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?; self.record(declaration); Ok(()) } @@ -1592,18 +1570,6 @@ impl FeatureInstallContext<'_> { } } - pub fn notifications(&mut self) -> FeatureNotificationSink<'_> { - FeatureNotificationSink { - report: self.report, - } - } - - pub fn alerts(&mut self) -> FeatureAlertSink<'_> { - FeatureAlertSink { - report: self.report, - } - } - pub fn diagnostics(&mut self) -> FeatureDiagnosticSink<'_> { FeatureDiagnosticSink { report: self.report, diff --git a/crates/worker/src/feature/background.rs b/crates/worker/src/feature/background.rs index 76254e5e..695e45d3 100644 --- a/crates/worker/src/feature/background.rs +++ b/crates/worker/src/feature/background.rs @@ -12,7 +12,9 @@ use std::time::Duration; use async_trait::async_trait; use serde::{Deserialize, Serialize}; +#[cfg(test)] use tokio::sync::Notify; +use tokio::sync::watch; use tokio::task::JoinHandle; use super::{BackgroundTaskDeclaration, FeatureId, FeatureInstallError}; @@ -22,6 +24,7 @@ const MAX_TASK_CONCURRENCY: u16 = 64; const MAX_TASK_ATTEMPTS: u16 = 16; const MAX_TASK_TIMEOUT_MS: u64 = 24 * 60 * 60 * 1_000; const MAX_RETAINED_DIAGNOSTICS: usize = 128; +const TASK_SETTLE_TIMEOUT_MS: u64 = 30_000; #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] @@ -107,31 +110,43 @@ pub struct BackgroundTaskContext { pub feature_id: FeatureId, pub task_name: String, pub execution_id: u64, + pub session_generation: u64, pub attempt: u16, } -#[derive(Clone, Default)] +#[derive(Clone)] pub struct BackgroundTaskCancellation { - cancelled: Arc, - notify: Arc, + sender: Arc>, +} + +impl Default for BackgroundTaskCancellation { + fn default() -> Self { + let (sender, _receiver) = watch::channel(false); + Self { + sender: Arc::new(sender), + } + } } impl BackgroundTaskCancellation { pub fn is_cancelled(&self) -> bool { - self.cancelled.load(Ordering::Acquire) + *self.sender.borrow() } pub async fn cancelled(&self) { - if self.is_cancelled() { + let mut receiver = self.sender.subscribe(); + if *receiver.borrow_and_update() { return; } - self.notify.notified().await; + while receiver.changed().await.is_ok() { + if *receiver.borrow_and_update() { + return; + } + } } fn cancel(&self) { - if !self.cancelled.swap(true, Ordering::AcqRel) { - self.notify.notify_one(); - } + self.sender.send_replace(true); } } @@ -200,6 +215,7 @@ impl FeatureBackgroundTaskRegistryBuilder { running: Mutex::new(BTreeMap::new()), diagnostics: Mutex::new(Vec::new()), next_execution_id: AtomicU64::new(1), + session_generation: AtomicU64::new(1), accepting: AtomicBool::new(true), }), } @@ -218,6 +234,7 @@ struct RegistryInner { running: Mutex>, diagnostics: Mutex>, next_execution_id: AtomicU64, + session_generation: AtomicU64, accepting: AtomicBool, } @@ -280,6 +297,7 @@ pub enum BackgroundTaskOutcome { TimedOut, Failed(HookError), JoinFailed, + StaleGenerationDiscarded, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -330,6 +348,7 @@ impl FeatureBackgroundTaskRegistry { } let execution_id = self.inner.next_execution_id.fetch_add(1, Ordering::Relaxed); + let session_generation = self.inner.session_generation.load(Ordering::Acquire); let cancellation = BackgroundTaskCancellation::default(); let task_cancellation = cancellation.clone(); let task = Arc::clone(®istration.task); @@ -347,10 +366,17 @@ impl FeatureBackgroundTaskRegistry { task_feature_id.clone(), task_task_name.clone(), execution_id, + session_generation, task_cancellation, ) .await; if let Some(inner) = weak_inner.upgrade() { + let outcome = + if inner.session_generation.load(Ordering::Acquire) == session_generation { + outcome + } else { + BackgroundTaskOutcome::StaleGenerationDiscarded + }; let mut diagnostics = inner.diagnostics.lock().expect("diagnostics poisoned"); diagnostics.push(BackgroundTaskDiagnostic { execution_id, @@ -405,7 +431,9 @@ impl FeatureBackgroundTaskRegistry { } pub async fn before_session_rewrite(&self) -> Result<(), HookError> { - self.settle(false).await + self.settle(false).await?; + self.inner.session_generation.fetch_add(1, Ordering::AcqRel); + Ok(()) } pub async fn shutdown(&self) -> Result<(), HookError> { @@ -469,17 +497,31 @@ impl FeatureBackgroundTaskRegistry { } } + let settle_deadline = + tokio::time::Instant::now() + Duration::from_millis(TASK_SETTLE_TIMEOUT_MS); let mut join_failed = false; - for (execution_id, feature_id, task_name, handle) in waiting { - if handle.await.is_err() { - join_failed = true; + for (execution_id, feature_id, task_name, mut handle) in waiting { + let outcome = match tokio::time::timeout_at(settle_deadline, &mut handle).await { + Ok(Ok(())) => None, + Ok(Err(_)) => { + join_failed = true; + Some(BackgroundTaskOutcome::JoinFailed) + } + Err(_) => { + join_failed = true; + handle.abort(); + let _ = handle.await; + Some(BackgroundTaskOutcome::TimedOut) + } + }; + if let Some(outcome) = outcome { let mut diagnostics = self.inner.diagnostics.lock().expect("diagnostics poisoned"); diagnostics.push(BackgroundTaskDiagnostic { execution_id, feature_id, task_name, attempts: 0, - outcome: BackgroundTaskOutcome::JoinFailed, + outcome, }); if diagnostics.len() > MAX_RETAINED_DIAGNOSTICS { let remove = diagnostics.len() - MAX_RETAINED_DIAGNOSTICS; @@ -490,7 +532,7 @@ impl FeatureBackgroundTaskRegistry { if join_failed { return Err(HookError::new( HookErrorCategory::Internal, - "feature background task join failed", + "feature background task settlement failed or exceeded its host deadline", )); } Ok(()) @@ -504,6 +546,7 @@ async fn execute_task( feature_id: FeatureId, task_name: String, execution_id: u64, + session_generation: u64, cancellation: BackgroundTaskCancellation, ) -> (u16, BackgroundTaskOutcome) { let (max_attempts, delay_ms) = match spec.retry { @@ -513,22 +556,24 @@ async fn execute_task( delay_ms, } => (max_attempts, delay_ms), }; + let deadline = tokio::time::Instant::now() + Duration::from_millis(spec.timeout_ms); for attempt in 1..=max_attempts { if cancellation.is_cancelled() { return (attempt, BackgroundTaskOutcome::Cancelled); } + if tokio::time::Instant::now() >= deadline { + return (attempt, BackgroundTaskOutcome::TimedOut); + } let context = BackgroundTaskContext { invocation: invocation.clone(), feature_id: feature_id.clone(), task_name: task_name.clone(), execution_id, + session_generation, attempt, }; - let result = tokio::time::timeout( - Duration::from_millis(spec.timeout_ms), - task.run(context, cancellation.clone()), - ) - .await; + let result = + tokio::time::timeout_at(deadline, task.run(context, cancellation.clone())).await; match result { Ok(Ok(())) => return (attempt, BackgroundTaskOutcome::Completed), Ok(Err(_error)) if cancellation.is_cancelled() => { @@ -541,8 +586,16 @@ async fn execute_task( Err(_) => return (attempt, BackgroundTaskOutcome::TimedOut), } if delay_ms > 0 { + let retry_at = std::cmp::min( + deadline, + tokio::time::Instant::now() + Duration::from_millis(delay_ms), + ); tokio::select! { - () = tokio::time::sleep(Duration::from_millis(delay_ms)) => {} + () = tokio::time::sleep_until(retry_at) => { + if tokio::time::Instant::now() >= deadline { + return (attempt, BackgroundTaskOutcome::TimedOut); + } + } () = cancellation.cancelled() => { return (attempt, BackgroundTaskOutcome::Cancelled); } @@ -684,6 +737,98 @@ mod tests { ); } + #[tokio::test] + async fn cancellation_observed_when_cancel_races_with_wait_registration() { + for _ in 0..100 { + let cancellation = BackgroundTaskCancellation::default(); + let waiter = cancellation.clone(); + let task = tokio::spawn(async move { + waiter.cancelled().await; + }); + cancellation.cancel(); + tokio::time::timeout(Duration::from_millis(100), task) + .await + .expect("watch-backed cancellation must not lose the transition") + .unwrap(); + } + } + + struct AlwaysFails { + calls: Arc, + } + + #[async_trait] + impl FeatureBackgroundTask for AlwaysFails { + async fn run( + &self, + _context: BackgroundTaskContext, + _cancellation: BackgroundTaskCancellation, + ) -> Result<(), HookError> { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(HookError::new(HookErrorCategory::Dependency, "retry")) + } + } + + #[tokio::test] + async fn retry_delay_consumes_one_total_execution_deadline() { + let feature = FeatureId::builtin("deadline-test"); + let declaration = BackgroundTaskDeclaration::worker_managed("deadline", "deadline"); + let calls = Arc::new(AtomicUsize::new(0)); + let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); + let mut spec = BackgroundTaskSpec::single_flight(declaration, Duration::from_millis(20)); + spec.shutdown = BackgroundTaskShutdownPolicy::Wait; + spec.retry = BackgroundTaskRetryPolicy::Bounded { + max_attempts: 3, + delay_ms: 100, + }; + builder + .register( + feature.clone(), + spec, + AlwaysFails { + calls: Arc::clone(&calls), + }, + ) + .unwrap(); + let registry = builder.build(); + registry.start(&feature, "deadline", invocation()).unwrap(); + registry.shutdown().await.unwrap(); + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!( + registry.diagnostics()[0].outcome, + BackgroundTaskOutcome::TimedOut + ); + } + + #[tokio::test] + async fn rewrite_joins_old_tasks_before_advancing_session_generation() { + let feature = FeatureId::builtin("generation-test"); + let declaration = BackgroundTaskDeclaration::worker_managed("generation", "generation"); + let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); + builder + .register( + feature.clone(), + BackgroundTaskSpec::single_flight(declaration, Duration::from_secs(1)), + WaitForCancellation, + ) + .unwrap(); + let registry = builder.build(); + registry + .start(&feature, "generation", invocation()) + .unwrap(); + assert_eq!(registry.inner.session_generation.load(Ordering::Acquire), 1); + + registry.before_session_rewrite().await.unwrap(); + + assert_eq!(registry.inner.session_generation.load(Ordering::Acquire), 2); + assert_eq!(registry.diagnostics().len(), 1); + assert_eq!( + registry.diagnostics()[0].outcome, + BackgroundTaskOutcome::Cancelled + ); + } + #[tokio::test] async fn block_policy_fences_rewrite_without_detaching_the_task() { let feature = FeatureId::builtin("rewrite-test"); diff --git a/crates/worker/src/hook.rs b/crates/worker/src/hook.rs index b326697a..1bebebfb 100644 --- a/crates/worker/src/hook.rs +++ b/crates/worker/src/hook.rs @@ -18,7 +18,6 @@ use std::ops::Deref; use std::sync::{Arc, Mutex}; -use agen::HistoryEntry; use agen::interceptor::{ PostToolAction, PreRequestAction, PreToolAction, PromptAction, TurnEndAction, }; @@ -29,7 +28,7 @@ use serde_json::Value; use session_store::{SystemItem, SystemReminder}; use thiserror::Error; -use crate::session_history::SessionHistoryMetadata; +use crate::SessionEntryRef; const HOOK_DIAGNOSTIC_MAX_BYTES: usize = 1_024; @@ -79,6 +78,39 @@ pub enum HookFailurePolicy { AttentionRequired, } +const DEFAULT_HOOK_TIMEOUT_MS: u64 = 30_000; +const MAX_HOOK_TIMEOUT_MS: u64 = 120_000; + +/// Host-enforced execution and failure budget for one hook registration. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct HookExecutionPolicy { + pub failure: HookFailurePolicy, + pub timeout_ms: u64, +} + +impl HookExecutionPolicy { + pub const fn new(failure: HookFailurePolicy, timeout_ms: u64) -> Self { + Self { + failure, + timeout_ms, + } + } + + pub const fn fail_closed() -> Self { + Self::new(HookFailurePolicy::FailClosed, DEFAULT_HOOK_TIMEOUT_MS) + } + + fn validate(self) -> Result { + if self.timeout_ms == 0 || self.timeout_ms > MAX_HOOK_TIMEOUT_MS { + return Err(HookError::new( + HookErrorCategory::InvalidInput, + format!("hook timeout_ms must be within 1..={MAX_HOOK_TIMEOUT_MS}"), + )); + } + Ok(self) + } +} + fn bounded_utf8(mut value: String, max_bytes: usize) -> String { if value.len() <= max_bytes { return value; @@ -421,6 +453,16 @@ pub trait Hook: Send + Sync { // Hook Registry // ============================================================================= +/// Bounded identity-only view of the committed history visible at a lifecycle +/// boundary. Hook context never carries history payloads; a Feature that needs +/// content must use a separately granted session-exploration service. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct HookHistoryRange { + pub first_entry_ref: Option, + pub last_entry_ref: Option, + pub entry_count: usize, +} + /// Stable provenance attached to every Worker lifecycle callback. #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct HookInvocationContext { @@ -452,8 +494,7 @@ pub struct RunExitContext { pub struct RunCommittedContext { pub invocation: HookInvocationContext, pub exit: RunCommittedExit, - pub committed_history: Vec>, - pub committed_history_len: usize, + pub committed_history: HookHistoryRange, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -468,8 +509,7 @@ pub enum SessionRewriteKind { pub struct BeforeSessionRewriteContext { pub invocation: HookInvocationContext, pub kind: SessionRewriteKind, - pub current_history: Vec>, - pub current_history_len: usize, + pub current_history: HookHistoryRange, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -511,20 +551,28 @@ impl HookEventKind for WorkerStopping { pub(crate) struct RegisteredHook { owner: String, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: Box>, } impl RegisteredHook { pub(crate) async fn call(&self, input: &E::Input) -> Result { - self.hook - .call(input) - .await - .map_err(|source| HookExecutionError { - owner: self.owner.clone(), - policy: self.policy, - source, - }) + let result = tokio::time::timeout( + std::time::Duration::from_millis(self.policy.timeout_ms), + self.hook.call(input), + ) + .await + .unwrap_or_else(|_| { + Err(HookError::new( + HookErrorCategory::Timeout, + "hook exceeded its host-enforced execution deadline", + )) + }); + result.map_err(|source| HookExecutionError { + owner: self.owner.clone(), + policy: self.policy.failure, + source, + }) } pub(crate) async fn call_optional( @@ -567,20 +615,23 @@ pub struct HookRegistryBuilder { macro_rules! add_hook_methods { ($default:ident, $named:ident, $field:ident, $event:ty) => { pub fn $default(&mut self, hook: impl Hook<$event> + 'static) { - self.$named("worker.host", HookFailurePolicy::FailClosed, hook); + self.$named("worker.host", HookExecutionPolicy::fail_closed(), hook) + .expect("host hook policy is valid"); } pub fn $named( &mut self, owner: impl Into, - policy: HookFailurePolicy, + policy: HookExecutionPolicy, hook: impl Hook<$event> + 'static, - ) { + ) -> Result<(), HookError> { + let policy = policy.validate()?; self.$field.push(RegisteredHook { owner: owner.into(), policy, hook: Box::new(hook), }); + Ok(()) } }; } @@ -706,6 +757,17 @@ impl HookRegistry { } } + pub(crate) fn record_chain_timeout(&self, lifecycle: &str) { + self.record_diagnostic(HookExecutionError { + owner: format!("worker.{lifecycle}"), + policy: HookFailurePolicy::AttentionRequired, + source: HookError::new( + HookErrorCategory::Timeout, + "hook chain exceeded its host-enforced lifecycle deadline", + ), + }); + } + pub fn diagnostics(&self) -> Vec { self.diagnostics .lock() @@ -899,28 +961,31 @@ mod tests { call_id: None, }, kind: SessionRewriteKind::Compact, - current_history: Vec::new(), - current_history_len: 8, + current_history: HookHistoryRange::default(), } } #[tokio::test] async fn rewrite_denials_are_resolved_by_owner_not_registration_order() { let mut builder = HookRegistryBuilder::new(); - builder.add_named_before_session_rewrite( - "z-feature", - HookFailurePolicy::FailClosed, - RewriteHook { - action: BeforeSessionRewriteAction::Deny("z denied".into()), - }, - ); - builder.add_named_before_session_rewrite( - "a-feature", - HookFailurePolicy::FailClosed, - RewriteHook { - action: BeforeSessionRewriteAction::Deny("a denied".into()), - }, - ); + builder + .add_named_before_session_rewrite( + "z-feature", + HookExecutionPolicy::fail_closed(), + RewriteHook { + action: BeforeSessionRewriteAction::Deny("z denied".into()), + }, + ) + .unwrap(); + builder + .add_named_before_session_rewrite( + "a-feature", + HookExecutionPolicy::fail_closed(), + RewriteHook { + action: BeforeSessionRewriteAction::Deny("a denied".into()), + }, + ) + .unwrap(); assert_eq!( builder @@ -935,11 +1000,13 @@ mod tests { #[tokio::test] async fn hook_failure_policy_is_applied_at_the_registry_boundary() { let mut fail_open = HookRegistryBuilder::new(); - fail_open.add_named_before_session_rewrite( - "feature", - HookFailurePolicy::FailOpenWithDiagnostic, - FailingRewriteHook, - ); + fail_open + .add_named_before_session_rewrite( + "feature", + HookExecutionPolicy::new(HookFailurePolicy::FailOpenWithDiagnostic, 30_000), + FailingRewriteHook, + ) + .unwrap(); let fail_open = fail_open.build(); assert_eq!( fail_open @@ -951,11 +1018,13 @@ mod tests { assert_eq!(fail_open.diagnostics().len(), 1); let mut fail_closed = HookRegistryBuilder::new(); - fail_closed.add_named_before_session_rewrite( - "feature", - HookFailurePolicy::FailClosed, - FailingRewriteHook, - ); + fail_closed + .add_named_before_session_rewrite( + "feature", + HookExecutionPolicy::fail_closed(), + FailingRewriteHook, + ) + .unwrap(); let error = fail_closed .build() .before_session_rewrite(&rewrite_context()) @@ -964,6 +1033,54 @@ mod tests { assert_eq!(error.source.category, HookErrorCategory::Dependency); } + struct NeverReturns; + + #[async_trait] + impl Hook for NeverReturns { + async fn call( + &self, + _input: &BeforeSessionRewriteContext, + ) -> Result { + std::future::pending().await + } + } + + #[tokio::test] + async fn hook_execution_deadline_cancels_a_non_cooperative_callback() { + let mut builder = HookRegistryBuilder::new(); + builder + .add_named_before_session_rewrite( + "feature", + HookExecutionPolicy::new(HookFailurePolicy::FailClosed, 10), + NeverReturns, + ) + .unwrap(); + let registry = builder.build(); + + let error = tokio::time::timeout( + std::time::Duration::from_millis(100), + registry.before_session_rewrite(&rewrite_context()), + ) + .await + .expect("host hook deadline must terminate the callback") + .unwrap_err(); + assert_eq!(error.source.category, HookErrorCategory::Timeout); + assert_eq!(registry.diagnostics().len(), 1); + } + + #[test] + fn invalid_hook_execution_budget_is_rejected_before_registration() { + let mut builder = HookRegistryBuilder::new(); + let error = builder + .add_named_before_session_rewrite( + "feature", + HookExecutionPolicy::new(HookFailurePolicy::FailClosed, 0), + NeverReturns, + ) + .unwrap_err(); + assert_eq!(error.category, HookErrorCategory::InvalidInput); + } + #[test] fn hook_diagnostics_are_utf8_bounded() { let error = HookError::new(HookErrorCategory::Internal, "界".repeat(1_000)); diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index 0e5c5b40..30bf4e12 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -30,9 +30,9 @@ use crate::compact::usage_tracker::UsageTracker; use session_store::SystemItem; use crate::hook::{ - HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookPromptAction, HookRegistry, - HookTurnEndAction, PreRequestContext, PreRequestInfo, PromptSubmitInfo, SystemItemAppendHandle, - ToolCallSummary, ToolResultSummary, TurnEndInfo, + HookEventKind, HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookPromptAction, + HookRegistry, HookTurnEndAction, PreRequestContext, PreRequestInfo, PromptSubmitInfo, + RegisteredHook, SystemItemAppendHandle, ToolCallSummary, ToolResultSummary, TurnEndInfo, }; use crate::ipc::notify_buffer::{NotifyBuffer, build_system_item_with_provenance}; use crate::prompt::catalog::PromptCatalog; @@ -43,6 +43,32 @@ use agen::token_counter::total_tokens; /// Maximum number of bytes copied into `TurnEndInfo::final_text_preview`. const FINAL_TEXT_PREVIEW_LIMIT: usize = 512; +const INLINE_HOOK_CHAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); + +async fn call_hook_before_deadline( + hook: &RegisteredHook, + input: &E::Input, + deadline: tokio::time::Instant, +) -> Result, InterceptorError> { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + return Err(InterceptorError::new( + InterceptorErrorCategory::Policy, + "Worker hook chain exceeded its host-enforced lifecycle deadline", + )); + } + tokio::time::timeout(remaining, hook.call_optional(input)) + .await + .map_err(|_| { + InterceptorError::new( + InterceptorErrorCategory::Policy, + "Worker hook chain exceeded its host-enforced lifecycle deadline", + ) + })? + .map_err(|error| { + InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) + }) +} pub(crate) struct WorkerInterceptor { registry: Arc, @@ -246,12 +272,10 @@ impl Interceptor for WorkerInterceptor { input_text: extract_message_text(item).unwrap_or_default(), turn_index, }; + let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT; let mut cancellations = Vec::new(); for hook in &self.registry.on_prompt_submit { - let Some(action) = hook.call_optional(&info).await.map_err(|error| { - InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) - })? - else { + let Some(action) = call_hook_before_deadline(hook, &info, deadline).await? else { continue; }; if let HookPromptAction::Cancel(reason) = action { @@ -354,12 +378,11 @@ impl Interceptor for WorkerInterceptor { .as_ref() .map(|_| SystemItemAppendHandle::new(Arc::clone(&pending_hook_system_items))); let hook_context = PreRequestContext::new(info, system_item_sink); + let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT; let mut cancellations = Vec::new(); let mut should_yield = false; for hook in &self.registry.pre_llm_request { - let Some(action) = hook.call_optional(&hook_context).await.map_err(|error| { - InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) - })? + let Some(action) = call_hook_before_deadline(hook, &hook_context, deadline).await? else { continue; }; @@ -431,14 +454,12 @@ impl Interceptor for WorkerInterceptor { tool_name: info.call.name.clone(), arguments: info.call.input.clone(), }; + let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT; let mut aborts = Vec::new(); let mut should_pause = false; let mut denials = Vec::new(); for hook in &self.registry.pre_tool_call { - let Some(action) = hook.call_optional(&summary).await.map_err(|error| { - InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) - })? - else { + let Some(action) = call_hook_before_deadline(hook, &summary, deadline).await? else { continue; }; @@ -479,12 +500,10 @@ impl Interceptor for WorkerInterceptor { attachments: Vec::new(), }, }; + let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT; let mut aborts = Vec::new(); for hook in &self.registry.post_tool_call { - let Some(action) = hook.call_optional(&summary).await.map_err(|error| { - InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) - })? - else { + let Some(action) = call_hook_before_deadline(hook, &summary, deadline).await? else { continue; }; @@ -516,12 +535,10 @@ impl Interceptor for WorkerInterceptor { tool_calls_count: self.tool_calls_this_turn.load(Ordering::Relaxed), final_text_preview, }; + let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT; let mut should_pause = false; for hook in &self.registry.on_turn_end { - let Some(action) = hook.call_optional(&info).await.map_err(|error| { - InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string()) - })? - else { + let Some(action) = call_hook_before_deadline(hook, &info, deadline).await? else { continue; }; if matches!(action, HookTurnEndAction::Pause) { diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 7db1f37d..d5f7f8d1 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -50,6 +50,7 @@ pub use prompt::system::{SystemPromptContext, SystemPromptError, SystemPromptTem pub use protocol::{ErrorCode, Event, Method, TurnResult, WorkerStatus}; pub use runtime::dir::RuntimeDir; pub use segment_log_sink::SegmentLogSink; +pub use session_capture::SessionEntryRef; pub use session_history::{ SessionHistoryDerivation, SessionHistoryEntryId, SessionHistoryMetadata, WorkerHistoryProvenance, WorkerSubjectSnapshot, diff --git a/crates/worker/src/session_capture.rs b/crates/worker/src/session_capture.rs index 19026fb8..92f14518 100644 --- a/crates/worker/src/session_capture.rs +++ b/crates/worker/src/session_capture.rs @@ -23,14 +23,14 @@ const OVERVIEW_ANCHOR_STRIDE: usize = 8; #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(transparent)] -pub(crate) struct SessionEntryRef(String); +pub struct SessionEntryRef(String); impl SessionEntryRef { - pub(crate) fn from_history_entry_id(entry_id: &crate::SessionHistoryEntryId) -> Self { + pub fn from_history_entry_id(entry_id: &crate::SessionHistoryEntryId) -> Self { Self(format!("E{}", entry_id.0)) } - pub(crate) fn parse(value: &str) -> Option { + pub fn parse(value: &str) -> Option { let suffix = value.strip_prefix('E')?; if suffix.is_empty() || suffix.len() > 64 @@ -43,7 +43,7 @@ impl SessionEntryRef { Some(Self(value.to_string())) } - pub(crate) fn as_str(&self) -> &str { + pub fn as_str(&self) -> &str { &self.0 } diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 582a6a81..6b944d39 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -51,10 +51,10 @@ use crate::feature::{ FeatureRegistryInstallReport, dedupe_instruction_contributions, }; use crate::hook::{ - BeforeSessionRewriteAction, BeforeSessionRewriteContext, Hook, HookInvocationContext, - HookRegistry, HookRegistryBuilder, OnPromptSubmit, OnTurnEnd, PostToolCall, PreLlmRequest, - PreToolCall, RunCommittedContext, RunCommittedExit, RunExitContext, SessionRewriteKind, - WorkerStoppingContext, + BeforeSessionRewriteAction, BeforeSessionRewriteContext, Hook, HookHistoryRange, + HookInvocationContext, HookRegistry, HookRegistryBuilder, OnPromptSubmit, OnTurnEnd, + PostToolCall, PreLlmRequest, PreToolCall, RunCommittedContext, RunCommittedExit, + RunExitContext, SessionRewriteKind, WorkerStoppingContext, }; use crate::in_flight::InFlightEvents; use crate::internal_worker::{ @@ -66,6 +66,7 @@ const COMPACTION_EXTENSION_DOMAIN: &str = "yoi.compaction"; const LARGE_PASTE_INLINE_MAX_BYTES: usize = 32 * 1024; const WORKER_ORCHESTRATION_INSTRUCTION_ID: &str = "worker.orchestration"; const WORKER_ORCHESTRATION_PROMPT_REF: &str = "common.worker_orchestration"; +const FEATURE_HOOK_CHAIN_TIMEOUT: Duration = Duration::from_secs(30); fn hook_run_exit(exit: &EngineRunExit) -> RunCommittedExit { match exit { @@ -1281,8 +1282,20 @@ impl Worker { invocation: self.hook_invocation_context(None), reason: reason.into(), }; - if let Err(error) = hooks.on_worker_stopping(&context).await { - tracing::warn!(error = %error, "worker-stopping hook requires attention"); + match tokio::time::timeout( + FEATURE_HOOK_CHAIN_TIMEOUT, + hooks.on_worker_stopping(&context), + ) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + tracing::warn!(error = %error, "worker-stopping hook requires attention"); + } + Err(_) => { + hooks.record_chain_timeout("worker-stopping"); + tracing::warn!("worker-stopping hook chain exceeded its host deadline"); + } } } if let Err(error) = self.feature_background_tasks.shutdown().await { @@ -1860,6 +1873,19 @@ impl Worker { self.engine.as_ref().expect("worker taken during run") } + fn hook_history_range(&self) -> HookHistoryRange { + let entries = self.session.history().entries(); + HookHistoryRange { + first_entry_ref: entries.first().map(|entry| { + crate::SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id) + }), + last_entry_ref: entries.last().map(|entry| { + crate::SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id) + }), + entry_count: entries.len(), + } + } + fn hook_invocation_context(&self, run_id: Option) -> HookInvocationContext { HookInvocationContext { workspace_id: self @@ -3444,8 +3470,17 @@ impl Worker { exit: hook_exit, history_len: self.session.history().len(), }; - if let Err(error) = hooks.on_run_exit(&context).await { - tracing::warn!(error = %error, "run-exit hook failed; preserving terminal commit"); + match tokio::time::timeout(FEATURE_HOOK_CHAIN_TIMEOUT, hooks.on_run_exit(&context)) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + tracing::warn!(error = %error, "run-exit hook failed; preserving terminal commit"); + } + Err(_) => { + hooks.record_chain_timeout("run-exit"); + tracing::warn!("run-exit hook chain exceeded its host deadline"); + } } } if matches!(&result, EngineRunExit::Interrupted(_)) { @@ -3457,11 +3492,19 @@ impl Worker { let context = RunCommittedContext { invocation: committed_invocation.clone(), exit: hook_exit, - committed_history: self.session.history().entries().to_vec(), - committed_history_len: self.session.history().len(), + committed_history: self.hook_history_range(), }; - if let Err(error) = hooks.on_run_committed(&context).await { - tracing::warn!(error = %error, "run-committed hook requires attention"); + match tokio::time::timeout(FEATURE_HOOK_CHAIN_TIMEOUT, hooks.on_run_committed(&context)) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + tracing::warn!(error = %error, "run-committed hook requires attention"); + } + Err(_) => { + hooks.record_chain_timeout("run-committed"); + tracing::warn!("run-committed hook chain exceeded its host deadline"); + } } } if let Err(error) = self @@ -3659,14 +3702,21 @@ impl Worker { let context = BeforeSessionRewriteContext { invocation: self.hook_invocation_context(None), kind, - current_history: self.session.history().entries().to_vec(), - current_history_len: self.session.history().len(), + current_history: self.hook_history_range(), }; - match hooks - .before_session_rewrite(&context) - .await - .map_err(|error| WorkerError::FeatureLifecycle(error.to_string()))? - { + let action = tokio::time::timeout( + FEATURE_HOOK_CHAIN_TIMEOUT, + hooks.before_session_rewrite(&context), + ) + .await + .map_err(|_| { + hooks.record_chain_timeout("before-session-rewrite"); + WorkerError::FeatureLifecycle( + "session-rewrite hook chain exceeded its host deadline".to_string(), + ) + })? + .map_err(|error| WorkerError::FeatureLifecycle(error.to_string()))?; + match action { BeforeSessionRewriteAction::Continue => {} BeforeSessionRewriteAction::Deny(reason) => { return Err(WorkerError::FeatureLifecycle(reason)); From f1dc90621ccb8704dbb0ded8fdc3322059b9c624 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 16:13:35 +0900 Subject: [PATCH 11/22] fix: fence feature task shutdown races --- crates/worker/src/feature/background.rs | 124 +++++++++++++++++++++--- 1 file changed, 112 insertions(+), 12 deletions(-) diff --git a/crates/worker/src/feature/background.rs b/crates/worker/src/feature/background.rs index 695e45d3..b18ac7db 100644 --- a/crates/worker/src/feature/background.rs +++ b/crates/worker/src/feature/background.rs @@ -104,13 +104,55 @@ impl BackgroundTaskSpec { } } +#[derive(Clone)] +pub struct BackgroundTaskGenerationFence { + expected: u64, + current: Arc, +} + +impl BackgroundTaskGenerationFence { + pub fn generation(&self) -> u64 { + self.expected + } + + pub fn ensure_current(&self) -> Result<(), HookError> { + if self.current.load(Ordering::Acquire) == self.expected { + Ok(()) + } else { + Err(HookError::new( + HookErrorCategory::Cancelled, + "background task belongs to a stale Session generation", + )) + } + } +} + +impl std::fmt::Debug for BackgroundTaskGenerationFence { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("BackgroundTaskGenerationFence") + .field("expected", &self.expected) + .finish_non_exhaustive() + } +} + +impl PartialEq for BackgroundTaskGenerationFence { + fn eq(&self, other: &Self) -> bool { + self.expected == other.expected && Arc::ptr_eq(&self.current, &other.current) + } +} + +impl Eq for BackgroundTaskGenerationFence {} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct BackgroundTaskContext { pub invocation: HookInvocationContext, pub feature_id: FeatureId, pub task_name: String, pub execution_id: u64, - pub session_generation: u64, + /// Token that must be checked by any separately granted mutation handle + /// immediately before committing output derived by this task. + pub generation_fence: BackgroundTaskGenerationFence, pub attempt: u16, } @@ -215,7 +257,7 @@ impl FeatureBackgroundTaskRegistryBuilder { running: Mutex::new(BTreeMap::new()), diagnostics: Mutex::new(Vec::new()), next_execution_id: AtomicU64::new(1), - session_generation: AtomicU64::new(1), + session_generation: Arc::new(AtomicU64::new(1)), accepting: AtomicBool::new(true), }), } @@ -234,7 +276,7 @@ struct RegistryInner { running: Mutex>, diagnostics: Mutex>, next_execution_id: AtomicU64, - session_generation: AtomicU64, + session_generation: Arc, accepting: AtomicBool, } @@ -338,6 +380,15 @@ impl FeatureBackgroundTaskRegistry { .running .lock() .expect("background tasks poisoned"); + // `shutdown()` flips accepting before taking this same lock. Recheck + // while holding the spawn/drain serialization boundary so a starter + // that passed the optimistic check cannot publish a task after drain. + if !self.inner.accepting.load(Ordering::Acquire) { + return Err(HookError::new( + HookErrorCategory::ScopeDisposed, + "background task scope is stopping", + )); + } running.retain(|_, running| !running.handle.is_finished()); let active = running .values() @@ -349,6 +400,10 @@ impl FeatureBackgroundTaskRegistry { let execution_id = self.inner.next_execution_id.fetch_add(1, Ordering::Relaxed); let session_generation = self.inner.session_generation.load(Ordering::Acquire); + let generation_fence = BackgroundTaskGenerationFence { + expected: session_generation, + current: Arc::clone(&self.inner.session_generation), + }; let cancellation = BackgroundTaskCancellation::default(); let task_cancellation = cancellation.clone(); let task = Arc::clone(®istration.task); @@ -366,17 +421,16 @@ impl FeatureBackgroundTaskRegistry { task_feature_id.clone(), task_task_name.clone(), execution_id, - session_generation, + generation_fence.clone(), task_cancellation, ) .await; if let Some(inner) = weak_inner.upgrade() { - let outcome = - if inner.session_generation.load(Ordering::Acquire) == session_generation { - outcome - } else { - BackgroundTaskOutcome::StaleGenerationDiscarded - }; + let outcome = if generation_fence.ensure_current().is_ok() { + outcome + } else { + BackgroundTaskOutcome::StaleGenerationDiscarded + }; let mut diagnostics = inner.diagnostics.lock().expect("diagnostics poisoned"); diagnostics.push(BackgroundTaskDiagnostic { execution_id, @@ -546,7 +600,7 @@ async fn execute_task( feature_id: FeatureId, task_name: String, execution_id: u64, - session_generation: u64, + generation_fence: BackgroundTaskGenerationFence, cancellation: BackgroundTaskCancellation, ) -> (u16, BackgroundTaskOutcome) { let (max_attempts, delay_ms) = match spec.retry { @@ -569,7 +623,7 @@ async fn execute_task( feature_id: feature_id.clone(), task_name: task_name.clone(), execution_id, - session_generation, + generation_fence: generation_fence.clone(), attempt, }; let result = @@ -737,6 +791,44 @@ mod tests { ); } + #[tokio::test] + async fn concurrent_start_and_shutdown_leave_no_task_after_drain() { + for _ in 0..100 { + let feature = FeatureId::builtin("shutdown-race"); + let declaration = BackgroundTaskDeclaration::worker_managed("race", "race"); + let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); + builder + .register( + feature.clone(), + BackgroundTaskSpec::single_flight(declaration, Duration::from_secs(1)), + WaitForCancellation, + ) + .unwrap(); + let registry = builder.build(); + let gate = Arc::new(std::sync::Barrier::new(2)); + let start_registry = registry.clone(); + let start_feature = feature.clone(); + let start_gate = Arc::clone(&gate); + let starter = tokio::task::spawn_blocking(move || { + start_gate.wait(); + start_registry.start(&start_feature, "race", invocation()) + }); + + gate.wait(); + registry.shutdown().await.unwrap(); + let start_result = starter.await.unwrap(); + assert!(matches!( + start_result, + Ok(BackgroundTaskStart::Started { .. }) + | Err(HookError { + category: HookErrorCategory::ScopeDisposed, + .. + }) + )); + assert!(registry.inner.running.lock().unwrap().is_empty()); + } + } + #[tokio::test] async fn cancellation_observed_when_cancel_races_with_wait_registration() { for _ in 0..100 { @@ -814,6 +906,10 @@ mod tests { ) .unwrap(); let registry = builder.build(); + let stale_fence = BackgroundTaskGenerationFence { + expected: registry.inner.session_generation.load(Ordering::Acquire), + current: Arc::clone(®istry.inner.session_generation), + }; registry .start(&feature, "generation", invocation()) .unwrap(); @@ -822,6 +918,10 @@ mod tests { registry.before_session_rewrite().await.unwrap(); assert_eq!(registry.inner.session_generation.load(Ordering::Acquire), 2); + assert_eq!( + stale_fence.ensure_current().unwrap_err().category, + HookErrorCategory::Cancelled + ); assert_eq!(registry.diagnostics().len(), 1); assert_eq!( registry.diagnostics()[0].outcome, From 60a5495ccdc631e8c6c2d712cd04a4e81409b0cc Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 16:34:15 +0900 Subject: [PATCH 12/22] fix: hold feature task barrier across rewrites --- crates/worker/src/feature/background.rs | 183 ++++++++++++++++++++++-- crates/worker/src/worker.rs | 20 +-- 2 files changed, 181 insertions(+), 22 deletions(-) diff --git a/crates/worker/src/feature/background.rs b/crates/worker/src/feature/background.rs index b18ac7db..abc18e6d 100644 --- a/crates/worker/src/feature/background.rs +++ b/crates/worker/src/feature/background.rs @@ -6,7 +6,7 @@ //! feature was explicitly granted at install time. use std::collections::BTreeMap; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU8, AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -25,6 +25,9 @@ const MAX_TASK_ATTEMPTS: u16 = 16; const MAX_TASK_TIMEOUT_MS: u64 = 24 * 60 * 60 * 1_000; const MAX_RETAINED_DIAGNOSTICS: usize = 128; const TASK_SETTLE_TIMEOUT_MS: u64 = 30_000; +const TASK_SCOPE_ACCEPTING: u8 = 0; +const TASK_SCOPE_REWRITING: u8 = 1; +const TASK_SCOPE_STOPPING: u8 = 2; #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] @@ -258,7 +261,7 @@ impl FeatureBackgroundTaskRegistryBuilder { diagnostics: Mutex::new(Vec::new()), next_execution_id: AtomicU64::new(1), session_generation: Arc::new(AtomicU64::new(1)), - accepting: AtomicBool::new(true), + lifecycle: AtomicU8::new(TASK_SCOPE_ACCEPTING), }), } } @@ -277,7 +280,7 @@ struct RegistryInner { diagnostics: Mutex>, next_execution_id: AtomicU64, session_generation: Arc, - accepting: AtomicBool, + lifecycle: AtomicU8, } impl Drop for RegistryInner { @@ -297,6 +300,31 @@ pub struct FeatureBackgroundTaskRegistry { inner: Arc, } +/// Holds the task registry quiescent across the entire Session rewrite. Drop +/// reopens starts only when shutdown has not moved the scope to Stopping. +pub struct BackgroundTaskRewriteGuard { + inner: Arc, +} + +impl std::fmt::Debug for BackgroundTaskRewriteGuard { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("BackgroundTaskRewriteGuard") + .finish_non_exhaustive() + } +} + +impl Drop for BackgroundTaskRewriteGuard { + fn drop(&mut self) { + let _ = self.inner.lifecycle.compare_exchange( + TASK_SCOPE_REWRITING, + TASK_SCOPE_ACCEPTING, + Ordering::AcqRel, + Ordering::Acquire, + ); + } +} + impl Default for FeatureBackgroundTaskRegistry { fn default() -> Self { FeatureBackgroundTaskRegistryBuilder::default().build() @@ -361,7 +389,7 @@ impl FeatureBackgroundTaskRegistry { task_name: &str, invocation: HookInvocationContext, ) -> Result { - if !self.inner.accepting.load(Ordering::Acquire) { + if self.inner.lifecycle.load(Ordering::Acquire) != TASK_SCOPE_ACCEPTING { return Err(HookError::new( HookErrorCategory::ScopeDisposed, "background task scope is stopping", @@ -380,10 +408,11 @@ impl FeatureBackgroundTaskRegistry { .running .lock() .expect("background tasks poisoned"); - // `shutdown()` flips accepting before taking this same lock. Recheck - // while holding the spawn/drain serialization boundary so a starter - // that passed the optimistic check cannot publish a task after drain. - if !self.inner.accepting.load(Ordering::Acquire) { + // `shutdown()` and `begin_session_rewrite()` change lifecycle before + // taking this same lock. Recheck while holding the spawn/drain + // serialization boundary so a starter that passed the optimistic + // check cannot publish after either barrier. + if self.inner.lifecycle.load(Ordering::Acquire) != TASK_SCOPE_ACCEPTING { return Err(HookError::new( HookErrorCategory::ScopeDisposed, "background task scope is stopping", @@ -484,14 +513,45 @@ impl FeatureBackgroundTaskRegistry { .clone() } - pub async fn before_session_rewrite(&self) -> Result<(), HookError> { + pub async fn begin_session_rewrite(&self) -> Result { + match self.inner.lifecycle.compare_exchange( + TASK_SCOPE_ACCEPTING, + TASK_SCOPE_REWRITING, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => {} + Err(TASK_SCOPE_REWRITING) => { + return Err(HookError::new( + HookErrorCategory::Dependency, + "another Session rewrite already owns the background task barrier", + )); + } + Err(_) => { + return Err(HookError::new( + HookErrorCategory::ScopeDisposed, + "background task scope is stopping", + )); + } + } + let guard = BackgroundTaskRewriteGuard { + inner: Arc::clone(&self.inner), + }; self.settle(false).await?; + if self.inner.lifecycle.load(Ordering::Acquire) != TASK_SCOPE_REWRITING { + return Err(HookError::new( + HookErrorCategory::ScopeDisposed, + "background task scope stopped during Session rewrite preparation", + )); + } self.inner.session_generation.fetch_add(1, Ordering::AcqRel); - Ok(()) + Ok(guard) } pub async fn shutdown(&self) -> Result<(), HookError> { - self.inner.accepting.store(false, Ordering::Release); + self.inner + .lifecycle + .store(TASK_SCOPE_STOPPING, Ordering::Release); self.settle(true).await } @@ -668,7 +728,7 @@ async fn execute_task( #[cfg(test)] mod tests { use super::*; - use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; fn invocation() -> HookInvocationContext { HookInvocationContext { @@ -893,6 +953,73 @@ mod tests { ); } + struct StartsTaskFromRewriteHook { + registry: FeatureBackgroundTaskRegistry, + feature: FeatureId, + rejected: Arc, + } + + #[async_trait] + impl crate::hook::Hook for StartsTaskFromRewriteHook { + async fn call( + &self, + _input: &crate::hook::BeforeSessionRewriteContext, + ) -> Result { + let error = self + .registry + .start(&self.feature, "hook-start", invocation()) + .unwrap_err(); + self.rejected.store( + error.category == HookErrorCategory::ScopeDisposed, + Ordering::Release, + ); + Ok(crate::hook::BeforeSessionRewriteAction::Continue) + } + } + + #[tokio::test] + async fn rewrite_hook_cannot_start_task_while_quiescence_guard_is_held() { + let feature = FeatureId::builtin("hook-start-test"); + let declaration = BackgroundTaskDeclaration::worker_managed("hook-start", "hook-start"); + let mut task_builder = FeatureBackgroundTaskRegistryBuilder::default(); + task_builder + .register( + feature.clone(), + BackgroundTaskSpec::single_flight(declaration, Duration::from_secs(1)), + WaitForCancellation, + ) + .unwrap(); + let tasks = task_builder.build(); + let rejected = Arc::new(AtomicBool::new(false)); + let mut hook_builder = crate::hook::HookRegistryBuilder::new(); + hook_builder + .add_named_before_session_rewrite( + "hook-start-test", + crate::hook::HookExecutionPolicy::fail_closed(), + StartsTaskFromRewriteHook { + registry: tasks.clone(), + feature, + rejected: Arc::clone(&rejected), + }, + ) + .unwrap(); + let hooks = hook_builder.build(); + let guard = tasks.begin_session_rewrite().await.unwrap(); + let context = crate::hook::BeforeSessionRewriteContext { + invocation: invocation(), + kind: crate::hook::SessionRewriteKind::Compact, + current_history: crate::hook::HookHistoryRange::default(), + }; + + assert_eq!( + hooks.before_session_rewrite(&context).await.unwrap(), + crate::hook::BeforeSessionRewriteAction::Continue + ); + assert!(rejected.load(Ordering::Acquire)); + drop(guard); + tasks.shutdown().await.unwrap(); + } + #[tokio::test] async fn rewrite_joins_old_tasks_before_advancing_session_generation() { let feature = FeatureId::builtin("generation-test"); @@ -915,9 +1042,16 @@ mod tests { .unwrap(); assert_eq!(registry.inner.session_generation.load(Ordering::Acquire), 1); - registry.before_session_rewrite().await.unwrap(); + let rewrite_guard = registry.begin_session_rewrite().await.unwrap(); assert_eq!(registry.inner.session_generation.load(Ordering::Acquire), 2); + assert!(matches!( + registry.start(&feature, "generation", invocation()), + Err(HookError { + category: HookErrorCategory::ScopeDisposed, + .. + }) + )); assert_eq!( stale_fence.ensure_current().unwrap_err().category, HookErrorCategory::Cancelled @@ -927,6 +1061,27 @@ mod tests { registry.diagnostics()[0].outcome, BackgroundTaskOutcome::Cancelled ); + drop(rewrite_guard); + assert!(matches!( + registry + .start(&feature, "generation", invocation()) + .unwrap(), + BackgroundTaskStart::Started { .. } + )); + registry.shutdown().await.unwrap(); + } + + #[tokio::test] + async fn shutdown_during_rewrite_prevents_guard_drop_from_reopening_starts() { + let registry = FeatureBackgroundTaskRegistry::default(); + let guard = registry.begin_session_rewrite().await.unwrap(); + registry.shutdown().await.unwrap(); + drop(guard); + + let error = registry + .start(&FeatureId::builtin("missing"), "missing", invocation()) + .unwrap_err(); + assert_eq!(error.category, HookErrorCategory::ScopeDisposed); } #[tokio::test] @@ -942,7 +1097,7 @@ mod tests { let registry = builder.build(); registry.start(&feature, "rewrite", invocation()).unwrap(); - let error = registry.before_session_rewrite().await.unwrap_err(); + let error = registry.begin_session_rewrite().await.unwrap_err(); assert_eq!(error.category, HookErrorCategory::Dependency); registry.shutdown().await.unwrap(); assert_eq!(registry.diagnostics().len(), 1); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 6b944d39..d69afe93 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -40,7 +40,7 @@ use manifest::{ use crate::compact::state::CompactState; use crate::compact::usage_tracker::UsageTracker; -use crate::feature::background::FeatureBackgroundTaskRegistry; +use crate::feature::background::{BackgroundTaskRewriteGuard, FeatureBackgroundTaskRegistry}; use crate::feature::builtin::memory::WorkspaceMemoryBackendError; use crate::feature::builtin::{ MemoryExtractFeature, MemoryExtractState, SessionExploreFeature, SessionExploreState, @@ -2031,7 +2031,8 @@ impl Worker { target: RewindTargetId, expected_head_entries: usize, ) -> Result { - self.prepare_session_rewrite(SessionRewriteKind::Rewind) + let _rewrite_guard = self + .prepare_session_rewrite(SessionRewriteKind::Rewind) .await .map_err(|error| RewindError::Invalid(error.to_string()))?; let loc = self.segment_state.location(); @@ -3393,7 +3394,8 @@ impl Worker { // state up to that turn). The new SegmentStart replaces the mirror // and is broadcast through the sink so existing subscribers reset // their view. - self.prepare_session_rewrite(SessionRewriteKind::Fork) + let _rewrite_guard = self + .prepare_session_rewrite(SessionRewriteKind::Fork) .await?; let w = self.engine.as_ref().unwrap(); let fork_segment_id = session_store::new_segment_id(); @@ -3693,9 +3695,10 @@ impl Worker { async fn prepare_session_rewrite( &mut self, kind: SessionRewriteKind, - ) -> Result<(), WorkerError> { - self.feature_background_tasks - .before_session_rewrite() + ) -> Result { + let rewrite_guard = self + .feature_background_tasks + .begin_session_rewrite() .await .map_err(|error| WorkerError::FeatureLifecycle(error.to_string()))?; if let Some(hooks) = self.hook_registry.clone() { @@ -3723,7 +3726,7 @@ impl Worker { } } } - Ok(()) + Ok(rewrite_guard) } pub async fn manual_compact(&mut self) -> Result { @@ -3941,7 +3944,8 @@ impl Worker { /// Runs one parent-owned observable compaction service and returns the new /// Segment ID. Lifecycle revisions are committed before they are broadcast. pub async fn compact(&mut self, retained_tokens: u64) -> Result { - self.prepare_session_rewrite(SessionRewriteKind::Compact) + let _rewrite_guard = self + .prepare_session_rewrite(SessionRewriteKind::Compact) .await?; let mut lifecycle = CompactionLifecycle { schema_version: 2, From fb13e53cb5fa9ef3990e50929fe075a41b9e8220 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 16:53:35 +0900 Subject: [PATCH 13/22] test: cover start race at rewrite barrier --- crates/worker/src/feature/background.rs | 57 +++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/crates/worker/src/feature/background.rs b/crates/worker/src/feature/background.rs index abc18e6d..976734b8 100644 --- a/crates/worker/src/feature/background.rs +++ b/crates/worker/src/feature/background.rs @@ -262,6 +262,8 @@ impl FeatureBackgroundTaskRegistryBuilder { next_execution_id: AtomicU64::new(1), session_generation: Arc::new(AtomicU64::new(1)), lifecycle: AtomicU8::new(TASK_SCOPE_ACCEPTING), + #[cfg(test)] + start_race_probe: Mutex::new(None), }), } } @@ -274,6 +276,12 @@ struct RunningTask { handle: JoinHandle<()>, } +#[cfg(test)] +struct StartRaceProbe { + optimistic_check_passed: Arc, + resume: Arc, +} + struct RegistryInner { registrations: BTreeMap<(FeatureId, String), Registration>, running: Mutex>, @@ -281,6 +289,8 @@ struct RegistryInner { next_execution_id: AtomicU64, session_generation: Arc, lifecycle: AtomicU8, + #[cfg(test)] + start_race_probe: Mutex>, } impl Drop for RegistryInner { @@ -395,6 +405,17 @@ impl FeatureBackgroundTaskRegistry { "background task scope is stopping", )); } + #[cfg(test)] + if let Some(probe) = self + .inner + .start_race_probe + .lock() + .expect("start race probe poisoned") + .take() + { + probe.optimistic_check_passed.wait(); + probe.resume.wait(); + } let key = (feature_id.clone(), task_name.to_string()); let registration = self.inner.registrations.get(&key).ok_or_else(|| { HookError::new( @@ -889,6 +910,42 @@ mod tests { } } + #[tokio::test] + async fn starter_paused_after_optimistic_check_is_rejected_by_rewrite_barrier() { + let feature = FeatureId::builtin("rewrite-race"); + let declaration = BackgroundTaskDeclaration::worker_managed("race", "race"); + let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); + builder + .register( + feature.clone(), + BackgroundTaskSpec::single_flight(declaration, Duration::from_secs(1)), + WaitForCancellation, + ) + .unwrap(); + let registry = builder.build(); + let optimistic_check_passed = Arc::new(std::sync::Barrier::new(2)); + let resume = Arc::new(std::sync::Barrier::new(2)); + *registry.inner.start_race_probe.lock().unwrap() = Some(StartRaceProbe { + optimistic_check_passed: Arc::clone(&optimistic_check_passed), + resume: Arc::clone(&resume), + }); + let start_registry = registry.clone(); + let starter = tokio::task::spawn_blocking(move || { + start_registry.start(&feature, "race", invocation()) + }); + + optimistic_check_passed.wait(); + let rewrite_guard = registry.begin_session_rewrite().await.unwrap(); + resume.wait(); + let error = starter.await.unwrap().unwrap_err(); + + assert_eq!(error.category, HookErrorCategory::ScopeDisposed); + assert!(registry.inner.running.lock().unwrap().is_empty()); + drop(rewrite_guard); + assert!(registry.inner.running.lock().unwrap().is_empty()); + registry.shutdown().await.unwrap(); + } + #[tokio::test] async fn cancellation_observed_when_cancel_races_with_wait_registration() { for _ in 0..100 { From 33d98868c37a79de48aebc7f47513802e7915f85 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 18:44:50 +0900 Subject: [PATCH 14/22] refactor: move memory extraction into lifecycle feature --- crates/worker/src/controller.rs | 38 +- crates/worker/src/feature.rs | 1 + crates/worker/src/feature/builtin.rs | 5 +- .../src/feature/builtin/memory_lifecycle.rs | 775 ++++++++++ ...ry_extract.rs => memory_staging_output.rs} | 36 +- crates/worker/src/feature/session.rs | 111 ++ crates/worker/src/internal_worker.rs | 8 +- crates/worker/src/worker.rs | 1327 ++--------------- crates/worker/tests/compact_events_test.rs | 238 --- 9 files changed, 1083 insertions(+), 1456 deletions(-) create mode 100644 crates/worker/src/feature/builtin/memory_lifecycle.rs rename crates/worker/src/feature/builtin/{memory_extract.rs => memory_staging_output.rs} (94%) create mode 100644 crates/worker/src/feature/session.rs diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 5e848f2c..43f28539 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -197,8 +197,8 @@ async fn finish_controller_run( { // history / user_segments are no longer mirrored on WorkerSharedState — // clients reconstruct them from `Event::Snapshot` + live - // `Event::Entry` deliveries driven by the session-log sink. We - // flip the status and kick post-run memory jobs here. + // `Event::Entry` deliveries driven by the session-log sink. The + // lifecycle hook/task registry observes the terminal commit separately. // // In-flight blocks are run-local streaming state, not durable transcript. // Any block not cleared by a committed AssistantItem must be discarded at @@ -206,7 +206,6 @@ async fn finish_controller_run( // partial text/tool arguments after newer entries. worker.clear_in_flight_events(); set_controller_status(shared_state, runtime_dir, working_event_tx, new_status).await; - worker.spawn_post_run_memory_jobs(); } /// Pending turn launch staged by an event handler for the next outer-loop @@ -995,6 +994,34 @@ where let worker_enabled = feature_config.worker.enabled; let sub_worker_enabled = feature_config.sub_worker.enabled; let mut feature_registry = FeatureRegistryBuilder::new(); + if feature_config.memory.enabled { + let config = memory_config.clone().ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "[feature.memory].enabled = true requires a [memory] configuration section", + ) + })?; + let workspace_client = worker.workspace_client_handle(); + if !workspace_client.is_available() || workspace_client.workspace_id().is_none() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Memory extraction requires Backend Workspace API authority", + )); + } + feature_registry.add_module( + crate::feature::builtin::memory_lifecycle::MemoryExtractionLifecycleFeature::new( + config, + worker.committed_session_capture_handle(), + worker.session_extension_handle(), + workspace_client, + spawner_manifest.clone(), + worker.llm_client_handle(), + prompts.clone(), + spawner_workspace_context.clone(), + worker.working_event_sender(), + ), + ); + } if sub_worker_enabled && !worker_enabled { feature_registry.add_module( crate::feature::builtin::manage_worker::sub_worker_control_feature( @@ -1717,11 +1744,6 @@ async fn controller_loop( // 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. - worker.wait_for_memory_jobs().await; - if let Some(session) = worker.workdir_session() && let Err(error) = session.close().await { diff --git a/crates/worker/src/feature.rs b/crates/worker/src/feature.rs index e91d44a6..be935209 100644 --- a/crates/worker/src/feature.rs +++ b/crates/worker/src/feature.rs @@ -2212,6 +2212,7 @@ pub mod background; pub mod builtin; pub mod mcp; pub mod plugin; +pub(crate) mod session; #[cfg(test)] mod tests { diff --git a/crates/worker/src/feature/builtin.rs b/crates/worker/src/feature/builtin.rs index 2f1be0a7..407e20dd 100644 --- a/crates/worker/src/feature/builtin.rs +++ b/crates/worker/src/feature/builtin.rs @@ -8,7 +8,8 @@ pub mod flow_transition; pub mod manage_workdir; pub mod manage_worker; pub mod memory; -pub mod memory_extract; +pub(crate) mod memory_lifecycle; +pub mod memory_staging_output; pub mod merge_request; pub mod objective; pub mod orchestration; @@ -19,8 +20,6 @@ pub mod ticket; pub mod worker_observation; pub mod workspace_worker_discovery; -pub(crate) use memory_extract::{MemoryExtractFeature, MemoryExtractState, render_extract_input}; -pub(crate) use session_explore::{SessionExploreFeature, SessionExploreState}; pub use task::{TaskFeature, task_tools_feature}; pub use ticket::{ TicketFeature, TicketFeatureAccess, ticket_tools_feature, ticket_tools_feature_with_access, diff --git a/crates/worker/src/feature/builtin/memory_lifecycle.rs b/crates/worker/src/feature/builtin/memory_lifecycle.rs new file mode 100644 index 00000000..30c14c91 --- /dev/null +++ b/crates/worker/src/feature/builtin/memory_lifecycle.rs @@ -0,0 +1,775 @@ +use std::sync::Arc; +use std::time::Duration; + +use agen::llm_client::LlmClient; +use arc_swap::ArcSwap; +use async_trait::async_trait; +use memory::extract; +use memory::schema::SourceRef; +use tokio::sync::broadcast; + +use crate::PromptCatalog; +use crate::Scope; +use crate::WorkerRunResult; +use crate::feature::background::{ + BackgroundTaskCancellation, BackgroundTaskContext, BackgroundTaskSpec, BackgroundTaskTrigger, + FeatureBackgroundTask, +}; +use crate::feature::builtin::memory_staging_output::{ + MemoryStagingOutputFeature, MemoryStagingOutputState, render_extract_input, +}; +use crate::feature::builtin::session_explore::{SessionExploreFeature, SessionExploreState}; +use crate::feature::session::{ + CommittedSessionCapture, CommittedSessionCaptureHandle, SessionExtensionHandle, +}; +use crate::feature::{ + BackgroundTaskDeclaration, FeatureDescriptor, FeatureInstallContext, FeatureInstallError, + FeatureModule, FeatureRegistryBuilder, +}; +use crate::hook::{HookError, HookErrorCategory}; +use crate::internal_worker::{ + InternalWorkerAuthority, InternalWorkerError, InternalWorkerIdentity, InternalWorkerResult, + InternalWorkerSpec, run_internal_worker_with_cancel_sender, +}; +use crate::session_capture::SessionCapture; +use crate::worker::{WorkerFilesystemAuthority, WorkerWorkspaceContext, WorkspaceClient}; +use agen::token_counter::total_tokens_at; +use manifest::WorkerManifest; +use protocol::Event; + +const TASK_NAME: &str = "memory-extraction"; +const TASK_TIMEOUT: Duration = Duration::from_secs(300); + +/// Parent-Worker lifecycle Feature that observes committed runs and schedules +/// bounded extraction work. It owns the Memory pointer, audit, restricted +/// Internal Worker, and staging disposition; Worker core owns only generic +/// hook/task/session plumbing. +#[derive(Clone)] +pub(crate) struct MemoryExtractionLifecycleFeature { + task: MemoryExtractionTask, +} + +#[derive(Clone)] +struct MemoryExtractionTask { + config: manifest::MemoryConfig, + capture: CommittedSessionCaptureHandle, + extensions: SessionExtensionHandle, + workspace_client: Arc, + manifest: WorkerManifest, + client: Box, + prompts: Arc>, + workspace_context: WorkerWorkspaceContext, + event_tx: Option>, +} + +impl MemoryExtractionLifecycleFeature { + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + config: manifest::MemoryConfig, + capture: CommittedSessionCaptureHandle, + extensions: SessionExtensionHandle, + workspace_client: Arc, + manifest: WorkerManifest, + client: Box, + prompts: Arc>, + workspace_context: WorkerWorkspaceContext, + event_tx: Option>, + ) -> Self { + Self { + task: MemoryExtractionTask { + config, + capture, + extensions, + workspace_client, + manifest, + client, + prompts, + workspace_context, + event_tx, + }, + } + } +} + +impl FeatureModule for MemoryExtractionLifecycleFeature { + fn descriptor(&self) -> FeatureDescriptor { + FeatureDescriptor::builtin("memory-extraction-lifecycle", "Memory Extraction Lifecycle") + .with_description( + "Observes terminal committed runs and schedules bounded Memory extraction.", + ) + .with_background_task(BackgroundTaskDeclaration::worker_managed( + TASK_NAME, + "Extract provenance-preserving Memory candidates after committed runs.", + )) + } + + fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> { + context + .background_tasks() + .register(memory_extraction_task_spec(), self.task.clone()) + } +} + +fn memory_extraction_task_spec() -> BackgroundTaskSpec { + let declaration = BackgroundTaskDeclaration::worker_managed( + TASK_NAME, + "Extract provenance-preserving Memory candidates after committed runs.", + ); + let mut spec = BackgroundTaskSpec::single_flight(declaration, TASK_TIMEOUT); + spec.trigger = BackgroundTaskTrigger::RunCommitted; + spec +} + +#[async_trait] +impl FeatureBackgroundTask for MemoryExtractionTask { + async fn run( + &self, + context: BackgroundTaskContext, + cancellation: BackgroundTaskCancellation, + ) -> Result<(), HookError> { + context.generation_fence.ensure_current()?; + let capture = self.capture.capture().map_err(hook_internal)?; + let pointer = extract_pointer(&capture)?; + if !extraction_threshold_reached(&capture, pointer.as_ref(), &self.config) { + return Ok(()); + } + + let history_start = pointer + .as_ref() + .map(|pointer| pointer.processed_through_history_len) + .unwrap_or(0) + .min(capture.history.len()); + let history_end = capture.history.len(); + if history_start >= history_end || capture.entry_count == 0 { + return Ok(()); + } + let view = SessionCapture::from_history_entries( + capture.segment_id.clone(), + capture.history[history_start..history_end].to_vec(), + ); + let start_entry = pointer + .as_ref() + .map(|pointer| pointer.processed_through_entry + 1) + .unwrap_or(0); + let source = SourceRef { + segment_id: capture.segment_id.clone(), + range: [start_entry as u64, (capture.entry_count - 1) as u64], + }; + let audit = WorkerAuditBase::new( + memory::audit::AuditWorker::MemoryExtract, + memory::audit::AuditTrigger::TokenThreshold, + self.config + .extract_model + .as_ref() + .or(Some(&self.manifest.model)) + .map(model_audit_from_manifest), + ) + .with_memory_settings(&self.config); + let extract_audit_base = memory::audit::ExtractAudit { + session_id: Some(capture.session_id.clone()), + segment_id: Some(capture.segment_id.clone()), + entry_range: Some([start_entry as u64, (capture.entry_count - 1) as u64]), + history_range: Some([history_start as u64, history_end as u64]), + ..Default::default() + }; + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Started, + "token_threshold_reached", + None, + Some(extract_audit_base.clone()), + None, + ) + .await; + let output_state = MemoryStagingOutputState::new( + view.clone(), + Arc::clone(&self.workspace_client), + source, + audit.run_id.to_string(), + ); + let client = if let Some(model) = self.config.extract_model.as_ref() { + match crate::model_client::build_client(model) { + Ok(client) => client, + Err(error) => { + self.record_preparation_failure(&audit, &extract_audit_base, error.to_string()) + .await; + return Ok(()); + } + } + } else { + self.client.clone_boxed() + }; + let Some(memory_language) = self + .config + .workspace_settings() + .map(|snapshot| snapshot.language) + else { + self.record_preparation_failure( + &audit, + &extract_audit_base, + "Memory extraction requires a bound Workspace Memory settings snapshot", + ) + .await; + return Ok(()); + }; + let system_prompt = match self + .prompts + .load_full() + .memory_extract_system(&memory_language) + { + Ok(prompt) => prompt, + Err(error) => { + self.record_preparation_failure(&audit, &extract_audit_base, error.to_string()) + .await; + return Ok(()); + } + }; + let mut manifest = self.manifest.clone(); + if let Some(model) = self.config.extract_model.clone() { + manifest.model = model; + } + + let cancel_observer = move |sender: tokio::sync::mpsc::Sender<()>| { + tokio::spawn(async move { + cancellation.cancelled().await; + let _ = sender.send(()).await; + }); + }; + let features = FeatureRegistryBuilder::new() + .with_module(SessionExploreFeature::new(SessionExploreState::new( + view.clone(), + ))) + .with_module(MemoryStagingOutputFeature::new(output_state.clone())); + let result = run_internal_worker_with_cancel_sender( + InternalWorkerSpec { + identity: InternalWorkerIdentity { + kind: "memory-extract", + run_id: audit.run_id, + }, + manifest, + client, + system_prompt, + input: render_extract_input(&view), + cache_key: Some(capture.segment_id.clone()), + max_turns: self + .config + .extract_worker_max_turns + .or(manifest::defaults::MEMORY_EXTRACT_WORKER_MAX_TURNS), + engine_configurator: None, + features, + required_tools: &[ + "ShowOverview", + "SearchEntries", + "ReadEntry", + "StageMemoryCandidate", + "FinishMemoryExtraction", + ], + authority: InternalWorkerAuthority { + workspace: self.workspace_context.clone(), + filesystem: WorkerFilesystemAuthority::None, + scope: Scope::empty(), + workdir_session: None, + }, + }, + cancel_observer, + ) + .await; + + let usage_event = match &result { + Ok(run) => { + tracing::debug!( + worker_kind = run.identity.kind, + run_id = %run.identity.run_id, + history_entries = run.history_entries, + "memory extraction Internal Worker completed" + ); + run.usage.as_ref() + } + Err(error) => error.usage.as_ref(), + }; + let usage_audit = usage_event.map(|event| memory::audit::UsageAudit { + input_tokens: event.input_tokens, + output_tokens: event.output_tokens, + total_tokens: event.total_tokens, + cache_read_input_tokens: event.cache_read_input_tokens, + cache_creation_input_tokens: event.cache_creation_input_tokens, + }); + let staging_ids = output_state.staged(); + let pointer_staging_id = staging_ids.first().cloned().unwrap_or_default(); + let extract_audit = Some(memory::audit::ExtractAudit { + staging_count: staging_ids.len(), + staging_paths: staging_ids, + ..extract_audit_base + }); + + match extraction_disposition(&result, output_state.is_finished()) { + ExtractionDisposition::Cancelled(reason) => { + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Cancelled, + reason, + usage_audit, + extract_audit, + None, + ) + .await; + return Ok(()); + } + ExtractionDisposition::Failed(reason) => { + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Failed, + reason, + usage_audit, + extract_audit, + None, + ) + .await; + return Ok(()); + } + ExtractionDisposition::Completed => {} + } + + context.generation_fence.ensure_current()?; + let next_pointer = memory::ExtractPointerPayload { + processed_through_entry: capture.entry_count - 1, + processed_through_history_len: capture.history.len(), + staging_id: pointer_staging_id, + }; + let payload = serde_json::to_value(&next_pointer).map_err(hook_internal)?; + if !self + .extensions + .append_if_current(&capture.location(), extract::EXTRACT_DOMAIN, payload) + .map_err(hook_internal)? + { + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Cancelled, + "session changed before memory-extract pointer commit", + usage_audit, + extract_audit, + None, + ) + .await; + return Ok(()); + } + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Completed, + "memory-extract completed", + usage_audit, + extract_audit, + None, + ) + .await; + Ok(()) + } +} + +impl MemoryExtractionTask { + async fn record_preparation_failure( + &self, + audit: &WorkerAuditBase, + extract: &memory::audit::ExtractAudit, + reason: impl Into, + ) { + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Failed, + reason, + None, + Some(extract.clone()), + None, + ) + .await; + } +} + +#[derive(Debug, PartialEq, Eq)] +enum ExtractionDisposition { + Completed, + Failed(String), + Cancelled(String), +} + +fn extraction_disposition( + result: &Result, + finish_called: bool, +) -> ExtractionDisposition { + match result { + Err(error) => { + // Preserve the Internal Worker result's immutable identity/history + // evidence for diagnostics even though the public audit reason is + // intentionally bounded to the typed source error. + tracing::debug!( + worker_kind = error.identity.kind, + run_id = %error.identity.run_id, + history_entries = error.history_entries, + "memory extraction Internal Worker failed" + ); + ExtractionDisposition::Failed(error.source.to_string()) + } + Ok(run) => match &run.lifecycle { + WorkerRunResult::RolledBack => { + ExtractionDisposition::Cancelled("memory-extract cancelled".to_string()) + } + WorkerRunResult::Interrupted { message, .. } => { + ExtractionDisposition::Failed(message.clone()) + } + WorkerRunResult::Finished | WorkerRunResult::Paused | WorkerRunResult::LimitReached + if finish_called => + { + ExtractionDisposition::Completed + } + WorkerRunResult::Finished | WorkerRunResult::Paused | WorkerRunResult::LimitReached => { + ExtractionDisposition::Failed( + "memory-extract did not call FinishMemoryExtraction".to_string(), + ) + } + }, + } +} + +fn now_millis() -> u64 { + use std::time::{SystemTime, UNIX_EPOCH}; + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .try_into() + .unwrap_or(u64::MAX) +} + +fn hook_internal(error: impl std::fmt::Display) -> HookError { + HookError::new(HookErrorCategory::Internal, error.to_string()) +} + +fn extract_pointer( + capture: &CommittedSessionCapture, +) -> Result, HookError> { + let pointer = memory::extract::fold_pointer(&capture.extensions); + if pointer.is_none() + && capture + .extensions + .iter() + .any(|(domain, _)| domain == extract::EXTRACT_DOMAIN) + { + return Err(hook_internal( + "latest committed Memory extraction pointer is malformed", + )); + } + Ok(pointer) +} + +fn extraction_threshold_reached( + capture: &CommittedSessionCapture, + pointer: Option<&memory::ExtractPointerPayload>, + config: &manifest::MemoryConfig, +) -> bool { + if capture.history.is_empty() { + return false; + } + let history_pointer = pointer + .map(|pointer| pointer.processed_through_history_len) + .unwrap_or(0) + .min(capture.history.len()); + let items = capture + .history + .iter() + .map(|entry| entry.item.clone()) + .collect::>(); + let current = total_tokens_at(&items, &capture.usage_history, capture.history.len()).tokens; + let baseline = total_tokens_at(&items, &capture.usage_history, history_pointer).tokens; + let Some(threshold) = config.extract_threshold.filter(|threshold| *threshold > 0) else { + return false; + }; + current.saturating_sub(baseline) >= threshold +} + +#[derive(Clone)] +struct WorkerAuditBase { + run_id: uuid::Uuid, + worker: memory::audit::AuditWorker, + trigger: memory::audit::AuditTrigger, + memory_settings: Option, + model: Option, +} + +impl WorkerAuditBase { + fn new( + worker: memory::audit::AuditWorker, + trigger: memory::audit::AuditTrigger, + model: Option, + ) -> Self { + Self { + run_id: uuid::Uuid::now_v7(), + worker, + trigger, + memory_settings: None, + model, + } + } + + fn with_memory_settings(mut self, config: &manifest::MemoryConfig) -> Self { + self.memory_settings = + config + .workspace_settings() + .map(|snapshot| memory::audit::MemorySettingsAudit { + workspace_id: snapshot.workspace_id, + settings_revision: snapshot.settings_revision, + language: snapshot.language, + }); + self + } + + async fn emit( + &self, + workspace_client: &dyn WorkspaceClient, + event_tx: Option<&broadcast::Sender>, + status: memory::audit::WorkerLifecycleStatus, + reason: impl Into, + usage: Option, + extract: Option, + consolidation: Option, + ) { + let reason = reason.into(); + let payload = memory::audit::WorkerLifecycleAudit { + run_id: self.run_id, + worker: self.worker.clone(), + status, + trigger: self.trigger, + reason: reason.clone(), + memory_settings: self.memory_settings.clone(), + model: self.model.clone(), + usage, + extract, + consolidation, + }; + let _ = workspace_client + .execute_memory_backend_operation(memory::backend::MemoryBackendOperation::AppendAudit( + memory::backend::MemoryAppendAuditOperation { + event: memory::audit::AuditEvent::new( + memory::audit::AuditPayload::WorkerLifecycle(payload), + ), + }, + )) + .await; + if let Some(tx) = event_tx { + let _ = tx.send(Event::MemoryWorker(protocol::MemoryWorkerEvent { + worker: self.worker.label().to_string(), + status: status.label().to_string(), + run_id: self.run_id.to_string(), + trigger: self.trigger.label().to_string(), + reason: reason.clone(), + message: format!( + "memory {} {}: {reason}", + self.worker.label(), + status.label() + ), + timestamp_ms: now_millis() as i64, + })); + } + } +} + +fn model_audit_from_manifest(model: &manifest::ModelManifest) -> memory::audit::ModelAudit { + memory::audit::ModelAudit { + ref_: model.ref_.clone(), + scheme: model.scheme.map(|scheme| format!("{scheme:?}")), + model_id: model.model_id.clone(), + } +} + +#[cfg(test)] +mod tests { + use agen::{HistoryEntry, Item, UsageRecord}; + + use super::*; + use crate::feature::background::{BackgroundTaskRewritePolicy, BackgroundTaskShutdownPolicy}; + use crate::session_history::SessionHistoryMetadata; + + fn internal_result( + lifecycle: WorkerRunResult, + ) -> Result { + Ok(InternalWorkerResult { + usage: None, + identity: InternalWorkerIdentity { + kind: "memory-extract", + run_id: uuid::Uuid::now_v7(), + }, + lifecycle, + history_entries: 1, + }) + } + + fn capture(history_len: usize, input_total_tokens: u64) -> CommittedSessionCapture { + CommittedSessionCapture { + session_id: "session-1".to_string(), + segment_id: "segment-1".to_string(), + session_revision: history_len.try_into().unwrap(), + entry_count: history_len, + history: (0..history_len) + .map(|index| HistoryEntry { + item: Item::user_message(format!("message-{index}")), + annotation: SessionHistoryMetadata::legacy_unknown(), + }) + .collect(), + usage_history: vec![UsageRecord { + history_len, + input_total_tokens, + cache_read_tokens: 0, + cache_write_tokens: 0, + output_tokens: 0, + }], + extensions: Vec::new(), + } + } + + #[test] + fn normal_and_empty_extraction_require_explicit_finish() { + let result = internal_result(WorkerRunResult::Finished); + assert_eq!( + extraction_disposition(&result, true), + ExtractionDisposition::Completed + ); + // `finish_called = true` with no staged ids is the explicit empty + // extraction outcome. Missing Finish is a failed extraction. + assert!(matches!( + extraction_disposition(&result, false), + ExtractionDisposition::Failed(reason) + if reason.contains("FinishMemoryExtraction") + )); + } + + #[test] + fn failed_and_pre_ai_cancelled_extraction_never_reach_pointer_commit() { + let failed = internal_result(WorkerRunResult::Interrupted { + code: crate::ErrorCode::Internal, + message: "provider failed".to_string(), + }); + assert!(matches!( + extraction_disposition(&failed, true), + ExtractionDisposition::Failed(_) + )); + let cancelled = internal_result(WorkerRunResult::RolledBack); + assert_eq!( + extraction_disposition(&cancelled, true), + ExtractionDisposition::Cancelled("memory-extract cancelled".to_string()) + ); + } + + #[test] + fn task_scope_cancels_and_joins_before_rewrite_and_shutdown() { + let spec = memory_extraction_task_spec(); + assert_eq!(spec.trigger, BackgroundTaskTrigger::RunCommitted); + assert_eq!(spec.max_concurrency, 1); + assert_eq!(spec.rewrite, BackgroundTaskRewritePolicy::CancelAndWait); + assert_eq!(spec.shutdown, BackgroundTaskShutdownPolicy::CancelAndWait); + } + + #[test] + fn threshold_uses_committed_usage_after_pointer() { + let capture = capture(2, 250); + let mut config = manifest::MemoryConfig::default(); + config.extract_threshold = Some(1); + assert!(extraction_threshold_reached( + &capture, + Some(&memory::ExtractPointerPayload { + processed_through_entry: 0, + processed_through_history_len: 1, + staging_id: "staging-1".to_string(), + }), + &config + )); + } + + #[test] + fn pointer_folds_latest_committed_extraction_extension() { + let mut capture = capture(2, 250); + let first = memory::ExtractPointerPayload { + processed_through_entry: 1, + processed_through_history_len: 1, + staging_id: "staging-1".to_string(), + }; + let latest = memory::ExtractPointerPayload { + processed_through_entry: 3, + processed_through_history_len: 2, + staging_id: "staging-2".to_string(), + }; + capture.extensions = vec![ + ( + extract::EXTRACT_DOMAIN.to_string(), + serde_json::to_value(&first).unwrap(), + ), + ("other.feature".to_string(), serde_json::json!({})), + ( + extract::EXTRACT_DOMAIN.to_string(), + serde_json::to_value(&latest).unwrap(), + ), + ]; + assert_eq!(extract_pointer(&capture).unwrap(), Some(latest)); + } + + #[test] + fn malformed_latest_pointer_fails_closed_instead_of_using_older_pointer() { + let mut capture = capture(2, 250); + capture.extensions = vec![ + ( + extract::EXTRACT_DOMAIN.to_string(), + serde_json::to_value(memory::ExtractPointerPayload { + processed_through_entry: 1, + processed_through_history_len: 1, + staging_id: "staging-1".to_string(), + }) + .unwrap(), + ), + ( + extract::EXTRACT_DOMAIN.to_string(), + serde_json::json!({"invalid": true}), + ), + ]; + assert!(extract_pointer(&capture).is_err()); + } + + #[test] + fn worker_core_no_longer_owns_memory_extraction_scheduler() { + let worker_source = include_str!("../../worker.rs"); + for removed in [ + "spawn_post_run_memory_jobs", + "run_extract_once_with_cancel_observer", + "consolidation_in_flight", + "extract_in_flight", + "memory_task:", + ] { + assert!( + !worker_source.contains(removed), + "Worker core still contains removed extraction scheduler symbol {removed}" + ); + } + let controller_source = include_str!("../../controller.rs"); + assert!(controller_source.contains("if feature_config.memory.enabled")); + assert!(controller_source.contains("MemoryExtractionLifecycleFeature::new")); + let internal_worker_source = include_str!("../../internal_worker.rs"); + assert!(!internal_worker_source.contains("manifest.memory = None")); + } + + #[test] + fn empty_capture_never_schedules_extraction() { + let capture = capture(0, 500); + let mut config = manifest::MemoryConfig::default(); + config.extract_threshold = Some(1); + assert!(!extraction_threshold_reached(&capture, None, &config)); + } +} diff --git a/crates/worker/src/feature/builtin/memory_extract.rs b/crates/worker/src/feature/builtin/memory_staging_output.rs similarity index 94% rename from crates/worker/src/feature/builtin/memory_extract.rs rename to crates/worker/src/feature/builtin/memory_staging_output.rs index 023dfab0..4ccc90af 100644 --- a/crates/worker/src/feature/builtin/memory_extract.rs +++ b/crates/worker/src/feature/builtin/memory_staging_output.rs @@ -28,7 +28,7 @@ const FINISH_DESCRIPTION: &str = "Finish Memory extraction after validating the number of candidates staged during this run."; #[derive(Clone)] -pub(crate) struct MemoryExtractState { +pub(crate) struct MemoryStagingOutputState { view: Arc, workspace_client: Arc, source: SourceRef, @@ -37,7 +37,7 @@ pub(crate) struct MemoryExtractState { finished: Arc>>, } -impl MemoryExtractState { +impl MemoryStagingOutputState { pub(crate) fn new( view: SessionCapture, workspace_client: Arc, @@ -70,22 +70,20 @@ impl MemoryExtractState { } #[derive(Clone)] -pub(crate) struct MemoryExtractFeature { - state: MemoryExtractState, +pub(crate) struct MemoryStagingOutputFeature { + state: MemoryStagingOutputState, } -impl MemoryExtractFeature { - pub(crate) fn new(state: MemoryExtractState) -> Self { +impl MemoryStagingOutputFeature { + pub(crate) fn new(state: MemoryStagingOutputState) -> Self { Self { state } } } -impl FeatureModule for MemoryExtractFeature { +impl FeatureModule for MemoryStagingOutputFeature { fn descriptor(&self) -> FeatureDescriptor { - FeatureDescriptor::builtin("memory-extract", "Memory Extract") - .with_description( - "Memory staging and extraction completion, independent from session exploration.", - ) + FeatureDescriptor::builtin("memory-staging-output", "Memory Staging Output") + .with_description("Restricted Memory staging output for an extraction Internal Worker.") .with_tool(ToolDeclaration::new( "StageMemoryCandidate", STAGE_DESCRIPTION, @@ -109,7 +107,7 @@ impl FeatureModule for MemoryExtractFeature { } } -fn stage_definition(state: MemoryExtractState) -> ToolDefinition { +fn stage_definition(state: MemoryStagingOutputState) -> ToolDefinition { Arc::new(move || { let schema = serde_json::to_value(schemars::schema_for!(StageMemoryCandidateParams)) .unwrap_or_else(|_| serde_json::json!({})); @@ -123,7 +121,7 @@ fn stage_definition(state: MemoryExtractState) -> ToolDefinition { }) } -fn finish_definition(state: MemoryExtractState) -> ToolDefinition { +fn finish_definition(state: MemoryStagingOutputState) -> ToolDefinition { Arc::new(move || { let schema = serde_json::to_value(schemars::schema_for!(FinishMemoryExtractionParams)) .unwrap_or_else(|_| serde_json::json!({})); @@ -157,7 +155,7 @@ struct FinishMemoryExtractionParams { } struct StageMemoryCandidateTool { - state: MemoryExtractState, + state: MemoryStagingOutputState, } #[async_trait] @@ -252,7 +250,7 @@ impl Tool for StageMemoryCandidateTool { } struct FinishMemoryExtractionTool { - state: MemoryExtractState, + state: MemoryStagingOutputState, } #[async_trait] @@ -431,8 +429,8 @@ mod tests { use super::*; - fn state() -> MemoryExtractState { - MemoryExtractState::new( + fn state() -> MemoryStagingOutputState { + MemoryStagingOutputState::new( SessionCapture::new("segment-1", vec![Item::user_message("durable decision")]), crate::worker::marker_workspace_client(None, "test-backend"), SourceRef { @@ -445,8 +443,8 @@ mod tests { #[test] fn memory_extract_declares_only_memory_mutation_tools() { - let descriptor = MemoryExtractFeature::new(state()).descriptor(); - assert_eq!(descriptor.id.as_str(), "builtin:memory-extract"); + let descriptor = MemoryStagingOutputFeature::new(state()).descriptor(); + assert_eq!(descriptor.id.as_str(), "builtin:memory-staging-output"); assert_eq!( descriptor .tools diff --git a/crates/worker/src/feature/session.rs b/crates/worker/src/feature/session.rs new file mode 100644 index 00000000..97345b10 --- /dev/null +++ b/crates/worker/src/feature/session.rs @@ -0,0 +1,111 @@ +use std::sync::Arc; + +use agen::{HistoryEntry, UsageRecord}; +use serde_json::Value; + +use crate::session_history::SessionHistoryMetadata; + +/// Immutable projection of one durably committed session-log location. +/// +/// Feature code receives this value only after the host has committed the +/// terminal run record. The projection deliberately carries annotated history +/// rather than the public flattened transcript so provenance-sensitive +/// features can construct their own bounded views. +#[derive(Clone)] +pub(crate) struct CommittedSessionCapture { + pub(crate) session_id: String, + pub(crate) segment_id: String, + /// Monotonic committed-log revision for the captured Segment. + pub(crate) session_revision: u64, + pub(crate) entry_count: usize, + pub(crate) history: Vec>, + pub(crate) usage_history: Vec, + pub(crate) extensions: Vec<(String, Value)>, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct CommittedSessionLocation { + pub(crate) session_id: String, + pub(crate) segment_id: String, + /// Monotonic committed-log revision for the captured Segment. + pub(crate) session_revision: u64, + pub(crate) entry_count: usize, +} + +impl CommittedSessionCapture { + pub(crate) fn location(&self) -> CommittedSessionLocation { + CommittedSessionLocation { + session_id: self.session_id.clone(), + segment_id: self.segment_id.clone(), + session_revision: self.session_revision, + entry_count: self.entry_count, + } + } +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum FeatureSessionError { + #[error("read committed session failed: {0}")] + Capture(String), + #[error("append session extension failed: {0}")] + Extension(String), +} + +#[derive(Clone)] +pub(crate) struct CommittedSessionCaptureHandle { + capture: Arc< + dyn Fn() -> Result + Send + Sync + 'static, + >, +} + +impl CommittedSessionCaptureHandle { + pub(crate) fn new( + capture: impl Fn() -> Result + + Send + + Sync + + 'static, + ) -> Self { + Self { + capture: Arc::new(capture), + } + } + + pub(crate) fn capture(&self) -> Result { + (self.capture)() + } +} + +#[derive(Clone)] +pub(crate) struct SessionExtensionHandle { + append: Arc< + dyn Fn(&CommittedSessionLocation, &str, Value) -> Result + + Send + + Sync + + 'static, + >, +} + +impl SessionExtensionHandle { + pub(crate) fn new( + append: impl Fn(&CommittedSessionLocation, &str, Value) -> Result + + Send + + Sync + + 'static, + ) -> Self { + Self { + append: Arc::new(append), + } + } + + /// Appends an extension only while the committed session is still at the + /// exact location captured by the feature. `Ok(false)` is a stale-write + /// fence, not an I/O failure. + pub(crate) fn append_if_current( + &self, + expected: &CommittedSessionLocation, + domain: &str, + payload: Value, + ) -> Result { + (self.append)(expected, domain, payload) + } +} diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index b69ac42c..c81787c5 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -123,14 +123,14 @@ where // Internal identities are run-scoped and never enter the public Runtime Worker catalog. manifest.worker.name = format!("internal-{}-{}", identity.kind, identity.run_id); - // Internal jobs only receive features supplied below. A parent manifest must not accidentally - // grant its normal public tool surface or recursively schedule memory work. + // Internal jobs only receive the explicitly supplied Feature set below. A + // parent manifest cannot accidentally grant its normal public tool surface + // or recursively schedule Feature-owned background work. manifest.feature = Default::default(); manifest.plugins = Default::default(); manifest.mcp = Default::default(); manifest.skills = None; manifest.compaction = None; - manifest.memory = None; let last_usage = Arc::new(Mutex::new(None::)); let usage_slot = last_usage.clone(); @@ -548,7 +548,6 @@ pub(crate) async fn spawn_internal_worker_session( authority, } = spec; manifest.worker.name = format!("internal-{}-{}", identity.kind, identity.run_id); - manifest.memory = None; let last_usage = Arc::new(Mutex::new(None::)); let usage_slot = last_usage.clone(); @@ -649,7 +648,6 @@ pub(crate) fn prepare_internal_worker_from_spec( manifest.mcp = Default::default(); manifest.skills = None; manifest.compaction = None; - manifest.memory = None; let mut engine = Engine::<_, agen::state::Mutable, crate::SessionHistoryMetadata>::new_annotated(client) diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index d69afe93..c5fcb23b 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -1,7 +1,7 @@ #[cfg(test)] use std::path::Path; use std::path::PathBuf; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -42,9 +42,10 @@ use crate::compact::state::CompactState; use crate::compact::usage_tracker::UsageTracker; use crate::feature::background::{BackgroundTaskRewriteGuard, FeatureBackgroundTaskRegistry}; use crate::feature::builtin::memory::WorkspaceMemoryBackendError; -use crate::feature::builtin::{ - MemoryExtractFeature, MemoryExtractState, SessionExploreFeature, SessionExploreState, - TaskFeature, WorkerObservationProvider, render_extract_input, +use crate::feature::builtin::{TaskFeature, WorkerObservationProvider}; +use crate::feature::session::{ + CommittedSessionCapture, CommittedSessionCaptureHandle, FeatureSessionError, + SessionExtensionHandle, }; use crate::feature::{ FeatureInstructionDeclaration, FeatureInstructionId, FeatureRegistryBuilder, @@ -59,7 +60,7 @@ use crate::hook::{ use crate::in_flight::InFlightEvents; use crate::internal_worker::{ InternalWorkerAuthority, InternalWorkerIdentity, InternalWorkerSpec, InternalWorkerVisibility, - prepare_internal_worker_from_spec, run_internal_worker, run_internal_worker_with_cancel_sender, + prepare_internal_worker_from_spec, }; const COMPACTION_EXTENSION_DOMAIN: &str = "yoi.compaction"; @@ -103,7 +104,6 @@ use protocol::{ }; use tokio::net::UnixStream; use tokio::sync::broadcast; -use tokio::task::JoinHandle; use workdir::{ LocalWorkdirSession, ReadOnlyWorkdirSession, WorkdirSessionCapabilities, WorkdirSessionHandle, }; @@ -772,6 +772,7 @@ where pub struct SegmentState { location: ArcSwap, entries_written: AtomicUsize, + append_lock: Mutex<()>, } impl SegmentState { @@ -782,6 +783,7 @@ impl SegmentState { segment_id, }), entries_written: AtomicUsize::new(entries_written), + append_lock: Mutex::new(()), }) } @@ -930,6 +932,15 @@ where /// mirror push → broadcast. The Store owns physical write ordering and /// partial-write recovery; publication happens only after it returns Ok. pub fn append_entry(&self, entry: LogEntry) -> Result<(), StoreError> { + let _append_guard = self + .state + .append_lock + .lock() + .expect("segment append lock poisoned"); + self.append_entry_locked(entry) + } + + fn append_entry_locked(&self, entry: LogEntry) -> Result<(), StoreError> { let loc = self.state.location(); self.store.append(loc.session_id, loc.segment_id, &entry)?; self.state.increment_entries(); @@ -1114,9 +1125,6 @@ pub struct Worker { /// Cloned into local WorkdirSession providers used by builtin tools, fs_view, /// and compaction so updates propagate at the next permission check. scope: SharedScope, - /// Filesystem authority this Worker may pass to spawned children. Direct tools - /// 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>, @@ -1221,28 +1229,6 @@ pub struct Worker { /// context from the workspace Memory document. Internal disposable /// workers disable this so resident memory exposure is opt-in per Worker. inject_resident_summary: bool, - /// When true (default), the system-prompt assembler may append resident - /// resident context. This is intentionally independent from - /// summary residency: each section has its own gate. - /// extract (memory.extract) reentry guard. `true` while an extract - /// worker is running; subsequent triggers are skipped per spec - /// (`docs/plan/memory.md` §Extract 並走防止). `Arc` so - /// the flag survives across `try_post_run_extract` calls without a - /// `&mut self` race. - extract_in_flight: Arc, - /// consolidation (memory.consolidation) in-process reentry guard. - consolidation_in_flight: Arc, - /// Last completed extract boundary. `None` means no extract has - /// run yet on this session — next extract starts from entry 0. - /// Restored from `RestoredState.extensions` on `restore`, updated - /// after each successful extract via `save_extension`. - extract_pointer: Arc>>, - /// extract/consolidation memory job running outside the controller method loop. - /// The task owns the extract/consolidate worker execution and is joined - /// at shutdown. A single slot is enough: extract/consolidation implementations loop - /// until thresholds fall below their trigger points, and concurrent - /// triggers are coalesced by skipping when this handle is still active. - memory_task: Option>, /// Typed user submissions in submit order. K-th entry corresponds to /// the K-th `Item::user_message` in `worker.history()` (modulo seed /// history loaded via `AnnotatedSegmentStart.history`, whose original segments @@ -1302,82 +1288,6 @@ impl Worker { 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 - { - tracing::warn!(error = %e, "Post-run memory task join failed"); - } - } -} - -impl Worker { - fn clone_for_memory_task(&self) -> Self { - // The cloned Worker's worker exists only as a snapshot for the memory - // task: `run_extract_once` reads `worker.history()`, and the - // extract/consolidate workers are built fresh inside their own - // methods using `worker.client()` as fallback when no override - // model is configured. system_prompt / request_config / cache_key - // are unused on this path, so we deliberately skip copying them. - let source_worker = self.engine.as_ref().expect("worker present"); - let worker = Engine::::new_annotated( - source_worker.client().clone(), - ); - Self { - manifest: self.manifest.clone(), - engine: Some(worker), - session: self.session.clone(), - last_run_interrupted: false, - store: self.store.clone(), - worker_metadata_writer: None, - segment_state: self.segment_state.clone(), - filesystem_authority: self.filesystem_authority.clone(), - workdir_session: self.workdir_session.clone(), - workspace_context: self.workspace_context.clone(), - flow_runtime_state: self.flow_runtime_state.clone(), - flow_feature_enabled: self.flow_feature_enabled, - 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()), - metrics_tracker: Arc::new(crate::compact::metrics_tracker::MetricsTracker::new()), - usage_history: self.usage_history.clone(), - tracker: None, - task_feature: self.task_feature.clone(), - worker_observation_provider: None, - system_prompt_template: None, - feature_instructions: self.feature_instructions.clone(), - alerter: self.alerter.clone(), - working_event_tx: self.working_event_tx.clone(), - internal_worker_registry: self.internal_worker_registry.clone(), - in_flight: self.in_flight.clone(), - ai_activity_counter: self.ai_activity_counter.clone(), - pending_notifies: NotifyBuffer::new(), - pending_attachments: Arc::new(Mutex::new(Vec::::new())), - pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())), - scope_allocation: None, - callback_socket: None, - runtime_ticket_role: None, - prompts: self.prompts.clone(), - inject_resident_summary: self.inject_resident_summary, - extract_in_flight: self.extract_in_flight.clone(), - consolidation_in_flight: self.consolidation_in_flight.clone(), - extract_pointer: self.extract_pointer.clone(), - memory_task: None, - user_segments: self.user_segments.clone(), - // The memory-task clone never appends to the session log - // (it only reads `worker.history()`), so a fresh sink is - // fine — nothing observes its broadcast. - sink: SegmentLogSink::new(), - history_persistence_wired: false, - log_writer: None, - } - } } impl Worker { @@ -1467,28 +1377,6 @@ impl Worker { } } -impl Worker { - pub fn spawn_post_run_memory_jobs(&mut self) { - // Drop a finished prior handle so we can spawn a fresh task. - // If the prior task is still running, coalesce by skipping — - // extract/consolidation implementations re-evaluate thresholds on completion. - self.cleanup_finished_memory_task(); - if self.memory_task.is_some() { - return; - } - - let mut worker = self.clone_for_memory_task(); - self.memory_task = Some(tokio::spawn(async move { - if let Err(e) = worker.try_post_run_extract().await { - tracing::warn!(error = %e, "Post-run memory extract task error"); - } - if let Err(e) = worker.try_post_run_consolidate().await { - tracing::warn!(error = %e, "Post-run memory consolidate task error"); - } - })); - } -} - impl Worker { /// Create a new Worker from a pre-built Engine and store. /// @@ -1515,8 +1403,7 @@ impl Worker { let session_id = session_store::new_session_id(); let segment_id = session_store::new_segment_id(); let prompts = Arc::new(ArcSwap::from(PromptCatalog::builtins_only()?)); - let delegation_scope = - DelegationScope::from_config(&manifest.delegation_scope).map_err(WorkerError::Scope)?; + DelegationScope::from_config(&manifest.delegation_scope).map_err(WorkerError::Scope)?; let scope = SharedScope::new(scope); let workdir_session = workdir_session_from_authority(&filesystem_authority, &scope); let mut worker = Self { @@ -1533,7 +1420,6 @@ impl Worker { flow_runtime_state: Arc::new(Mutex::new(None)), flow_feature_enabled: false, scope, - delegation_scope, hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), @@ -1560,10 +1446,6 @@ impl Worker { runtime_ticket_role: None, prompts, inject_resident_summary: true, - extract_in_flight: Arc::new(AtomicBool::new(false)), - consolidation_in_flight: Arc::new(AtomicBool::new(false)), - extract_pointer: Arc::new(Mutex::new(None)), - memory_task: None, user_segments: Vec::new(), sink: SegmentLogSink::new(), history_persistence_wired: false, @@ -2062,7 +1944,6 @@ impl Worker { let retained = entries[..truncate_entries].to_vec(); let tool_side_effect_warning = suffix_has_tool_side_effects(&entries[truncate_entries..]); let state = segment_log::collect_state(&retained); - let extract_pointer = memory::extract::fold_pointer(&state.extensions); let summary = RewindSummary { truncated_to_entries: truncate_entries, discarded_entries: entries.len().saturating_sub(truncate_entries), @@ -2093,10 +1974,6 @@ impl Worker { .pending_attachments .lock() .expect("pending_attachments poisoned") = Vec::new(); - *self - .extract_pointer - .lock() - .expect("extract_pointer poisoned") = extract_pointer; Ok(RewindAppliedState { entries: retained, @@ -2197,21 +2074,6 @@ impl Worker { &self.user_segments } - pub fn extract_pointer(&self) -> Option { - self.extract_pointer - .lock() - .expect("extract_pointer poisoned") - .clone() - } - - /// Test/diagnostic handle to the consolidation in-flight guard. Production - /// callers do not need this; tests use it to assert that the reentry - /// guard skips an in-progress consolidation without losing data. - #[doc(hidden)] - pub fn consolidation_in_flight_handle(&self) -> Arc { - self.consolidation_in_flight.clone() - } - /// Shared handle to the cumulative Usage history. /// /// Callbacks that need live access to the latest measurements (e.g. @@ -2228,6 +2090,90 @@ impl Worker { self.usage_history.clone() } + /// Narrow read-only handle used by lifecycle Features after terminal run + /// persistence. The capture is rebuilt from the committed log under the + /// same append fence used by log writers. + pub(crate) fn committed_session_capture_handle(&self) -> CommittedSessionCaptureHandle + where + St: Clone + Send + Sync + 'static, + { + let store = self.store.clone(); + let state = Arc::clone(&self.segment_state); + CommittedSessionCaptureHandle::new(move || { + let _append_guard = state + .append_lock + .lock() + .expect("segment append lock poisoned"); + let location = state.location(); + let entries = store + .read_all(location.session_id, location.segment_id) + .map_err(|error| FeatureSessionError::Capture(error.to_string()))?; + let restored = segment_log::collect_state(&entries); + let history = + restore_history_entries(location.session_id, location.segment_id, &entries) + .map_err(|error| FeatureSessionError::Capture(error.to_string()))?; + Ok(CommittedSessionCapture { + session_id: location.session_id.to_string(), + segment_id: location.segment_id.to_string(), + session_revision: entries.len().try_into().unwrap_or(u64::MAX), + entry_count: entries.len(), + history, + usage_history: restored.usage_history, + extensions: restored.extensions.into_iter().collect(), + }) + }) + } + + /// Narrow fenced append authority for Feature-owned Session extensions. + pub(crate) fn session_extension_handle(&self) -> SessionExtensionHandle + where + St: Clone + Send + Sync + 'static, + { + let writer = self.log_writer_handle(); + SessionExtensionHandle::new(move |expected, domain, payload| { + let _append_guard = writer + .state + .append_lock + .lock() + .expect("segment append lock poisoned"); + let location = writer.state.location(); + if location.session_id.to_string() != expected.session_id + || location.segment_id.to_string() != expected.segment_id + || u64::try_from(writer.state.entries_written()).unwrap_or(u64::MAX) + != expected.session_revision + || writer.state.entries_written() != expected.entry_count + { + return Ok(false); + } + writer + .append_entry_locked(LogEntry::Extension { + ts: segment_log::now_millis(), + domain: domain.to_string(), + payload, + }) + .map_err(|error| FeatureSessionError::Extension(error.to_string()))?; + Ok(true) + }) + } + + pub(crate) fn llm_client_handle(&self) -> Box + where + C: Clone, + { + Box::new( + (*self + .engine + .as_ref() + .expect("worker taken during run") + .client()) + .clone(), + ) + } + + pub(crate) fn working_event_sender(&self) -> Option> { + self.working_event_tx.clone() + } + /// Handle to the per-LLM-request `UsageTracker`. /// /// Sibling modules (e.g. the prune observer) clone this `Arc` to stash @@ -2601,24 +2547,6 @@ impl Worker { self.run(vec![Segment::text(s)]).await } - /// Drop the prior memory_task handle if it has finished. Keep it if - /// still running so callers can decide whether to wait or coalesce. - fn cleanup_finished_memory_task(&mut self) { - if self.memory_task.as_ref().is_some_and(|h| h.is_finished()) { - self.memory_task = None; - } - } - - /// Wait for the in-flight memory task (if any) to finish. Used before - /// compact rewrites history (extract reads the same history). - async fn join_memory_task(&mut self) { - if let Some(handle) = self.memory_task.take() - && let Err(e) = handle.await - { - tracing::warn!(error = %e, "Memory task join failed"); - } - } - /// Whether `try_pre_run_compact` would actually compact. The same /// check is duplicated inside `try_pre_run_compact` itself for /// defensive reasons; this is the gate for joining the memory task @@ -2642,11 +2570,8 @@ impl Worker { self.refresh_prompt_projection_for_future_operations()?; self.ensure_interceptor_installed(); self.ensure_system_prompt_materialized().await?; - self.cleanup_finished_memory_task(); self.ensure_segment_head().await?; - if self.should_pre_run_compact() { - self.join_memory_task().await; - } + if self.should_pre_run_compact() {} self.try_pre_run_compact().await; Ok(()) } @@ -3744,7 +3669,6 @@ impl Worker { } self.ensure_interceptor_installed(); - self.cleanup_finished_memory_task(); self.ensure_segment_head().await?; let state = self.compact_state.clone(); @@ -3770,7 +3694,6 @@ impl Worker { return Ok(ManualCompactResult::Skipped { message }); } - self.join_memory_task().await; match self.compact(retained).await { Ok(new_segment_id) => { info!(new_segment_id = %new_segment_id, "Manual compaction succeeded"); @@ -4578,26 +4501,12 @@ impl Worker { // compact layout guarantees history[0] is the summary. worker.set_cache_anchor(Some(0)); // Re-key the OpenAI Responses prompt cache namespace to the new - // segment_id so post-compact turns share a key with extract / - // consolidate workers running in the same session. + // segment_id so post-compact turns use the rewritten session namespace. worker.set_cache_key(Some(new_segment_id.to_string())); self.usage_history .lock() .expect("usage_history poisoned") .clear(); - // Reset extract pointer alongside usage_history: the compacted - // session has a fresh log with no `LogEntry::Extension` entries - // yet, so a cold restore here would set extract_pointer to None - // via fold_pointer. The in-memory pointer must match — otherwise - // `tokens_added_since(old_history_len)` would treat the new - // (shorter) history as if it had already been processed, and - // extract would stop firing for the rest of the process's - // lifetime. - *self - .extract_pointer - .lock() - .expect("extract_pointer poisoned") = None; - Ok((new_segment_id, summary_text)) } @@ -4615,736 +4524,15 @@ impl Worker { let worker = self.engine.as_ref().expect("worker taken during run"); Ok(worker.client().clone_boxed()) } - - /// Build the LlmClient for the extract (memory.extract) Engine. - /// - /// Uses `memory.extract_model` from manifest if set, otherwise clones - /// the main client. - fn build_extractor_client( - &self, - memory_cfg: &manifest::MemoryConfig, - ) -> Result, WorkerError> { - if let Some(ref m) = memory_cfg.extract_model { - let client = crate::model_client::build_client(m)?; - return Ok(client); - } - let worker = self.engine.as_ref().expect("worker taken during run"); - Ok(worker.client().clone_boxed()) - } - - /// pointer 以降に増えたプロンプト全長の推定。extract trigger が - /// 閾値判定に使う。 - /// - /// `total_tokens_at(now) - total_tokens_at(pointer)` の差分で、 - /// compact と同じ accounting (measured / interpolated / extrapolated) - /// に乗る。`history_len_pointer == 0` は「未抽出」扱いで現プロンプト - /// 全長そのものが返る。 - /// - /// 素朴な `usage_history.input_total_tokens` の合計は使わない: - /// `input_total_tokens` は **送信時の prompt prefix 全長** であって - /// 増分ではないので、長い turn 内の連続 LLM call では super-set を - /// 何度も足し込んでしまい実消費の数倍に膨らむ。 - fn tokens_added_since(&self, history_len_pointer: usize) -> u64 { - let now = self.history().len(); - let total_now = self.total_tokens_at(now).tokens; - let total_at_pointer = self.total_tokens_at(history_len_pointer).tokens; - total_now.saturating_sub(total_at_pointer) - } - - /// extract (memory.extract) post-run trigger. - /// - /// Called by the Controller before spawning the background memory task so - /// the extract worker sees a stable session-log entry range while compact - /// is deferred until the next turn starts. Best-effort: failures are - /// logged but not propagated. - /// - /// Behaviour follows `docs/plan/memory.md` §Extract 並走防止: - /// in-flight 中の trigger は skip し、完了時点で閾値再評価する - /// (the loop below). Pending state is not retained — the - /// re-evaluation happens naturally because the in-memory pointer - /// has advanced. - pub async fn try_post_run_extract(&mut self) -> Result<(), WorkerError> { - let Some(memory_cfg) = self.manifest.memory.clone() else { - return Ok(()); - }; - // `Some(0)` means disabled, same as `None`. Otherwise the - // `tokens_since >= 0` comparison would fire on every post-run. - let Some(threshold) = memory_cfg.extract_threshold.filter(|n| *n > 0) else { - let model = memory_cfg - .extract_model - .as_ref() - .unwrap_or(&self.manifest.model); - WorkerAuditBase::new( - memory::audit::AuditWorker::MemoryExtract, - memory::audit::AuditTrigger::TokenThreshold, - Some(model_audit_from_manifest(model)), - ) - .with_memory_settings(&memory_cfg) - .emit( - self.workspace_client(), - self.working_event_tx.as_ref(), - memory::audit::WorkerLifecycleStatus::Skipped, - "extract_threshold_disabled", - None, - None, - None, - ) - .await; - return Ok(()); - }; - - loop { - // CAS the in-flight flag. If another task is already running - // an extract for this Worker, skip per spec. - if self - .extract_in_flight - .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) - .is_err() - { - let model = memory_cfg - .extract_model - .as_ref() - .unwrap_or(&self.manifest.model); - WorkerAuditBase::new( - memory::audit::AuditWorker::MemoryExtract, - memory::audit::AuditTrigger::TokenThreshold, - Some(model_audit_from_manifest(model)), - ) - .with_memory_settings(&memory_cfg) - .emit( - self.workspace_client(), - self.working_event_tx.as_ref(), - memory::audit::WorkerLifecycleStatus::Skipped, - "extract_already_in_flight", - None, - None, - None, - ) - .await; - return Ok(()); - } - let result = self.run_extract_once(&memory_cfg, threshold).await; - self.extract_in_flight.store(false, Ordering::Release); - - match result { - Ok(ExtractDecision::Skipped) => return Ok(()), - Ok(ExtractDecision::Completed) => { - // Re-evaluate threshold against the newly advanced - // pointer. In the current synchronous architecture - // this normally exits via Skipped on the next pass, - // but the loop is forward-looking for the case - // where new activity piles up while extract runs. - continue; - } - Err(e) => { - tracing::warn!(error = %e, "extract failed"); - self.alert( - AlertLevel::Warn, - AlertSource::Worker, - format!("memory extract failed: {e}"), - ); - return Ok(()); - } - } - } - } - - /// Single extract iteration: snapshot pointer, decide whether to - /// fire, run the worker if so, persist results and the new pointer. - async fn run_extract_once( - &mut self, - memory_cfg: &manifest::MemoryConfig, - threshold: u64, - ) -> Result { - self.run_extract_once_with_cancel_observer(memory_cfg, threshold, None) - .await - } - - async fn run_extract_once_with_cancel_observer( - &mut self, - memory_cfg: &manifest::MemoryConfig, - threshold: u64, - cancel_observer: Option) + Send + 'static>>, - ) -> Result { - use memory::extract; - - let model = memory_cfg - .extract_model - .as_ref() - .unwrap_or(&self.manifest.model); - let audit = WorkerAuditBase::new( - memory::audit::AuditWorker::MemoryExtract, - memory::audit::AuditTrigger::TokenThreshold, - Some(model_audit_from_manifest(model)), - ) - .with_memory_settings(memory_cfg); - let working_event_tx = self.working_event_tx.as_ref(); - - let pointer_snapshot = self - .extract_pointer - .lock() - .expect("extract_pointer poisoned") - .clone(); - let processed_history_len = pointer_snapshot - .as_ref() - .map(|p| p.processed_through_history_len) - .unwrap_or(0); - - let tokens_since = self.tokens_added_since(processed_history_len); - if tokens_since < threshold { - audit.emit( - self.workspace_client(), - working_event_tx, - memory::audit::WorkerLifecycleStatus::Skipped, - format!( - "token_threshold_not_reached tokens_since={tokens_since} threshold={threshold}" - ), - None, - None, - None, - ).await; - return Ok(ExtractDecision::Skipped); - } - - let current_history_len = self.session.history().len(); - if current_history_len <= processed_history_len { - audit - .emit( - self.workspace_client(), - working_event_tx, - memory::audit::WorkerLifecycleStatus::Skipped, - "no_new_history_items", - None, - Some(memory::audit::ExtractAudit { - history_range: Some([ - processed_history_len as u64, - current_history_len as u64, - ]), - ..Default::default() - }), - None, - ) - .await; - return Ok(ExtractDecision::Skipped); - } - - // Read the session log to get the current entry count. This is - // the boundary for the source.range end_entry. Called once per - // extract, on a small local file. - let entries_now = self - .store - .read_all(self.session_id(), self.segment_id())? - .len(); - if entries_now == 0 { - audit - .emit( - self.workspace_client(), - working_event_tx, - memory::audit::WorkerLifecycleStatus::Skipped, - "empty_segment_log", - None, - None, - None, - ) - .await; - return Ok(ExtractDecision::Skipped); - } - let end_entry = entries_now - 1; - let start_entry = pointer_snapshot - .as_ref() - .map(|p| p.processed_through_entry + 1) - .unwrap_or(0); - if start_entry > end_entry { - audit - .emit( - self.workspace_client(), - working_event_tx, - memory::audit::WorkerLifecycleStatus::Skipped, - "no_new_segment_entries", - None, - Some(memory::audit::ExtractAudit { - session_id: Some(self.session_id().to_string()), - segment_id: Some(self.segment_id().to_string()), - entry_range: Some([start_entry as u64, end_entry as u64]), - history_range: Some([ - processed_history_len as u64, - current_history_len as u64, - ]), - ..Default::default() - }), - None, - ) - .await; - return Ok(ExtractDecision::Skipped); - } - - let extract_audit_base = memory::audit::ExtractAudit { - session_id: Some(self.session_id().to_string()), - segment_id: Some(self.segment_id().to_string()), - entry_range: Some([start_entry as u64, end_entry as u64]), - history_range: Some([processed_history_len as u64, current_history_len as u64]), - ..Default::default() - }; - audit - .emit( - self.workspace_client(), - working_event_tx, - memory::audit::WorkerLifecycleStatus::Started, - format!( - "token_threshold_reached tokens_since={tokens_since} threshold={threshold}" - ), - None, - Some(extract_audit_base.clone()), - None, - ) - .await; - - let entries_to_extract = - self.session.history().entries()[processed_history_len..current_history_len].to_vec(); - - let extract_worker_max_turns = memory_cfg - .extract_worker_max_turns - .or(manifest::defaults::MEMORY_EXTRACT_WORKER_MAX_TURNS); - - let client = match self.build_extractor_client(memory_cfg) { - Ok(client) => client, - Err(err) => { - audit - .emit( - self.workspace_client(), - working_event_tx, - memory::audit::WorkerLifecycleStatus::Failed, - format!("client_build_failed: {err}"), - None, - Some(extract_audit_base), - None, - ) - .await; - return Err(err); - } - }; - let memory_language = memory_language(memory_cfg)?; - let extract_system_prompt = match self - .prompts - .load_full() - .memory_extract_system(&memory_language) - { - Ok(prompt) => prompt, - Err(err) => { - audit - .emit( - self.workspace_client(), - working_event_tx, - memory::audit::WorkerLifecycleStatus::Failed, - format!("prompt_render_failed: {err}"), - None, - Some(extract_audit_base), - None, - ) - .await; - return Err(WorkerError::PromptCatalog(err)); - } - }; - let source_segment_id = self.segment_state.segment_id(); - let source = memory::schema::SourceRef { - segment_id: source_segment_id.to_string(), - range: [start_entry as u64, end_entry as u64], - }; - let session_view = crate::session_capture::SessionCapture::from_history_entries( - source_segment_id.to_string(), - entries_to_extract, - ); - let session_explore_state = SessionExploreState::new(session_view.clone()); - let memory_extract_state = MemoryExtractState::new( - session_view, - self.workspace_client_handle(), - source, - audit.run_id.to_string(), - ); - let input_text = render_extract_input(session_explore_state.view()); - let features = FeatureRegistryBuilder::new() - .with_module(SessionExploreFeature::new(session_explore_state.clone())) - .with_module(MemoryExtractFeature::new(memory_extract_state.clone())); - let mut internal_manifest = self.manifest.clone(); - internal_manifest.model = model.clone(); - let internal_spec = InternalWorkerSpec { - identity: InternalWorkerIdentity { - kind: "memory-extract", - run_id: audit.run_id, - }, - manifest: internal_manifest, - client, - system_prompt: extract_system_prompt, - input: input_text, - cache_key: Some(self.segment_id().to_string()), - max_turns: extract_worker_max_turns, - engine_configurator: None, - features, - required_tools: &[ - "ShowOverview", - "SearchEntries", - "ReadEntry", - "StageMemoryCandidate", - "FinishMemoryExtraction", - ], - authority: InternalWorkerAuthority { - workspace: self.workspace_context.clone(), - filesystem: WorkerFilesystemAuthority::None, - scope: Scope::empty(), - workdir_session: None, - }, - }; - let internal_result = match cancel_observer { - Some(observer) => run_internal_worker_with_cancel_sender(internal_spec, observer).await, - None => run_internal_worker(internal_spec).await, - }; - let usage = match internal_result { - Ok(result) => { - tracing::debug!( - internal_worker_kind = result.identity.kind, - internal_worker_run_id = %result.identity.run_id, - history_entries = result.history_entries, - lifecycle = ?result.lifecycle, - "internal Worker execution completed" - ); - let usage = result.usage.as_ref().map(usage_audit_from_event); - if let Some(error) = extract_internal_worker_lifecycle_error(&result.lifecycle) { - audit - .emit( - self.workspace_client(), - working_event_tx, - memory::audit::WorkerLifecycleStatus::Cancelled, - "worker_cancelled: internal Worker run rolled back before AI output", - usage, - Some(extract_audit_base), - None, - ) - .await; - return Err(error); - } - usage - } - Err(err) => { - tracing::debug!( - internal_worker_kind = err.identity.kind, - internal_worker_run_id = %err.identity.run_id, - history_entries = err.history_entries, - "internal Worker execution failed" - ); - let usage = err.usage.as_ref().map(usage_audit_from_event); - audit - .emit( - self.workspace_client(), - working_event_tx, - lifecycle_status_for_worker_error(&err.source), - format!("worker_failed: {}", err.source), - usage, - Some(extract_audit_base), - None, - ) - .await; - return Err(err.source); - } - }; - - let staging_results = memory_extract_state.staged(); - if !memory_extract_state.is_finished() { - tracing::warn!( - staged_count = staging_results.len(), - "extract worker did not call FinishMemoryExtraction; advancing pointer with staged output" - ); - } - let staging_id = staging_results.first().cloned().unwrap_or_default(); - - let pointer_payload = extract::ExtractPointerPayload { - processed_through_entry: end_entry, - processed_through_history_len: current_history_len, - staging_id: staging_id.clone(), - }; - let payload_value = serde_json::to_value(&pointer_payload) - .expect("ExtractPointerPayload is always JSON-serializable"); - self.commit_entry(LogEntry::Extension { - ts: segment_log::now_millis(), - domain: extract::EXTRACT_DOMAIN.into(), - payload: payload_value, - })?; - - *self - .extract_pointer - .lock() - .expect("extract_pointer poisoned") = Some(pointer_payload); - - let mut extract_audit = extract_audit_base; - extract_audit.staging_count = staging_results.len(); - for id in &staging_results { - extract_audit.staging_ids.push(id.clone()); - } - let reason = if staging_id.is_empty() { - "completed_no_staging_output" - } else { - "completed_staging_written" - }; - audit - .emit( - self.workspace_client(), - working_event_tx, - memory::audit::WorkerLifecycleStatus::Completed, - reason, - usage, - Some(extract_audit), - None, - ) - .await; - - Ok(ExtractDecision::Completed) - } - - /// Request Backend-managed Memory staging consolidation after a Worker turn. - /// - /// Worker has no local Workspace memory authority. It only asks the Backend - /// Workspace to notify or spawn the dedicated consolidater Worker. - pub async fn try_post_run_consolidate(&mut self) -> Result<(), WorkerError> { - let Some(memory_cfg) = self.manifest.memory.clone() else { - return Ok(()); - }; - let model = memory_cfg - .consolidation_model - .as_ref() - .unwrap_or(&self.manifest.model); - let files_threshold = memory_cfg.consolidation_threshold_files.filter(|n| *n > 0); - let bytes_threshold = memory_cfg.consolidation_threshold_bytes.filter(|n| *n > 0); - if files_threshold.is_none() && bytes_threshold.is_none() { - WorkerAuditBase::new( - memory::audit::AuditWorker::MemoryConsolidation, - memory::audit::AuditTrigger::StagingBacklog, - Some(model_audit_from_manifest(model)), - ) - .with_memory_settings(&memory_cfg) - .emit( - self.workspace_client(), - self.working_event_tx.as_ref(), - memory::audit::WorkerLifecycleStatus::Skipped, - "consolidation_threshold_disabled", - None, - None, - None, - ) - .await; - return Ok(()); - } - - match self - .workspace_client() - .request_memory_staging_consolidation( - memory::backend::MemoryConsolidateStagingOperation { - force: false, - threshold_files: files_threshold, - threshold_bytes: bytes_threshold, - }, - ) - .await - { - Ok(output) => { - tracing::debug!( - status = output.status.as_str(), - summary = output.summary.as_str(), - "requested backend memory staging consolidation" - ); - } - Err(error) => { - tracing::warn!( - error = %error, - "failed to request backend memory staging consolidation" - ); - WorkerAuditBase::new( - memory::audit::AuditWorker::MemoryConsolidation, - memory::audit::AuditTrigger::StagingBacklog, - Some(model_audit_from_manifest(model)), - ) - .with_memory_settings(&memory_cfg) - .emit( - self.workspace_client(), - self.working_event_tx.as_ref(), - memory::audit::WorkerLifecycleStatus::Skipped, - "consolidation_backend_operation_failed", - None, - None, - None, - ) - .await; - } - } - Ok(()) - } } -fn extract_internal_worker_lifecycle_error(lifecycle: &WorkerRunResult) -> Option { - match lifecycle { - WorkerRunResult::RolledBack => Some(WorkerError::Engine(EngineError::Cancelled)), - WorkerRunResult::Interrupted { message, .. } => { - Some(WorkerError::Engine(EngineError::Aborted(message.clone()))) - } - WorkerRunResult::Finished | WorkerRunResult::Paused | WorkerRunResult::LimitReached => None, - } -} - -fn lifecycle_status_for_worker_error(err: &WorkerError) -> memory::audit::WorkerLifecycleStatus { - if matches!(err, WorkerError::Engine(EngineError::Cancelled)) { - memory::audit::WorkerLifecycleStatus::Cancelled - } else { - memory::audit::WorkerLifecycleStatus::Failed - } -} - -fn usage_audit_from_event( - event: &agen::llm_client::event::UsageEvent, -) -> memory::audit::UsageAudit { - memory::audit::UsageAudit { - input_tokens: event.input_tokens, - output_tokens: event.output_tokens, - total_tokens: event.total_tokens, - cache_read_input_tokens: event.cache_read_input_tokens, - cache_creation_input_tokens: event.cache_creation_input_tokens, - } -} - -fn model_audit_from_manifest(model: &manifest::ModelManifest) -> memory::audit::ModelAudit { - memory::audit::ModelAudit { - ref_: model.ref_.clone(), - scheme: model.scheme.map(|scheme| format!("{scheme:?}")), - model_id: model.model_id.clone(), - } -} - -fn emit_memory_worker_event( - working_event_tx: Option<&broadcast::Sender>, - run_id: uuid::Uuid, - worker: memory::audit::AuditWorker, - status: memory::audit::WorkerLifecycleStatus, - trigger: memory::audit::AuditTrigger, - reason: &str, -) { - let Some(working_event_tx) = working_event_tx else { - return; - }; - let message = format!("memory {} {}: {reason}", worker.label(), status.label()); - let _ = working_event_tx.send(Event::MemoryWorker(protocol::MemoryWorkerEvent { - worker: worker.label().to_string(), - status: status.label().to_string(), - run_id: run_id.to_string(), - trigger: trigger.label().to_string(), - reason: reason.to_string(), - message, - timestamp_ms: segment_log::now_millis() as i64, - })); -} - -#[derive(Debug, Clone)] -struct WorkerAuditBase { - run_id: uuid::Uuid, - worker: memory::audit::AuditWorker, - trigger: memory::audit::AuditTrigger, - model: Option, - memory_settings: Option, -} - -impl WorkerAuditBase { - fn new( - worker: memory::audit::AuditWorker, - trigger: memory::audit::AuditTrigger, - model: Option, - ) -> Self { - Self { - run_id: uuid::Uuid::now_v7(), - worker, - trigger, - model, - memory_settings: None, - } - } - - fn with_memory_settings(mut self, memory_config: &manifest::MemoryConfig) -> Self { - self.memory_settings = - memory_config - .workspace_settings() - .map(|snapshot| memory::audit::MemorySettingsAudit { - workspace_id: snapshot.workspace_id, - settings_revision: snapshot.settings_revision, - language: snapshot.language, - }); - self - } - - async fn emit( - &self, - workspace_client: &dyn WorkspaceClient, - working_event_tx: Option<&broadcast::Sender>, - status: memory::audit::WorkerLifecycleStatus, - reason: impl Into, - usage: Option, - extract: Option, - consolidation: Option, - ) { - let reason = reason.into(); - let payload = memory::audit::WorkerLifecycleAudit { - run_id: self.run_id, - worker: self.worker, - status, - trigger: self.trigger, - reason: reason.clone(), - memory_settings: self.memory_settings.clone(), - model: self.model.clone(), - usage, - extract, - consolidation, - }; - let _ = workspace_client - .execute_memory_backend_operation(memory::backend::MemoryBackendOperation::AppendAudit( - memory::backend::MemoryAppendAuditOperation { - event: memory::audit::AuditEvent::new( - memory::audit::AuditPayload::WorkerLifecycle(payload), - ), - }, - )) - .await; - if should_emit_memory_worker_event(self.worker, status, &reason) { - emit_memory_worker_event( - working_event_tx, - self.run_id, - self.worker, - status, - self.trigger, - &reason, - ); - } - } -} - -fn should_emit_memory_worker_event( - worker: memory::audit::AuditWorker, - status: memory::audit::WorkerLifecycleStatus, - reason: &str, -) -> bool { - if worker == memory::audit::AuditWorker::MemoryConsolidation - && status == memory::audit::WorkerLifecycleStatus::Skipped - { - return !is_idle_consolidation_skip_reason(reason); - } - true -} - -fn is_idle_consolidation_skip_reason(reason: &str) -> bool { - reason == "no_staging_entries" - || reason == "consolidation_threshold_disabled" - || reason.starts_with("threshold_not_reached") -} - -fn memory_language(cfg: &manifest::MemoryConfig) -> Result { - cfg.workspace_settings() +fn memory_language(config: &manifest::MemoryConfig) -> Result { + config + .workspace_settings() .map(|snapshot| snapshot.language) .ok_or_else(|| { WorkerError::InvalidState( - "Memory operation requires a bound Workspace Memory settings snapshot".to_string(), + "Memory is enabled without a bound Workspace Memory settings snapshot".to_string(), ) }) } @@ -5358,15 +4546,6 @@ fn worker_language(cfg: &manifest::EngineManifest) -> &str { } } -/// Outcome of a single extract iteration. Internal to -/// `try_post_run_extract` / `run_extract_once`. -enum ExtractDecision { - /// Threshold not reached, or no items to extract. - Skipped, - /// Extract ran and pointer advanced. Caller re-evaluates threshold. - Completed, -} - impl Worker, St> where St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, @@ -5480,7 +4659,6 @@ where flow_runtime_state: Arc::new(Mutex::new(None)), flow_feature_enabled: false, scope, - delegation_scope: common.delegation_scope, hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), @@ -5507,10 +4685,6 @@ where runtime_ticket_role: None, prompts: common.prompts, inject_resident_summary: true, - extract_in_flight: Arc::new(AtomicBool::new(false)), - consolidation_in_flight: Arc::new(AtomicBool::new(false)), - extract_pointer: Arc::new(Mutex::new(None)), - memory_task: None, user_segments: Vec::new(), sink: SegmentLogSink::new(), history_persistence_wired: false, @@ -5566,7 +4740,6 @@ where flow_runtime_state: Arc::new(Mutex::new(None)), flow_feature_enabled: false, scope, - delegation_scope: common.delegation_scope, hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), @@ -5593,10 +4766,6 @@ where runtime_ticket_role: None, prompts: common.prompts, inject_resident_summary: true, - extract_in_flight: Arc::new(AtomicBool::new(false)), - consolidation_in_flight: Arc::new(AtomicBool::new(false)), - extract_pointer: Arc::new(Mutex::new(None)), - memory_task: None, user_segments: Vec::new(), sink: SegmentLogSink::new(), history_persistence_wired: false, @@ -5687,7 +4856,6 @@ where flow_runtime_state: Arc::new(Mutex::new(None)), flow_feature_enabled: false, scope, - delegation_scope: common.delegation_scope, hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), @@ -5714,10 +4882,6 @@ where runtime_ticket_role: None, prompts: common.prompts, inject_resident_summary: true, - extract_in_flight: Arc::new(AtomicBool::new(false)), - consolidation_in_flight: Arc::new(AtomicBool::new(false)), - extract_pointer: Arc::new(Mutex::new(None)), - memory_task: None, user_segments: Vec::new(), sink: SegmentLogSink::new(), history_persistence_wired: false, @@ -6041,7 +5205,6 @@ where worker.set_cache_anchor(Some(0)); } - let extract_pointer = memory::extract::fold_pointer(&state.extensions); let task_feature = TaskFeature::from_history(&state.history); let worker_metadata_writer = Some(worker_metadata_writer_for_store(&store)); let scope = SharedScope::new(common.scope); @@ -6063,7 +5226,6 @@ where )?)), flow_feature_enabled: false, scope, - delegation_scope: common.delegation_scope, hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), @@ -6092,10 +5254,6 @@ where runtime_ticket_role: None, prompts: common.prompts, inject_resident_summary: true, - extract_in_flight: Arc::new(AtomicBool::new(false)), - consolidation_in_flight: Arc::new(AtomicBool::new(false)), - extract_pointer: Arc::new(Mutex::new(extract_pointer)), - memory_task: None, user_segments: state.user_segments, // Seed the mirror with the entries we just replayed so a // late-attaching client sees the full prefix without an @@ -6891,7 +6049,6 @@ struct WorkerCommon { filesystem_authority: WorkerFilesystemAuthority, workspace_context: WorkerWorkspaceContext, scope: Scope, - delegation_scope: DelegationScope, client: Box, prompts: Arc>, system_prompt_template: Option, @@ -7053,8 +6210,7 @@ fn prepare_worker_common_from_scope( }); } } - let delegation_scope = - DelegationScope::from_config(&manifest.delegation_scope).map_err(WorkerError::Scope)?; + DelegationScope::from_config(&manifest.delegation_scope).map_err(WorkerError::Scope)?; let client = match model_client { Some(client) => client, @@ -7074,7 +6230,6 @@ fn prepare_worker_common_from_scope( filesystem_authority, workspace_context, scope, - delegation_scope, client, prompts, system_prompt_template, @@ -7542,45 +6697,6 @@ permission = "read" } } -#[cfg(test)] -mod memory_worker_event_tests { - use super::*; - - #[test] - fn suppresses_idle_consolidation_skip_worker_events() { - assert!(!should_emit_memory_worker_event( - memory::audit::AuditWorker::MemoryConsolidation, - memory::audit::WorkerLifecycleStatus::Skipped, - "no_staging_entries", - )); - assert!(!should_emit_memory_worker_event( - memory::audit::AuditWorker::MemoryConsolidation, - memory::audit::WorkerLifecycleStatus::Skipped, - "threshold_not_reached files=1 bytes=64 min_files=2 min_bytes=1048576", - )); - assert!(!should_emit_memory_worker_event( - memory::audit::AuditWorker::MemoryConsolidation, - memory::audit::WorkerLifecycleStatus::Skipped, - "consolidation_threshold_disabled", - )); - assert!(should_emit_memory_worker_event( - memory::audit::AuditWorker::MemoryConsolidation, - memory::audit::WorkerLifecycleStatus::Skipped, - "no_valid_staging_entries invalid=1", - )); - assert!(should_emit_memory_worker_event( - memory::audit::AuditWorker::MemoryConsolidation, - memory::audit::WorkerLifecycleStatus::Completed, - "completed", - )); - assert!(should_emit_memory_worker_event( - memory::audit::AuditWorker::MemoryExtract, - memory::audit::WorkerLifecycleStatus::Skipped, - "threshold_not_reached files=1", - )); - } -} - #[cfg(test)] mod build_summary_prompt_tests { use super::*; @@ -7711,77 +6827,11 @@ mod build_summary_prompt_tests { assert!(prompt.contains("[1 Assistant] done")); } - #[derive(Clone)] - struct CancelBeforeAiExtractClient { - cancel_tx: Arc>>>, - } - - #[async_trait] - impl LlmClient for CancelBeforeAiExtractClient { - async fn stream( - &self, - _request: agen::llm_client::Request, - ) -> Result< - std::pin::Pin< - Box< - dyn futures::Stream< - Item = Result< - agen::llm_client::event::Event, - agen::llm_client::ClientError, - >, - > + Send, - >, - >, - agen::llm_client::ClientError, - > { - let tx = self - .cancel_tx - .lock() - .expect("cancel sender lock") - .clone() - .expect("extract caller must install the Internal Worker cancel sender"); - tx.send(()).await.expect("cancel Internal Worker"); - Ok(Box::pin(futures::stream::pending())) - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } - } - #[derive(Debug, Default)] struct RecordingAuditWorkspaceClient { requests: Mutex>, } - impl RecordingAuditWorkspaceClient { - fn lifecycle_audits(&self) -> Vec { - self.requests - .lock() - .expect("recorded workspace requests lock") - .iter() - .filter_map(|request| { - let operation: memory::backend::MemoryBackendOperation = serde_json::from_str( - request - .body - .as_deref() - .expect("memory backend operation body"), - ) - .expect("memory backend operation"); - match operation { - memory::backend::MemoryBackendOperation::AppendAudit(operation) => { - match operation.event.payload { - memory::audit::AuditPayload::WorkerLifecycle(audit) => Some(audit), - _ => None, - } - } - _ => None, - } - }) - .collect() - } - } - impl WorkspaceClient for RecordingAuditWorkspaceClient { fn workspace_id(&self) -> Option<&str> { Some("workspace-test") @@ -9439,154 +8489,65 @@ mod build_summary_prompt_tests { } #[tokio::test] - async fn cancelled_internal_extract_does_not_commit_pointer_or_completed_audit() { + async fn feature_session_extension_is_fenced_by_exact_committed_location() { let dir = tempfile::tempdir().unwrap(); + let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap(); let cwd = dir.path().join("workspace"); std::fs::create_dir_all(&cwd).unwrap(); - let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap(); - let cancel_tx = Arc::new(Mutex::new(None)); - let client = CancelBeforeAiExtractClient { - cancel_tx: cancel_tx.clone(), - }; - let audit_client = Arc::new(RecordingAuditWorkspaceClient::default()); - let mut manifest = minimal_manifest(); - manifest.memory = Some(manifest::MemoryConfig { - extract_threshold: Some(1), - workspace_id: Some("workspace-test".to_string()), - settings_revision: Some(1), - language: Some("English".to_string()), - ..Default::default() - }); - let memory_config = manifest.memory.clone().unwrap(); let mut worker = Worker::new( - manifest, - Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(client), + minimal_manifest(), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), store, - WorkerWorkspaceContext::with_client( - Some(WorkspaceId::new("workspace-test").unwrap()), - audit_client.clone(), - ), + WorkerWorkspaceContext::no_workspace(), WorkerFilesystemAuthority::local(cwd.clone(), cwd.clone()), Scope::writable(&cwd).unwrap(), ) .await .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.", - ); - worker.set_history_for_test(vec![evidence.clone()]); + + let capture_handle = worker.committed_session_capture_handle(); + let extension_handle = worker.session_extension_handle(); + let stale = capture_handle.capture().unwrap(); worker - .commit_entry(LogEntry::AnnotatedUserInput { + .commit_entry(LogEntry::TurnEnd { ts: segment_log::now_millis(), - extensions: vec![], - history: vec![crate::session_history::test_logged_history_entry( - evidence.clone(), - )], - segments: vec![text_segment( - "The cancellation regression must leave this evidence available for retry.", - )], + turn_count: 1, }) .unwrap(); - worker - .usage_history - .lock() - .expect("usage history lock") - .push(UsageRecord { - history_len: 1, - input_total_tokens: 100, - cache_read_tokens: 0, - cache_write_tokens: 0, - output_tokens: 0, - }); + assert!( + !extension_handle + .append_if_current( + &stale.location(), + "test.feature", + serde_json::json!({"revision": "stale"}), + ) + .unwrap() + ); - let entries_before = worker + let current = capture_handle.capture().unwrap(); + assert!( + extension_handle + .append_if_current( + ¤t.location(), + "test.feature", + serde_json::json!({"revision": "current"}), + ) + .unwrap() + ); + let entries = worker .store .read_all(worker.session_id(), worker.segment_id()) .unwrap(); - assert!( - worker - .extract_pointer - .lock() - .expect("extract pointer lock") - .is_none() - ); - - let cancel_tx_for_extract = cancel_tx.clone(); - let error = match worker - .run_extract_once_with_cancel_observer( - &memory_config, - 1, - Some(Box::new(move |cancel_sender| { - *cancel_tx_for_extract - .lock() - .expect("cancel sender slot lock") = Some(cancel_sender); - })), - ) - .await - { - Err(error) => error, - Ok(_) => panic!("pre-AI cancellation must not complete extraction"), - }; - - assert!(matches!(error, WorkerError::Engine(EngineError::Cancelled))); - assert!( - worker - .extract_pointer - .lock() - .expect("extract pointer lock") - .is_none() - ); - assert_eq!(worker.history(), &[evidence]); - - let entries_after = worker - .store - .read_all(worker.session_id(), worker.segment_id()) - .unwrap(); - assert_eq!(entries_after.len(), entries_before.len()); - assert!(!entries_after.iter().any(|entry| matches!( - entry, - LogEntry::Extension { domain, .. } if domain == memory::extract::EXTRACT_DOMAIN - ))); - - let audits = audit_client.lifecycle_audits(); - assert_eq!(audits.len(), 2); - assert_eq!(audits[0].run_id, audits[1].run_id); - assert_eq!(audits[0].worker, memory::audit::AuditWorker::MemoryExtract); - assert!(audits.iter().all(|audit| { - audit.memory_settings - == Some(memory::audit::MemorySettingsAudit { - workspace_id: "workspace-test".to_string(), - settings_revision: 1, - language: "English".to_string(), - }) - })); assert_eq!( - audits.iter().map(|audit| audit.status).collect::>(), - vec![ - memory::audit::WorkerLifecycleStatus::Started, - memory::audit::WorkerLifecycleStatus::Cancelled, - ] - ); - assert!( - !audits + entries .iter() - .any(|audit| { audit.status == memory::audit::WorkerLifecycleStatus::Completed }) + .filter(|entry| matches!(entry, LogEntry::Extension { domain, .. } if domain == "test.feature")) + .count(), + 1 ); } - #[test] - fn successful_internal_extract_lifecycles_enter_the_commit_path() { - for lifecycle in [ - WorkerRunResult::Finished, - WorkerRunResult::Paused, - WorkerRunResult::LimitReached, - ] { - assert!(extract_internal_worker_lifecycle_error(&lifecycle).is_none()); - } - } - fn minimal_manifest() -> WorkerManifest { let toml_str = r#" [worker] diff --git a/crates/worker/tests/compact_events_test.rs b/crates/worker/tests/compact_events_test.rs index afef80bb..47fbd351 100644 --- a/crates/worker/tests/compact_events_test.rs +++ b/crates/worker/tests/compact_events_test.rs @@ -578,138 +578,6 @@ async fn mid_turn_compact_success_broadcasts_start_and_done() { assert_eq!(new_id_in_event, Some(worker.segment_id())); } -/// Regression: `Worker::compact()` must reset the in-memory -/// `extract_pointer` so extract keeps firing on the new compacted -/// session. -/// -/// Without the reset, the pointer's `processed_through_history_len` -/// holds the old (typically large) item count, while the new compacted -/// session starts with a much shorter history (`[summary, ...]`). -/// `cumulative_input_tokens_since` would then filter every new -/// usage record out (their `history_len` is below the stale pointer) -/// and extract would never re-fire for the rest of the process. -const EXTRACT_PLUS_COMPACT_MANIFEST: &str = r#" -[worker] -name = "test-worker" -pwd = "./" - -[model] -scheme = "anthropic" -model_id = "test-model" - -[engine] -max_tokens = 100 - -[memory] -workspace_id = "test-workspace" -settings_revision = 1 -language = "English" -extract_threshold = 1 - -[compaction] -compact_threshold = 1 -compact_retained_tokens = 0 - -[[scope.allow]] -target = "./" -permission = "write" -"#; - -fn finish_memory_extraction_tool_use_events(call_id: &str) -> Vec { - let input = serde_json::json!({ - "staged_count": 0, - "no_candidates_reason": "test run has no durable candidates" - }) - .to_string(); - vec![ - LlmEvent::tool_use_start(0, call_id, "FinishMemoryExtraction"), - LlmEvent::tool_input_delta(0, input), - LlmEvent::tool_use_stop(0), - LlmEvent::Status(StatusEvent { - status: ResponseStatus::Completed, - }), - ] -} - -#[tokio::test] -async fn compact_resets_extract_pointer_so_extract_can_fire_again() { - // Mock LLM responses, in call order: - // [0] first run with usage(1000) so extract threshold (=1) fires. - // [1] extract worker invokes FinishMemoryExtraction with empty output. - // [2] extract worker closes after the tool result. - // [3] compact worker invokes write_summary. - // [4] compact worker closes after the tool result. - let client = MockClient::new(vec![ - text_events_with_usage("hi", 1000), - finish_memory_extraction_tool_use_events("ec1"), - single_text_events("done"), - write_summary_tool_use_events("sc1", "summary"), - single_text_events("done"), - ]); - let mut worker = make_worker_with_manifest(EXTRACT_PLUS_COMPACT_MANIFEST, client).await; - - worker.run_text("first").await.unwrap(); - - // extract fires; pointer becomes Some. - worker.try_post_run_extract().await.unwrap(); - assert!( - worker.extract_pointer().is_some(), - "extract_pointer should be Some after a successful extract" - ); - - // Compact runs. Without the fix the in-memory pointer would still - // reference the old Segment's history_len. - worker.try_pre_run_compact().await; - assert!( - worker.extract_pointer().is_none(), - "extract_pointer must be reset to None after compact (matches cold-restore on the new Segment)" - ); -} - -/// `extract_threshold = 0` is treated as "disabled" — without this, a -/// raw `>=` comparison against `tokens_since` would fire extract on -/// every post-run regardless of activity. Mirrors the consolidation -/// zero-threshold convention so users have a single way to opt out -/// without removing the `[memory]` section. -const EXTRACT_THRESHOLD_ZERO_MANIFEST: &str = r#" -[worker] -name = "test-worker" -pwd = "./" - -[model] -scheme = "anthropic" -model_id = "test-model" - -[engine] -max_tokens = 100 - -[memory] -extract_threshold = 0 - -[[scope.allow]] -target = "./" -permission = "write" -"#; - -#[tokio::test] -async fn extract_threshold_zero_is_disabled() { - // Mock provides exactly one response — the first run. If extract - // were treated as "fire on any change" because of `tokens_since >= 0`, - // it would call into the extract worker and exhaust the mock. - let client = MockClient::new(vec![text_events_with_usage("hi", 1000)]); - let mut worker = make_worker_with_manifest(EXTRACT_THRESHOLD_ZERO_MANIFEST, client).await; - - worker.run_text("first").await.unwrap(); - worker - .try_post_run_extract() - .await - .expect("extract_threshold=0 must skip silently, not fail"); - assert!( - worker.extract_pointer().is_none(), - "no extract should have run — pointer must remain None" - ); -} - #[tokio::test] async fn pre_run_compact_failure_broadcasts_start_and_failed() { // Only the first run has a response. Compaction will run the @@ -746,112 +614,6 @@ async fn pre_run_compact_failure_broadcasts_start_and_failed() { ); } -// --------------------------------------------------------------------------- -// Detached post-run memory jobs (`spawn_post_run_memory_jobs` / -// `wait_for_memory_jobs`). Covers the detach round-trip and the structural -// invariant that the cloned memory-task Worker shares `SegmentState` with the -// source Worker, so that `save_extension` from the background extract does not -// leave the next turn's `save_user_input` looking at a stale session pointer. - -const EXTRACT_NO_COMPACT_MANIFEST: &str = r#" -[worker] -name = "test-worker" -pwd = "./" - -[model] -scheme = "anthropic" -model_id = "test-model" - -[engine] -max_tokens = 100 - -[memory] -workspace_id = "test-workspace" -settings_revision = 1 -language = "English" -extract_threshold = 1 - -[[scope.allow]] -target = "./" -permission = "write" -"#; - -#[tokio::test] -async fn extract_large_unprocessed_range_does_not_abort_on_input_occupancy() { - let client = MockClient::new(vec![ - text_events_with_usage("recorded", 1000), - finish_memory_extraction_tool_use_events("ec-large"), - single_text_events("done"), - ]); - let mut worker = make_worker_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await; - - let large_request = format!("remember this large slice: {}", "x ".repeat(200_000)); - worker.run_text(&large_request).await.unwrap(); - - worker.try_post_run_extract().await.expect( - "large unprocessed extract ranges must reach the extract worker, not abort locally", - ); - assert!( - worker.extract_pointer().is_some(), - "successful extract should advance the pointer even when the input range is large" - ); -} - -#[tokio::test] -async fn spawn_and_wait_drives_extract_to_completion() { - let client = MockClient::new(vec![ - text_events_with_usage("hi", 1000), - finish_memory_extraction_tool_use_events("ec1"), - single_text_events("done"), - ]); - let mut worker = make_worker_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await; - - worker.run_text("first").await.unwrap(); - assert!( - worker.extract_pointer().is_none(), - "extract has not run yet — pointer must be None" - ); - - worker.spawn_post_run_memory_jobs(); - worker.wait_for_memory_jobs().await; - - assert!( - worker.extract_pointer().is_some(), - "spawn + wait must complete extract; pointer should be set" - ); -} - -#[tokio::test] -async fn detached_extract_does_not_fork_session_log() { - // Source worker and the cloned memory-task worker share `SegmentState` via - // `Arc<_>`. The detached extract advances the entry tally through - // `save_extension`; the next `run` must see that same tally so - // `ensure_head_or_fork` does not spawn a new session. - let client = MockClient::new(vec![ - text_events_with_usage("hi", 1000), - finish_memory_extraction_tool_use_events("ec1"), - single_text_events("done"), - text_events_with_usage("ok", 1000), - ]); - let mut worker = make_worker_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await; - - worker.run_text("first").await.unwrap(); - let session_before = worker.segment_id(); - - worker.spawn_post_run_memory_jobs(); - worker.wait_for_memory_jobs().await; - - worker.run_text("second").await.unwrap(); - let session_after = worker.segment_id(); - - assert_eq!( - session_before, session_after, - "detached extract's save_extension and the next turn's save_user_input \ - must share the entry tally through SegmentState — a fork here means the \ - clone carried its own counter" - ); -} - #[tokio::test] async fn controller_compact_method_emits_start_and_done() { let client = MockClient::new(vec![ From 27e5df106f86b40356e95ef6db76f6d9f0fe71fc Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 19:16:01 +0900 Subject: [PATCH 15/22] fix: complete memory lifecycle feature boundaries --- crates/worker/src/controller.rs | 10 +- .../src/feature/builtin/memory_lifecycle.rs | 304 +++++++++++++++--- crates/worker/src/feature/session.rs | 8 + crates/worker/src/worker.rs | 40 ++- 4 files changed, 312 insertions(+), 50 deletions(-) diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 43f28539..ddef4480 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -994,13 +994,7 @@ where let worker_enabled = feature_config.worker.enabled; let sub_worker_enabled = feature_config.sub_worker.enabled; let mut feature_registry = FeatureRegistryBuilder::new(); - if feature_config.memory.enabled { - let config = memory_config.clone().ok_or_else(|| { - std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "[feature.memory].enabled = true requires a [memory] configuration section", - ) - })?; + if let Some(config) = memory_config.clone() { let workspace_client = worker.workspace_client_handle(); if !workspace_client.is_available() || workspace_client.workspace_id().is_none() { return Err(std::io::Error::new( @@ -1009,7 +1003,7 @@ where )); } feature_registry.add_module( - crate::feature::builtin::memory_lifecycle::MemoryExtractionLifecycleFeature::new( + crate::feature::builtin::memory_lifecycle::MemoryLifecycleFeature::new( config, worker.committed_session_capture_handle(), worker.session_extension_handle(), diff --git a/crates/worker/src/feature/builtin/memory_lifecycle.rs b/crates/worker/src/feature/builtin/memory_lifecycle.rs index 30c14c91..4d8db2f4 100644 --- a/crates/worker/src/feature/builtin/memory_lifecycle.rs +++ b/crates/worker/src/feature/builtin/memory_lifecycle.rs @@ -20,7 +20,8 @@ use crate::feature::builtin::memory_staging_output::{ }; use crate::feature::builtin::session_explore::{SessionExploreFeature, SessionExploreState}; use crate::feature::session::{ - CommittedSessionCapture, CommittedSessionCaptureHandle, SessionExtensionHandle, + CommittedRunExit, CommittedSessionCapture, CommittedSessionCaptureHandle, + SessionExtensionHandle, }; use crate::feature::{ BackgroundTaskDeclaration, FeatureDescriptor, FeatureInstallContext, FeatureInstallError, @@ -37,7 +38,7 @@ use agen::token_counter::total_tokens_at; use manifest::WorkerManifest; use protocol::Event; -const TASK_NAME: &str = "memory-extraction"; +const TASK_NAME: &str = "memory-lifecycle"; const TASK_TIMEOUT: Duration = Duration::from_secs(300); /// Parent-Worker lifecycle Feature that observes committed runs and schedules @@ -45,12 +46,12 @@ const TASK_TIMEOUT: Duration = Duration::from_secs(300); /// Internal Worker, and staging disposition; Worker core owns only generic /// hook/task/session plumbing. #[derive(Clone)] -pub(crate) struct MemoryExtractionLifecycleFeature { - task: MemoryExtractionTask, +pub(crate) struct MemoryLifecycleFeature { + task: MemoryLifecycleTask, } #[derive(Clone)] -struct MemoryExtractionTask { +struct MemoryLifecycleTask { config: manifest::MemoryConfig, capture: CommittedSessionCaptureHandle, extensions: SessionExtensionHandle, @@ -62,7 +63,7 @@ struct MemoryExtractionTask { event_tx: Option>, } -impl MemoryExtractionLifecycleFeature { +impl MemoryLifecycleFeature { #[allow(clippy::too_many_arguments)] pub(crate) fn new( config: manifest::MemoryConfig, @@ -76,7 +77,7 @@ impl MemoryExtractionLifecycleFeature { event_tx: Option>, ) -> Self { Self { - task: MemoryExtractionTask { + task: MemoryLifecycleTask { config, capture, extensions, @@ -91,46 +92,133 @@ impl MemoryExtractionLifecycleFeature { } } -impl FeatureModule for MemoryExtractionLifecycleFeature { +impl FeatureModule for MemoryLifecycleFeature { fn descriptor(&self) -> FeatureDescriptor { - FeatureDescriptor::builtin("memory-extraction-lifecycle", "Memory Extraction Lifecycle") + FeatureDescriptor::builtin("memory-lifecycle", "Memory Lifecycle") .with_description( - "Observes terminal committed runs and schedules bounded Memory extraction.", + "Observes terminal committed runs and schedules bounded Memory extraction and Backend consolidation requests.", ) .with_background_task(BackgroundTaskDeclaration::worker_managed( TASK_NAME, - "Extract provenance-preserving Memory candidates after committed runs.", + "Extract provenance-preserving Memory candidates and request Backend consolidation after committed runs.", )) } fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> { context .background_tasks() - .register(memory_extraction_task_spec(), self.task.clone()) + .register(memory_lifecycle_task_spec(), self.task.clone()) } } -fn memory_extraction_task_spec() -> BackgroundTaskSpec { +fn memory_lifecycle_task_spec() -> BackgroundTaskSpec { let declaration = BackgroundTaskDeclaration::worker_managed( TASK_NAME, - "Extract provenance-preserving Memory candidates after committed runs.", + "Extract provenance-preserving Memory candidates and request Backend consolidation after committed runs.", ); let mut spec = BackgroundTaskSpec::single_flight(declaration, TASK_TIMEOUT); spec.trigger = BackgroundTaskTrigger::RunCommitted; spec } -#[async_trait] -impl FeatureBackgroundTask for MemoryExtractionTask { - async fn run( +impl MemoryLifecycleTask { + async fn run_extraction( &self, context: BackgroundTaskContext, cancellation: BackgroundTaskCancellation, ) -> Result<(), HookError> { context.generation_fence.ensure_current()?; - let capture = self.capture.capture().map_err(hook_internal)?; - let pointer = extract_pointer(&capture)?; - if !extraction_threshold_reached(&capture, pointer.as_ref(), &self.config) { + let audit = WorkerAuditBase::new( + memory::audit::AuditWorker::MemoryExtract, + memory::audit::AuditTrigger::TokenThreshold, + self.config + .extract_model + .as_ref() + .or(Some(&self.manifest.model)) + .map(model_audit_from_manifest), + ) + .with_memory_settings(&self.config); + let capture = match self.capture.capture() { + Ok(capture) => capture, + Err(error) => { + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Failed, + format!("committed_session_capture_failed: {error}"), + None, + None, + None, + ) + .await; + return Ok(()); + } + }; + if !extraction_run_eligible(capture.run_exit) { + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Skipped, + format!("parent_run_not_finished: {:?}", capture.run_exit), + None, + None, + None, + ) + .await; + return Ok(()); + } + let pointer = match extract_pointer(&capture) { + Ok(pointer) => pointer, + Err(error) => { + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Failed, + format!("extract_pointer_invalid: {error}"), + None, + None, + None, + ) + .await; + return Ok(()); + } + }; + let Some(threshold) = self + .config + .extract_threshold + .filter(|threshold| *threshold > 0) + else { + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Skipped, + "token_threshold_disabled", + None, + None, + None, + ) + .await; + return Ok(()); + }; + let tokens_since = tokens_since_pointer(&capture, pointer.as_ref()); + if tokens_since < threshold { + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Skipped, + format!( + "token_threshold_not_reached tokens_since={tokens_since} threshold={threshold}" + ), + None, + None, + None, + ) + .await; return Ok(()); } @@ -141,6 +229,22 @@ impl FeatureBackgroundTask for MemoryExtractionTask { .min(capture.history.len()); let history_end = capture.history.len(); if history_start >= history_end || capture.entry_count == 0 { + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Skipped, + "no_new_committed_session_entries", + None, + Some(memory::audit::ExtractAudit { + session_id: Some(capture.session_id.clone()), + segment_id: Some(capture.segment_id.clone()), + history_range: Some([history_start as u64, history_end as u64]), + ..Default::default() + }), + None, + ) + .await; return Ok(()); } let view = SessionCapture::from_history_entries( @@ -155,16 +259,6 @@ impl FeatureBackgroundTask for MemoryExtractionTask { segment_id: capture.segment_id.clone(), range: [start_entry as u64, (capture.entry_count - 1) as u64], }; - let audit = WorkerAuditBase::new( - memory::audit::AuditWorker::MemoryExtract, - memory::audit::AuditTrigger::TokenThreshold, - self.config - .extract_model - .as_ref() - .or(Some(&self.manifest.model)) - .map(model_audit_from_manifest), - ) - .with_memory_settings(&self.config); let extract_audit_base = memory::audit::ExtractAudit { session_id: Some(capture.session_id.clone()), segment_id: Some(capture.segment_id.clone()), @@ -376,7 +470,86 @@ impl FeatureBackgroundTask for MemoryExtractionTask { } } -impl MemoryExtractionTask { +#[async_trait] +impl FeatureBackgroundTask for MemoryLifecycleTask { + async fn run( + &self, + context: BackgroundTaskContext, + cancellation: BackgroundTaskCancellation, + ) -> Result<(), HookError> { + let extraction = self + .run_extraction(context.clone(), cancellation.clone()) + .await; + if !cancellation.is_cancelled() { + context.generation_fence.ensure_current()?; + self.request_consolidation().await; + } + extraction + } +} + +impl MemoryLifecycleTask { + async fn request_consolidation(&self) { + let audit = WorkerAuditBase::new( + memory::audit::AuditWorker::MemoryConsolidation, + memory::audit::AuditTrigger::StagingBacklog, + self.config + .consolidation_model + .as_ref() + .or(Some(&self.manifest.model)) + .map(model_audit_from_manifest), + ) + .with_memory_settings(&self.config); + let Some((threshold_files, threshold_bytes)) = consolidation_thresholds(&self.config) + else { + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Skipped, + "consolidation_threshold_disabled", + None, + None, + None, + ) + .await; + return; + }; + match self + .workspace_client + .request_memory_staging_consolidation( + memory::backend::MemoryConsolidateStagingOperation { + force: false, + threshold_files, + threshold_bytes, + }, + ) + .await + { + Ok(output) => { + tracing::debug!( + status = output.status.as_str(), + summary = output.summary.as_str(), + "requested Backend Memory staging consolidation" + ); + } + Err(error) => { + tracing::warn!(%error, "request Backend Memory staging consolidation failed"); + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Skipped, + "consolidation_backend_operation_failed", + None, + None, + None, + ) + .await; + } + } + } + async fn record_preparation_failure( &self, audit: &WorkerAuditBase, @@ -473,14 +646,30 @@ fn extract_pointer( Ok(pointer) } -fn extraction_threshold_reached( +fn consolidation_thresholds( + config: &manifest::MemoryConfig, +) -> Option<(Option, Option)> { + let threshold_files = config + .consolidation_threshold_files + .filter(|threshold| *threshold > 0); + let threshold_bytes = config + .consolidation_threshold_bytes + .filter(|threshold| *threshold > 0); + if threshold_files.is_none() && threshold_bytes.is_none() { + None + } else { + Some((threshold_files, threshold_bytes)) + } +} + +fn extraction_run_eligible(exit: CommittedRunExit) -> bool { + exit == CommittedRunExit::Finished +} + +fn tokens_since_pointer( capture: &CommittedSessionCapture, pointer: Option<&memory::ExtractPointerPayload>, - config: &manifest::MemoryConfig, -) -> bool { - if capture.history.is_empty() { - return false; - } +) -> u64 { let history_pointer = pointer .map(|pointer| pointer.processed_through_history_len) .unwrap_or(0) @@ -492,10 +681,22 @@ fn extraction_threshold_reached( .collect::>(); let current = total_tokens_at(&items, &capture.usage_history, capture.history.len()).tokens; let baseline = total_tokens_at(&items, &capture.usage_history, history_pointer).tokens; + current.saturating_sub(baseline) +} + +#[cfg(test)] +fn extraction_threshold_reached( + capture: &CommittedSessionCapture, + pointer: Option<&memory::ExtractPointerPayload>, + config: &manifest::MemoryConfig, +) -> bool { + if capture.history.is_empty() { + return false; + } let Some(threshold) = config.extract_threshold.filter(|threshold| *threshold > 0) else { return false; }; - current.saturating_sub(baseline) >= threshold + tokens_since_pointer(capture, pointer) >= threshold } #[derive(Clone)] @@ -620,6 +821,7 @@ mod tests { segment_id: "segment-1".to_string(), session_revision: history_len.try_into().unwrap(), entry_count: history_len, + run_exit: CommittedRunExit::Finished, history: (0..history_len) .map(|index| HistoryEntry { item: Item::user_message(format!("message-{index}")), @@ -637,6 +839,24 @@ mod tests { } } + #[test] + fn consolidation_thresholds_enable_backend_request_on_either_limit() { + let mut config = manifest::MemoryConfig::default(); + assert_eq!(consolidation_thresholds(&config), None); + config.consolidation_threshold_files = Some(3); + assert_eq!(consolidation_thresholds(&config), Some((Some(3), None))); + config.consolidation_threshold_files = None; + config.consolidation_threshold_bytes = Some(4096); + assert_eq!(consolidation_thresholds(&config), Some((None, Some(4096)))); + } + + #[test] + fn interrupted_parent_run_is_not_extraction_eligible() { + assert!(extraction_run_eligible(CommittedRunExit::Finished)); + assert!(!extraction_run_eligible(CommittedRunExit::NonFinal)); + assert!(!extraction_run_eligible(CommittedRunExit::Interrupted)); + } + #[test] fn normal_and_empty_extraction_require_explicit_finish() { let result = internal_result(WorkerRunResult::Finished); @@ -672,7 +892,7 @@ mod tests { #[test] fn task_scope_cancels_and_joins_before_rewrite_and_shutdown() { - let spec = memory_extraction_task_spec(); + let spec = memory_lifecycle_task_spec(); assert_eq!(spec.trigger, BackgroundTaskTrigger::RunCommitted); assert_eq!(spec.max_concurrency, 1); assert_eq!(spec.rewrite, BackgroundTaskRewritePolicy::CancelAndWait); @@ -759,8 +979,10 @@ mod tests { ); } let controller_source = include_str!("../../controller.rs"); - assert!(controller_source.contains("if feature_config.memory.enabled")); - assert!(controller_source.contains("MemoryExtractionLifecycleFeature::new")); + assert!(controller_source.contains("if let Some(config) = memory_config.clone()")); + assert!(controller_source.contains("MemoryLifecycleFeature::new")); + let lifecycle_source = include_str!("memory_lifecycle.rs"); + assert!(lifecycle_source.contains("request_memory_staging_consolidation")); let internal_worker_source = include_str!("../../internal_worker.rs"); assert!(!internal_worker_source.contains("manifest.memory = None")); } diff --git a/crates/worker/src/feature/session.rs b/crates/worker/src/feature/session.rs index 97345b10..0729f0af 100644 --- a/crates/worker/src/feature/session.rs +++ b/crates/worker/src/feature/session.rs @@ -5,6 +5,13 @@ use serde_json::Value; use crate::session_history::SessionHistoryMetadata; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum CommittedRunExit { + Finished, + NonFinal, + Interrupted, +} + /// Immutable projection of one durably committed session-log location. /// /// Feature code receives this value only after the host has committed the @@ -18,6 +25,7 @@ pub(crate) struct CommittedSessionCapture { /// Monotonic committed-log revision for the captured Segment. pub(crate) session_revision: u64, pub(crate) entry_count: usize, + pub(crate) run_exit: CommittedRunExit, pub(crate) history: Vec>, pub(crate) usage_history: Vec, pub(crate) extensions: Vec<(String, Value)>, diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index c5fcb23b..4b74f647 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -44,7 +44,7 @@ use crate::feature::background::{BackgroundTaskRewriteGuard, FeatureBackgroundTa use crate::feature::builtin::memory::WorkspaceMemoryBackendError; use crate::feature::builtin::{TaskFeature, WorkerObservationProvider}; use crate::feature::session::{ - CommittedSessionCapture, CommittedSessionCaptureHandle, FeatureSessionError, + CommittedRunExit, CommittedSessionCapture, CommittedSessionCaptureHandle, FeatureSessionError, SessionExtensionHandle, }; use crate::feature::{ @@ -2108,6 +2108,19 @@ impl Worker { let entries = store .read_all(location.session_id, location.segment_id) .map_err(|error| FeatureSessionError::Capture(error.to_string()))?; + let run_exit = entries + .iter() + .rev() + .find_map(|entry| match entry { + LogEntry::RunCompleted { + result: EngineResult::Finished, + .. + } => Some(CommittedRunExit::Finished), + LogEntry::RunCompleted { .. } => Some(CommittedRunExit::NonFinal), + LogEntry::RunErrored { .. } => Some(CommittedRunExit::Interrupted), + _ => None, + }) + .unwrap_or(CommittedRunExit::NonFinal); let restored = segment_log::collect_state(&entries); let history = restore_history_entries(location.session_id, location.segment_id, &entries) @@ -2117,6 +2130,7 @@ impl Worker { segment_id: location.segment_id.to_string(), session_revision: entries.len().try_into().unwrap_or(u64::MAX), entry_count: entries.len(), + run_exit, history, usage_history: restored.usage_history, extensions: restored.extensions.into_iter().collect(), @@ -8546,6 +8560,30 @@ mod build_summary_prompt_tests { .count(), 1 ); + + worker + .commit_entry(LogEntry::RunErrored { + ts: segment_log::now_millis(), + interrupted: true, + message: "cancelled".to_string(), + }) + .unwrap(); + assert_eq!( + capture_handle.capture().unwrap().run_exit, + CommittedRunExit::Interrupted + ); + worker + .commit_entry(LogEntry::RunCompleted { + ts: segment_log::now_millis(), + interrupted: false, + result: EngineResult::Finished, + active_run_turn_count: None, + }) + .unwrap(); + assert_eq!( + capture_handle.capture().unwrap().run_exit, + CommittedRunExit::Finished + ); } fn minimal_manifest() -> WorkerManifest { From 532d078720a094fcb81b9a5ce2a948ad956fe339 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 19:41:04 +0900 Subject: [PATCH 16/22] test: exercise memory lifecycle background behavior --- .../src/feature/builtin/memory_lifecycle.rs | 280 ++++++++++++++++++ 1 file changed, 280 insertions(+) diff --git a/crates/worker/src/feature/builtin/memory_lifecycle.rs b/crates/worker/src/feature/builtin/memory_lifecycle.rs index 4d8db2f4..59de7caf 100644 --- a/crates/worker/src/feature/builtin/memory_lifecycle.rs +++ b/crates/worker/src/feature/builtin/memory_lifecycle.rs @@ -795,12 +795,292 @@ fn model_audit_from_manifest(model: &manifest::ModelManifest) -> memory::audit:: #[cfg(test)] mod tests { + use std::pin::Pin; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex}; + + use agen::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent}; + use agen::llm_client::{ClientError, Request}; use agen::{HistoryEntry, Item, UsageRecord}; + use futures::Stream; use super::*; + use crate::feature::background::FeatureBackgroundTaskRegistryBuilder; use crate::feature::background::{BackgroundTaskRewritePolicy, BackgroundTaskShutdownPolicy}; + use crate::feature::session::CommittedSessionLocation; + use crate::hook::HookInvocationContext; use crate::session_history::SessionHistoryMetadata; + #[derive(Debug, Default)] + struct RecordingWorkspaceClient { + requests: Mutex>, + } + + impl WorkspaceClient for RecordingWorkspaceClient { + fn workspace_id(&self) -> Option<&str> { + Some("workspace-1") + } + + fn kind(&self) -> &str { + "memory-lifecycle-test" + } + + fn is_available(&self) -> bool { + true + } + + fn execute( + &self, + request: crate::worker::WorkspaceRequest, + ) -> Result { + self.requests.lock().unwrap().push(request); + Err(crate::worker::WorkspaceClientError::Unavailable( + "recording client".to_string(), + )) + } + } + + #[derive(Clone)] + struct ScriptClient { + responses: Arc>>, + calls: Arc, + } + + impl ScriptClient { + fn new(responses: Vec>) -> Self { + Self { + responses: Arc::new(responses), + calls: Arc::new(AtomicUsize::new(0)), + } + } + } + + #[async_trait] + impl LlmClient for ScriptClient { + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + async fn stream( + &self, + _request: Request, + ) -> Result> + Send>>, ClientError> + { + let index = self.calls.fetch_add(1, Ordering::SeqCst); + let events = self.responses.get(index).cloned().ok_or_else(|| { + ClientError::Config("memory lifecycle test client exhausted".to_string()) + })?; + Ok(Box::pin(futures::stream::iter(events.into_iter().map(Ok)))) + } + } + + fn finish_empty_events(call_id: &str) -> Vec { + vec![ + LlmEvent::tool_use_start(0, call_id, "FinishMemoryExtraction"), + LlmEvent::tool_input_delta( + 0, + serde_json::json!({ + "staged_count": 0, + "no_candidates_reason": "no durable candidates" + }) + .to_string(), + ), + LlmEvent::tool_use_stop(0), + LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] + } + + fn completed_events() -> Vec { + vec![ + LlmEvent::text_block_start(0), + LlmEvent::text_delta(0, "done"), + LlmEvent::text_block_stop(0, None), + LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] + } + + fn test_manifest() -> WorkerManifest { + WorkerManifest::from_toml( + r#" +[worker] +name = "memory-lifecycle-test" +scope = "main" + +[model] +scheme = "anthropic" +model_id = "test-model" + +[engine] + +[[scope.allow]] +target = "/memory-lifecycle-test" +permission = "write" +"#, + ) + .unwrap() + } + + fn test_config() -> manifest::MemoryConfig { + let mut config = manifest::MemoryConfig { + extract_threshold: Some(1), + ..Default::default() + }; + config.bind_workspace_settings(&manifest::WorkspaceMemorySettingsSnapshot { + workspace_id: "workspace-1".to_string(), + settings_revision: 1, + language: "English".to_string(), + }); + config + } + + fn test_task( + capture: CommittedSessionCapture, + client: ScriptClient, + extension_writes: Arc>>, + event_tx: broadcast::Sender, + workspace_client: Arc, + ) -> MemoryLifecycleTask { + let capture_handle = CommittedSessionCaptureHandle::new(move || Ok(capture.clone())); + let extensions = SessionExtensionHandle::new(move |location, domain, payload| { + extension_writes + .lock() + .unwrap() + .push((location.clone(), domain.to_string(), payload)); + Ok(true) + }); + MemoryLifecycleTask { + config: test_config(), + capture: capture_handle, + extensions, + workspace_client, + manifest: test_manifest(), + client: Box::new(client), + prompts: Arc::new(ArcSwap::from(PromptCatalog::builtins_only().unwrap())), + workspace_context: WorkerWorkspaceContext::no_workspace(), + event_tx: Some(event_tx), + } + } + + async fn run_background_task(task: MemoryLifecycleTask) { + let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); + builder + .register( + crate::feature::FeatureId::builtin("memory-lifecycle"), + memory_lifecycle_task_spec(), + task, + ) + .unwrap(); + let registry = builder.build(); + registry + .start_run_committed(HookInvocationContext { + worker_id: "worker-1".to_string(), + session_id: "session-1".to_string(), + session_revision: 1, + run_id: Some("run-1".to_string()), + ..Default::default() + }) + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if !registry.diagnostics().is_empty() { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("memory lifecycle background task should finish"); + registry.shutdown().await.unwrap(); + } + + #[tokio::test] + async fn run_committed_background_task_finishes_empty_extraction_and_commits_pointer() { + let client = ScriptClient::new(vec![finish_empty_events("finish-1"), completed_events()]); + let calls = Arc::clone(&client.calls); + let extension_writes = Arc::new(Mutex::new(Vec::new())); + let (event_tx, _) = broadcast::channel(16); + let workspace_client: Arc = + Arc::new(RecordingWorkspaceClient::default()); + run_background_task(test_task( + capture(2, 250), + client, + Arc::clone(&extension_writes), + event_tx, + workspace_client, + )) + .await; + + assert_eq!(calls.load(Ordering::SeqCst), 2); + let writes = extension_writes.lock().unwrap(); + assert_eq!(writes.len(), 1); + assert_eq!(writes[0].1, extract::EXTRACT_DOMAIN); + let pointer: memory::ExtractPointerPayload = + serde_json::from_value(writes[0].2.clone()).unwrap(); + assert_eq!(pointer.processed_through_history_len, 2); + assert_eq!(pointer.processed_through_entry, 1); + } + + #[tokio::test] + async fn interrupted_committed_run_skips_internal_worker_and_pointer_commit() { + let client = ScriptClient::new(Vec::new()); + let calls = Arc::clone(&client.calls); + let extension_writes = Arc::new(Mutex::new(Vec::new())); + let (event_tx, mut event_rx) = broadcast::channel(16); + let workspace_client: Arc = + Arc::new(RecordingWorkspaceClient::default()); + let mut interrupted = capture(2, 250); + interrupted.run_exit = CommittedRunExit::Interrupted; + run_background_task(test_task( + interrupted, + client, + Arc::clone(&extension_writes), + event_tx, + workspace_client, + )) + .await; + + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert!(extension_writes.lock().unwrap().is_empty()); + let events = std::iter::from_fn(|| event_rx.try_recv().ok()).collect::>(); + assert!(events.iter().any(|event| { + matches!(event, Event::MemoryWorker(event) if event.reason.contains("parent_run_not_finished")) + })); + } + + #[tokio::test] + async fn lifecycle_task_requests_backend_consolidation_from_configured_threshold() { + let client = ScriptClient::new(Vec::new()); + let extension_writes = Arc::new(Mutex::new(Vec::new())); + let (event_tx, _) = broadcast::channel(16); + let workspace_client = Arc::new(RecordingWorkspaceClient::default()); + let mut interrupted = capture(2, 250); + interrupted.run_exit = CommittedRunExit::Interrupted; + let mut task = test_task( + interrupted, + client, + extension_writes, + event_tx, + workspace_client.clone(), + ); + task.config.consolidation_threshold_files = Some(3); + run_background_task(task).await; + + let requests = workspace_client.requests.lock().unwrap(); + assert!( + requests.iter().any(|request| { + request.path.contains("memory") + && request.body.as_deref().is_some_and(|body| { + body.contains("\"threshold_files\":3") && body.contains("\"force\":false") + }) + }), + "recorded requests: {requests:?}" + ); + } + fn internal_result( lifecycle: WorkerRunResult, ) -> Result { From d1f5661881f766b95249db4be3b23ea78f9fffbe Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 20:11:29 +0900 Subject: [PATCH 17/22] test: cover memory lifecycle outcomes --- crates/worker/src/controller.rs | 68 ++++- .../src/feature/builtin/memory_lifecycle.rs | 244 +++++++++++++++++- 2 files changed, 301 insertions(+), 11 deletions(-) diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index ddef4480..2d3b44f4 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -901,6 +901,21 @@ pub(crate) fn wire_event_bridges_on_engine( // per-item commit channel is wired at the top of this function. } +fn add_memory_lifecycle_if_configured( + registry: &mut FeatureRegistryBuilder, + config: Option, + build: impl FnOnce(manifest::MemoryConfig) -> std::io::Result, +) -> std::io::Result +where + M: crate::feature::FeatureModule + 'static, +{ + let Some(config) = config else { + return Ok(false); + }; + registry.add_module(build(config)?); + Ok(true) +} + /// Register the builtin file-manipulation tools, optional memory tools, /// and the Worker-orchestration tools (SubWorkerSpawn + comm) on the Worker's /// Engine. Returns the WorkdirSession handle used to attach a `WorkerFsView` to @@ -994,7 +1009,7 @@ where let worker_enabled = feature_config.worker.enabled; let sub_worker_enabled = feature_config.sub_worker.enabled; let mut feature_registry = FeatureRegistryBuilder::new(); - if let Some(config) = memory_config.clone() { + add_memory_lifecycle_if_configured(&mut feature_registry, memory_config.clone(), |config| { let workspace_client = worker.workspace_client_handle(); if !workspace_client.is_available() || workspace_client.workspace_id().is_none() { return Err(std::io::Error::new( @@ -1002,7 +1017,7 @@ where "Memory extraction requires Backend Workspace API authority", )); } - feature_registry.add_module( + Ok( crate::feature::builtin::memory_lifecycle::MemoryLifecycleFeature::new( config, worker.committed_session_capture_handle(), @@ -1014,8 +1029,8 @@ where spawner_workspace_context.clone(), worker.working_event_sender(), ), - ); - } + ) + })?; if sub_worker_enabled && !worker_enabled { feature_registry.add_module( crate::feature::builtin::manage_worker::sub_worker_control_feature( @@ -2135,6 +2150,51 @@ mod tests { use tempfile::TempDir; use tokio::net::UnixListener; + #[test] + fn memory_lifecycle_registration_depends_only_on_memory_config_presence() { + #[derive(Clone)] + struct TestMemoryLifecycleModule; + + impl crate::feature::FeatureModule for TestMemoryLifecycleModule { + fn descriptor(&self) -> crate::feature::FeatureDescriptor { + crate::feature::FeatureDescriptor::builtin( + "test-memory-lifecycle", + "Test Memory Lifecycle", + ) + } + + fn install( + &self, + _context: &mut crate::feature::FeatureInstallContext<'_>, + ) -> Result<(), crate::feature::FeatureInstallError> { + Ok(()) + } + } + + let mut registry = FeatureRegistryBuilder::new(); + let configured = std::cell::Cell::new(false); + let installed = add_memory_lifecycle_if_configured( + &mut registry, + Some(manifest::MemoryConfig::default()), + |_| { + configured.set(true); + Ok(TestMemoryLifecycleModule) + }, + ) + .unwrap(); + assert!(installed); + assert!(configured.get()); + + let mut registry = FeatureRegistryBuilder::new(); + let installed = add_memory_lifecycle_if_configured::( + &mut registry, + None, + |_| panic!("disabled Memory must not construct its lifecycle Feature"), + ) + .unwrap(); + assert!(!installed); + } + #[test] fn image_attachment_gate_requires_vision_and_supported_openai_scheme() { let openai = manifest::ModelManifest { diff --git a/crates/worker/src/feature/builtin/memory_lifecycle.rs b/crates/worker/src/feature/builtin/memory_lifecycle.rs index 59de7caf..b7083428 100644 --- a/crates/worker/src/feature/builtin/memory_lifecycle.rs +++ b/crates/worker/src/feature/builtin/memory_lifecycle.rs @@ -833,13 +833,69 @@ mod tests { &self, request: crate::worker::WorkspaceRequest, ) -> Result { + let is_stage_candidate = request + .body + .as_deref() + .is_some_and(|body| body.contains("stage_candidate")); + let is_append_audit = request + .body + .as_deref() + .is_some_and(|body| body.contains("append_audit")); self.requests.lock().unwrap().push(request); + if is_stage_candidate { + return Ok(crate::worker::WorkspaceResponse { + status: 200, + body: serde_json::to_string(&memory::backend::MemoryBackendHttpResponse::Ok { + result: memory::backend::MemoryBackendOperationResult::StagingWritten( + memory::backend::MemoryStagingWriteOutput { + staging_count: 1, + staging_ids: vec!["candidate-1".to_string()], + }, + ), + }) + .unwrap(), + }); + } + if is_append_audit { + return Ok(crate::worker::WorkspaceResponse { + status: 200, + body: serde_json::to_string(&memory::backend::MemoryBackendHttpResponse::Ok { + result: memory::backend::MemoryBackendOperationResult::Acknowledged( + memory::backend::MemoryBackendAckOutput { + summary: "audit recorded".to_string(), + }, + ), + }) + .unwrap(), + }); + } Err(crate::worker::WorkspaceClientError::Unavailable( "recording client".to_string(), )) } } + #[derive(Clone)] + struct PendingClient { + calls: Arc, + } + + #[async_trait] + impl LlmClient for PendingClient { + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + async fn stream( + &self, + _request: Request, + ) -> Result> + Send>>, ClientError> + { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(Box::pin(futures::stream::pending())) + } + } + #[derive(Clone)] struct ScriptClient { responses: Arc>>, @@ -874,6 +930,40 @@ mod tests { } } + fn stage_candidate_events(call_id: &str, entry_ref: &str) -> Vec { + vec![ + LlmEvent::tool_use_start(0, call_id, "StageMemoryCandidate"), + LlmEvent::tool_input_delta( + 0, + serde_json::json!({ + "kind": "decision", + "claim": "Keep lifecycle work feature-owned.", + "why_useful": "Prevents Worker core coupling.", + "entry_refs": [entry_ref] + }) + .to_string(), + ), + LlmEvent::tool_use_stop(0), + LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] + } + + fn finish_events(call_id: &str, staged_count: usize) -> Vec { + vec![ + LlmEvent::tool_use_start(0, call_id, "FinishMemoryExtraction"), + LlmEvent::tool_input_delta( + 0, + serde_json::json!({"staged_count": staged_count}).to_string(), + ), + LlmEvent::tool_use_stop(0), + LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] + } + fn finish_empty_events(call_id: &str) -> Vec { vec![ LlmEvent::tool_use_start(0, call_id, "FinishMemoryExtraction"), @@ -939,7 +1029,7 @@ permission = "write" fn test_task( capture: CommittedSessionCapture, - client: ScriptClient, + client: Box, extension_writes: Arc>>, event_tx: broadcast::Sender, workspace_client: Arc, @@ -958,14 +1048,16 @@ permission = "write" extensions, workspace_client, manifest: test_manifest(), - client: Box::new(client), + client, prompts: Arc::new(ArcSwap::from(PromptCatalog::builtins_only().unwrap())), workspace_context: WorkerWorkspaceContext::no_workspace(), event_tx: Some(event_tx), } } - async fn run_background_task(task: MemoryLifecycleTask) { + fn start_background_task( + task: MemoryLifecycleTask, + ) -> crate::feature::background::FeatureBackgroundTaskRegistry { let mut builder = FeatureBackgroundTaskRegistryBuilder::default(); builder .register( @@ -984,6 +1076,11 @@ permission = "write" ..Default::default() }) .unwrap(); + registry + } + + async fn run_background_task(task: MemoryLifecycleTask) { + let registry = start_background_task(task); tokio::time::timeout(Duration::from_secs(5), async { loop { if !registry.diagnostics().is_empty() { @@ -1007,7 +1104,7 @@ permission = "write" Arc::new(RecordingWorkspaceClient::default()); run_background_task(test_task( capture(2, 250), - client, + Box::new(client), Arc::clone(&extension_writes), event_tx, workspace_client, @@ -1024,6 +1121,139 @@ permission = "write" assert_eq!(pointer.processed_through_entry, 1); } + #[tokio::test] + async fn run_committed_background_task_stages_non_empty_extraction_and_commits_pointer() { + let source = capture(2, 250); + let entry_ref = + SessionCapture::from_history_entries(source.segment_id.clone(), source.history.clone()) + .overview()[0] + .id + .to_string(); + let client = ScriptClient::new(vec![ + stage_candidate_events("stage-1", &entry_ref), + finish_events("finish-1", 1), + completed_events(), + ]); + let calls = Arc::clone(&client.calls); + let extension_writes = Arc::new(Mutex::new(Vec::new())); + let (event_tx, mut event_rx) = broadcast::channel(64); + let workspace_client = Arc::new(RecordingWorkspaceClient::default()); + run_background_task(test_task( + source, + Box::new(client), + Arc::clone(&extension_writes), + event_tx, + workspace_client.clone(), + )) + .await; + + assert_eq!(calls.load(Ordering::SeqCst), 3); + let writes = extension_writes.lock().unwrap(); + assert_eq!( + writes.len(), + 1, + "recorded requests: {:?}; events: {:?}", + workspace_client.requests.lock().unwrap(), + std::iter::from_fn(|| event_rx.try_recv().ok()).collect::>() + ); + let pointer: memory::ExtractPointerPayload = + serde_json::from_value(writes[0].2.clone()).unwrap(); + assert_eq!(pointer.staging_id, "candidate-1"); + assert!( + workspace_client + .requests + .lock() + .unwrap() + .iter() + .any(|request| { + request + .body + .as_deref() + .is_some_and(|body| body.contains("stage_candidate")) + }) + ); + } + + #[tokio::test] + async fn failed_extraction_emits_failure_event_and_durable_audit_without_pointer() { + let client = ScriptClient::new(Vec::new()); + let calls = Arc::clone(&client.calls); + let extension_writes = Arc::new(Mutex::new(Vec::new())); + let (event_tx, mut event_rx) = broadcast::channel(16); + let workspace_client = Arc::new(RecordingWorkspaceClient::default()); + run_background_task(test_task( + capture(2, 250), + Box::new(client), + Arc::clone(&extension_writes), + event_tx, + workspace_client.clone(), + )) + .await; + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert!(extension_writes.lock().unwrap().is_empty()); + let events = std::iter::from_fn(|| event_rx.try_recv().ok()).collect::>(); + assert!(events.iter().any(|event| { + matches!(event, Event::MemoryWorker(event) if event.status == "failed") + })); + let requests = workspace_client.requests.lock().unwrap(); + assert!( + requests.iter().any(|request| { + request.body.as_deref().is_some_and(|body| { + body.contains("append_audit") + && body.contains("worker_lifecycle") + && body.contains("failed") + }) + }), + "recorded requests: {requests:?}" + ); + } + + #[tokio::test] + async fn rewrite_barrier_cancels_active_extraction_and_emits_cancelled_without_pointer() { + let calls = Arc::new(AtomicUsize::new(0)); + let client = PendingClient { + calls: Arc::clone(&calls), + }; + let extension_writes = Arc::new(Mutex::new(Vec::new())); + let (event_tx, mut event_rx) = broadcast::channel(16); + let workspace_client = Arc::new(RecordingWorkspaceClient::default()); + let registry = start_background_task(test_task( + capture(2, 250), + Box::new(client), + Arc::clone(&extension_writes), + event_tx, + workspace_client.clone(), + )); + tokio::time::timeout(Duration::from_secs(5), async { + while calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("extraction child should reach its first provider request"); + let rewrite_guard = registry.begin_session_rewrite().await.unwrap(); + drop(rewrite_guard); + registry.shutdown().await.unwrap(); + + assert!(extension_writes.lock().unwrap().is_empty()); + let events = std::iter::from_fn(|| event_rx.try_recv().ok()).collect::>(); + assert!(events.iter().any(|event| { + matches!(event, Event::MemoryWorker(event) if event.status == "cancelled") + })); + let requests = workspace_client.requests.lock().unwrap(); + assert!( + requests.iter().any(|request| { + request.body.as_deref().is_some_and(|body| { + body.contains("append_audit") + && body.contains("worker_lifecycle") + && body.contains("cancelled") + }) + }), + "recorded requests: {requests:?}" + ); + } + #[tokio::test] async fn interrupted_committed_run_skips_internal_worker_and_pointer_commit() { let client = ScriptClient::new(Vec::new()); @@ -1036,7 +1266,7 @@ permission = "write" interrupted.run_exit = CommittedRunExit::Interrupted; run_background_task(test_task( interrupted, - client, + Box::new(client), Arc::clone(&extension_writes), event_tx, workspace_client, @@ -1061,7 +1291,7 @@ permission = "write" interrupted.run_exit = CommittedRunExit::Interrupted; let mut task = test_task( interrupted, - client, + Box::new(client), extension_writes, event_tx, workspace_client.clone(), @@ -1259,7 +1489,7 @@ permission = "write" ); } let controller_source = include_str!("../../controller.rs"); - assert!(controller_source.contains("if let Some(config) = memory_config.clone()")); + assert!(controller_source.contains("add_memory_lifecycle_if_configured")); assert!(controller_source.contains("MemoryLifecycleFeature::new")); let lifecycle_source = include_str!("memory_lifecycle.rs"); assert!(lifecycle_source.contains("request_memory_staging_consolidation")); From 5ee77698db5ef875b23c9bd12b985dc3d07ae2d6 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 20:53:27 +0900 Subject: [PATCH 18/22] fix: gate workspace memory lifecycle authority --- crates/worker/src/controller.rs | 104 +++++++++++++++++++-------- crates/worker/src/internal_worker.rs | 3 + crates/worker/src/worker.rs | 16 +++++ 3 files changed, 93 insertions(+), 30 deletions(-) diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 2d3b44f4..8355f7ee 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -904,6 +904,7 @@ pub(crate) fn wire_event_bridges_on_engine( fn add_memory_lifecycle_if_configured( registry: &mut FeatureRegistryBuilder, config: Option, + workspace_bound: bool, build: impl FnOnce(manifest::MemoryConfig) -> std::io::Result, ) -> std::io::Result where @@ -912,6 +913,15 @@ where let Some(config) = config else { return Ok(false); }; + if config.workspace_settings().is_none() { + if workspace_bound { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Workspace-bound Memory requires a Backend-authored settings snapshot", + )); + } + return Ok(false); + } registry.add_module(build(config)?); Ok(true) } @@ -1009,28 +1019,36 @@ where let worker_enabled = feature_config.worker.enabled; let sub_worker_enabled = feature_config.sub_worker.enabled; let mut feature_registry = FeatureRegistryBuilder::new(); - add_memory_lifecycle_if_configured(&mut feature_registry, memory_config.clone(), |config| { - let workspace_client = worker.workspace_client_handle(); - if !workspace_client.is_available() || workspace_client.workspace_id().is_none() { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "Memory extraction requires Backend Workspace API authority", - )); - } - Ok( - crate::feature::builtin::memory_lifecycle::MemoryLifecycleFeature::new( - config, - worker.committed_session_capture_handle(), - worker.session_extension_handle(), - workspace_client, - spawner_manifest.clone(), - worker.llm_client_handle(), - prompts.clone(), - spawner_workspace_context.clone(), - worker.working_event_sender(), - ), - ) - })?; + add_memory_lifecycle_if_configured( + &mut feature_registry, + worker + .manifest_lifecycle_features_enabled() + .then(|| memory_config.clone()) + .flatten(), + spawner_workspace_context.workspace_id().is_some(), + |config| { + let workspace_client = worker.workspace_client_handle(); + if !workspace_client.is_available() || workspace_client.workspace_id().is_none() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Memory extraction requires Backend Workspace API authority", + )); + } + Ok( + crate::feature::builtin::memory_lifecycle::MemoryLifecycleFeature::new( + config, + worker.committed_session_capture_handle(), + worker.session_extension_handle(), + workspace_client, + spawner_manifest.clone(), + worker.llm_client_handle(), + prompts.clone(), + spawner_workspace_context.clone(), + worker.working_event_sender(), + ), + ) + }, + )?; if sub_worker_enabled && !worker_enabled { feature_registry.add_module( crate::feature::builtin::manage_worker::sub_worker_control_feature( @@ -2151,7 +2169,7 @@ mod tests { use tokio::net::UnixListener; #[test] - fn memory_lifecycle_registration_depends_only_on_memory_config_presence() { + fn memory_lifecycle_registration_requires_bound_workspace_memory_config() { #[derive(Clone)] struct TestMemoryLifecycleModule; @@ -2173,15 +2191,18 @@ mod tests { let mut registry = FeatureRegistryBuilder::new(); let configured = std::cell::Cell::new(false); - let installed = add_memory_lifecycle_if_configured( - &mut registry, - Some(manifest::MemoryConfig::default()), - |_| { + let mut memory_config = manifest::MemoryConfig::default(); + memory_config.bind_workspace_settings(&manifest::WorkspaceMemorySettingsSnapshot { + workspace_id: "workspace-1".to_string(), + settings_revision: 1, + language: "English".to_string(), + }); + let installed = + add_memory_lifecycle_if_configured(&mut registry, Some(memory_config), true, |_| { configured.set(true); Ok(TestMemoryLifecycleModule) - }, - ) - .unwrap(); + }) + .unwrap(); assert!(installed); assert!(configured.get()); @@ -2189,10 +2210,33 @@ mod tests { let installed = add_memory_lifecycle_if_configured::( &mut registry, None, + false, |_| panic!("disabled Memory must not construct its lifecycle Feature"), ) .unwrap(); assert!(!installed); + + let installed = add_memory_lifecycle_if_configured::( + &mut registry, + Some(manifest::MemoryConfig::default()), + false, + |_| panic!("Memory without a Backend-authored settings snapshot must stay disabled"), + ) + .unwrap(); + assert!(!installed); + + let error = add_memory_lifecycle_if_configured::( + &mut registry, + Some(manifest::MemoryConfig::default()), + true, + |_| panic!("invalid Workspace Memory config must fail before Feature construction"), + ) + .unwrap_err(); + assert!( + error + .to_string() + .contains("Backend-authored settings snapshot") + ); } #[test] diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index c81787c5..ec7ef586 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -164,6 +164,7 @@ where identity: identity.clone(), history_entries: 0, })?; + worker.disable_manifest_lifecycle_features(); if let Some(session) = inherited_workdir_session { worker.bind_workdir_session(Some(session)); } @@ -578,6 +579,7 @@ pub(crate) async fn spawn_internal_worker_session( .map_err(|source| InternalWorkerSessionError::Build { message: source.to_string(), })?; + worker.disable_manifest_lifecycle_features(); if let Some(session) = inherited_workdir_session { worker.bind_workdir_session(Some(session)); } @@ -671,6 +673,7 @@ pub(crate) fn prepare_internal_worker_from_spec( .map_err(|source| InternalWorkerSessionError::Build { message: source.to_string(), })?; + worker.disable_manifest_lifecycle_features(); if let Some(session) = inherited_workdir_session { worker.bind_workdir_session(Some(session)); } diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 4b74f647..483f508a 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -1130,6 +1130,9 @@ pub struct Worker { hook_registry: Option>, /// Executable background tasks registered by successfully installed features. feature_background_tasks: FeatureBackgroundTaskRegistry, + /// Internal Workers install an explicit Feature composition and disable + /// manifest-derived lifecycle Features before controller startup. + manifest_lifecycle_features_enabled: bool, interceptor_installed: bool, /// Shared compaction state (present when threshold is configured). compact_state: Option>, @@ -1423,6 +1426,7 @@ impl Worker { hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), + manifest_lifecycle_features_enabled: true, interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -1800,6 +1804,14 @@ impl Worker { self.session.replace_history(entries); } + pub(crate) fn disable_manifest_lifecycle_features(&mut self) { + self.manifest_lifecycle_features_enabled = false; + } + + pub(crate) fn manifest_lifecycle_features_enabled(&self) -> bool { + self.manifest_lifecycle_features_enabled + } + /// Install enabled feature modules into the Worker host surfaces. pub fn install_features( &mut self, @@ -4676,6 +4688,7 @@ where hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), + manifest_lifecycle_features_enabled: true, interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -4757,6 +4770,7 @@ where hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), + manifest_lifecycle_features_enabled: false, interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -4873,6 +4887,7 @@ where hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), + manifest_lifecycle_features_enabled: true, interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), @@ -5243,6 +5258,7 @@ where hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), + manifest_lifecycle_features_enabled: true, interceptor_installed: false, compact_state: None, usage_tracker: Arc::new(UsageTracker::new()), From 1e674d70c2b7e2c5febfab2b97e1c2ab8f9a23de Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 22:55:37 +0900 Subject: [PATCH 19/22] refactor: unify Memory feature configuration authority --- crates/manifest/src/config.rs | 208 ++++--- crates/manifest/src/defaults.rs | 2 +- crates/manifest/src/lib.rs | 527 +++++++++++++----- crates/manifest/src/profile.rs | 70 +-- crates/memory/src/backend.rs | 18 +- crates/memory/src/consolidate/staging.rs | 3 +- crates/memory/src/workspace.rs | 41 +- crates/tui/src/setup_model.rs | 7 +- crates/worker-runtime/src/fs_store.rs | 24 +- crates/worker-runtime/src/worker_backend.rs | 26 +- crates/worker/src/controller.rs | 162 +++--- crates/worker/src/feature/builtin/memory.rs | 56 ++ .../src/feature/builtin/memory_lifecycle.rs | 127 ++--- crates/worker/src/spawn/tool.rs | 62 ++- crates/worker/src/worker.rs | 132 ++--- crates/workspace-server/src/server.rs | 33 +- docs/manifest.toml | 67 +-- resources/profiles/base.dcdl | 15 +- resources/profiles/default.dcdl | 2 +- resources/profiles/memory-consolidation.dcdl | 2 +- 20 files changed, 947 insertions(+), 637 deletions(-) diff --git a/crates/manifest/src/config.rs b/crates/manifest/src/config.rs index f7d535be..6a6e6e96 100644 --- a/crates/manifest/src/config.rs +++ b/crates/manifest/src/config.rs @@ -18,8 +18,9 @@ use crate::model::{AuthRef, ModelManifest, ReasoningControl}; use crate::plugin::PluginConfig; use crate::{ CompactionConfig, EngineManifest, FeatureConfig, FeatureFlagConfig, FileUploadLimits, - McpConfig, McpEnvValue, McpStdioCwdPolicy, MemoryConfig, MemoryFeatureConfig, - MergeRequestFeatureConfig, ScopeConfig, SessionConfig, SkillsConfig, TicketFeatureConfig, + McpConfig, McpEnvValue, McpStdioCwdPolicy, MemoryExtractionProfileConfig, + MemoryFeatureProfileConfig, MemoryResidentProfileConfig, MergeRequestFeatureConfig, + ResolvedMemoryFeatureConfig, ScopeConfig, SessionConfig, SkillsConfig, TicketFeatureConfig, ToolOutputLimits, ToolPermissionConfig, ToolPermissionRule, WebConfig, WorkerFeatureConfig, WorkerManifest, WorkerMeta, }; @@ -67,9 +68,6 @@ pub struct WorkerManifestConfig { /// First-class web tool opt-in. See [`WebConfig`]. #[serde(default)] pub web: Option, - /// Memory subsystem opt-in. See [`MemoryConfig`]. - #[serde(default)] - pub memory: Option, /// External Agent Skills directories. See [`crate::SkillsConfig`]. #[serde(default)] pub skills: Option, @@ -193,18 +191,72 @@ impl From for WorkerFeatureConfig { } #[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct MemoryFeatureConfigPartial { #[serde(default)] pub enabled: Option, #[serde(default)] - pub staging: Option, + pub staging_tools: Option, + #[serde(default)] + pub resident: Option, + #[serde(default)] + pub extraction: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MemoryResidentProfileConfigPartial { + #[serde(default)] + pub inject_summary: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MemoryExtractionProfileConfigPartial { + #[serde(default)] + pub enabled: Option, + #[serde(default)] + pub model: Option, + #[serde(default)] + pub threshold: Option, + #[serde(default)] + pub worker_max_turns: Option, } impl MemoryFeatureConfigPartial { fn merge(self, other: Self) -> Self { Self { enabled: other.enabled.or(self.enabled), - staging: other.staging.or(self.staging), + staging_tools: other.staging_tools.or(self.staging_tools), + resident: merge_option( + self.resident, + other.resident, + MemoryResidentProfileConfigPartial::merge, + ), + extraction: merge_option( + self.extraction, + other.extraction, + MemoryExtractionProfileConfigPartial::merge, + ), + } + } +} + +impl MemoryResidentProfileConfigPartial { + fn merge(self, other: Self) -> Self { + Self { + inject_summary: other.inject_summary.or(self.inject_summary), + } + } +} + +impl MemoryExtractionProfileConfigPartial { + fn merge(self, other: Self) -> Self { + Self { + enabled: other.enabled.or(self.enabled), + model: other.model.or(self.model), + threshold: other.threshold.or(self.threshold), + worker_max_turns: other.worker_max_turns.or(self.worker_max_turns), } } } @@ -259,7 +311,7 @@ impl From for FeatureConfig { task: value.task.map(FeatureFlagConfig::from).unwrap_or_default(), memory: value .memory - .map(MemoryFeatureConfig::from) + .map(ResolvedMemoryFeatureConfig::from) .unwrap_or_default(), web: value.web.map(FeatureFlagConfig::from).unwrap_or_default(), image: value.image.map(FeatureFlagConfig::from).unwrap_or_default(), @@ -329,20 +381,45 @@ impl From for WorkerFeatureConfigPartial { } } -impl From for MemoryFeatureConfig { +impl From for ResolvedMemoryFeatureConfig { fn from(value: MemoryFeatureConfigPartial) -> Self { + let resident = value.resident.unwrap_or_default(); + let extraction = value.extraction.unwrap_or_default(); Self { - enabled: value.enabled.unwrap_or_default(), - staging: value.staging.unwrap_or_default(), + profile: MemoryFeatureProfileConfig { + enabled: value.enabled.unwrap_or_default(), + staging_tools: value.staging_tools.unwrap_or_default(), + resident: MemoryResidentProfileConfig { + inject_summary: resident.inject_summary.unwrap_or(true), + }, + extraction: MemoryExtractionProfileConfig { + enabled: extraction.enabled.unwrap_or(true), + model: extraction.model, + threshold: extraction.threshold.or(Some(50_000)), + worker_max_turns: extraction + .worker_max_turns + .or(defaults::MEMORY_EXTRACT_WORKER_MAX_TURNS), + }, + }, + workspace_settings: None, } } } -impl From for MemoryFeatureConfigPartial { - fn from(value: MemoryFeatureConfig) -> Self { +impl From for MemoryFeatureConfigPartial { + fn from(value: ResolvedMemoryFeatureConfig) -> Self { Self { - enabled: Some(value.enabled), - staging: Some(value.staging), + enabled: Some(value.profile.enabled), + staging_tools: Some(value.profile.staging_tools), + resident: Some(MemoryResidentProfileConfigPartial { + inject_summary: Some(value.profile.resident.inject_summary), + }), + extraction: Some(MemoryExtractionProfileConfigPartial { + enabled: Some(value.profile.extraction.enabled), + model: value.profile.extraction.model, + threshold: value.profile.extraction.threshold, + worker_max_turns: value.profile.extraction.worker_max_turns, + }), } } } @@ -543,13 +620,9 @@ pub(crate) fn reject_removed_manifest_fields(s: &str) -> Result<(), toml::de::Er (removed; use compaction.prune_protected_tokens)", )); } - if value - .get("memory") - .and_then(toml::Value::as_table) - .is_some_and(|table| table.contains_key("extract_worker_max_input_tokens")) - { + if value.get("memory").is_some() { return Err(toml::de::Error::custom( - "unknown field in manifest: memory.extract_worker_max_input_tokens (removed)", + "unknown field in manifest: memory (removed; configure feature.memory)", )); } if value @@ -633,11 +706,6 @@ impl WorkerManifestConfig { for rule in &mut self.delegation_scope.deny { rule.target = join_if_relative(base, &rule.target); } - if let Some(ref mut memory) = self.memory - && let Some(ref mut root) = memory.workspace_root - { - *root = join_if_relative(base, root); - } if let Some(ref mut compaction) = self.compaction && let Some(ref mut cp) = compaction.model { @@ -682,7 +750,6 @@ impl WorkerManifestConfig { CompactionConfigPartial::merge, ), web: merge_option(self.web, upper.web, WebConfig::merge), - memory: merge_option(self.memory, upper.memory, MemoryConfig::merge), skills: merge_option(self.skills, upper.skills, SkillsConfig::merge), } } @@ -754,32 +821,6 @@ impl crate::WebFetchConfig { } } -impl MemoryConfig { - fn merge(self, upper: Self) -> Self { - Self { - workspace_root: upper.workspace_root.or(self.workspace_root), - query_result_limit: upper.query_result_limit.or(self.query_result_limit), - query_excerpt_lines: upper.query_excerpt_lines.or(self.query_excerpt_lines), - inject_summary: upper.inject_summary.or(self.inject_summary), - workspace_id: upper.workspace_id.or(self.workspace_id), - settings_revision: upper.settings_revision.or(self.settings_revision), - language: upper.language.or(self.language), - extract_model: upper.extract_model.or(self.extract_model), - extract_threshold: upper.extract_threshold.or(self.extract_threshold), - extract_worker_max_turns: upper - .extract_worker_max_turns - .or(self.extract_worker_max_turns), - consolidation_model: upper.consolidation_model.or(self.consolidation_model), - consolidation_threshold_files: upper - .consolidation_threshold_files - .or(self.consolidation_threshold_files), - consolidation_threshold_bytes: upper - .consolidation_threshold_bytes - .or(self.consolidation_threshold_bytes), - } - } -} - impl WorkerMetaConfig { fn merge(self, upper: Self) -> Self { Self { @@ -1223,7 +1264,6 @@ impl TryFrom for WorkerManifest { mcp: cfg.mcp, compaction, web: cfg.web, - memory: cfg.memory, skills: cfg.skills, profile: None, }) @@ -1271,7 +1311,6 @@ mod tests { session: None, compaction: None, web: None, - memory: None, skills: None, } } @@ -1846,29 +1885,46 @@ prune_protected_turns = 3 } #[test] - fn from_toml_rejects_removed_extract_worker_max_input_tokens_field() { - let bad = r#" -[memory] -extract_worker_max_input_tokens = 30000 -"#; - let err = WorkerManifestConfig::from_toml(bad).unwrap_err(); - assert!( - err.to_string() - .contains("memory.extract_worker_max_input_tokens"), - "unexpected error: {err}" - ); + fn from_toml_accepts_memory_extraction_settings_only_under_feature_memory() { + let cfg = WorkerManifestConfig::from_toml( + r#" +[feature.memory] +enabled = true +staging_tools = false + +[feature.memory.resident] +inject_summary = false + +[feature.memory.extraction] +enabled = true +threshold = 42000 +worker_max_turns = 2 +"#, + ) + .unwrap(); + let memory = cfg.feature.memory.unwrap(); + assert_eq!(memory.enabled, Some(true)); + assert_eq!(memory.staging_tools, Some(false)); + assert_eq!(memory.resident.unwrap().inject_summary, Some(false)); + let extraction = memory.extraction.unwrap(); + assert_eq!(extraction.enabled, Some(true)); + assert_eq!(extraction.threshold, Some(42_000)); + assert_eq!(extraction.worker_max_turns, Some(2)); } #[test] - fn from_toml_accepts_extract_worker_max_turns() { - let cfg = WorkerManifestConfig::from_toml( + fn from_toml_rejects_legacy_top_level_memory_authority() { + let err = WorkerManifestConfig::from_toml( r#" [memory] extract_worker_max_turns = 2 "#, ) - .unwrap(); - assert_eq!(cfg.memory.unwrap().extract_worker_max_turns, Some(2)); + .unwrap_err(); + assert!( + err.to_string().contains("memory"), + "unexpected error: {err}" + ); } #[test] @@ -1948,7 +2004,7 @@ worker_max_turns = 7 fn feature_flags_default_disabled_in_resolved_manifest() { let manifest: WorkerManifest = minimal_valid().try_into().unwrap(); assert!(!manifest.feature.task.enabled); - assert!(!manifest.feature.memory.enabled); + assert!(!manifest.feature.memory.profile.enabled); assert!(!manifest.feature.web.enabled); assert!(!manifest.feature.sub_worker.enabled); assert!(!manifest.feature.objective.enabled); @@ -2025,8 +2081,8 @@ enabled = false } ); assert!(!manifest.feature.orchestration.enabled); - assert!(!manifest.feature.memory.enabled); - assert!(!manifest.feature.memory.staging); + assert!(!manifest.feature.memory.profile.enabled); + assert!(!manifest.feature.memory.profile.staging_tools); assert!(!manifest.feature.objective.enabled); } @@ -2074,7 +2130,7 @@ readiness_check = true enabled = true [feature.memory] -staging = true +staging_tools = true [feature.manage_workdir] enabled = true @@ -2111,8 +2167,8 @@ enabled = true }) .try_into() .unwrap(); - assert!(manifest.feature.memory.enabled); - assert!(manifest.feature.memory.staging); + assert!(manifest.feature.memory.profile.enabled); + assert!(manifest.feature.memory.profile.staging_tools); assert!(manifest.feature.manage_workdir.enabled); assert!(manifest.feature.ticket.enabled); assert!(!manifest.feature.ticket.authoring); diff --git a/crates/manifest/src/defaults.rs b/crates/manifest/src/defaults.rs index 70327972..88aeafab 100644 --- a/crates/manifest/src/defaults.rs +++ b/crates/manifest/src/defaults.rs @@ -93,5 +93,5 @@ pub const COMPACT_RESULT_CONTEXT_MAX_TOKENS: u64 = 60_000; pub const COMPACT_DEFAULT_REFERENCE_COUNT: usize = 5; /// Optional maximum extract-worker tool-loop depth. `None` means unlimited. -/// See [`crate::MemoryConfig::extract_worker_max_turns`]. +/// See [`crate::MemoryExtractionProfileConfig::worker_max_turns`]. pub const MEMORY_EXTRACT_WORKER_MAX_TURNS: Option = Some(8); diff --git a/crates/manifest/src/lib.rs b/crates/manifest/src/lib.rs index 66535082..11c6c2f7 100644 --- a/crates/manifest/src/lib.rs +++ b/crates/manifest/src/lib.rs @@ -47,6 +47,7 @@ use serde::{Deserialize, Serialize}; /// part of the manifest — it is the process's `std::env::current_dir()` /// at construction time. #[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct WorkerManifest { pub worker: WorkerMeta, pub model: ModelManifest, @@ -80,11 +81,6 @@ pub struct WorkerManifest { pub mcp: McpConfig, #[serde(default)] pub compaction: Option, - /// Memory subsystem configuration. Presence of `[memory]` configures memory - /// storage, extraction, consolidation, and resident injection, but memory - /// tools are surfaced only when `[feature.memory].enabled = true`. - #[serde(default)] - pub memory: Option, /// First-class web tools configuration. Network access remains fail-closed /// under this config; WebSearch/WebFetch schemas are surfaced only when /// `[feature.web].enabled = true`. @@ -109,12 +105,12 @@ pub struct WorkerManifest { /// profile/config data only: they do not carry runtime Worker names, sockets, /// sessions, secrets, or resolved host state. Tool registration still applies /// the normal scope, host-authority, backend, memory, and network checks. -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct FeatureConfig { #[serde(default)] pub task: FeatureFlagConfig, #[serde(default)] - pub memory: MemoryFeatureConfig, + pub memory: ResolvedMemoryFeatureConfig, #[serde(default)] pub web: FeatureFlagConfig, #[serde(default)] @@ -147,7 +143,7 @@ impl Default for FeatureConfig { fn default() -> Self { Self { task: FeatureFlagConfig::disabled(), - memory: MemoryFeatureConfig::disabled(), + memory: ResolvedMemoryFeatureConfig::default(), web: FeatureFlagConfig::disabled(), image: FeatureFlagConfig::disabled(), sub_worker: FeatureFlagConfig::disabled(), @@ -222,34 +218,117 @@ const fn default_true() -> bool { true } -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] -pub struct MemoryFeatureConfig { - #[serde(default)] +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(default, deny_unknown_fields)] +pub struct MemoryFeatureProfileConfig { pub enabled: bool, /// Exposes Memory staging queue tools in addition to normal Memory CRUD/query tools. - #[serde(default)] - pub staging: bool, + pub staging_tools: bool, + pub resident: MemoryResidentProfileConfig, + pub extraction: MemoryExtractionProfileConfig, } -impl MemoryFeatureConfig { - pub const fn disabled() -> Self { - Self { - enabled: false, - staging: false, - } +impl MemoryFeatureProfileConfig { + pub fn disabled() -> Self { + Self::default() } - pub const fn enabled() -> Self { + pub fn enabled() -> Self { Self { enabled: true, - staging: false, + ..Self::default() } } } -impl Default for MemoryFeatureConfig { +impl Default for MemoryFeatureProfileConfig { fn default() -> Self { - Self::disabled() + Self { + enabled: false, + staging_tools: false, + resident: MemoryResidentProfileConfig::default(), + extraction: MemoryExtractionProfileConfig::default(), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(default, deny_unknown_fields)] +pub struct MemoryResidentProfileConfig { + pub inject_summary: bool, +} + +impl Default for MemoryResidentProfileConfig { + fn default() -> Self { + Self { + inject_summary: true, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(default, deny_unknown_fields)] +pub struct MemoryExtractionProfileConfig { + pub enabled: bool, + pub model: Option, + pub threshold: Option, + pub worker_max_turns: Option, +} + +impl Default for MemoryExtractionProfileConfig { + fn default() -> Self { + Self { + enabled: true, + model: None, + threshold: Some(50_000), + worker_max_turns: defaults::MEMORY_EXTRACT_WORKER_MAX_TURNS, + } + } +} + +/// Immutable Memory execution configuration persisted in a resolved Worker Manifest. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] +#[serde(default, deny_unknown_fields)] +pub struct ResolvedMemoryFeatureConfig { + pub profile: MemoryFeatureProfileConfig, + #[serde(skip_serializing_if = "Option::is_none")] + pub workspace_settings: Option, +} + +impl ResolvedMemoryFeatureConfig { + pub fn enabled(&self) -> bool { + self.profile.enabled + } + + pub fn bind_workspace_settings( + &mut self, + settings: WorkspaceMemorySettingsSnapshot, + ) -> Result<(), &'static str> { + if !self.profile.enabled { + if self.workspace_settings.is_some() { + return Err("disabled Memory feature must not carry Workspace settings"); + } + return Ok(()); + } + if self.workspace_settings.is_some() { + return Err("memory Workspace settings are already bound"); + } + self.workspace_settings = Some(settings); + Ok(()) + } + + pub fn workspace_settings(&self) -> Option { + self.workspace_settings.clone() + } + + pub fn validate_execution(&self) -> Result<(), &'static str> { + if self.profile.enabled && self.workspace_settings.is_none() { + return Err("enabled Memory feature requires trusted Workspace settings"); + } + if !self.profile.enabled && self.workspace_settings.is_some() { + return Err("disabled Memory feature must not carry Workspace settings"); + } + Ok(()) } } @@ -484,98 +563,6 @@ pub struct WorkspaceMemorySettingsSnapshot { pub language: String, } -/// Memory subsystem configuration. Presence in the manifest enables -/// memory; `workspace_root` pins the memory workspace explicitly. When it -/// is absent, memory resolution searches upward from the Worker's pwd for a -/// `.yoi/memory` marker rather than treating `.yoi` project records alone -/// as a memory root. -/// -/// All fields are `Option`; defaults are applied at the consumer -/// (`.unwrap_or(defaults::...)`). This keeps cascade `merge` simple -/// (`upper.x.or(self.x)`) without a separate partial/resolved split. -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -pub struct MemoryConfig { - /// Override for the memory workspace root. When `None`, consumers resolve - /// the root from their default path and ancestor `.yoi/memory` markers. - /// When set, must be an absolute path. - #[serde(default)] - pub workspace_root: Option, - /// Maximum number of records returned by `MemoryQuery` / - /// `MemoryQuery` per call. `None` ⇒ tool default (20). - #[serde(default)] - pub query_result_limit: Option, - /// Lines of context before and after each match in query excerpts. - /// Ignored when the request omits `query`. `None` ⇒ tool default (3). - #[serde(default)] - pub query_excerpt_lines: Option, - /// Whether the body of `memory/summary.md` is exposed in the resident - /// system-prompt section. `None` ⇒ enabled. - #[serde(default)] - pub inject_summary: Option, - /// Workspace that owns the bound Memory settings revision. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub workspace_id: Option, - /// Monotonic revision of the bound Workspace Memory settings. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub settings_revision: Option, - /// Language from the bound Workspace Memory settings revision. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub language: Option, - /// Optional model for the extract worker. When `None`, - /// the main engine model is cloned via `clone_boxed()`. Lightweight - /// reasoning-capable models (Haiku / 4o-mini / Flash class) are - /// recommended. - #[serde(default)] - pub extract_model: Option, - /// Cumulative input-token threshold (since the last extract pointer) - /// that triggers an extract run. `None` disables the extract trigger - /// entirely; memory tools and resident injection still work, only - /// the auto-extract trigger is dormant. - #[serde(default)] - pub extract_threshold: Option, - /// Optional maximum extract-worker tool-loop depth. `None` leaves - /// the worker unlimited; the default bounds runaway short-context - /// loops. Falls through to - /// [`defaults::MEMORY_EXTRACT_WORKER_MAX_TURNS`] when unset. - #[serde(default)] - pub extract_worker_max_turns: Option, - /// Optional model for the consolidation worker. When - /// `None`, the main engine model is cloned via `clone_boxed()`. - /// Reasoning-class models are recommended. - #[serde(default)] - pub consolidation_model: Option, - /// Consolidation trigger: file-count threshold of `_staging/`. The - /// consolidation run fires when the staging directory has at least - /// this many entries. Either threshold reaching its limit fires - /// consolidation (logical OR). `None` for both thresholds ⇒ - /// consolidation disabled. - #[serde(default)] - pub consolidation_threshold_files: Option, - /// Consolidation trigger: byte-size threshold across all `_staging/` - /// entries. Either threshold reaching its limit fires consolidation. - /// `None` for both thresholds ⇒ consolidation disabled. - #[serde(default)] - pub consolidation_threshold_bytes: Option, -} - -impl MemoryConfig { - /// Replace any untrusted manifest values with a trusted Workspace snapshot. - pub fn bind_workspace_settings(&mut self, snapshot: &WorkspaceMemorySettingsSnapshot) { - self.workspace_id = Some(snapshot.workspace_id.clone()); - self.settings_revision = Some(snapshot.settings_revision); - self.language = Some(snapshot.language.clone()); - } - - /// Return the complete bound Workspace settings snapshot, if every field is present. - pub fn workspace_settings(&self) -> Option { - Some(WorkspaceMemorySettingsSnapshot { - workspace_id: self.workspace_id.clone()?, - settings_revision: self.settings_revision?, - language: self.language.clone()?, - }) - } -} - /// Worker metadata. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct WorkerMeta { @@ -941,6 +928,167 @@ impl WorkerManifest { } } +const RESOLVED_MANIFEST_SNAPSHOT_SCHEMA_VERSION: u64 = 2; + +/// Serialize a resolved Worker Manifest for durable Worker-specific storage. +pub fn write_persisted_worker_manifest_snapshot( + manifest: &WorkerManifest, +) -> Result { + Ok(serde_json::json!({ + "schema_version": RESOLVED_MANIFEST_SNAPSHOT_SCHEMA_VERSION, + "manifest": serde_json::to_value(manifest)?, + })) +} + +/// Read a durable resolved Worker Manifest through the versioned compatibility +/// boundary. Runtime code must not deserialize persisted snapshots directly. +pub fn read_persisted_worker_manifest_snapshot( + snapshot: serde_json::Value, +) -> Result { + let object = snapshot.as_object().ok_or_else(|| { + serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "resolved Worker manifest snapshot must be an object", + )) + })?; + if let Some(version) = object.get("schema_version") { + let version = version.as_u64().ok_or_else(|| { + serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "resolved Worker manifest snapshot schema_version must be an integer", + )) + })?; + if version != RESOLVED_MANIFEST_SNAPSHOT_SCHEMA_VERSION { + return Err(serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("unsupported resolved Worker manifest snapshot schema version {version}"), + ))); + } + if object.len() != 2 { + return Err(serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "resolved Worker manifest snapshot contains unknown fields", + ))); + } + let manifest = object.get("manifest").cloned().ok_or_else(|| { + serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "resolved Worker manifest snapshot is missing manifest", + )) + })?; + return serde_json::from_value(manifest); + } + + migrate_legacy_resolved_manifest_snapshot(snapshot) +} + +fn migrate_legacy_resolved_manifest_snapshot( + mut snapshot: serde_json::Value, +) -> Result { + let root = snapshot.as_object_mut().ok_or_else(|| { + serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "legacy resolved Worker manifest snapshot must be an object", + )) + })?; + let legacy_memory = root.remove("memory"); + let feature = root + .entry("feature") + .or_insert_with(|| serde_json::json!({})) + .as_object_mut() + .ok_or_else(|| { + serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "legacy resolved Worker manifest feature must be an object", + )) + })?; + let legacy_feature_memory = feature + .remove("memory") + .unwrap_or_else(|| serde_json::json!({})); + let legacy_feature_memory = legacy_feature_memory.as_object().ok_or_else(|| { + serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "legacy resolved Worker manifest feature.memory must be an object", + )) + })?; + if legacy_feature_memory + .keys() + .any(|key| !matches!(key.as_str(), "enabled" | "staging")) + { + return Err(serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "legacy resolved Worker manifest mixes old and new Memory configuration", + ))); + } + let enabled = legacy_feature_memory + .get("enabled") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + let staging_tools = legacy_feature_memory + .get("staging") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + + let legacy_memory = legacy_memory.unwrap_or_else(|| serde_json::json!({})); + let legacy_memory = legacy_memory.as_object().ok_or_else(|| { + serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "legacy resolved Worker manifest memory must be an object", + )) + })?; + let workspace_id = legacy_memory.get("workspace_id").cloned(); + let settings_revision = legacy_memory.get("settings_revision").cloned(); + let language = legacy_memory.get("language").cloned(); + let workspace_settings = match (workspace_id, settings_revision, language) { + (Some(workspace_id), Some(settings_revision), Some(language)) => Some(serde_json::json!({ + "workspace_id": workspace_id, + "settings_revision": settings_revision, + "language": language, + })), + (None, None, None) => None, + _ => { + return Err(serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "legacy resolved Worker manifest contains a partial Memory settings snapshot", + ))); + } + }; + let extraction_threshold = legacy_memory + .get("extract_threshold") + .cloned() + .unwrap_or(serde_json::Value::Null); + let extraction_enabled = !extraction_threshold.is_null(); + let mut resolved = serde_json::json!({ + "profile": { + "enabled": enabled, + "staging_tools": staging_tools, + "resident": { + "inject_summary": legacy_memory + .get("inject_summary") + .and_then(serde_json::Value::as_bool) + .unwrap_or(true), + }, + "extraction": { + "enabled": extraction_enabled, + "model": legacy_memory.get("extract_model").cloned().unwrap_or(serde_json::Value::Null), + "threshold": extraction_threshold, + "worker_max_turns": legacy_memory + .get("extract_worker_max_turns") + .cloned() + .unwrap_or(serde_json::Value::Null), + }, + }, + }); + if let Some(workspace_settings) = workspace_settings { + resolved + .as_object_mut() + .expect("resolved Memory config is an object") + .insert("workspace_settings".to_string(), workspace_settings); + } + feature.insert("memory".to_string(), resolved); + serde_json::from_value(snapshot) +} + #[cfg(test)] mod tests { use super::*; @@ -1246,36 +1394,129 @@ model_id = "claude-sonnet-4-20250514" } #[test] - fn omitted_memory_is_none() { + fn omitted_memory_feature_is_disabled() { let manifest = WorkerManifest::from_toml(MINIMAL_REQUIRED).unwrap(); - assert!(manifest.memory.is_none()); + assert!(!manifest.feature.memory.profile.enabled); + assert!(manifest.feature.memory.workspace_settings.is_none()); } #[test] - fn empty_memory_section_enables_with_default_root() { - let toml = format!("{MINIMAL_REQUIRED}\n[memory]\n"); + fn resolved_memory_feature_requires_nested_profile_and_trusted_snapshot() { + let toml = format!( + "{MINIMAL_REQUIRED}\n\ + [feature.memory.profile]\n\ + enabled = true\n\ + staging_tools = false\n\n\ + [feature.memory.profile.resident]\n\ + inject_summary = false\n\n\ + [feature.memory.profile.extraction]\n\ + enabled = true\n\ + threshold = 42000\n\ + worker_max_turns = 2\n\n\ + [feature.memory.workspace_settings]\n\ + workspace_id = \"workspace-1\"\n\ + settings_revision = 7\n\ + language = \"日本語\"\n" + ); let manifest = WorkerManifest::from_toml(&toml).unwrap(); - let mem = manifest.memory.expect("memory section parsed"); - assert!(mem.workspace_root.is_none()); - assert_eq!(mem.inject_summary, None); - } - - #[test] - fn memory_section_with_inject_summary_false() { - let toml = format!("{MINIMAL_REQUIRED}\n[memory]\ninject_summary = false\n"); - let manifest = WorkerManifest::from_toml(&toml).unwrap(); - let mem = manifest.memory.unwrap(); - assert_eq!(mem.inject_summary, Some(false)); - } - - #[test] - fn memory_section_with_explicit_root() { - let toml = format!("{MINIMAL_REQUIRED}\n[memory]\nworkspace_root = \"/some/where\"\n"); - let manifest = WorkerManifest::from_toml(&toml).unwrap(); - let mem = manifest.memory.unwrap(); + assert!(manifest.feature.memory.profile.enabled); + assert!(!manifest.feature.memory.profile.resident.inject_summary); assert_eq!( - mem.workspace_root.unwrap(), - std::path::PathBuf::from("/some/where") + manifest.feature.memory.profile.extraction.threshold, + Some(42_000) + ); + assert_eq!( + manifest + .feature + .memory + .workspace_settings() + .unwrap() + .language, + "日本語" + ); + } + + #[test] + fn resolved_memory_execution_validation_fails_closed() { + let snapshot = WorkspaceMemorySettingsSnapshot { + workspace_id: "workspace-1".to_string(), + settings_revision: 1, + language: "English".to_string(), + }; + let mut enabled = ResolvedMemoryFeatureConfig::default(); + enabled.profile.enabled = true; + assert!(enabled.validate_execution().is_err()); + enabled.bind_workspace_settings(snapshot.clone()).unwrap(); + assert!(enabled.validate_execution().is_ok()); + + let mut disabled = ResolvedMemoryFeatureConfig::default(); + disabled.workspace_settings = Some(snapshot.clone()); + assert!(disabled.validate_execution().is_err()); + assert!(disabled.bind_workspace_settings(snapshot).is_err()); + } + + #[test] + fn current_manifest_rejects_legacy_top_level_memory_authority() { + let toml = format!("{MINIMAL_REQUIRED}\n[memory]\nlanguage = \"Japanese\"\n"); + assert!(WorkerManifest::from_toml(&toml).is_err()); + } + + #[test] + fn persisted_manifest_adapter_migrates_legacy_memory_authority() { + let mut manifest = + serde_json::to_value(WorkerManifest::from_toml(MINIMAL_REQUIRED).unwrap()).unwrap(); + manifest["feature"]["memory"] = serde_json::json!({ + "enabled": true, + "staging": true, + }); + manifest["memory"] = serde_json::json!({ + "workspace_root": "/discarded", + "query_result_limit": 999, + "inject_summary": false, + "workspace_id": "workspace-1", + "settings_revision": 9, + "language": "Français", + "extract_threshold": 1234, + "extract_worker_max_turns": 3, + "consolidation_threshold_files": 99, + }); + + let migrated = read_persisted_worker_manifest_snapshot(manifest).unwrap(); + assert!(migrated.feature.memory.profile.enabled); + assert!(migrated.feature.memory.profile.staging_tools); + assert!(!migrated.feature.memory.profile.resident.inject_summary); + assert_eq!( + migrated.feature.memory.profile.extraction.threshold, + Some(1234) + ); + assert_eq!( + migrated + .feature + .memory + .workspace_settings() + .unwrap() + .language, + "Français" + ); + let current = write_persisted_worker_manifest_snapshot(&migrated).unwrap(); + assert_eq!(current["schema_version"], 2); + assert!(current["manifest"].get("memory").is_none()); + } + + #[test] + fn persisted_manifest_adapter_rejects_mixed_or_future_authority() { + let manifest = + serde_json::to_value(WorkerManifest::from_toml(MINIMAL_REQUIRED).unwrap()).unwrap(); + let mut mixed = manifest.clone(); + mixed["feature"]["memory"] = serde_json::json!({ "enabled": true, "profile": {} }); + mixed["memory"] = serde_json::json!({}); + assert!(read_persisted_worker_manifest_snapshot(mixed).is_err()); + assert!( + read_persisted_worker_manifest_snapshot(serde_json::json!({ + "schema_version": 3, + "manifest": manifest, + })) + .is_err() ); } @@ -1291,14 +1532,6 @@ model_id = "claude-sonnet-4-20250514" )); } - #[test] - fn memory_section_with_language() { - let toml = format!("{MINIMAL_REQUIRED}\n[memory]\nlanguage = \"Japanese\"\n"); - let manifest = WorkerManifest::from_toml(&toml).unwrap(); - let mem = manifest.memory.unwrap(); - assert_eq!(mem.language.as_deref(), Some("Japanese")); - } - #[test] fn reject_unknown_scheme() { let toml = diff --git a/crates/manifest/src/profile.rs b/crates/manifest/src/profile.rs index 2c8e639f..40bc0e1a 100644 --- a/crates/manifest/src/profile.rs +++ b/crates/manifest/src/profile.rs @@ -20,9 +20,9 @@ use crate::config::{ use crate::model::{AuthRef, ModelManifest}; use crate::plugin::PluginConfig; use crate::{ - EngineManifestConfig, McpConfig, McpStdioCwdPolicy, MemoryConfig, Permission, ResolveError, - ScopeConfig, ScopeRule, SkillsConfig, WebConfig, WorkerManifest, WorkerManifestConfig, - WorkerMetaConfig, paths, + EngineManifestConfig, McpConfig, McpStdioCwdPolicy, Permission, ResolveError, ScopeConfig, + ScopeRule, SkillsConfig, WebConfig, WorkerManifest, WorkerManifestConfig, WorkerMetaConfig, + paths, }; const PROFILE_FORMAT_V1: &str = "yoi.profile.v1"; @@ -185,7 +185,7 @@ pub fn validate_profile_execution_target( if feature.manage_workdir.enabled { requirements.insert(WorkspaceAuthorityRequirement::ManageWorkdir); } - if feature.memory.enabled || feature.memory.staging { + if feature.memory.profile.enabled || feature.memory.profile.staging_tools { requirements.insert(WorkspaceAuthorityRequirement::Memory); } if feature.merge_request.show @@ -642,7 +642,6 @@ fn resolve_profile_value( mcp: profile.mcp, compaction, web: profile.web, - memory: profile.memory.map(Into::into), skills: profile.skills, }; let config = @@ -663,51 +662,6 @@ fn resolve_profile_value( }) } -#[derive(Debug, Default, Deserialize)] -#[serde(deny_unknown_fields)] -struct ProfileMemoryConfig { - #[serde(default)] - workspace_root: Option, - #[serde(default)] - query_result_limit: Option, - #[serde(default)] - query_excerpt_lines: Option, - #[serde(default)] - inject_summary: Option, - #[serde(default)] - extract_model: Option, - #[serde(default)] - extract_threshold: Option, - #[serde(default)] - extract_worker_max_turns: Option, - #[serde(default)] - consolidation_model: Option, - #[serde(default)] - consolidation_threshold_files: Option, - #[serde(default)] - consolidation_threshold_bytes: Option, -} - -impl From for MemoryConfig { - fn from(profile: ProfileMemoryConfig) -> Self { - Self { - workspace_root: profile.workspace_root, - query_result_limit: profile.query_result_limit, - query_excerpt_lines: profile.query_excerpt_lines, - inject_summary: profile.inject_summary, - workspace_id: None, - settings_revision: None, - language: None, - extract_model: profile.extract_model, - extract_threshold: profile.extract_threshold, - extract_worker_max_turns: profile.extract_worker_max_turns, - consolidation_model: profile.consolidation_model, - consolidation_threshold_files: profile.consolidation_threshold_files, - consolidation_threshold_bytes: profile.consolidation_threshold_bytes, - } - } -} - #[derive(Debug, Default, Deserialize)] #[serde(deny_unknown_fields)] struct ProfileConfig { @@ -738,8 +692,6 @@ struct ProfileConfig { #[serde(default)] web: Option, #[serde(default)] - memory: Option, - #[serde(default)] skills: Option, } @@ -940,12 +892,6 @@ fn validate_profile_paths(profile: &ProfileConfig) -> Result<(), ProfileError> { .map_err(|source| ProfileError::ProfileDeserialize { source })?; reject_absolute_auth_file(&model.auth, "compaction.model.auth.file")?; } - if let Some(memory) = &profile.memory - && let Some(root) = &memory.workspace_root - && root.is_absolute() - { - return Err(ProfileError::InvalidProfile("field `memory.workspace_root` is a resolved path and is not allowed in reusable Profiles".into())); - } if let Some(skills) = &profile.skills { for dir in &skills.directories { if dir.is_absolute() { @@ -1299,7 +1245,9 @@ mod tests { ("settings_revision", serde_json::json!(2)), ("language", serde_json::json!("Japanese")), ] { - let artifact = serde_json::json!({ "memory": { (field): value } }); + let artifact = serde_json::json!({ + "feature": { "memory": { (field): value } } + }); let error = resolve_profile_artifact_value( artifact, ProfileSource::Registry { @@ -1351,7 +1299,7 @@ mod tests { assert!(resolved.manifest.delegation_scope.allow.iter().any(|rule| { rule.permission == protocol::Permission::Write && rule.target == tmp.path() })); - assert!(!resolved.manifest.feature.memory.enabled); + assert!(!resolved.manifest.feature.memory.profile.enabled); assert!(!resolved.manifest.feature.ticket.enabled); assert!(!resolved.manifest.feature.objective.enabled); assert!(!resolved.manifest.feature.flow.enabled); @@ -1630,7 +1578,7 @@ enabled = false .unwrap(); assert_eq!(resolved.manifest.worker.name, "runtime-worker"); assert!(resolved.manifest.feature.task.enabled); - assert!(!resolved.manifest.feature.memory.enabled); + assert!(!resolved.manifest.feature.memory.profile.enabled); assert!(resolved.manifest.feature.web.enabled); assert!(resolved.manifest.feature.sub_worker.enabled); assert!(resolved.manifest.feature.ticket.enabled); diff --git a/crates/memory/src/backend.rs b/crates/memory/src/backend.rs index df8c548b..66a5b73a 100644 --- a/crates/memory/src/backend.rs +++ b/crates/memory/src/backend.rs @@ -152,13 +152,10 @@ pub enum MemoryStagingAffectedMemoryOperation { } #[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct MemoryConsolidateStagingOperation { #[serde(default)] pub force: bool, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub threshold_files: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub threshold_bytes: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -450,10 +447,21 @@ mod tests { use super::*; use crate::extract::{CandidateKind, ExtractedCandidate}; + #[test] + fn consolidation_operation_rejects_caller_owned_thresholds() { + let error = + serde_json::from_value::(serde_json::json!({ + "force": false, + "threshold_files": 1, + })) + .unwrap_err(); + assert!(error.to_string().contains("threshold_files")); + } + #[test] fn staging_list_read_close_records_reason_and_deletes_candidate() { let temp = tempfile::tempdir().unwrap(); - let layout = WorkspaceLayout::resolve(&manifest::MemoryConfig::default(), temp.path()); + let layout = WorkspaceLayout::resolve(temp.path()); let source = SourceRef { segment_id: "segment-1".into(), range: [0, 1], diff --git a/crates/memory/src/consolidate/staging.rs b/crates/memory/src/consolidate/staging.rs index 66ce4218..6449d63a 100644 --- a/crates/memory/src/consolidate/staging.rs +++ b/crates/memory/src/consolidate/staging.rs @@ -21,8 +21,7 @@ pub struct StagingEntry { pub id: Uuid, pub path: PathBuf, pub record: StagingRecord, - /// このファイルのバイト長。閾値判定 (`consolidation_threshold_bytes`) - /// に使う。 + /// このファイルのバイト長。Backendのconsolidation閾値判定に使用する。 pub bytes: u64, } diff --git a/crates/memory/src/workspace.rs b/crates/memory/src/workspace.rs index f4058994..3daa21b8 100644 --- a/crates/memory/src/workspace.rs +++ b/crates/memory/src/workspace.rs @@ -70,24 +70,12 @@ impl WorkspaceLayout { Self { root: root.into() } } - /// Resolve a layout from a `MemoryConfig`. + /// Resolve a layout from the nearest Memory marker. /// - /// An explicit `memory.workspace_root` is honored exactly. Without an - /// explicit root, resolution searches `default_root` and its ancestors for - /// the nearest `.yoi/memory` directory. This keeps child worktrees that - /// contain `.yoi` project records such as tickets from - /// becoming independent memory roots merely because they contain `.yoi`. - /// - /// If no memory marker exists, this falls back to `default_root` because - /// existing call sites require a concrete layout. That fallback is a - /// no-marker compatibility path, not a `.yoi` marker interpretation; it - /// must not be used as evidence that `.yoi` alone enables repo-local - /// memory. - pub fn resolve(cfg: &manifest::MemoryConfig, default_root: &Path) -> Self { - if let Some(root) = &cfg.workspace_root { - return Self::new(root.clone()); - } - + /// Resolution searches `default_root` and its ancestors for the nearest + /// `.yoi/memory` directory. This legacy local-storage helper owns its path + /// policy directly; resolved Worker Manifests do not carry storage paths. + pub fn resolve(default_root: &Path) -> Self { let root = find_memory_marker_root(default_root).unwrap_or_else(|| default_root.to_path_buf()); Self::new(root) @@ -335,16 +323,6 @@ mod tests { assert!(matches!(err, LintError::InvalidPath(_))); } - #[test] - fn resolve_uses_workspace_root_when_set() { - let cfg = manifest::MemoryConfig { - workspace_root: Some(PathBuf::from("/explicit")), - ..Default::default() - }; - let layout = WorkspaceLayout::resolve(&cfg, Path::new("/fallback")); - assert_eq!(layout.root(), Path::new("/explicit")); - } - #[test] fn resolve_selects_nearest_ancestor_memory_marker_when_workspace_root_missing() { let tmp = TempDir::new().unwrap(); @@ -353,8 +331,7 @@ mod tests { std::fs::create_dir_all(workspace.join(".yoi/memory")).unwrap(); std::fs::create_dir_all(&child).unwrap(); - let cfg = manifest::MemoryConfig::default(); - let layout = WorkspaceLayout::resolve(&cfg, &child); + let layout = WorkspaceLayout::resolve(&child); assert_eq!(layout.root(), workspace.as_path()); } @@ -366,8 +343,7 @@ mod tests { std::fs::create_dir_all(workspace.join(".yoi/memory")).unwrap(); std::fs::create_dir_all(child.join(".yoi/tickets")).unwrap(); - let cfg = manifest::MemoryConfig::default(); - let layout = WorkspaceLayout::resolve(&cfg, &child); + let layout = WorkspaceLayout::resolve(&child); assert_eq!(layout.root(), workspace.as_path()); } @@ -381,8 +357,7 @@ mod tests { assert_eq!(find_memory_marker_root(&child), None); - let cfg = manifest::MemoryConfig::default(); - let layout = WorkspaceLayout::resolve(&cfg, &child); + let layout = WorkspaceLayout::resolve(&child); assert_eq!(layout.root(), child.as_path()); } } diff --git a/crates/tui/src/setup_model.rs b/crates/tui/src/setup_model.rs index 951b434b..7bf733d2 100644 --- a/crates/tui/src/setup_model.rs +++ b/crates/tui/src/setup_model.rs @@ -228,7 +228,7 @@ worker_context_max_tokens = 100000 enabled = true [feature.memory] -enabled = true +enabled = false [feature.web] enabled = true @@ -241,11 +241,6 @@ enabled = true authoring = true thread = true -[memory] -extract_threshold = 50000 -consolidation_threshold_files = 5 -consolidation_threshold_bytes = 50000 - [web] enabled = true diff --git a/crates/worker-runtime/src/fs_store.rs b/crates/worker-runtime/src/fs_store.rs index 120fb558..fdb33c1d 100644 --- a/crates/worker-runtime/src/fs_store.rs +++ b/crates/worker-runtime/src/fs_store.rs @@ -759,8 +759,8 @@ fn migrate_worker_aggregate_document( .get_mut("resolved_manifest_snapshot") .filter(|snapshot| !snapshot.is_null()) { - let manifest: manifest::WorkerManifest = - serde_json::from_value(snapshot.clone()).map_err(|error| { + let mut manifest = manifest::read_persisted_worker_manifest_snapshot(snapshot.clone()) + .map_err(|error| { runtime_store_corrupt( metadata_path, format!("decode Worker aggregate resolved manifest snapshot: {error}"), @@ -775,20 +775,14 @@ fn migrate_worker_aggregate_document( ), )); } - snapshot - .as_object_mut() - .and_then(|manifest| manifest.get_mut("worker")) - .and_then(serde_json::Value::as_object_mut) - .ok_or_else(|| { + manifest.worker.name = expected_name.clone(); + *snapshot = + manifest::write_persisted_worker_manifest_snapshot(&manifest).map_err(|error| { runtime_store_corrupt( metadata_path, - "Worker aggregate resolved manifest is missing worker metadata".to_string(), + format!("encode migrated Worker aggregate resolved manifest: {error}"), ) - })? - .insert( - "name".to_string(), - serde_json::Value::String(expected_name.clone()), - ); + })?; } metadata.insert( "worker_name".to_string(), @@ -809,8 +803,8 @@ fn migrate_worker_aggregate_document( )); } if let Some(snapshot) = metadata.resolved_manifest_snapshot { - let manifest: manifest::WorkerManifest = - serde_json::from_value(snapshot).map_err(|error| { + let manifest = + manifest::read_persisted_worker_manifest_snapshot(snapshot).map_err(|error| { runtime_store_corrupt( metadata_path, format!("decode migrated Worker aggregate resolved manifest: {error}"), diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 15d34050..95d0c941 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -746,9 +746,15 @@ fn bind_workspace_memory_settings( )); } manifest + .feature .memory - .get_or_insert_with(manifest::MemoryConfig::default) - .bind_workspace_settings(snapshot); + .bind_workspace_settings(snapshot.clone()) + .map_err(str::to_string)?; + manifest + .feature + .memory + .validate_execution() + .map_err(str::to_string)?; Ok(()) } @@ -759,10 +765,18 @@ fn validate_worker_memory_settings( let Some(expected) = request.memory_settings.as_ref() else { return Ok(()); }; - let actual = manifest + manifest + .feature .memory - .as_ref() - .and_then(manifest::MemoryConfig::workspace_settings) + .validate_execution() + .map_err(str::to_string)?; + if !manifest.feature.memory.profile.enabled { + return Ok(()); + } + let actual = manifest + .feature + .memory + .workspace_settings() .ok_or_else(|| { "Workspace Worker restored without its bound Memory settings snapshot".to_string() })?; @@ -3165,7 +3179,7 @@ mod tests { Some(session_store::WorkerActiveSegmentRef::pending_segment( session_id, )), - Some(serde_json::to_value(&manifest).unwrap()), + Some(manifest::write_persisted_worker_manifest_snapshot(&manifest).unwrap()), ) .unwrap(); diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 8355f7ee..fb93f785 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -901,27 +901,39 @@ pub(crate) fn wire_event_bridges_on_engine( // per-item commit channel is wired at the top of this function. } -fn add_memory_lifecycle_if_configured( +fn add_memory_tools_if_configured( registry: &mut FeatureRegistryBuilder, - config: Option, - workspace_bound: bool, - build: impl FnOnce(manifest::MemoryConfig) -> std::io::Result, + config: &manifest::ResolvedMemoryFeatureConfig, + build: impl FnOnce() -> std::io::Result, ) -> std::io::Result where M: crate::feature::FeatureModule + 'static, { - let Some(config) = config else { - return Ok(false); - }; - if config.workspace_settings().is_none() { - if workspace_bound { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "Workspace-bound Memory requires a Backend-authored settings snapshot", - )); - } + if !config.profile.enabled { return Ok(false); } + config + .validate_execution() + .map_err(|message| std::io::Error::new(std::io::ErrorKind::InvalidInput, message))?; + registry.add_module(build()?); + Ok(true) +} + +fn add_memory_lifecycle_if_configured( + registry: &mut FeatureRegistryBuilder, + config: manifest::ResolvedMemoryFeatureConfig, + lifecycle_enabled: bool, + build: impl FnOnce(manifest::ResolvedMemoryFeatureConfig) -> std::io::Result, +) -> std::io::Result +where + M: crate::feature::FeatureModule + 'static, +{ + if !lifecycle_enabled || !config.profile.enabled || !config.profile.extraction.enabled { + return Ok(false); + } + config + .validate_execution() + .map_err(|message| std::io::Error::new(std::io::ErrorKind::InvalidInput, message))?; registry.add_module(build(config)?); Ok(true) } @@ -962,7 +974,7 @@ where let local_filesystem = worker.local_working_directory().cloned(); let local_workspace_root = local_filesystem.as_ref().map(|local| local.root.clone()); let task_feature = worker.task_feature(); - let memory_config = worker.manifest().memory.clone(); + let memory_config = feature_config.memory.clone(); let web_config = worker.manifest().web.clone(); let mcp_config = worker.manifest().mcp.clone(); let spawner_name = worker.manifest().worker.name.clone(); @@ -1019,13 +1031,23 @@ where let worker_enabled = feature_config.worker.enabled; let sub_worker_enabled = feature_config.sub_worker.enabled; let mut feature_registry = FeatureRegistryBuilder::new(); + add_memory_tools_if_configured(&mut feature_registry, &memory_config, || { + let workspace_client = worker.workspace_client_handle(); + if !workspace_client.is_available() || workspace_client.workspace_id().is_none() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Memory tools require Backend Workspace API authority", + )); + } + Ok(crate::feature::builtin::memory::MemoryToolsFeature::new( + workspace_client, + memory_config.profile.staging_tools, + )) + })?; add_memory_lifecycle_if_configured( &mut feature_registry, - worker - .manifest_lifecycle_features_enabled() - .then(|| memory_config.clone()) - .flatten(), - spawner_workspace_context.workspace_id().is_some(), + memory_config.clone(), + worker.manifest_lifecycle_features_enabled(), |config| { let workspace_client = worker.workspace_client_handle(); if !workspace_client.is_available() || workspace_client.workspace_id().is_none() { @@ -1202,38 +1224,6 @@ where } } - // Memory tools require explicit feature exposure. Workspace memory access - // is authority-bound to the Backend Workspace API; the Worker must not - // register local filesystem memory tools even when it has local cwd/root - // authority for shell/file tools. - if feature_config.memory.enabled { - let _mem = memory_config.as_ref().ok_or_else(|| { - std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "[feature.memory].enabled = true requires a [memory] configuration section", - ) - })?; - if workspace_client.is_available() && workspace_client.workspace_id().is_some() { - let definitions = if feature_config.memory.staging { - crate::feature::builtin::memory::workspace_http_memory_consolidation_tools( - workspace_client.clone(), - ) - } else { - crate::feature::builtin::memory::workspace_http_memory_tools( - workspace_client.clone(), - ) - }; - for definition in definitions { - engine.register_tool(definition); - } - } else { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "memory tools require Backend Workspace API authority", - )); - } - } - let mut observation_providers: Vec< Arc, > = Vec::new(); @@ -2169,7 +2159,7 @@ mod tests { use tokio::net::UnixListener; #[test] - fn memory_lifecycle_registration_requires_bound_workspace_memory_config() { + fn memory_feature_registration_requires_bound_workspace_memory_config() { #[derive(Clone)] struct TestMemoryLifecycleModule; @@ -2189,16 +2179,43 @@ mod tests { } } + let mut registry = FeatureRegistryBuilder::new(); + let installed = add_memory_tools_if_configured::( + &mut registry, + &manifest::ResolvedMemoryFeatureConfig::default(), + || panic!("disabled Memory must not construct its tools Feature"), + ) + .unwrap(); + assert!(!installed); + + let mut missing_snapshot = manifest::ResolvedMemoryFeatureConfig::default(); + missing_snapshot.profile.enabled = true; + let error = add_memory_tools_if_configured::( + &mut registry, + &missing_snapshot, + || panic!("invalid Memory config must fail before tools Feature construction"), + ) + .unwrap_err(); + assert!( + error + .to_string() + .contains("requires trusted Workspace settings") + ); + let mut registry = FeatureRegistryBuilder::new(); let configured = std::cell::Cell::new(false); - let mut memory_config = manifest::MemoryConfig::default(); - memory_config.bind_workspace_settings(&manifest::WorkspaceMemorySettingsSnapshot { - workspace_id: "workspace-1".to_string(), - settings_revision: 1, - language: "English".to_string(), - }); + let mut memory_config = manifest::ResolvedMemoryFeatureConfig::default(); + memory_config.profile.enabled = true; + memory_config.profile.extraction.enabled = true; + memory_config + .bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot { + workspace_id: "workspace-1".to_string(), + settings_revision: 1, + language: "English".to_string(), + }) + .unwrap(); let installed = - add_memory_lifecycle_if_configured(&mut registry, Some(memory_config), true, |_| { + add_memory_lifecycle_if_configured(&mut registry, memory_config, true, |_| { configured.set(true); Ok(TestMemoryLifecycleModule) }) @@ -2209,25 +2226,38 @@ mod tests { let mut registry = FeatureRegistryBuilder::new(); let installed = add_memory_lifecycle_if_configured::( &mut registry, - None, - false, + manifest::ResolvedMemoryFeatureConfig::default(), + true, |_| panic!("disabled Memory must not construct its lifecycle Feature"), ) .unwrap(); assert!(!installed); + let mut lifecycle_disabled = manifest::ResolvedMemoryFeatureConfig::default(); + lifecycle_disabled.profile.enabled = true; + lifecycle_disabled.profile.extraction.enabled = true; + lifecycle_disabled + .bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot { + workspace_id: "workspace-1".to_string(), + settings_revision: 1, + language: "English".to_string(), + }) + .unwrap(); let installed = add_memory_lifecycle_if_configured::( &mut registry, - Some(manifest::MemoryConfig::default()), + lifecycle_disabled, false, - |_| panic!("Memory without a Backend-authored settings snapshot must stay disabled"), + |_| panic!("disabled lifecycle must not construct its Feature"), ) .unwrap(); assert!(!installed); + let mut missing_snapshot = manifest::ResolvedMemoryFeatureConfig::default(); + missing_snapshot.profile.enabled = true; + missing_snapshot.profile.extraction.enabled = true; let error = add_memory_lifecycle_if_configured::( &mut registry, - Some(manifest::MemoryConfig::default()), + missing_snapshot, true, |_| panic!("invalid Workspace Memory config must fail before Feature construction"), ) @@ -2235,7 +2265,7 @@ mod tests { assert!( error .to_string() - .contains("Backend-authored settings snapshot") + .contains("requires trusted Workspace settings") ); } diff --git a/crates/worker/src/feature/builtin/memory.rs b/crates/worker/src/feature/builtin/memory.rs index 595eebc1..62403d1a 100644 --- a/crates/worker/src/feature/builtin/memory.rs +++ b/crates/worker/src/feature/builtin/memory.rs @@ -18,6 +18,10 @@ use schemars::JsonSchema; use serde::de::DeserializeOwned; use serde_json::json; +use crate::feature::{ + FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureModule, ToolContribution, + ToolDeclaration, +}; use crate::worker::{ WorkspaceClient, WorkspaceClientError, WorkspaceRequest, WorkspaceRequestMethod, }; @@ -338,6 +342,44 @@ fn query_schema() -> serde_json::Value { }) } +#[derive(Clone)] +pub struct MemoryToolsFeature { + tools: Vec, +} + +impl MemoryToolsFeature { + pub fn new(client: Arc, staging_tools: bool) -> Self { + let tools = if staging_tools { + workspace_http_memory_consolidation_tools(client) + } else { + workspace_http_memory_tools(client) + }; + Self { tools } + } +} + +impl FeatureModule for MemoryToolsFeature { + fn descriptor(&self) -> FeatureDescriptor { + let mut descriptor = FeatureDescriptor::builtin("memory", "Memory") + .with_description("Workspace Memory document, query, and staging tools."); + for tool in &self.tools { + let (meta, _) = tool(); + descriptor = descriptor.with_tool(ToolDeclaration::new(meta.name, meta.description)); + } + descriptor + } + + fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> { + for tool in &self.tools { + let (meta, _) = tool(); + context + .tools() + .register(ToolContribution::new(meta.name, tool.clone()))?; + } + Ok(()) + } +} + #[cfg(test)] mod tests { use super::*; @@ -368,6 +410,20 @@ mod tests { .input_schema } + #[test] + fn memory_feature_owns_normal_and_staging_tool_surfaces() { + let normal = MemoryToolsFeature::new(test_client(), false); + let normal_names = tool_names(normal.tools); + assert!(normal_names.contains(&"MemoryQuery".to_string())); + assert!(!normal_names.contains(&"MemoryStagingList".to_string())); + + let staging = MemoryToolsFeature::new(test_client(), true); + assert_eq!(staging.descriptor().id.as_str(), "builtin:memory"); + let staging_names = tool_names(staging.tools); + assert!(staging_names.contains(&"MemoryQuery".to_string())); + assert!(staging_names.contains(&"MemoryStagingList".to_string())); + } + #[test] fn normal_workspace_memory_tools_do_not_include_staging_tools() { let names = tool_names(workspace_http_memory_tools(test_client())); diff --git a/crates/worker/src/feature/builtin/memory_lifecycle.rs b/crates/worker/src/feature/builtin/memory_lifecycle.rs index b7083428..9c601263 100644 --- a/crates/worker/src/feature/builtin/memory_lifecycle.rs +++ b/crates/worker/src/feature/builtin/memory_lifecycle.rs @@ -52,7 +52,7 @@ pub(crate) struct MemoryLifecycleFeature { #[derive(Clone)] struct MemoryLifecycleTask { - config: manifest::MemoryConfig, + config: manifest::ResolvedMemoryFeatureConfig, capture: CommittedSessionCaptureHandle, extensions: SessionExtensionHandle, workspace_client: Arc, @@ -66,7 +66,7 @@ struct MemoryLifecycleTask { impl MemoryLifecycleFeature { #[allow(clippy::too_many_arguments)] pub(crate) fn new( - config: manifest::MemoryConfig, + config: manifest::ResolvedMemoryFeatureConfig, capture: CommittedSessionCaptureHandle, extensions: SessionExtensionHandle, workspace_client: Arc, @@ -132,7 +132,9 @@ impl MemoryLifecycleTask { memory::audit::AuditWorker::MemoryExtract, memory::audit::AuditTrigger::TokenThreshold, self.config - .extract_model + .profile + .extraction + .model .as_ref() .or(Some(&self.manifest.model)) .map(model_audit_from_manifest), @@ -188,7 +190,9 @@ impl MemoryLifecycleTask { }; let Some(threshold) = self .config - .extract_threshold + .profile + .extraction + .threshold .filter(|threshold| *threshold > 0) else { audit @@ -283,7 +287,7 @@ impl MemoryLifecycleTask { source, audit.run_id.to_string(), ); - let client = if let Some(model) = self.config.extract_model.as_ref() { + let client = if let Some(model) = self.config.profile.extraction.model.as_ref() { match crate::model_client::build_client(model) { Ok(client) => client, Err(error) => { @@ -321,7 +325,7 @@ impl MemoryLifecycleTask { } }; let mut manifest = self.manifest.clone(); - if let Some(model) = self.config.extract_model.clone() { + if let Some(model) = self.config.profile.extraction.model.clone() { manifest.model = model; } @@ -349,7 +353,9 @@ impl MemoryLifecycleTask { cache_key: Some(capture.segment_id.clone()), max_turns: self .config - .extract_worker_max_turns + .profile + .extraction + .worker_max_turns .or(manifest::defaults::MEMORY_EXTRACT_WORKER_MAX_TURNS), engine_configurator: None, features, @@ -493,36 +499,13 @@ impl MemoryLifecycleTask { let audit = WorkerAuditBase::new( memory::audit::AuditWorker::MemoryConsolidation, memory::audit::AuditTrigger::StagingBacklog, - self.config - .consolidation_model - .as_ref() - .or(Some(&self.manifest.model)) - .map(model_audit_from_manifest), + Some(model_audit_from_manifest(&self.manifest.model)), ) .with_memory_settings(&self.config); - let Some((threshold_files, threshold_bytes)) = consolidation_thresholds(&self.config) - else { - audit - .emit( - self.workspace_client.as_ref(), - self.event_tx.as_ref(), - memory::audit::WorkerLifecycleStatus::Skipped, - "consolidation_threshold_disabled", - None, - None, - None, - ) - .await; - return; - }; match self .workspace_client .request_memory_staging_consolidation( - memory::backend::MemoryConsolidateStagingOperation { - force: false, - threshold_files, - threshold_bytes, - }, + memory::backend::MemoryConsolidateStagingOperation { force: false }, ) .await { @@ -646,22 +629,6 @@ fn extract_pointer( Ok(pointer) } -fn consolidation_thresholds( - config: &manifest::MemoryConfig, -) -> Option<(Option, Option)> { - let threshold_files = config - .consolidation_threshold_files - .filter(|threshold| *threshold > 0); - let threshold_bytes = config - .consolidation_threshold_bytes - .filter(|threshold| *threshold > 0); - if threshold_files.is_none() && threshold_bytes.is_none() { - None - } else { - Some((threshold_files, threshold_bytes)) - } -} - fn extraction_run_eligible(exit: CommittedRunExit) -> bool { exit == CommittedRunExit::Finished } @@ -688,12 +655,17 @@ fn tokens_since_pointer( fn extraction_threshold_reached( capture: &CommittedSessionCapture, pointer: Option<&memory::ExtractPointerPayload>, - config: &manifest::MemoryConfig, + config: &manifest::ResolvedMemoryFeatureConfig, ) -> bool { if capture.history.is_empty() { return false; } - let Some(threshold) = config.extract_threshold.filter(|threshold| *threshold > 0) else { + let Some(threshold) = config + .profile + .extraction + .threshold + .filter(|threshold| *threshold > 0) + else { return false; }; tokens_since_pointer(capture, pointer) >= threshold @@ -723,7 +695,7 @@ impl WorkerAuditBase { } } - fn with_memory_settings(mut self, config: &manifest::MemoryConfig) -> Self { + fn with_memory_settings(mut self, config: &manifest::ResolvedMemoryFeatureConfig) -> Self { self.memory_settings = config .workspace_settings() @@ -1014,16 +986,18 @@ permission = "write" .unwrap() } - fn test_config() -> manifest::MemoryConfig { - let mut config = manifest::MemoryConfig { - extract_threshold: Some(1), - ..Default::default() - }; - config.bind_workspace_settings(&manifest::WorkspaceMemorySettingsSnapshot { - workspace_id: "workspace-1".to_string(), - settings_revision: 1, - language: "English".to_string(), - }); + fn test_config() -> manifest::ResolvedMemoryFeatureConfig { + let mut config = manifest::ResolvedMemoryFeatureConfig::default(); + config.profile.enabled = true; + config.profile.extraction.enabled = true; + config.profile.extraction.threshold = Some(1); + config + .bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot { + workspace_id: "workspace-1".to_string(), + settings_revision: 1, + language: "English".to_string(), + }) + .unwrap(); config } @@ -1282,30 +1256,30 @@ permission = "write" } #[tokio::test] - async fn lifecycle_task_requests_backend_consolidation_from_configured_threshold() { + async fn lifecycle_task_requests_backend_owned_consolidation_eligibility() { let client = ScriptClient::new(Vec::new()); let extension_writes = Arc::new(Mutex::new(Vec::new())); let (event_tx, _) = broadcast::channel(16); let workspace_client = Arc::new(RecordingWorkspaceClient::default()); let mut interrupted = capture(2, 250); interrupted.run_exit = CommittedRunExit::Interrupted; - let mut task = test_task( + let task = test_task( interrupted, Box::new(client), extension_writes, event_tx, workspace_client.clone(), ); - task.config.consolidation_threshold_files = Some(3); run_background_task(task).await; let requests = workspace_client.requests.lock().unwrap(); assert!( requests.iter().any(|request| { request.path.contains("memory") - && request.body.as_deref().is_some_and(|body| { - body.contains("\"threshold_files\":3") && body.contains("\"force\":false") - }) + && request + .body + .as_deref() + .is_some_and(|body| body == "{\"force\":false}") }), "recorded requests: {requests:?}" ); @@ -1349,17 +1323,6 @@ permission = "write" } } - #[test] - fn consolidation_thresholds_enable_backend_request_on_either_limit() { - let mut config = manifest::MemoryConfig::default(); - assert_eq!(consolidation_thresholds(&config), None); - config.consolidation_threshold_files = Some(3); - assert_eq!(consolidation_thresholds(&config), Some((Some(3), None))); - config.consolidation_threshold_files = None; - config.consolidation_threshold_bytes = Some(4096); - assert_eq!(consolidation_thresholds(&config), Some((None, Some(4096)))); - } - #[test] fn interrupted_parent_run_is_not_extraction_eligible() { assert!(extraction_run_eligible(CommittedRunExit::Finished)); @@ -1412,8 +1375,8 @@ permission = "write" #[test] fn threshold_uses_committed_usage_after_pointer() { let capture = capture(2, 250); - let mut config = manifest::MemoryConfig::default(); - config.extract_threshold = Some(1); + let mut config = manifest::ResolvedMemoryFeatureConfig::default(); + config.profile.extraction.threshold = Some(1); assert!(extraction_threshold_reached( &capture, Some(&memory::ExtractPointerPayload { @@ -1500,8 +1463,8 @@ permission = "write" #[test] fn empty_capture_never_schedules_extraction() { let capture = capture(0, 500); - let mut config = manifest::MemoryConfig::default(); - config.extract_threshold = Some(1); + let mut config = manifest::ResolvedMemoryFeatureConfig::default(); + config.profile.extraction.threshold = Some(1); assert!(!extraction_threshold_reached(&capture, None, &config)); } } diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index 8c3ee231..c2498874 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -425,6 +425,8 @@ impl Tool for SubWorkerSpawnTool { WorkerManifestConfig::resolution_defaults().merge(child_config), ) .map_err(|error| ToolError::ExecutionFailed(format!("resolve child manifest: {error}")))?; + bind_child_memory_settings(&self.spawner_manifest, &mut child_manifest) + .map_err(ToolError::ExecutionFailed)?; // Delegated children stay bound to their scoped session and cannot use // Workspace attachment tools to replace it with parent-level authority. child_manifest.feature.manage_workdir.enabled = false; @@ -827,6 +829,33 @@ fn profile_error_with_available(error: ProfileError, available: &AvailableProfil ) } +fn bind_child_memory_settings( + parent: &manifest::WorkerManifest, + child: &mut manifest::WorkerManifest, +) -> Result<(), String> { + if !child.feature.memory.profile.enabled { + return child + .feature + .memory + .validate_execution() + .map_err(str::to_string); + } + let workspace_settings = parent.feature.memory.workspace_settings().ok_or_else(|| { + "enabled child Memory feature requires the parent's trusted Workspace settings snapshot" + .to_string() + })?; + child + .feature + .memory + .bind_workspace_settings(workspace_settings) + .map_err(str::to_string)?; + child + .feature + .memory + .validate_execution() + .map_err(str::to_string) +} + fn manifest_to_reusable_config(manifest: &WorkerManifest) -> WorkerManifestConfig { WorkerManifestConfig { worker: WorkerMetaConfig { @@ -894,7 +923,6 @@ fn manifest_to_reusable_config(manifest: &WorkerManifest) -> WorkerManifestConfi model: c.model.clone(), }), web: manifest.web.clone(), - memory: manifest.memory.clone(), skills: manifest.skills.clone(), } } @@ -1091,10 +1119,7 @@ enabled = true thread = true [feature.memory] -enabled = true - -[memory] -extract_threshold = 4000 +enabled = false "#; #[tokio::test] @@ -1526,6 +1551,33 @@ extract_threshold = 4000 .unwrap() } + #[test] + fn child_memory_inherits_only_the_parents_trusted_settings_snapshot() { + let temp = tempfile::tempdir().unwrap(); + let mut parent = parent_manifest(temp.path(), None); + parent.feature.memory.profile.enabled = true; + parent + .feature + .memory + .bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot { + workspace_id: "workspace-1".to_string(), + settings_revision: 4, + language: "日本語".to_string(), + }) + .unwrap(); + let mut child = parent.clone(); + child.feature.memory.workspace_settings = None; + + bind_child_memory_settings(&parent, &mut child).unwrap(); + assert_eq!( + child.feature.memory.workspace_settings(), + parent.feature.memory.workspace_settings() + ); + + child.feature.memory.profile.enabled = false; + assert!(bind_child_memory_settings(&parent, &mut child).is_err()); + } + fn write_project_profile_registry( project: &Path, default: Option<&str>, diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 483f508a..76888eff 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -2478,11 +2478,10 @@ impl Worker { ) }); if is_memory_consolidation { - let memory_config = self.manifest.memory.as_ref().ok_or_else(|| { - WorkerError::InvalidState( - "Memory consolidation Worker has no Memory configuration".to_string(), - ) - })?; + let memory_config = &self.manifest.feature.memory; + memory_config + .validate_execution() + .map_err(|message| WorkerError::InvalidState(message.to_string()))?; let language = memory_language(memory_config)?; let rendered = self .prompts @@ -2517,11 +2516,8 @@ impl Worker { } } let inject_summary = self.inject_resident_summary - && self - .manifest - .memory - .as_ref() - .is_some_and(|m| m.inject_summary.unwrap_or(true)); + && self.manifest.feature.memory.profile.enabled + && self.manifest.feature.memory.profile.resident.inject_summary; let resident_summary: Option = if inject_summary { match self.resident_summary_from_workspace_authority().await { Ok(summary) => summary, @@ -4552,7 +4548,7 @@ impl Worker { } } -fn memory_language(config: &manifest::MemoryConfig) -> Result { +fn memory_language(config: &manifest::ResolvedMemoryFeatureConfig) -> Result { config .workspace_settings() .map(|snapshot| snapshot.language) @@ -5426,7 +5422,8 @@ fn worker_metadata_for_manifest( metadata = metadata.with_workspace_root(local_workspace_root.to_path_buf()); } if should_persist_resolved_manifest_snapshot(manifest) { - metadata.resolved_manifest_snapshot = serde_json::to_value(manifest).ok(); + metadata.resolved_manifest_snapshot = + manifest::write_persisted_worker_manifest_snapshot(manifest).ok(); } metadata } @@ -5439,10 +5436,20 @@ fn validate_workspace_memory_snapshot( let Some(workspace_id) = workspace_context.workspace_id() else { return Ok(()); }; - let snapshot = manifest + manifest + .feature .memory - .as_ref() - .and_then(manifest::MemoryConfig::workspace_settings) + .validate_execution() + .map_err(|message| { + WorkerError::InvalidState(format!("Workspace Worker {worker_name}: {message}")) + })?; + if !manifest.feature.memory.profile.enabled { + return Ok(()); + } + let snapshot = manifest + .feature + .memory + .workspace_settings() .ok_or_else(|| { WorkerError::InvalidState(format!( "Workspace Worker {worker_name} has no complete persisted Memory settings snapshot" @@ -5468,11 +5475,7 @@ fn validate_workspace_memory_snapshot( fn should_persist_resolved_manifest_snapshot(manifest: &WorkerManifest) -> bool { manifest.profile.is_some() || manifest.plugins.has_resolved_plan() - || manifest - .memory - .as_ref() - .and_then(manifest::MemoryConfig::workspace_settings) - .is_some() + || manifest.feature.memory.workspace_settings.is_some() } fn restore_manifest_from_worker_metadata_snapshot( @@ -5481,12 +5484,14 @@ fn restore_manifest_from_worker_metadata_snapshot( fallback: WorkerManifest, ) -> Result { match snapshot { - Some(snapshot) => serde_json::from_value(snapshot).map_err(|source| { - WorkerError::WorkerMetadataManifestSnapshot { - worker_name: worker_name.to_string(), - source, - } - }), + Some(snapshot) => { + manifest::read_persisted_worker_manifest_snapshot(snapshot).map_err(|source| { + WorkerError::WorkerMetadataManifestSnapshot { + worker_name: worker_name.to_string(), + source, + } + }) + } None => Ok(fallback), } } @@ -6198,11 +6203,6 @@ fn prepare_worker_common_with_context_and_model_client( WorkerFilesystemAuthority::Local(LocalWorkingDirectory { root, cwd }) } }; - let mut scope_config = scope_config; - if let (Some(mem), Some(local)) = (manifest.memory.as_ref(), filesystem_authority.as_local()) { - let layout = memory::WorkspaceLayout::resolve(mem, &local.root); - scope_config.deny.extend(memory::deny_write_rules(&layout)); - } let scope = if scope_config.allow.is_empty() && filesystem_authority.as_local().is_none() { Scope::empty() } else { @@ -6292,8 +6292,7 @@ mod spawned_context_tests { std::fs::create_dir_all(&workspace_root).unwrap(); std::fs::create_dir_all(&cwd).unwrap(); - let mut manifest = minimal_manifest_for_context_test(&workspace_root, &cwd); - manifest.memory = Some(manifest::MemoryConfig::default()); + let manifest = minimal_manifest_for_context_test(&workspace_root, &cwd); let common = prepare_worker_common_with_context( &manifest, &PromptCatalogSource::builtins_only(), @@ -6327,8 +6326,7 @@ mod spawned_context_tests { let workspace_root = tmp.path().join("workspace-root"); let cwd = workspace_root.join("nested"); std::fs::create_dir_all(&cwd).unwrap(); - let mut manifest = minimal_manifest_for_context_test(&workspace_root, &cwd); - manifest.memory = Some(manifest::MemoryConfig::default()); + let manifest = minimal_manifest_for_context_test(&workspace_root, &cwd); let loader = PromptCatalogSource::builtins_only(); let workspace_id = WorkspaceId::new("ws-api-only").unwrap(); let common = prepare_worker_common_with_context( @@ -6535,7 +6533,7 @@ permission = "write" let restored = restore_manifest_from_worker_metadata_snapshot( "restore-scope", - Some(serde_json::to_value(&saved).unwrap()), + Some(manifest::write_persisted_worker_manifest_snapshot(&saved).unwrap()), current, ) .unwrap(); @@ -6590,24 +6588,26 @@ permission = "read" "#, ) .unwrap(); - manifest.memory = Some(manifest::MemoryConfig::default()); - manifest.memory.as_mut().unwrap().bind_workspace_settings( - &manifest::WorkspaceMemorySettingsSnapshot { + manifest.feature.memory.profile.enabled = true; + manifest + .feature + .memory + .bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot { workspace_id: "workspace-a".to_string(), settings_revision: 7, language: "Japanese".to_string(), - }, - ); + }) + .unwrap(); let metadata = worker_metadata_for_manifest(&manifest, None, None, None); - let restored: WorkerManifest = serde_json::from_value( + let restored = manifest::read_persisted_worker_manifest_snapshot( metadata .resolved_manifest_snapshot .expect("Memory settings require a resolved manifest snapshot"), ) .unwrap(); assert_eq!( - restored.memory.unwrap().workspace_settings(), + restored.feature.memory.workspace_settings(), Some(manifest::WorkspaceMemorySettingsSnapshot { workspace_id: "workspace-a".to_string(), settings_revision: 7, @@ -6638,7 +6638,7 @@ permission = "read" ); let mut missing = manifest.clone(); - missing.memory.as_mut().unwrap().settings_revision = None; + missing.feature.memory.workspace_settings = None; assert!( validate_workspace_memory_snapshot( "memory-snapshot", @@ -6715,7 +6715,7 @@ permission = "read" let snapshot = metadata .resolved_manifest_snapshot .expect("plugin-resolved manifest should be snapshotted"); - let restored: WorkerManifest = serde_json::from_value(snapshot).unwrap(); + let restored = manifest::read_persisted_worker_manifest_snapshot(snapshot).unwrap(); assert!(restored.profile.is_none()); assert_eq!(restored.plugins.resolved.len(), 1); @@ -8203,13 +8203,16 @@ mod build_summary_prompt_tests { }, profile: None, }); - let mut memory = manifest::MemoryConfig::default(); - memory.bind_workspace_settings(&manifest::WorkspaceMemorySettingsSnapshot { - workspace_id: "workspace-test".to_string(), - settings_revision: 3, - language: "Japanese".to_string(), - }); - manifest.memory = Some(memory); + let mut memory = manifest::ResolvedMemoryFeatureConfig::default(); + memory.profile.enabled = true; + memory + .bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot { + workspace_id: "workspace-test".to_string(), + settings_revision: 3, + language: "Japanese".to_string(), + }) + .unwrap(); + manifest.feature.memory = memory; let mut worker = Worker::new( manifest, Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), @@ -8235,7 +8238,7 @@ mod build_summary_prompt_tests { async fn render_system_prompt_with_summary( summary_doc: Option<&str>, - memory_config: Option, + memory_config: Option, resident_injection: bool, ) -> String { render_system_prompt_with_resident_sections( @@ -8249,7 +8252,7 @@ mod build_summary_prompt_tests { async fn render_system_prompt_with_resident_sections( summary_doc: Option<&str>, - memory_config: Option, + memory_config: Option, gates: ResidentInjectionGates, _unused: bool, ) -> String { @@ -8258,12 +8261,15 @@ mod build_summary_prompt_tests { let cwd = dir.path().join("workspace"); std::fs::create_dir_all(&cwd).unwrap(); let mut manifest = minimal_manifest(); - manifest.memory = memory_config.clone(); + manifest.feature.memory = memory_config.clone().unwrap_or_default(); + if memory_config.is_some() { + manifest.feature.memory.profile.enabled = true; + } let scope = Scope::writable(&cwd).unwrap(); let authority = WorkerFilesystemAuthority::local(cwd.clone(), cwd.clone()); let workspace_context = if memory_config .as_ref() - .is_some_and(|cfg| cfg.inject_summary.unwrap_or(true)) + .is_some_and(|cfg| cfg.profile.resident.inject_summary) && gates.summary { stub_memory_backend_context(summary_doc.and_then(summary_content_for_backend)) @@ -8350,7 +8356,7 @@ mod build_summary_prompt_tests { async fn resident_summary_body_is_injected_without_frontmatter() { let rendered = render_system_prompt_with_summary( Some(&summary_doc("summary body for resident prompt\n")), - Some(manifest::MemoryConfig::default()), + Some(manifest::ResolvedMemoryFeatureConfig::default()), true, ) .await; @@ -8362,10 +8368,8 @@ mod build_summary_prompt_tests { #[tokio::test] async fn resident_summary_injection_can_be_disabled_by_manifest() { - let memory = manifest::MemoryConfig { - inject_summary: Some(false), - ..manifest::MemoryConfig::default() - }; + let mut memory = manifest::ResolvedMemoryFeatureConfig::default(); + memory.profile.resident.inject_summary = false; let rendered = render_system_prompt_with_summary( Some(&summary_doc("disabled summary body\n")), Some(memory), @@ -8377,7 +8381,7 @@ mod build_summary_prompt_tests { } #[tokio::test] - async fn resident_summary_is_absent_without_memory_config() { + async fn resident_summary_is_absent_when_memory_feature_is_disabled() { let rendered = render_system_prompt_with_summary( Some(&summary_doc("memory-disabled summary body\n")), None, @@ -8392,7 +8396,7 @@ mod build_summary_prompt_tests { async fn malformed_resident_summary_does_not_fail_render() { let rendered = render_system_prompt_with_summary( Some("---\nthis is not yaml: : :\n---\nbad summary body\n"), - Some(manifest::MemoryConfig::default()), + Some(manifest::ResolvedMemoryFeatureConfig::default()), true, ) .await; @@ -8405,7 +8409,7 @@ mod build_summary_prompt_tests { async fn resident_summary_gate_false_omits_only_summary() { let prompt = render_system_prompt_with_resident_sections( Some(&summary_doc("resident summary marker")), - Some(manifest::MemoryConfig::default()), + Some(manifest::ResolvedMemoryFeatureConfig::default()), ResidentInjectionGates { summary: false }, true, ) diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 584cb96c..3599f26d 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -7993,17 +7993,15 @@ fn start_memory_staging_consolidation( total_bytes, }); } - let reached_files = operation - .threshold_files - .is_some_and(|threshold| candidate_count >= threshold); - let reached_bytes = operation - .threshold_bytes - .is_some_and(|threshold| total_bytes >= threshold); + const CONSOLIDATION_THRESHOLD_FILES: usize = 5; + const CONSOLIDATION_THRESHOLD_BYTES: u64 = 50_000; + let reached_files = candidate_count >= CONSOLIDATION_THRESHOLD_FILES; + let reached_bytes = total_bytes >= CONSOLIDATION_THRESHOLD_BYTES; if !operation.force && !reached_files && !reached_bytes { return Ok(MemoryConsolidationOutput { status: "skipped_below_threshold".to_string(), summary: format!( - "Memory staging backlog has {candidate_count} candidate(s), {total_bytes} byte(s), below configured threshold." + "Memory staging backlog has {candidate_count} candidate(s), {total_bytes} byte(s), below Backend policy threshold." ), candidate_count, total_bytes, @@ -19921,11 +19919,7 @@ mod tests { let output = match start_memory_staging_consolidation( api, - MemoryConsolidateStagingOperation { - force: true, - threshold_files: None, - threshold_bytes: None, - }, + MemoryConsolidateStagingOperation { force: true }, ) { Ok(output) => output, Err(_) => panic!("unexpected ApiError from memory consolidation trigger"), @@ -19956,6 +19950,15 @@ mod tests { ) .unwrap(); + let below_threshold = start_memory_staging_consolidation( + api.clone(), + MemoryConsolidateStagingOperation { force: false }, + ) + .unwrap(); + assert_eq!(below_threshold.status, "skipped_below_threshold"); + assert_eq!(below_threshold.candidate_count, 1); + assert!(below_threshold.summary.contains("Backend policy threshold")); + let resolved_config_bundle = None; let existing = api .runtime @@ -20010,11 +20013,7 @@ mod tests { let second = match start_memory_staging_consolidation( api.clone(), - MemoryConsolidateStagingOperation { - force: true, - threshold_files: None, - threshold_bytes: None, - }, + MemoryConsolidateStagingOperation { force: true }, ) { Ok(output) => output, Err(_) => panic!("unexpected ApiError from second memory consolidation trigger"), diff --git a/docs/manifest.toml b/docs/manifest.toml index 28ce2e2f..2dc2386e 100644 --- a/docs/manifest.toml +++ b/docs/manifest.toml @@ -222,55 +222,38 @@ permission = "write" # # ref = "anthropic/claude-haiku-4-5" -# ===== [memory] ============================================================= -# Memory subsystem の opt-in。 -# - セクションが *ある* … memory tools (MemoryRead/Write/Edit) を登録、 -# `/memory/` と `/` -# の通常 write を Worker 自体に対して deny する。 -# - セクションが *無い* … 何も起きない (legacy 動作)。 -# `[memory]` だけ書いて中身を省略するのも有効 (全フィールド既定値で有効化)。 -# [memory] +# ===== [feature.memory] ====================================================== +# Memory は `feature.memory` だけを入口にする。resolved Worker Manifest では +# Profile由来の設定を `profile` に、Backend由来のWorkspace設定snapshotを +# `workspace_settings` に分離して保存する。`workspace_settings` はBackendだけが +# bindする信頼済み入力で、Profile・Browser・model入力から指定できない。 +# `profile.enabled = false` の場合、Memory tools、resident injection、extract、 +# consolidation requestをすべて無効にし、snapshotも保持しない。 # -# # 任意。デフォルト: Worker の pwd (構築時)。 -# # 必ず絶対パス (相対なら manifest base 起点で resolve)。 -# workspace_root = "/abs/path/to/workspace" +# [feature.memory.profile] +# enabled = true +# staging_tools = false # -# # 任意。デフォルト: tool 側既定 = 20。 -# # MemoryQuery / MemoryQuery が 1 回に返す最大件数。 -# query_result_limit = 20 +# [feature.memory.profile.resident] +# inject_summary = true # -# # 任意。デフォルト: tool 側既定 = 3。 -# # 各マッチ前後に表示するコンテキスト行数。`query` 省略時は無視。 -# query_excerpt_lines = 3 +# [feature.memory.profile.extraction] +# enabled = true +# threshold = 30000 +# worker_max_turns = 8 # -# # 任意。デフォルト: メインモデルを `clone_boxed()` で複製。 -# # extract ワーカーのモデル ([model] と同じ形式)。 -# # Haiku / 4o-mini / Flash クラスの軽量 reasoning モデル推奨。 -# # [memory.extract_model] +# # 任意。省略時はmain modelをcloneする。 +# # [feature.memory.profile.extraction.model] # # ref = "anthropic/claude-haiku-4-5" # -# # 任意。デフォルト: なし (extract 自動発火を完全停止)。 -# # 前回 extract pointer 以降の累積入力 token がこの値を超えると extract 起動。 -# # ※ memory tools と resident injection は extract_threshold が None でも動く。 -# extract_threshold = 30000 +# # Backendがresolved Manifestへbindする。手書き/Profile入力では指定しない。 +# # [feature.memory.workspace_settings] +# # workspace_id = "workspace-id" +# # settings_revision = 1 +# # language = "日本語" # -# # 任意。デフォルト: 8 (`defaults::MEMORY_EXTRACT_WORKER_MAX_TURNS`)。 -# # extract worker 自身の tool loop 上限。Rust config で None の場合のみ無制限。 -# extract_worker_max_turns = 8 -# -# # 任意。デフォルト: メインモデルを `clone_boxed()` で複製。 -# # consolidation ワーカーのモデル。reasoning クラス推奨。 -# # [memory.consolidation_model] -# # ref = "anthropic/claude-sonnet-4-6" -# -# # 任意。デフォルト: なし。 -# # `_staging/` のエントリ数がこの値以上で consolidation 発火 (files / bytes は OR)。 -# consolidation_threshold_files = 50 -# -# # 任意。デフォルト: なし。 -# # `_staging/` の総バイト数がこの値以上で consolidation 発火 (files / bytes は OR)。 -# # files / bytes の両方が None だと consolidation 完全無効。 -# consolidation_threshold_bytes = 1048576 +# Query結果/抜粋の上限とconsolidation eligibility/thresholdはoperation/Backend +# policyが所有し、通常Worker Manifestには含めない。legacy `[memory]` は拒否する。 # ===== [skills] ============================================================= diff --git a/resources/profiles/base.dcdl b/resources/profiles/base.dcdl index c91afed8..88db46c7 100644 --- a/resources/profiles/base.dcdl +++ b/resources/profiles/base.dcdl @@ -23,7 +23,14 @@ compaction = { feature = { task = { enabled = true; }; - memory = { enabled = true; }; + memory = { + enabled = true; + resident = { inject_summary = true; }; + extraction = { + enabled = true; + threshold = 50000; + }; + }; web = { enabled = true; }; image = { enabled = true; }; sub_worker = { enabled = false; }; @@ -40,12 +47,6 @@ feature = { }; }; -memory = { - extract_threshold = 50000; - consolidation_threshold_files = 5; - consolidation_threshold_bytes = 50000; -}; - web = { enabled = true; search = { diff --git a/resources/profiles/default.dcdl b/resources/profiles/default.dcdl index 3d3721c2..84f0f3c5 100644 --- a/resources/profiles/default.dcdl +++ b/resources/profiles/default.dcdl @@ -6,7 +6,7 @@ import "./base.dcdl" // { feature = { task = { enabled = true; }; - memory = { enabled = false; staging = false; }; + memory = { enabled = false; staging_tools = false; }; web = { enabled = true; }; image = { enabled = true; }; sub_worker = { enabled = true; }; diff --git a/resources/profiles/memory-consolidation.dcdl b/resources/profiles/memory-consolidation.dcdl index 689fb934..d1a0474e 100644 --- a/resources/profiles/memory-consolidation.dcdl +++ b/resources/profiles/memory-consolidation.dcdl @@ -5,7 +5,7 @@ import "./base.dcdl" // { feature = { task = { enabled = false; }; - memory = { enabled = true; staging = true; }; + memory = { enabled = true; staging_tools = true; }; web = { enabled = false; }; sub_worker = { enabled = false; }; worker = { enabled = false; }; From 12646b6ca0927671df4fd27b27ceb1f92823a86d Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 23:40:12 +0900 Subject: [PATCH 20/22] refactor: install Memory prompt contributions through Feature --- crates/manifest/src/lib.rs | 12 + crates/worker/src/controller.rs | 227 +++------------- crates/worker/src/feature/builtin/memory.rs | 210 +++++++++++++++ .../src/feature/builtin/memory_lifecycle.rs | 46 +++- crates/worker/src/worker.rs | 253 ++++-------------- docs/manifest.toml | 4 +- 6 files changed, 356 insertions(+), 396 deletions(-) diff --git a/crates/manifest/src/lib.rs b/crates/manifest/src/lib.rs index 11c6c2f7..0e892677 100644 --- a/crates/manifest/src/lib.rs +++ b/crates/manifest/src/lib.rs @@ -328,6 +328,12 @@ impl ResolvedMemoryFeatureConfig { if !self.profile.enabled && self.workspace_settings.is_some() { return Err("disabled Memory feature must not carry Workspace settings"); } + if let Some(settings) = &self.workspace_settings + && (settings.settings_revision == 0 + || !is_normalized_workspace_memory_language(&settings.language)) + { + return Err("Memory Workspace settings snapshot metadata is invalid"); + } Ok(()) } } @@ -918,6 +924,12 @@ impl Default for CompactionConfig { } impl WorkerManifest { + pub fn requires_persisted_execution_snapshot(&self) -> bool { + self.profile.is_some() + || self.plugins.has_resolved_plan() + || self.feature.memory.workspace_settings.is_some() + } + /// Parse a manifest from a TOML string. pub fn from_toml(s: &str) -> Result { config::reject_removed_manifest_fields(s)?; diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index fb93f785..4e431d3d 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -901,43 +901,6 @@ pub(crate) fn wire_event_bridges_on_engine( // per-item commit channel is wired at the top of this function. } -fn add_memory_tools_if_configured( - registry: &mut FeatureRegistryBuilder, - config: &manifest::ResolvedMemoryFeatureConfig, - build: impl FnOnce() -> std::io::Result, -) -> std::io::Result -where - M: crate::feature::FeatureModule + 'static, -{ - if !config.profile.enabled { - return Ok(false); - } - config - .validate_execution() - .map_err(|message| std::io::Error::new(std::io::ErrorKind::InvalidInput, message))?; - registry.add_module(build()?); - Ok(true) -} - -fn add_memory_lifecycle_if_configured( - registry: &mut FeatureRegistryBuilder, - config: manifest::ResolvedMemoryFeatureConfig, - lifecycle_enabled: bool, - build: impl FnOnce(manifest::ResolvedMemoryFeatureConfig) -> std::io::Result, -) -> std::io::Result -where - M: crate::feature::FeatureModule + 'static, -{ - if !lifecycle_enabled || !config.profile.enabled || !config.profile.extraction.enabled { - return Ok(false); - } - config - .validate_execution() - .map_err(|message| std::io::Error::new(std::io::ErrorKind::InvalidInput, message))?; - registry.add_module(build(config)?); - Ok(true) -} - /// Register the builtin file-manipulation tools, optional memory tools, /// and the Worker-orchestration tools (SubWorkerSpawn + comm) on the Worker's /// Engine. Returns the WorkdirSession handle used to attach a `WorkerFsView` to @@ -974,7 +937,6 @@ where let local_filesystem = worker.local_working_directory().cloned(); let local_workspace_root = local_filesystem.as_ref().map(|local| local.root.clone()); let task_feature = worker.task_feature(); - let memory_config = feature_config.memory.clone(); let web_config = worker.manifest().web.clone(); let mcp_config = worker.manifest().mcp.clone(); let spawner_name = worker.manifest().worker.name.clone(); @@ -1031,46 +993,41 @@ where let worker_enabled = feature_config.worker.enabled; let sub_worker_enabled = feature_config.sub_worker.enabled; let mut feature_registry = FeatureRegistryBuilder::new(); - add_memory_tools_if_configured(&mut feature_registry, &memory_config, || { - let workspace_client = worker.workspace_client_handle(); - if !workspace_client.is_available() || workspace_client.workspace_id().is_none() { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "Memory tools require Backend Workspace API authority", - )); - } - Ok(crate::feature::builtin::memory::MemoryToolsFeature::new( - workspace_client, - memory_config.profile.staging_tools, - )) - })?; - add_memory_lifecycle_if_configured( - &mut feature_registry, - memory_config.clone(), - worker.manifest_lifecycle_features_enabled(), - |config| { - let workspace_client = worker.workspace_client_handle(); - if !workspace_client.is_available() || workspace_client.workspace_id().is_none() { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "Memory extraction requires Backend Workspace API authority", - )); - } - Ok( - crate::feature::builtin::memory_lifecycle::MemoryLifecycleFeature::new( - config, - worker.committed_session_capture_handle(), - worker.session_extension_handle(), - workspace_client, - spawner_manifest.clone(), - worker.llm_client_handle(), - prompts.clone(), - spawner_workspace_context.clone(), - worker.working_event_sender(), - ), - ) - }, - )?; + let memory_install_plan = crate::feature::builtin::memory::MemoryFeatureInstallPlan::prepare( + worker.manifest(), + worker.workspace_client_handle(), + worker.prompts().load_full(), + ) + .await?; + let memory_prompt_contribution = memory_install_plan.as_ref().map(|plan| { + ( + plan.resident_summary.clone(), + plan.system_prompt_override.clone(), + ) + }); + let memory_lifecycle_config = memory_install_plan + .as_ref() + .map(|plan| plan.resolved_config.clone()); + if let Some(plan) = memory_install_plan { + feature_registry.add_module(plan.module); + } + if let Some(memory_config) = memory_lifecycle_config + && let Some(memory_lifecycle) = + crate::feature::builtin::memory_lifecycle::MemoryLifecycleFeature::from_resolved_config( + worker.manifest_lifecycle_features_enabled(), + memory_config, + worker.committed_session_capture_handle(), + worker.session_extension_handle(), + worker.workspace_client_handle(), + spawner_manifest.clone(), + worker.llm_client_handle(), + prompts.clone(), + spawner_workspace_context.clone(), + worker.working_event_sender(), + )? + { + feature_registry.add_module(memory_lifecycle); + } if sub_worker_enabled && !worker_enabled { feature_registry.add_module( crate::feature::builtin::manage_worker::sub_worker_control_feature( @@ -1279,6 +1236,9 @@ where ), )); } + if let Some((resident_summary, system_prompt_override)) = memory_prompt_contribution { + worker.install_system_prompt_contribution(resident_summary, system_prompt_override); + } if let Some(tracker) = tracker { worker.attach_tracker(tracker); } @@ -2158,117 +2118,6 @@ mod tests { use tempfile::TempDir; use tokio::net::UnixListener; - #[test] - fn memory_feature_registration_requires_bound_workspace_memory_config() { - #[derive(Clone)] - struct TestMemoryLifecycleModule; - - impl crate::feature::FeatureModule for TestMemoryLifecycleModule { - fn descriptor(&self) -> crate::feature::FeatureDescriptor { - crate::feature::FeatureDescriptor::builtin( - "test-memory-lifecycle", - "Test Memory Lifecycle", - ) - } - - fn install( - &self, - _context: &mut crate::feature::FeatureInstallContext<'_>, - ) -> Result<(), crate::feature::FeatureInstallError> { - Ok(()) - } - } - - let mut registry = FeatureRegistryBuilder::new(); - let installed = add_memory_tools_if_configured::( - &mut registry, - &manifest::ResolvedMemoryFeatureConfig::default(), - || panic!("disabled Memory must not construct its tools Feature"), - ) - .unwrap(); - assert!(!installed); - - let mut missing_snapshot = manifest::ResolvedMemoryFeatureConfig::default(); - missing_snapshot.profile.enabled = true; - let error = add_memory_tools_if_configured::( - &mut registry, - &missing_snapshot, - || panic!("invalid Memory config must fail before tools Feature construction"), - ) - .unwrap_err(); - assert!( - error - .to_string() - .contains("requires trusted Workspace settings") - ); - - let mut registry = FeatureRegistryBuilder::new(); - let configured = std::cell::Cell::new(false); - let mut memory_config = manifest::ResolvedMemoryFeatureConfig::default(); - memory_config.profile.enabled = true; - memory_config.profile.extraction.enabled = true; - memory_config - .bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot { - workspace_id: "workspace-1".to_string(), - settings_revision: 1, - language: "English".to_string(), - }) - .unwrap(); - let installed = - add_memory_lifecycle_if_configured(&mut registry, memory_config, true, |_| { - configured.set(true); - Ok(TestMemoryLifecycleModule) - }) - .unwrap(); - assert!(installed); - assert!(configured.get()); - - let mut registry = FeatureRegistryBuilder::new(); - let installed = add_memory_lifecycle_if_configured::( - &mut registry, - manifest::ResolvedMemoryFeatureConfig::default(), - true, - |_| panic!("disabled Memory must not construct its lifecycle Feature"), - ) - .unwrap(); - assert!(!installed); - - let mut lifecycle_disabled = manifest::ResolvedMemoryFeatureConfig::default(); - lifecycle_disabled.profile.enabled = true; - lifecycle_disabled.profile.extraction.enabled = true; - lifecycle_disabled - .bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot { - workspace_id: "workspace-1".to_string(), - settings_revision: 1, - language: "English".to_string(), - }) - .unwrap(); - let installed = add_memory_lifecycle_if_configured::( - &mut registry, - lifecycle_disabled, - false, - |_| panic!("disabled lifecycle must not construct its Feature"), - ) - .unwrap(); - assert!(!installed); - - let mut missing_snapshot = manifest::ResolvedMemoryFeatureConfig::default(); - missing_snapshot.profile.enabled = true; - missing_snapshot.profile.extraction.enabled = true; - let error = add_memory_lifecycle_if_configured::( - &mut registry, - missing_snapshot, - true, - |_| panic!("invalid Workspace Memory config must fail before Feature construction"), - ) - .unwrap_err(); - assert!( - error - .to_string() - .contains("requires trusted Workspace settings") - ); - } - #[test] fn image_attachment_gate_requires_vision_and_supported_openai_scheme() { let openai = manifest::ModelManifest { diff --git a/crates/worker/src/feature/builtin/memory.rs b/crates/worker/src/feature/builtin/memory.rs index 62403d1a..e04054b0 100644 --- a/crates/worker/src/feature/builtin/memory.rs +++ b/crates/worker/src/feature/builtin/memory.rs @@ -342,6 +342,113 @@ fn query_schema() -> serde_json::Value { }) } +pub struct MemoryFeatureInstallPlan { + pub module: MemoryToolsFeature, + pub resident_summary: Option, + pub system_prompt_override: Option, + pub(crate) resolved_config: manifest::ResolvedMemoryFeatureConfig, +} + +impl MemoryFeatureInstallPlan { + pub async fn prepare( + manifest: &manifest::WorkerManifest, + client: Arc, + prompts: Arc, + ) -> std::io::Result> { + Self::prepare_resolved( + manifest.feature.memory.clone(), + client, + prompts, + manifest.profile.clone(), + ) + .await + } + + async fn prepare_resolved( + config: manifest::ResolvedMemoryFeatureConfig, + client: Arc, + prompts: Arc, + profile: Option, + ) -> std::io::Result> { + let memory_consolidation_worker = profile.as_ref().is_some_and(|snapshot| { + matches!( + &snapshot.source, + manifest::ProfileSource::Registry { + source: manifest::ProfileRegistrySource::Builtin, + name, + .. + } if name == "memory-consolidation" + ) + }); + config + .validate_execution() + .map_err(|message| std::io::Error::new(std::io::ErrorKind::InvalidInput, message))?; + if !config.profile.enabled { + return Ok(None); + } + let workspace_id = client.workspace_id().ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Memory tools require Backend Workspace API authority", + ) + })?; + let settings = config + .workspace_settings() + .expect("validated enabled Memory config has Workspace settings"); + if settings.workspace_id != workspace_id { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "Memory settings belong to {} instead of {}", + settings.workspace_id, workspace_id + ), + )); + } + + let resident_summary = if config.profile.resident.inject_summary { + match client + .execute_memory_backend_operation( + memory::backend::MemoryBackendOperation::ResidentSummary( + memory::backend::MemoryResidentSummaryOperation::default(), + ), + ) + .await + { + Ok(memory::backend::MemoryBackendOperationResult::ToolOutput(output)) => { + output.content + } + Ok(other) => { + tracing::debug!(?other, "unexpected resident Memory Backend result"); + None + } + Err(error) => { + tracing::debug!(%error, "resident Memory summary unavailable"); + None + } + } + } else { + None + }; + let system_prompt_override = if memory_consolidation_worker { + let language = settings.language; + Some( + prompts + .memory_consolidation_system(&language) + .map_err(|error| std::io::Error::other(error.to_string()))?, + ) + } else { + None + }; + + Ok(Some(Self { + module: MemoryToolsFeature::new(client, config.profile.staging_tools), + resident_summary, + system_prompt_override, + resolved_config: config, + })) + } +} + #[derive(Clone)] pub struct MemoryToolsFeature { tools: Vec, @@ -392,6 +499,39 @@ mod tests { )) } + fn resident_client(content: &str) -> Arc { + use std::io::{Read, Write}; + use std::net::TcpListener; + + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + let content = content.to_string(); + std::thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + let mut request = [0_u8; 1024]; + let _ = stream.read(&mut request).unwrap(); + let body = serde_json::json!({ + "status": "ok", + "result": { + "kind": "tool_output", + "summary": "resident Memory summary collected", + "content": content, + } + }) + .to_string(); + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ); + stream.write_all(response.as_bytes()).unwrap(); + }); + Arc::new(crate::worker::TestWorkspaceHttpClient::new( + "workspace", + format!("http://{addr}"), + )) + } + fn tool_names(definitions: Vec) -> Vec { let mut names = definitions .into_iter() @@ -410,6 +550,76 @@ mod tests { .input_schema } + #[tokio::test] + async fn memory_install_plan_is_the_fail_closed_config_boundary() { + let prompts = crate::prompt::catalog::PromptCatalog::builtins_only().unwrap(); + let disabled = MemoryFeatureInstallPlan::prepare_resolved( + manifest::ResolvedMemoryFeatureConfig::default(), + test_client(), + prompts.clone(), + None, + ) + .await + .unwrap(); + assert!(disabled.is_none()); + + let mut enabled = manifest::ResolvedMemoryFeatureConfig::default(); + enabled.profile.enabled = true; + enabled.profile.resident.inject_summary = false; + assert!( + MemoryFeatureInstallPlan::prepare_resolved( + enabled.clone(), + test_client(), + prompts.clone(), + None, + ) + .await + .is_err() + ); + enabled + .bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot { + workspace_id: "workspace".to_string(), + settings_revision: 1, + language: "English".to_string(), + }) + .unwrap(); + let mut foreign = enabled.clone(); + foreign.workspace_settings.as_mut().unwrap().workspace_id = "other-workspace".to_string(); + assert!( + MemoryFeatureInstallPlan::prepare_resolved( + foreign, + test_client(), + prompts.clone(), + None, + ) + .await + .is_err() + ); + let plan = MemoryFeatureInstallPlan::prepare_resolved( + enabled.clone(), + test_client(), + prompts.clone(), + None, + ) + .await + .unwrap() + .unwrap(); + assert!(plan.resident_summary.is_none()); + assert!(plan.system_prompt_override.is_none()); + + enabled.profile.resident.inject_summary = true; + let plan = MemoryFeatureInstallPlan::prepare_resolved( + enabled, + resident_client("# Durable Memory"), + prompts, + None, + ) + .await + .unwrap() + .unwrap(); + assert_eq!(plan.resident_summary.as_deref(), Some("# Durable Memory")); + } + #[test] fn memory_feature_owns_normal_and_staging_tool_surfaces() { let normal = MemoryToolsFeature::new(test_client(), false); diff --git a/crates/worker/src/feature/builtin/memory_lifecycle.rs b/crates/worker/src/feature/builtin/memory_lifecycle.rs index 9c601263..f2d0caab 100644 --- a/crates/worker/src/feature/builtin/memory_lifecycle.rs +++ b/crates/worker/src/feature/builtin/memory_lifecycle.rs @@ -64,6 +64,44 @@ struct MemoryLifecycleTask { } impl MemoryLifecycleFeature { + #[allow(clippy::too_many_arguments)] + pub(crate) fn from_resolved_config( + lifecycle_enabled: bool, + config: manifest::ResolvedMemoryFeatureConfig, + capture: CommittedSessionCaptureHandle, + extensions: SessionExtensionHandle, + workspace_client: Arc, + manifest: WorkerManifest, + client: Box, + prompts: Arc>, + workspace_context: WorkerWorkspaceContext, + event_tx: Option>, + ) -> std::io::Result> { + if !lifecycle_enabled || !config.profile.enabled || !config.profile.extraction.enabled { + return Ok(None); + } + config + .validate_execution() + .map_err(|message| std::io::Error::new(std::io::ErrorKind::InvalidInput, message))?; + if !workspace_client.is_available() || workspace_client.workspace_id().is_none() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Memory extraction requires Backend Workspace API authority", + )); + } + Ok(Some(Self::new( + config, + capture, + extensions, + workspace_client, + manifest, + client, + prompts, + workspace_context, + event_tx, + ))) + } + #[allow(clippy::too_many_arguments)] pub(crate) fn new( config: manifest::ResolvedMemoryFeatureConfig, @@ -1452,8 +1490,12 @@ permission = "write" ); } let controller_source = include_str!("../../controller.rs"); - assert!(controller_source.contains("add_memory_lifecycle_if_configured")); - assert!(controller_source.contains("MemoryLifecycleFeature::new")); + assert!(controller_source.contains("MemoryFeatureInstallPlan::prepare")); + assert!(controller_source.contains("MemoryLifecycleFeature::from_resolved_config")); + let worker_production = worker_source.split("#[cfg(test)]").next().unwrap(); + let controller_production = controller_source.split("#[cfg(test)]").next().unwrap(); + assert!(!worker_production.contains(".feature.memory")); + assert!(!controller_production.contains(".feature.memory")); let lifecycle_source = include_str!("memory_lifecycle.rs"); assert!(lifecycle_source.contains("request_memory_staging_consolidation")); let internal_worker_source = include_str!("../../internal_worker.rs"); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 76888eff..d4a3b85c 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -1228,10 +1228,12 @@ pub struct Worker { /// [`Self::from_manifest`], or defaults to the builtin pack when a /// Worker is constructed through lower-level paths that have no loader. prompts: Arc>, - /// When true (default), the system-prompt assembler may append resident - /// context from the workspace Memory document. Internal disposable - /// workers disable this so resident memory exposure is opt-in per Worker. + /// Test/internal policy gate for installed resident prompt contributions. inject_resident_summary: bool, + /// Materialized resident prompt context installed by an enabled Feature. + feature_resident_summary: Option, + /// Complete system prompt replacement installed by an enabled Feature. + feature_system_prompt_override: Option, /// Typed user submissions in submit order. K-th entry corresponds to /// the K-th `Item::user_message` in `worker.history()` (modulo seed /// history loaded via `AnnotatedSegmentStart.history`, whose original segments @@ -1450,6 +1452,8 @@ impl Worker { runtime_ticket_role: None, prompts, inject_resident_summary: true, + feature_resident_summary: None, + feature_system_prompt_override: None, user_segments: Vec::new(), sink: SegmentLogSink::new(), history_persistence_wired: false, @@ -1487,11 +1491,20 @@ impl Worker { self.inject_resident_summary = enabled; } - /// Toggle workspace Memory document resident injection in the system prompt. + /// Internal/test gate for installed resident prompt contributions. pub fn set_resident_summary_injection(&mut self, enabled: bool) { self.inject_resident_summary = enabled; } + pub(crate) fn install_system_prompt_contribution( + &mut self, + resident_summary: Option, + system_prompt_override: Option, + ) { + self.feature_resident_summary = resident_summary; + self.feature_system_prompt_override = system_prompt_override; + } + pub fn prompts(&self) -> Arc> { Arc::clone(&self.prompts) } @@ -1625,25 +1638,6 @@ impl Worker { self.worker_observation_provider.clone() } - async fn resident_summary_from_workspace_authority( - &self, - ) -> Result, WorkerError> { - let result = self - .workspace_client() - .execute_memory_backend_operation( - memory::backend::MemoryBackendOperation::ResidentSummary( - memory::backend::MemoryResidentSummaryOperation::default(), - ), - ) - .await?; - match result { - memory::backend::MemoryBackendOperationResult::ToolOutput(output) => Ok(output.content), - other => Err(WorkerError::FeatureInstall(format!( - "unexpected memory backend result for resident summary: {other:?}" - ))), - } - } - /// Activate an Agent Skill through the Workspace backend/client and commit /// the returned SKILL.md body to history before it can influence an LLM run. /// @@ -2467,26 +2461,7 @@ impl Worker { let Some(template) = self.system_prompt_template.take() else { return Ok(()); }; - let is_memory_consolidation = self.manifest.profile.as_ref().is_some_and(|snapshot| { - matches!( - &snapshot.source, - manifest::ProfileSource::Registry { - source: manifest::ProfileRegistrySource::Builtin, - name, - .. - } if name == "memory-consolidation" - ) - }); - if is_memory_consolidation { - let memory_config = &self.manifest.feature.memory; - memory_config - .validate_execution() - .map_err(|message| WorkerError::InvalidState(message.to_string()))?; - let language = memory_language(memory_config)?; - let rendered = self - .prompts - .load_full() - .memory_consolidation_system(&language)?; + if let Some(rendered) = self.feature_system_prompt_override.take() { self.engine .as_mut() .expect("worker present") @@ -2515,20 +2490,10 @@ impl Worker { } } } - let inject_summary = self.inject_resident_summary - && self.manifest.feature.memory.profile.enabled - && self.manifest.feature.memory.profile.resident.inject_summary; - let resident_summary: Option = if inject_summary { - match self.resident_summary_from_workspace_authority().await { - Ok(summary) => summary, - Err(error) => { - tracing::debug!(%error, "resident memory summary unavailable"); - None - } - } - } else { - None - }; + let resident_summary = self + .inject_resident_summary + .then(|| self.feature_resident_summary.clone()) + .flatten(); let worker_language = worker_language(&self.manifest.engine); let scope_snapshot = self.scope.snapshot(); let cwd_for_prompt = self @@ -4548,17 +4513,6 @@ impl Worker { } } -fn memory_language(config: &manifest::ResolvedMemoryFeatureConfig) -> Result { - config - .workspace_settings() - .map(|snapshot| snapshot.language) - .ok_or_else(|| { - WorkerError::InvalidState( - "Memory is enabled without a bound Workspace Memory settings snapshot".to_string(), - ) - }) -} - fn worker_language(cfg: &manifest::EngineManifest) -> &str { let language = cfg.language.trim(); if language.is_empty() { @@ -4623,7 +4577,6 @@ where filesystem_authority: WorkerFilesystemAuthority, model_client: Option>, ) -> Result { - validate_workspace_memory_snapshot(&manifest.worker.name, &manifest, &workspace_context)?; let common = prepare_worker_common_with_context_and_model_client( &manifest, &loader, @@ -4708,6 +4661,8 @@ where runtime_ticket_role: None, prompts: common.prompts, inject_resident_summary: true, + feature_resident_summary: None, + feature_system_prompt_override: None, user_segments: Vec::new(), sink: SegmentLogSink::new(), history_persistence_wired: false, @@ -4790,6 +4745,8 @@ where runtime_ticket_role: None, prompts: common.prompts, inject_resident_summary: true, + feature_resident_summary: None, + feature_system_prompt_override: None, user_segments: Vec::new(), sink: SegmentLogSink::new(), history_persistence_wired: false, @@ -4836,7 +4793,6 @@ where workspace_context: WorkerWorkspaceContext, filesystem_authority: WorkerFilesystemAuthority, ) -> Result { - validate_workspace_memory_snapshot(&manifest.worker.name, &manifest, &workspace_context)?; let common = prepare_worker_common_with_context( &manifest, &loader, @@ -4907,6 +4863,8 @@ where runtime_ticket_role: None, prompts: common.prompts, inject_resident_summary: true, + feature_resident_summary: None, + feature_system_prompt_override: None, user_segments: Vec::new(), sink: SegmentLogSink::new(), history_persistence_wired: false, @@ -5280,6 +5238,8 @@ where runtime_ticket_role: None, prompts: common.prompts, inject_resident_summary: true, + feature_resident_summary: None, + feature_system_prompt_override: None, user_segments: state.user_segments, // Seed the mirror with the entries we just replayed so a // late-attaching client sees the full prefix without an @@ -5428,54 +5388,8 @@ fn worker_metadata_for_manifest( metadata } -fn validate_workspace_memory_snapshot( - worker_name: &str, - manifest: &WorkerManifest, - workspace_context: &WorkerWorkspaceContext, -) -> Result<(), WorkerError> { - let Some(workspace_id) = workspace_context.workspace_id() else { - return Ok(()); - }; - manifest - .feature - .memory - .validate_execution() - .map_err(|message| { - WorkerError::InvalidState(format!("Workspace Worker {worker_name}: {message}")) - })?; - if !manifest.feature.memory.profile.enabled { - return Ok(()); - } - let snapshot = manifest - .feature - .memory - .workspace_settings() - .ok_or_else(|| { - WorkerError::InvalidState(format!( - "Workspace Worker {worker_name} has no complete persisted Memory settings snapshot" - )) - })?; - if snapshot.workspace_id != workspace_id.as_str() { - return Err(WorkerError::InvalidState(format!( - "Workspace Worker {worker_name} Memory settings belong to {} instead of {}", - snapshot.workspace_id, - workspace_id.as_str() - ))); - } - if snapshot.settings_revision == 0 - || !manifest::is_normalized_workspace_memory_language(&snapshot.language) - { - return Err(WorkerError::InvalidState(format!( - "Workspace Worker {worker_name} has corrupt Memory settings snapshot metadata" - ))); - } - Ok(()) -} - fn should_persist_resolved_manifest_snapshot(manifest: &WorkerManifest) -> bool { - manifest.profile.is_some() - || manifest.plugins.has_resolved_plan() - || manifest.feature.memory.workspace_settings.is_some() + manifest.requires_persisted_execution_snapshot() } fn restore_manifest_from_worker_metadata_snapshot( @@ -6565,7 +6479,7 @@ permission = "write" } #[test] - fn workspace_memory_settings_snapshot_is_persisted_and_scope_checked() { + fn workspace_memory_settings_snapshot_is_persisted_through_versioned_adapter() { let mut manifest = WorkerManifest::from_toml( r#" [worker] @@ -6614,42 +6528,6 @@ permission = "read" language: "Japanese".to_string(), }) ); - assert!( - validate_workspace_memory_snapshot( - "memory-snapshot", - &manifest, - &WorkerWorkspaceContext::unavailable( - Some(WorkspaceId::new("workspace-a").unwrap()), - "test", - ) - ) - .is_ok() - ); - assert!( - validate_workspace_memory_snapshot( - "memory-snapshot", - &manifest, - &WorkerWorkspaceContext::unavailable( - Some(WorkspaceId::new("workspace-b").unwrap()), - "test", - ) - ) - .is_err() - ); - - let mut missing = manifest.clone(); - missing.feature.memory.workspace_settings = None; - assert!( - validate_workspace_memory_snapshot( - "memory-snapshot", - &missing, - &WorkerWorkspaceContext::unavailable( - Some(WorkspaceId::new("workspace-a").unwrap()), - "test", - ) - ) - .is_err() - ); } #[test] @@ -8230,6 +8108,12 @@ mod build_summary_prompt_tests { ) .unwrap(), ); + let prompt_override = worker + .prompts() + .load_full() + .memory_consolidation_system("Japanese") + .unwrap(); + worker.install_system_prompt_contribution(None, Some(prompt_override)); worker.ensure_system_prompt_materialized().await.unwrap(); let prompt = worker.engine().get_system_prompt().unwrap(); assert!(prompt.contains("`language`: `Japanese`")); @@ -8267,15 +8151,7 @@ mod build_summary_prompt_tests { } let scope = Scope::writable(&cwd).unwrap(); let authority = WorkerFilesystemAuthority::local(cwd.clone(), cwd.clone()); - let workspace_context = if memory_config - .as_ref() - .is_some_and(|cfg| cfg.profile.resident.inject_summary) - && gates.summary - { - stub_memory_backend_context(summary_doc.and_then(summary_content_for_backend)) - } else { - WorkerWorkspaceContext::local_filesystem(None) - }; + let workspace_context = WorkerWorkspaceContext::local_filesystem(None); let mut worker = Worker::new( manifest, Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), @@ -8287,6 +8163,16 @@ mod build_summary_prompt_tests { .await .unwrap(); worker.set_resident_memory_injection(gates.summary); + let resident_summary = if memory_config + .as_ref() + .is_some_and(|cfg| cfg.profile.resident.inject_summary) + && gates.summary + { + summary_doc.and_then(summary_content_for_backend) + } else { + None + }; + worker.install_system_prompt_contribution(resident_summary, None); let template = SystemPromptTemplate::parse( "default", crate::prompt::source::PromptCatalogSource::builtins_only(), @@ -8313,45 +8199,6 @@ mod build_summary_prompt_tests { Some(doc.to_string()) } - fn stub_memory_backend_context(content: Option) -> WorkerWorkspaceContext { - use std::io::{Read, Write}; - use std::net::TcpListener; - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let addr = listener.local_addr().unwrap(); - std::thread::spawn(move || { - let (mut stream, _) = listener.accept().unwrap(); - let mut buffer = [0_u8; 1024]; - let _ = stream.read(&mut buffer).unwrap(); - let body = serde_json::json!({ - "status": "ok", - "result": { - "kind": "tool_output", - "summary": if content.is_some() { - "resident memory summary collected" - } else { - "resident memory summary unavailable" - }, - "content": content, - } - }) - .to_string(); - let response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - body.len(), - body - ); - stream.write_all(response.as_bytes()).unwrap(); - }); - WorkerWorkspaceContext::with_client( - Some(WorkspaceId::new("test-memory").unwrap()), - Arc::new(TestWorkspaceHttpClient::new( - "test-memory", - format!("http://{addr}"), - )), - ) - } - #[tokio::test] async fn resident_summary_body_is_injected_without_frontmatter() { let rendered = render_system_prompt_with_summary( @@ -8367,7 +8214,7 @@ mod build_summary_prompt_tests { } #[tokio::test] - async fn resident_summary_injection_can_be_disabled_by_manifest() { + async fn resident_summary_injection_can_be_disabled_by_memory_feature() { let mut memory = manifest::ResolvedMemoryFeatureConfig::default(); memory.profile.resident.inject_summary = false; let rendered = render_system_prompt_with_summary( diff --git a/docs/manifest.toml b/docs/manifest.toml index 2dc2386e..2becb8bc 100644 --- a/docs/manifest.toml +++ b/docs/manifest.toml @@ -227,8 +227,8 @@ permission = "write" # Profile由来の設定を `profile` に、Backend由来のWorkspace設定snapshotを # `workspace_settings` に分離して保存する。`workspace_settings` はBackendだけが # bindする信頼済み入力で、Profile・Browser・model入力から指定できない。 -# `profile.enabled = false` の場合、Memory tools、resident injection、extract、 -# consolidation requestをすべて無効にし、snapshotも保持しない。 +# `profile.enabled = false` の場合、Memory tools、Feature prompt contributionによる +# resident injection、extract、consolidation requestをすべて無効にし、snapshotも保持しない。 # # [feature.memory.profile] # enabled = true From 4df277c81f21ad2fb311f86a1ede07f0836b2b2b Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 23:53:59 +0900 Subject: [PATCH 21/22] fix: preserve Memory consolidation enablement --- crates/manifest/src/config.rs | 43 ++++++++++-- crates/manifest/src/lib.rs | 68 ++++++++++++++++++- .../src/feature/builtin/memory_lifecycle.rs | 34 +++++++++- docs/manifest.toml | 3 + resources/profiles/base.dcdl | 1 + 5 files changed, 142 insertions(+), 7 deletions(-) diff --git a/crates/manifest/src/config.rs b/crates/manifest/src/config.rs index 6a6e6e96..d5ae7049 100644 --- a/crates/manifest/src/config.rs +++ b/crates/manifest/src/config.rs @@ -18,11 +18,11 @@ use crate::model::{AuthRef, ModelManifest, ReasoningControl}; use crate::plugin::PluginConfig; use crate::{ CompactionConfig, EngineManifest, FeatureConfig, FeatureFlagConfig, FileUploadLimits, - McpConfig, McpEnvValue, McpStdioCwdPolicy, MemoryExtractionProfileConfig, - MemoryFeatureProfileConfig, MemoryResidentProfileConfig, MergeRequestFeatureConfig, - ResolvedMemoryFeatureConfig, ScopeConfig, SessionConfig, SkillsConfig, TicketFeatureConfig, - ToolOutputLimits, ToolPermissionConfig, ToolPermissionRule, WebConfig, WorkerFeatureConfig, - WorkerManifest, WorkerMeta, + McpConfig, McpEnvValue, McpStdioCwdPolicy, MemoryConsolidationProfileConfig, + MemoryExtractionProfileConfig, MemoryFeatureProfileConfig, MemoryResidentProfileConfig, + MergeRequestFeatureConfig, ResolvedMemoryFeatureConfig, ScopeConfig, SessionConfig, + SkillsConfig, TicketFeatureConfig, ToolOutputLimits, ToolPermissionConfig, ToolPermissionRule, + WebConfig, WorkerFeatureConfig, WorkerManifest, WorkerMeta, }; /// Partial-form Worker manifest. Every field is optional; one or more @@ -201,6 +201,8 @@ pub struct MemoryFeatureConfigPartial { pub resident: Option, #[serde(default)] pub extraction: Option, + #[serde(default)] + pub consolidation: Option, } #[derive(Debug, Clone, Default, Serialize, Deserialize)] @@ -223,6 +225,13 @@ pub struct MemoryExtractionProfileConfigPartial { pub worker_max_turns: Option, } +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MemoryConsolidationProfileConfigPartial { + #[serde(default)] + pub request_enabled: Option, +} + impl MemoryFeatureConfigPartial { fn merge(self, other: Self) -> Self { Self { @@ -238,6 +247,11 @@ impl MemoryFeatureConfigPartial { other.extraction, MemoryExtractionProfileConfigPartial::merge, ), + consolidation: merge_option( + self.consolidation, + other.consolidation, + MemoryConsolidationProfileConfigPartial::merge, + ), } } } @@ -305,6 +319,14 @@ impl MergeRequestFeatureConfigPartial { } } +impl MemoryConsolidationProfileConfigPartial { + fn merge(self, other: Self) -> Self { + Self { + request_enabled: other.request_enabled.or(self.request_enabled), + } + } +} + impl From for FeatureConfig { fn from(value: FeatureConfigPartial) -> Self { Self { @@ -385,6 +407,7 @@ impl From for ResolvedMemoryFeatureConfig { fn from(value: MemoryFeatureConfigPartial) -> Self { let resident = value.resident.unwrap_or_default(); let extraction = value.extraction.unwrap_or_default(); + let consolidation = value.consolidation.unwrap_or_default(); Self { profile: MemoryFeatureProfileConfig { enabled: value.enabled.unwrap_or_default(), @@ -400,6 +423,9 @@ impl From for ResolvedMemoryFeatureConfig { .worker_max_turns .or(defaults::MEMORY_EXTRACT_WORKER_MAX_TURNS), }, + consolidation: MemoryConsolidationProfileConfig { + request_enabled: consolidation.request_enabled.unwrap_or(true), + }, }, workspace_settings: None, } @@ -420,6 +446,9 @@ impl From for MemoryFeatureConfigPartial { threshold: value.profile.extraction.threshold, worker_max_turns: value.profile.extraction.worker_max_turns, }), + consolidation: Some(MemoryConsolidationProfileConfigPartial { + request_enabled: Some(value.profile.consolidation.request_enabled), + }), } } } @@ -1899,6 +1928,9 @@ inject_summary = false enabled = true threshold = 42000 worker_max_turns = 2 + +[feature.memory.consolidation] +request_enabled = false "#, ) .unwrap(); @@ -1906,6 +1938,7 @@ worker_max_turns = 2 assert_eq!(memory.enabled, Some(true)); assert_eq!(memory.staging_tools, Some(false)); assert_eq!(memory.resident.unwrap().inject_summary, Some(false)); + assert_eq!(memory.consolidation.unwrap().request_enabled, Some(false)); let extraction = memory.extraction.unwrap(); assert_eq!(extraction.enabled, Some(true)); assert_eq!(extraction.threshold, Some(42_000)); diff --git a/crates/manifest/src/lib.rs b/crates/manifest/src/lib.rs index 0e892677..5ebdc72f 100644 --- a/crates/manifest/src/lib.rs +++ b/crates/manifest/src/lib.rs @@ -226,6 +226,7 @@ pub struct MemoryFeatureProfileConfig { pub staging_tools: bool, pub resident: MemoryResidentProfileConfig, pub extraction: MemoryExtractionProfileConfig, + pub consolidation: MemoryConsolidationProfileConfig, } impl MemoryFeatureProfileConfig { @@ -248,6 +249,7 @@ impl Default for MemoryFeatureProfileConfig { staging_tools: false, resident: MemoryResidentProfileConfig::default(), extraction: MemoryExtractionProfileConfig::default(), + consolidation: MemoryConsolidationProfileConfig::default(), } } } @@ -286,6 +288,20 @@ impl Default for MemoryExtractionProfileConfig { } } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(default, deny_unknown_fields)] +pub struct MemoryConsolidationProfileConfig { + pub request_enabled: bool, +} + +impl Default for MemoryConsolidationProfileConfig { + fn default() -> Self { + Self { + request_enabled: true, + } + } +} + /// Immutable Memory execution configuration persisted in a resolved Worker Manifest. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] #[serde(default, deny_unknown_fields)] @@ -1070,6 +1086,31 @@ fn migrate_legacy_resolved_manifest_snapshot( .cloned() .unwrap_or(serde_json::Value::Null); let extraction_enabled = !extraction_threshold.is_null(); + if legacy_memory + .get("consolidation_model") + .is_some_and(|model| !model.is_null()) + { + return Err(serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "legacy resolved Worker manifest uses a Worker-owned consolidation model that cannot be migrated to Backend authority", + ))); + } + let threshold_files = legacy_memory + .get("consolidation_threshold_files") + .and_then(serde_json::Value::as_u64); + let threshold_bytes = legacy_memory + .get("consolidation_threshold_bytes") + .and_then(serde_json::Value::as_u64); + let consolidation_enabled = match (threshold_files, threshold_bytes) { + (None, None) => false, + (Some(5), Some(50_000)) => true, + _ => { + return Err(serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "legacy resolved Worker manifest uses custom consolidation thresholds that cannot be migrated to Backend policy", + ))); + } + }; let mut resolved = serde_json::json!({ "profile": { "enabled": enabled, @@ -1089,6 +1130,9 @@ fn migrate_legacy_resolved_manifest_snapshot( .cloned() .unwrap_or(serde_json::Value::Null), }, + "consolidation": { + "request_enabled": consolidation_enabled, + }, }, }); if let Some(workspace_settings) = workspace_settings { @@ -1490,7 +1534,8 @@ model_id = "claude-sonnet-4-20250514" "language": "Français", "extract_threshold": 1234, "extract_worker_max_turns": 3, - "consolidation_threshold_files": 99, + "consolidation_threshold_files": 5, + "consolidation_threshold_bytes": 50000, }); let migrated = read_persisted_worker_manifest_snapshot(manifest).unwrap(); @@ -1501,6 +1546,14 @@ model_id = "claude-sonnet-4-20250514" migrated.feature.memory.profile.extraction.threshold, Some(1234) ); + assert!( + migrated + .feature + .memory + .profile + .consolidation + .request_enabled + ); assert_eq!( migrated .feature @@ -1523,6 +1576,19 @@ model_id = "claude-sonnet-4-20250514" mixed["feature"]["memory"] = serde_json::json!({ "enabled": true, "profile": {} }); mixed["memory"] = serde_json::json!({}); assert!(read_persisted_worker_manifest_snapshot(mixed).is_err()); + + let mut custom_policy = + serde_json::to_value(WorkerManifest::from_toml(MINIMAL_REQUIRED).unwrap()).unwrap(); + custom_policy["feature"]["memory"] = serde_json::json!({ "enabled": true }); + custom_policy["memory"] = serde_json::json!({ + "workspace_id": "workspace-1", + "settings_revision": 1, + "language": "English", + "consolidation_threshold_files": 99, + "consolidation_threshold_bytes": 50000, + }); + assert!(read_persisted_worker_manifest_snapshot(custom_policy).is_err()); + assert!( read_persisted_worker_manifest_snapshot(serde_json::json!({ "schema_version": 3, diff --git a/crates/worker/src/feature/builtin/memory_lifecycle.rs b/crates/worker/src/feature/builtin/memory_lifecycle.rs index f2d0caab..7349f842 100644 --- a/crates/worker/src/feature/builtin/memory_lifecycle.rs +++ b/crates/worker/src/feature/builtin/memory_lifecycle.rs @@ -526,7 +526,9 @@ impl FeatureBackgroundTask for MemoryLifecycleTask { .await; if !cancellation.is_cancelled() { context.generation_fence.ensure_current()?; - self.request_consolidation().await; + if self.config.profile.consolidation.request_enabled { + self.request_consolidation().await; + } } extraction } @@ -1323,6 +1325,36 @@ permission = "write" ); } + #[tokio::test] + async fn lifecycle_task_does_not_request_consolidation_when_profile_disables_it() { + let client = ScriptClient::new(Vec::new()); + let extension_writes = Arc::new(Mutex::new(Vec::new())); + let (event_tx, _) = broadcast::channel(16); + let workspace_client = Arc::new(RecordingWorkspaceClient::default()); + let mut interrupted = capture(2, 250); + interrupted.run_exit = CommittedRunExit::Interrupted; + let mut task = test_task( + interrupted, + Box::new(client), + extension_writes, + event_tx, + workspace_client.clone(), + ); + task.config.profile.consolidation.request_enabled = false; + run_background_task(task).await; + + let requests = workspace_client.requests.lock().unwrap(); + assert!( + !requests.iter().any(|request| { + request + .body + .as_deref() + .is_some_and(|body| body == "{\"force\":false}") + }), + "recorded requests: {requests:?}" + ); + } + fn internal_result( lifecycle: WorkerRunResult, ) -> Result { diff --git a/docs/manifest.toml b/docs/manifest.toml index 2becb8bc..e7f72fd3 100644 --- a/docs/manifest.toml +++ b/docs/manifest.toml @@ -246,6 +246,9 @@ permission = "write" # # [feature.memory.profile.extraction.model] # # ref = "anthropic/claude-haiku-4-5" # +# [feature.memory.profile.consolidation] +# request_enabled = true +# # # Backendがresolved Manifestへbindする。手書き/Profile入力では指定しない。 # # [feature.memory.workspace_settings] # # workspace_id = "workspace-id" diff --git a/resources/profiles/base.dcdl b/resources/profiles/base.dcdl index 88db46c7..c6cdc32c 100644 --- a/resources/profiles/base.dcdl +++ b/resources/profiles/base.dcdl @@ -30,6 +30,7 @@ feature = { enabled = true; threshold = 50000; }; + consolidation = { request_enabled = true; }; }; web = { enabled = true; }; image = { enabled = true; }; From aa96bbedbc7db7a0443ac19371a2f25637fa128e Mon Sep 17 00:00:00 2001 From: Hare Date: Sat, 5 Sep 2026 00:24:00 +0900 Subject: [PATCH 22/22] fix: harden Memory restore and scheduling --- crates/manifest/src/lib.rs | 178 +++++++++++------- crates/worker/src/feature/builtin/memory.rs | 42 +++++ .../src/feature/builtin/memory_lifecycle.rs | 46 ++++- 3 files changed, 197 insertions(+), 69 deletions(-) diff --git a/crates/manifest/src/lib.rs b/crates/manifest/src/lib.rs index 5ebdc72f..a7908702 100644 --- a/crates/manifest/src/lib.rs +++ b/crates/manifest/src/lib.rs @@ -956,6 +956,34 @@ impl WorkerManifest { } } +#[derive(Debug, Default, Deserialize)] +#[serde(default, deny_unknown_fields)] +struct LegacyMemoryFeatureConfig { + enabled: bool, + staging: bool, +} + +#[derive(Debug, Default, Deserialize)] +#[serde(default, deny_unknown_fields)] +struct LegacyMemoryConfig { + #[serde(rename = "workspace_root")] + _workspace_root: Option, + #[serde(rename = "query_result_limit")] + _query_result_limit: Option, + #[serde(rename = "query_excerpt_lines")] + _query_excerpt_lines: Option, + inject_summary: Option, + workspace_id: Option, + settings_revision: Option, + language: Option, + extract_model: Option, + extract_threshold: Option, + extract_worker_max_turns: Option, + consolidation_model: Option, + consolidation_threshold_files: Option, + consolidation_threshold_bytes: Option, +} + const RESOLVED_MANIFEST_SNAPSHOT_SCHEMA_VERSION: u64 = 2; /// Serialize a resolved Worker Manifest for durable Worker-specific storage. @@ -1004,12 +1032,37 @@ pub fn read_persisted_worker_manifest_snapshot( "resolved Worker manifest snapshot is missing manifest", )) })?; - return serde_json::from_value(manifest); + if manifest + .as_object() + .is_some_and(|manifest| manifest.contains_key("memory")) + { + return Err(serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "current resolved Worker manifest contains removed top-level memory authority", + ))); + } + return validate_persisted_worker_manifest(serde_json::from_value(manifest)?); } migrate_legacy_resolved_manifest_snapshot(snapshot) } +fn validate_persisted_worker_manifest( + manifest: WorkerManifest, +) -> Result { + manifest + .feature + .memory + .validate_execution() + .map_err(|message| { + serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + message, + )) + })?; + Ok(manifest) +} + fn migrate_legacy_resolved_manifest_snapshot( mut snapshot: serde_json::Value, ) -> Result { @@ -1030,44 +1083,21 @@ fn migrate_legacy_resolved_manifest_snapshot( "legacy resolved Worker manifest feature must be an object", )) })?; - let legacy_feature_memory = feature - .remove("memory") - .unwrap_or_else(|| serde_json::json!({})); - let legacy_feature_memory = legacy_feature_memory.as_object().ok_or_else(|| { - serde_json::Error::io(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "legacy resolved Worker manifest feature.memory must be an object", - )) - })?; - if legacy_feature_memory - .keys() - .any(|key| !matches!(key.as_str(), "enabled" | "staging")) - { - return Err(serde_json::Error::io(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "legacy resolved Worker manifest mixes old and new Memory configuration", - ))); - } - let enabled = legacy_feature_memory - .get("enabled") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false); - let staging_tools = legacy_feature_memory - .get("staging") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false); + let legacy_feature_memory: LegacyMemoryFeatureConfig = serde_json::from_value( + feature + .remove("memory") + .unwrap_or_else(|| serde_json::json!({})), + )?; + let enabled = legacy_feature_memory.enabled; + let staging_tools = legacy_feature_memory.staging; - let legacy_memory = legacy_memory.unwrap_or_else(|| serde_json::json!({})); - let legacy_memory = legacy_memory.as_object().ok_or_else(|| { - serde_json::Error::io(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "legacy resolved Worker manifest memory must be an object", - )) - })?; - let workspace_id = legacy_memory.get("workspace_id").cloned(); - let settings_revision = legacy_memory.get("settings_revision").cloned(); - let language = legacy_memory.get("language").cloned(); - let workspace_settings = match (workspace_id, settings_revision, language) { + let legacy_memory: LegacyMemoryConfig = + serde_json::from_value(legacy_memory.unwrap_or_else(|| serde_json::json!({})))?; + let mut workspace_settings = match ( + legacy_memory.workspace_id, + legacy_memory.settings_revision, + legacy_memory.language, + ) { (Some(workspace_id), Some(settings_revision), Some(language)) => Some(serde_json::json!({ "workspace_id": workspace_id, "settings_revision": settings_revision, @@ -1081,27 +1111,20 @@ fn migrate_legacy_resolved_manifest_snapshot( ))); } }; - let extraction_threshold = legacy_memory - .get("extract_threshold") - .cloned() - .unwrap_or(serde_json::Value::Null); - let extraction_enabled = !extraction_threshold.is_null(); - if legacy_memory - .get("consolidation_model") - .is_some_and(|model| !model.is_null()) - { + if !enabled { + workspace_settings = None; + } + let extraction_enabled = legacy_memory.extract_threshold.is_some(); + if legacy_memory.consolidation_model.is_some() { return Err(serde_json::Error::io(std::io::Error::new( std::io::ErrorKind::InvalidData, "legacy resolved Worker manifest uses a Worker-owned consolidation model that cannot be migrated to Backend authority", ))); } - let threshold_files = legacy_memory - .get("consolidation_threshold_files") - .and_then(serde_json::Value::as_u64); - let threshold_bytes = legacy_memory - .get("consolidation_threshold_bytes") - .and_then(serde_json::Value::as_u64); - let consolidation_enabled = match (threshold_files, threshold_bytes) { + let consolidation_enabled = match ( + legacy_memory.consolidation_threshold_files, + legacy_memory.consolidation_threshold_bytes, + ) { (None, None) => false, (Some(5), Some(50_000)) => true, _ => { @@ -1116,19 +1139,13 @@ fn migrate_legacy_resolved_manifest_snapshot( "enabled": enabled, "staging_tools": staging_tools, "resident": { - "inject_summary": legacy_memory - .get("inject_summary") - .and_then(serde_json::Value::as_bool) - .unwrap_or(true), + "inject_summary": legacy_memory.inject_summary.unwrap_or(true), }, "extraction": { "enabled": extraction_enabled, - "model": legacy_memory.get("extract_model").cloned().unwrap_or(serde_json::Value::Null), - "threshold": extraction_threshold, - "worker_max_turns": legacy_memory - .get("extract_worker_max_turns") - .cloned() - .unwrap_or(serde_json::Value::Null), + "model": serde_json::to_value(legacy_memory.extract_model)?, + "threshold": legacy_memory.extract_threshold, + "worker_max_turns": legacy_memory.extract_worker_max_turns, }, "consolidation": { "request_enabled": consolidation_enabled, @@ -1142,7 +1159,7 @@ fn migrate_legacy_resolved_manifest_snapshot( .insert("workspace_settings".to_string(), workspace_settings); } feature.insert("memory".to_string(), resolved); - serde_json::from_value(snapshot) + validate_persisted_worker_manifest(serde_json::from_value(snapshot)?) } #[cfg(test)] @@ -1566,6 +1583,18 @@ model_id = "claude-sonnet-4-20250514" let current = write_persisted_worker_manifest_snapshot(&migrated).unwrap(); assert_eq!(current["schema_version"], 2); assert!(current["manifest"].get("memory").is_none()); + + let mut disabled = + serde_json::to_value(WorkerManifest::from_toml(MINIMAL_REQUIRED).unwrap()).unwrap(); + disabled["feature"]["memory"] = serde_json::json!({ "enabled": false }); + disabled["memory"] = serde_json::json!({ + "workspace_id": "workspace-1", + "settings_revision": 9, + "language": "Français", + }); + let disabled = read_persisted_worker_manifest_snapshot(disabled).unwrap(); + assert!(!disabled.feature.memory.profile.enabled); + assert!(disabled.feature.memory.workspace_settings.is_none()); } #[test] @@ -1589,6 +1618,25 @@ model_id = "claude-sonnet-4-20250514" }); assert!(read_persisted_worker_manifest_snapshot(custom_policy).is_err()); + let current = WorkerManifest::from_toml(MINIMAL_REQUIRED).unwrap(); + let mut current = write_persisted_worker_manifest_snapshot(¤t).unwrap(); + current["manifest"]["memory"] = serde_json::json!({ + "workspace_id": "workspace-1", + "settings_revision": 1, + "language": "English", + }); + assert!(read_persisted_worker_manifest_snapshot(current).is_err()); + + let mut missing_settings = WorkerManifest::from_toml(MINIMAL_REQUIRED).unwrap(); + missing_settings.feature.memory.profile.enabled = true; + let missing_settings = write_persisted_worker_manifest_snapshot(&missing_settings).unwrap(); + assert!(read_persisted_worker_manifest_snapshot(missing_settings).is_err()); + + let mut malformed_legacy = manifest.clone(); + malformed_legacy["feature"]["memory"] = serde_json::json!({ "enabled": "yes" }); + malformed_legacy["memory"] = serde_json::json!({ "unknown": true }); + assert!(read_persisted_worker_manifest_snapshot(malformed_legacy).is_err()); + assert!( read_persisted_worker_manifest_snapshot(serde_json::json!({ "schema_version": 3, diff --git a/crates/worker/src/feature/builtin/memory.rs b/crates/worker/src/feature/builtin/memory.rs index e04054b0..33c5eff3 100644 --- a/crates/worker/src/feature/builtin/memory.rs +++ b/crates/worker/src/feature/builtin/memory.rs @@ -620,6 +620,48 @@ mod tests { assert_eq!(plan.resident_summary.as_deref(), Some("# Durable Memory")); } + #[tokio::test] + async fn memory_prompt_contribution_rereads_resident_summary_for_each_install() { + let prompts = crate::prompt::catalog::PromptCatalog::builtins_only().unwrap(); + let mut config = manifest::ResolvedMemoryFeatureConfig::default(); + config.profile.enabled = true; + config + .bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot { + workspace_id: "workspace".to_string(), + settings_revision: 1, + language: "English".to_string(), + }) + .unwrap(); + + let first = MemoryFeatureInstallPlan::prepare_resolved( + config.clone(), + resident_client("first resident summary"), + prompts.clone(), + None, + ) + .await + .unwrap() + .unwrap(); + let restored = MemoryFeatureInstallPlan::prepare_resolved( + config, + resident_client("updated resident summary"), + prompts, + None, + ) + .await + .unwrap() + .unwrap(); + + assert_eq!( + first.resident_summary.as_deref(), + Some("first resident summary") + ); + assert_eq!( + restored.resident_summary.as_deref(), + Some("updated resident summary") + ); + } + #[test] fn memory_feature_owns_normal_and_staging_tool_surfaces() { let normal = MemoryToolsFeature::new(test_client(), false); diff --git a/crates/worker/src/feature/builtin/memory_lifecycle.rs b/crates/worker/src/feature/builtin/memory_lifecycle.rs index 7349f842..13d501ea 100644 --- a/crates/worker/src/feature/builtin/memory_lifecycle.rs +++ b/crates/worker/src/feature/builtin/memory_lifecycle.rs @@ -77,7 +77,10 @@ impl MemoryLifecycleFeature { workspace_context: WorkerWorkspaceContext, event_tx: Option>, ) -> std::io::Result> { - if !lifecycle_enabled || !config.profile.enabled || !config.profile.extraction.enabled { + if !lifecycle_enabled + || !config.profile.enabled + || (!config.profile.extraction.enabled && !config.profile.consolidation.request_enabled) + { return Ok(None); } config @@ -521,9 +524,12 @@ impl FeatureBackgroundTask for MemoryLifecycleTask { context: BackgroundTaskContext, cancellation: BackgroundTaskCancellation, ) -> Result<(), HookError> { - let extraction = self - .run_extraction(context.clone(), cancellation.clone()) - .await; + let extraction = if self.config.profile.extraction.enabled { + self.run_extraction(context.clone(), cancellation.clone()) + .await + } else { + Ok(()) + }; if !cancellation.is_cancelled() { context.generation_fence.ensure_current()?; if self.config.profile.consolidation.request_enabled { @@ -1355,6 +1361,38 @@ permission = "write" ); } + #[tokio::test] + async fn lifecycle_task_requests_consolidation_when_extraction_is_disabled() { + let client = ScriptClient::new(Vec::new()); + let extension_writes = Arc::new(Mutex::new(Vec::new())); + let (event_tx, _) = broadcast::channel(16); + let workspace_client = Arc::new(RecordingWorkspaceClient::default()); + let mut interrupted = capture(2, 250); + interrupted.run_exit = CommittedRunExit::Interrupted; + let mut task = test_task( + interrupted, + Box::new(client), + extension_writes.clone(), + event_tx, + workspace_client.clone(), + ); + task.config.profile.extraction.enabled = false; + task.config.profile.consolidation.request_enabled = true; + run_background_task(task).await; + + assert!(extension_writes.lock().unwrap().is_empty()); + let requests = workspace_client.requests.lock().unwrap(); + assert!( + requests.iter().any(|request| { + request + .body + .as_deref() + .is_some_and(|body| body == "{\"force\":false}") + }), + "recorded requests: {requests:?}" + ); + } + fn internal_result( lifecycle: WorkerRunResult, ) -> Result {