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![