From 848a1e9a5b5c8d366617e0dda93d7f828e59941f Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 26 Oct 2026 09:16:37 +0900 Subject: [PATCH] feat: worker separate restored controller preparation from publication Prepare the real receiver, full snapshot and paired streams privately; retain activation effects and startup cleanup until the host decision. --- crates/worker/src/bootstrap.rs | 220 +++++++++++++ crates/worker/src/controller.rs | 473 ++++++++++++++++++++++++---- crates/worker/src/feature/mcp.rs | 308 ++++++++++++++++-- crates/worker/src/spawn/registry.rs | 426 ++++++++++++++++++++++--- crates/worker/src/worker.rs | 163 ++++++++-- 5 files changed, 1412 insertions(+), 178 deletions(-) diff --git a/crates/worker/src/bootstrap.rs b/crates/worker/src/bootstrap.rs index 00e6c5ddf..2a591cb3f 100644 --- a/crates/worker/src/bootstrap.rs +++ b/crates/worker/src/bootstrap.rs @@ -323,3 +323,223 @@ where } } } + +/// Host bindings for one unpublished restored Controller. These do not expose +/// Worker execution or permit arbitrary mutation during preparation. +pub struct RestoredControllerBindings { + pub backend_job: Option, + pub workdir_sessions: std::sync::Arc, + pub observation_provider: + Option>, + pub subjektiv_attached: bool, +} + +/// The same real receiver, snapshot sources and paired subscription retained +/// privately across the Runtime's durable restore decision. +pub struct UnpublishedRestoredController { + initialized: Option, St>>, + activation: Option, + streams: Option, + subjektiv_attached: bool, + effects_applied: bool, + activation_failed: bool, + closing: bool, +} + +pub struct ActivatedRestoredController { + pub controller: BootstrappedWorker, + pub streams: crate::ipc::protocol_session::WorkerProtocolSessionStreams, +} + +impl PreparedRestoredWorker +where + St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, +{ + pub fn workspace_client_handle( + &mut self, + ) -> Result, WorkerError> { + Ok(self.worker_mut()?.workspace_client_handle()) + } + + pub fn retain_startup_feature_cleanup( + &mut self, + ) -> Result { + Ok(self.worker_mut()?.retain_startup_feature_cleanup()) + } + + /// Initialize the real Controller without allocation, canonical restore + /// writes, attribution changes, registry reclamation or input consumption. + /// MCP discovery and Feature tasks belong to the retained startup cleanup. + pub async fn prepare_runtime_controller( + mut self, + run_dir: &std::path::Path, + transport: WorkerControllerTransport, + bindings: RestoredControllerBindings, + ) -> Result, WorkerBootstrapError> { + let worker = self.worker_mut().map_err(WorkerBootstrapError::Worker)?; + if let Some(binding) = bindings.backend_job { + worker.bind_backend_job(binding).map_err(|error| { + WorkerBootstrapError::Worker(WorkerError::FeatureLifecycle(error.into())) + })?; + } + worker + .workdir_sessions() + .close_all() + .await + .map_err(|error| { + WorkerBootstrapError::Worker(WorkerError::FeatureLifecycle(error.to_string())) + })?; + worker.bind_workdir_sessions(bindings.workdir_sessions); + worker.bind_worker_observation_provider(bindings.observation_provider); + let flow_enabled = worker.manifest().feature.flow.enabled; + if flow_enabled { + let report = worker + .install_runtime_flow_transition_feature() + .map_err(|error| { + WorkerBootstrapError::Worker(WorkerError::FeatureLifecycle(error)) + })?; + if report.reports.iter().any(|report| !report.installed) { + return Err(WorkerBootstrapError::Worker(WorkerError::FeatureLifecycle( + "Flow transition installation refused".into(), + ))); + } + } + let (worker, activation) = self + .into_controller_parts() + .map_err(WorkerBootstrapError::Worker)?; + let initialized = WorkerController::initialize_restored(worker, run_dir, transport) + .await + .map_err(|source| WorkerBootstrapError::Controller { + source, + cleanup_failed: true, + })?; + if flow_enabled { + initialized.handle.shared_state.enable_flow_transition(); + } + let streams = { + let _guard = initialized.handle.protocol_snapshot_publish_guard(); + crate::ipc::protocol_session::subscribe_worker_protocol_session(&initialized.handle) + }; + Ok(UnpublishedRestoredController { + initialized: Some(initialized), + activation: Some(activation), + streams: Some(streams), + subjektiv_attached: bindings.subjektiv_attached, + effects_applied: false, + activation_failed: false, + closing: false, + }) + } +} + +impl UnpublishedRestoredController +where + St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, +{ + pub fn snapshot(&self) -> &protocol::Event { + &self + .streams + .as_ref() + .expect("unpublished restore streams") + .snapshot_event + } + + /// Only publication failures after confirmed effects can retry on this + /// exact owner. Partial activation must be reconciled, never blindly rerun. + pub fn can_retry_publication(&self) -> bool { + self.initialized.is_some() + && self.effects_applied + && !self.activation_failed + && !self.closing + } + + /// Apply canonical restore effects after the host's durable decision, while + /// retaining the same private Controller for subsequent publication. + pub async fn apply_effects(&mut self) -> Result<(), WorkerBootstrapError> { + let initialized = self + .initialized + .as_mut() + .ok_or_else(|| WorkerBootstrapError::Worker(WorkerError::RestoreCandidateClosing))?; + if self.activation_failed || self.closing { + return Err(WorkerBootstrapError::Worker( + WorkerError::RestoreCandidateClosing, + )); + } + if !self.effects_applied { + self.activation_failed = true; + self.activation + .as_ref() + .expect("restore activation") + .activate(&mut initialized.worker) + .await + .map_err(WorkerBootstrapError::Worker)?; + initialized + .worker + .finalize_subjektiv_session_attribution( + self.subjektiv_attached, + crate::SubjektivSessionAttributionLifecycle::RestoredSession, + ) + .map_err(WorkerBootstrapError::Worker)?; + self.effects_applied = true; + self.activation_failed = false; + self.activation = None; + } + Ok(()) + } + + pub fn allocation_cleanup( + &self, + ) -> Option { + self.initialized + .as_ref() + .and_then(|initialized| initialized.worker.restore_allocation_cleanup_handle()) + } + + pub async fn activate(&mut self) -> Result { + self.apply_effects().await?; + let initialized = self + .initialized + .as_mut() + .expect("same initialized Controller"); + let publication = initialized.prepare_publication().await.map_err(|source| { + WorkerBootstrapError::Controller { + source, + cleanup_failed: true, + } + })?; + let initialized = self + .initialized + .take() + .expect("same initialized Controller"); + let (handle, shutdown, controller_task) = initialized.start(publication); + Ok(ActivatedRestoredController { + controller: BootstrappedWorker { + handle, + shutdown, + controller_task, + }, + streams: self.streams.take().expect("same paired subscription"), + }) + } + + pub async fn discard(&mut self) -> Result<(), WorkerError> { + self.closing = true; + if let Some(initialized) = self.initialized.as_mut() { + initialized + .worker + .retain_startup_feature_cleanup() + .shutdown() + .await + .map_err(|error| WorkerError::FeatureLifecycle(error.to_string()))?; + initialized + .worker + .workdir_sessions() + .close_all() + .await + .map_err(|error| WorkerError::FeatureLifecycle(error.to_string()))?; + } + self.streams = None; + self.initialized = None; + Ok(()) + } +} diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 6721753d6..c20d5d833 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -23,7 +23,7 @@ use crate::shutdown_after_idle::{ ShutdownAfterIdleRequest, TicketIntakeReadyShutdownHook, is_ticket_intake_role, take_shutdown_request_after_status, }; -use crate::spawn::registry::SpawnedWorkerRegistry; +use crate::spawn::registry::{PreparedSpawnedWorkerRegistryLoad, SpawnedWorkerRegistry}; use crate::spawn::tool::sub_worker_spawn_tool; use crate::worker::{ PendingSubmissionHandle, SystemItemCommitter, Worker, WorkerError, WorkerRunResult, @@ -688,6 +688,28 @@ pub enum WorkerControllerTransport { pub struct WorkerController; +// Own the real command receiver and all snapshot sources before starting any +// protocol listener or controller loop. Fresh startup may materialize the head; +// Runtime restore requires an already materialized source and defers canonical +// restore effects until its durable decision. Tool discovery can own external +// resources, so private initialization is not synonymous with side-effect-free. +pub(crate) struct InitializedWorkerController { + pub(crate) worker: Worker, + pub(crate) handle: WorkerHandle, + method_rx: mpsc::Receiver, + shutdown_tx: oneshot::Sender<()>, + shutdown_rx: ShutdownReceiver, + registry_reclaim: PreparedSpawnedWorkerRegistryLoad, + shutdown_after_idle: ShutdownAfterIdleRequest, + transport: WorkerControllerTransport, + command_events: Option, +} + +pub(crate) struct PreparedControllerPublication { + socket_server: Option, + command_observer: Option>, +} + impl WorkerController { /// Explicit recovery of an unmaterialized Session head. This installs the /// saved Feature composition needed to produce the seed, but creates no @@ -877,9 +899,8 @@ impl WorkerController { } async fn spawn_initialized( - mut worker: Worker, + worker: Worker, runtime_base: &Path, - runtime_managed: bool, runtime_run: Option<&Path>, transport: WorkerControllerTransport, @@ -888,6 +909,71 @@ impl WorkerController { C: LlmClient + Clone + 'static, St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, { + Self::initialize( + worker, + runtime_base, + runtime_managed, + runtime_run, + transport, + ) + .await? + .publish() + .await + } + + async fn initialize( + worker: Worker, + runtime_base: &Path, + runtime_managed: bool, + runtime_run: Option<&Path>, + transport: WorkerControllerTransport, + ) -> Result, std::io::Error> + where + C: LlmClient + Clone + 'static, + St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, + { + Self::initialize_inner( + worker, + runtime_base, + runtime_managed, + runtime_run, + transport, + true, + ) + .await + } + + pub(crate) async fn initialize_restored( + worker: Worker, St>, + run_dir: &Path, + transport: WorkerControllerTransport, + ) -> std::io::Result, St>> + where + St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, + { + let parent = run_dir + .parent() + .ok_or_else(|| std::io::Error::other("run path has no parent"))?; + Self::initialize_inner(worker, parent, true, Some(run_dir), transport, false).await + } + + async fn initialize_inner( + mut worker: Worker, + runtime_base: &Path, + runtime_managed: bool, + runtime_run: Option<&Path>, + transport: WorkerControllerTransport, + canonical_startup: bool, + ) -> Result, std::io::Error> + where + C: LlmClient + Clone + 'static, + St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, + { + if !canonical_startup && worker.needs_initial_session_head_materialization() { + return Err(std::io::Error::other( + "restore source requires operation-owned head recovery", + )); + } worker .validate_subjektiv_body_publication() .map_err(std::io::Error::other)?; @@ -913,22 +999,13 @@ impl WorkerController { }); let spawner_name = worker.manifest().worker.name.clone(); - let self_parent_socket = worker.callback_socket().cloned(); - let loaded_registry = SpawnedWorkerRegistry::load_from_worker_state_with_reclaim( + let registry_reclaim = SpawnedWorkerRegistry::prepare_from_worker_state_with_reclaim( runtime_dir.clone(), worker.store().clone(), spawner_name.clone(), Some(worker.scope().clone()), - ) - .await?; - let reclaimed_unreachable = loaded_registry.reclaimed_unreachable; - let spawned_registry = loaded_registry.registry; - if reclaimed_unreachable { - worker.push_notify( - "Restored Worker state contained unreachable delegated child Workers; their delegated write scopes were reclaimed before resume." - .to_string(), - ); - } + )?; + let spawned_registry = registry_reclaim.registry.clone(); // Hand the alerter to the Worker so internal operations (compaction, // AGENTS.md ingestion during the first turn) can emit user-facing @@ -1020,9 +1097,6 @@ impl WorkerController { if let Some(fs_for_view) = fs_for_view.as_ref() { shared_state.set_fs_view(crate::fs_view::WorkerFsView::new(fs_for_view.clone())); } - runtime_dir.write_manifest(&manifest_toml).await?; - runtime_dir.write_status(&shared_state).await?; - let artifact_store: Arc = Arc::new(worker.store().clone()); let session_id = worker.session_id(); let pending_activations = worker.pending_activation_state(); @@ -1047,25 +1121,88 @@ impl WorkerController { pending_activations, }; - // Do not start an observer before fallible head/recovery/runtime-file - // initialization. The only remaining fallible step is socket startup, - // whose error path must abort AND join this task before returning. - let command_observer = fs_for_view + let command_events = fs_for_view .as_ref() - .and_then(|session| wire_workdir_command_events(session, &in_flight)); - let socket_server = match transport { - WorkerControllerTransport::UnixSocket => match SocketServer::start(&handle).await { - Ok(server) => Some(server), - Err(error) => { - if let Some(observer) = command_observer { - observer.abort(); - let _ = observer.await; - } - return Err(error); - } - }, + .map(|session| PreparedWorkdirCommandEvents::new(session, &handle.in_flight)); + Ok(InitializedWorkerController { + worker, + handle, + method_rx, + shutdown_tx, + shutdown_rx, + registry_reclaim, + shutdown_after_idle, + transport, + command_events, + }) + } +} + +impl InitializedWorkerController +where + C: LlmClient + Clone + 'static, + St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, +{ + async fn publish( + mut self, + ) -> Result<(WorkerHandle, ShutdownReceiver, tokio::task::JoinHandle<()>), std::io::Error> { + let publication = self.prepare_publication().await?; + Ok(self.start(publication)) + } + + pub(crate) async fn prepare_publication( + &mut self, + ) -> std::io::Result { + let runtime_dir = &self.handle.runtime_dir; + runtime_dir + .write_manifest(&self.handle.shared_state.manifest_toml) + .await?; + runtime_dir.write_status(&self.handle.shared_state).await?; + let loaded_registry = self.registry_reclaim.apply().await?; + debug_assert!(Arc::ptr_eq( + &loaded_registry.registry, + &self.handle.spawned_registry + )); + if loaded_registry.reclaimed_unreachable { + self.worker.push_notify("Restored Worker state contained unreachable delegated child Workers; their delegated write scopes were reclaimed before resume.".to_string()); + } + let socket_server = match self.transport { + WorkerControllerTransport::UnixSocket => Some(SocketServer::start(&self.handle).await?), WorkerControllerTransport::InProcess => None, }; + let command_observer = self + .command_events + .take() + .and_then(PreparedWorkdirCommandEvents::start); + Ok(PreparedControllerPublication { + socket_server, + command_observer, + }) + } + + pub(crate) fn start( + self, + publication: PreparedControllerPublication, + ) -> (WorkerHandle, ShutdownReceiver, tokio::task::JoinHandle<()>) { + let Self { + mut worker, + handle, + method_rx, + shutdown_tx, + shutdown_rx, + shutdown_after_idle, + .. + } = self; + let PreparedControllerPublication { + socket_server, + command_observer, + } = publication; + let runtime_dir = handle.runtime_dir.clone(); + let shared_state = handle.shared_state.clone(); + let working_event_tx = handle.working_event_tx.clone(); + let spawned_registry = handle.spawned_registry.clone(); + let self_parent_socket = worker.callback_socket().cloned(); + let spawner_name = worker.manifest().worker.name.clone(); // === 5. controller_loop === // Clone cancel sender and notification buffer before moving worker @@ -1094,7 +1231,7 @@ impl WorkerController { command_observer, )); - Ok((handle, shutdown_rx, controller_task)) + (handle, shutdown_rx, controller_task) } } @@ -1102,30 +1239,51 @@ pub(crate) fn wire_workdir_command_events( session: &Arc, in_flight: &InFlightEvents, ) -> Option> { - in_flight.replace_command_snapshot(protocol_command_snapshots(session.as_ref())); - let Some(mut events) = session.subscribe_command_events() else { - return None; - }; - // Keep only a weak reference in the observer task. Holding the session - // strongly here would keep its broadcast sender alive forever and prevent - // the receiver from observing closure during Worker teardown. - let session = Arc::downgrade(session); - let in_flight = in_flight.clone(); - Some(tokio::spawn(async move { - loop { - match events.recv().await { - Ok(event) => in_flight.publish_command_event(protocol_command_event(event)), - Err(broadcast::error::RecvError::Lagged(_)) => { - let Some(session) = session.upgrade() else { - break; - }; - in_flight - .replace_command_snapshot(protocol_command_snapshots(session.as_ref())); - } - Err(broadcast::error::RecvError::Closed) => break, - } + PreparedWorkdirCommandEvents::new(session, in_flight).start() +} + +struct PreparedWorkdirCommandEvents { + session: std::sync::Weak, + in_flight: InFlightEvents, + events: Option>, +} + +impl PreparedWorkdirCommandEvents { + fn new(session: &Arc, in_flight: &InFlightEvents) -> Self { + // Subscribe before reading command state; hold the receiver privately + // across the durable decision instead of creating a publication gap. + let events = session.subscribe_command_events(); + in_flight.replace_command_snapshot(protocol_command_snapshots(session.as_ref())); + Self { + session: Arc::downgrade(session), + in_flight: in_flight.clone(), + events, } - })) + } + + fn start(self) -> Option> { + let Self { + session, + in_flight, + events, + } = self; + let mut events = events?; + Some(tokio::spawn(async move { + loop { + match events.recv().await { + Ok(event) => in_flight.publish_command_event(protocol_command_event(event)), + Err(broadcast::error::RecvError::Lagged(_)) => { + let Some(session) = session.upgrade() else { + break; + }; + in_flight + .replace_command_snapshot(protocol_command_snapshots(session.as_ref())); + } + Err(broadcast::error::RecvError::Closed) => break, + } + } + })) + } } fn protocol_command_snapshots(session: &dyn WorkdirSession) -> Vec { @@ -1851,9 +2009,17 @@ where .add_module(crate::feature::builtin::orchestration::orchestration_feature()); } if let Some(workspace_root) = local_workspace_root.as_ref() { - if let Some(module) = + let module = if let Some(cleanup) = worker.startup_feature_cleanup_handle() { + crate::feature::mcp::discover_stdio_tool_feature_with_cleanup( + &mcp_config, + workspace_root, + &cleanup.mcp_discovery(), + ) + .await? + } else { crate::feature::mcp::discover_stdio_tool_feature(&mcp_config, workspace_root).await - { + }; + if let Some(module) = module { feature_registry = feature_registry.with_module(module); } } @@ -4123,6 +4289,171 @@ mod tests { } } + // Component connection: real Worker/Store, initialized endpoint, paired + // protocol subscription and Controller loop. The model boundary is an + // empty non-network script; this test must not request any generation. + #[tokio::test] + async fn initialized_controller_publishes_the_same_endpoint_and_subscription() { + let root = TempDir::new().unwrap(); + let manifest = manifest::WorkerManifest::from_toml( + "[worker]\nname = 'p'\n[model]\nscheme = 'anthropic'\nmodel_id = 'test'\n[engine]\n[scope]\n", + ) + .unwrap(); + let client = ScriptedClient::new(vec![]); + let requests = client.requests.clone(); + let engine = + Engine::<_, agen::state::Mutable, crate::SessionHistoryMetadata>::new_annotated( + client, + ); + let store = session_store::CombinedStore::new( + session_store::FsStore::new(root.path().join("sessions")).unwrap(), + session_store::FsWorkerStore::new(root.path().join("workers")).unwrap(), + ); + let mut worker = Worker::new( + manifest, + engine, + store, + crate::worker::WorkerWorkspaceContext::local_filesystem(None), + crate::worker::WorkerFilesystemAuthority::None, + manifest::Scope::empty(), + ) + .await + .unwrap(); + worker.set_system_prompt_template( + crate::SystemPromptTemplate::parse( + "default", + crate::PromptCatalogSource::builtins_only(), + ) + .unwrap(), + ); + let prepared = WorkerController::initialize( + worker, + root.path(), + false, + None, + WorkerControllerTransport::UnixSocket, + ) + .await + .unwrap(); + let endpoint = prepared.handle.clone(); + let mut streams = + crate::ipc::protocol_session::subscribe_worker_protocol_session(&endpoint); + let entries = endpoint.committed_entries(); + assert!(matches!( + &streams.snapshot_event, + Event::Snapshot { greeting, state, session, .. } + if !greeting.tools.is_empty() + && state.state == WorkerState::Idle + && session == &session_store::public_snapshot::project_current_session_snapshot(&entries) + )); + assert!(!endpoint.protocol_is_closed()); + assert!(!endpoint.runtime_dir.socket_path().exists()); + assert!(!endpoint.runtime_dir.path().join("status.json").exists()); + assert!(!endpoint.runtime_dir.path().join("manifest.toml").exists()); + assert!(matches!( + streams.events.try_recv(), + Err(broadcast::error::TryRecvError::Empty) + )); + assert!(requests.lock().unwrap().is_empty()); + + let (handle, shutdown, task) = prepared.publish().await.unwrap(); + assert!(handle.same_controller(&endpoint)); + assert!(handle.runtime_dir.socket_path().exists()); + assert!(handle.runtime_dir.path().join("status.json").exists()); + handle.send(Method::ListPendingSubmissions).await.unwrap(); + let pending = tokio::time::timeout(Duration::from_secs(5), async { + loop { + if let Event::PendingSubmissionsChanged { pending } = + streams.events.recv().await.unwrap() + { + break pending; + } + } + }) + .await + .unwrap(); + assert_eq!(pending, protocol::PendingSubmissionsSnapshot::default()); + assert_eq!( + serde_json::to_value(handle.committed_entries()).unwrap(), + serde_json::to_value(entries).unwrap(), + ); + handle + .send(Method::Shutdown { + command: WorkerCommandEnvelope { command_id: 1 }, + }) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), shutdown) + .await + .unwrap() + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), task) + .await + .unwrap() + .unwrap(); + assert!(endpoint.protocol_is_closed()); + assert!(!endpoint.runtime_dir.socket_path().exists()); + assert!(requests.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn tool_initialization_failure_does_not_reclaim_saved_children() { + let root = TempDir::new().unwrap(); + let mut manifest = manifest::WorkerManifest::from_toml( + "[worker]\nname = 'p'\n[model]\nscheme = 'anthropic'\nmodel_id = 'test'\n[engine]\n[scope]\n", + ).unwrap(); + manifest.feature.ticket.enabled = true; + let client = ScriptedClient::new(vec![]); + let requests = client.requests.clone(); + let store = session_store::CombinedStore::new( + session_store::FsStore::new(root.path().join("sessions")).unwrap(), + session_store::FsWorkerStore::new(root.path().join("workers")).unwrap(), + ); + // New feature activation is refused by the real unavailable host; + // legacy child metadata must not be consumed before that refusal. + let worker = Worker::new( + manifest, + Engine::<_, agen::state::Mutable, crate::SessionHistoryMetadata>::new_annotated( + client, + ), + store.clone(), + crate::worker::WorkerWorkspaceContext::no_workspace(), + crate::worker::WorkerFilesystemAuthority::None, + manifest::Scope::empty(), + ) + .await + .unwrap(); + let mut metadata = session_store::WorkerMetadata::new( + "p", + Some(session_store::WorkerActiveSegmentRef { + session_id: worker.session_id(), + segment_id: Some(worker.segment_id()), + }), + ); + metadata + .spawned_children + .push(session_store::WorkerSpawnedChild { + worker_name: "legacy".into(), + socket_path: root.path().join("legacy.sock"), + callback_address: root.path().join("p/sock"), + scope_delegated: Vec::new(), + }); + store.write(&metadata).unwrap(); + let before = store.read_by_name("p").unwrap(); + let result = WorkerController::initialize( + worker, + root.path(), + false, + None, + WorkerControllerTransport::InProcess, + ) + .await; + assert!(result.is_err()); + assert_eq!(store.read_by_name("p").unwrap(), before); + assert!(!root.path().join("p/spawned_workers.json").exists()); + assert!(requests.lock().unwrap().is_empty()); + } + #[derive(Debug)] struct DriveWorkspaceHost { available: bool, @@ -4450,20 +4781,24 @@ mod tests { assert!(feature_install < materialize_gate); assert!(materialize_gate < materialize); assert!(materialize < exposure); - let observer = startup.find("let command_observer = fs_for_view").unwrap(); + let observer_preparation = startup + .find("PreparedWorkdirCommandEvents::new(session") + .unwrap(); let recover = startup.find(".recover_unfinished_compaction()").unwrap(); let files = startup - .find("runtime_dir.write_status(&shared_state).await?") + .find(".write_status(&self.handle.shared_state)") + .unwrap(); + let listener = startup + .find("Some(SocketServer::start(&self.handle).await?)") + .unwrap(); + let observer_start = startup + .find(".and_then(PreparedWorkdirCommandEvents::start)") .unwrap(); let loop_start = startup.find("tokio::spawn(controller_loop(").unwrap(); - assert!(recover < observer && files < observer && observer < loop_start); - let socket_startup = &startup[observer..loop_start]; - let abort = socket_startup.find("observer.abort()").unwrap(); - let join = socket_startup.find("observer.await").unwrap(); - let failure = socket_startup.find("return Err(error)").unwrap(); + assert!(recover < observer_preparation && observer_preparation < files); assert!( - abort < join && join < failure, - "startup cannot detach its observer on error" + files < listener && listener < observer_start && observer_start < loop_start, + "observer must not start before the last fallible listener preparation" ); assert!( startup[materialize_gate..materialize] diff --git a/crates/worker/src/feature/mcp.rs b/crates/worker/src/feature/mcp.rs index 259af6db3..a617acb72 100644 --- a/crates/worker/src/feature/mcp.rs +++ b/crates/worker/src/feature/mcp.rs @@ -8,6 +8,7 @@ //! the regular registry decides Tool-name collisions. use std::collections::BTreeMap; +use std::io; use std::path::Path; use std::sync::Arc; @@ -25,6 +26,7 @@ use mcp::stdio::{ }; use serde::Serialize; use serde_json::{Map, Value}; +use tokio::sync::Mutex; use super::{ FeatureDescriptor, FeatureDiagnostic, FeatureInstallContext, FeatureInstallError, @@ -52,8 +54,88 @@ const MAX_LIST_ITEMS: usize = 128; const MAX_RESOURCE_CONTENTS: usize = 16; const MAX_PROMPT_MESSAGES: usize = 32; +/// Owns discovery clients until process reaping is confirmed. Retain a clone in +/// startup cleanup *before* awaiting discovery. Stop/cancel discovery before +/// calling `shutdown`; the shared lock serializes discovery and cleanup, not +/// cancellation of an in-flight discovery task. +/// +/// Error or cancellation leaves the actual clients here for another shutdown +/// attempt. Dropping the final owner only provides best-effort process killing. +#[derive(Clone, Default)] +pub struct McpDiscoveryCleanup { + clients: Arc>>, +} + +impl McpDiscoveryCleanup { + pub fn new() -> Self { + Self::default() + } + + /// Try every retained client once, removing only confirmed shutdowns. An + /// error is not a cleanup receipt; retry this same owner, not discovery. + pub async fn shutdown(&self) -> io::Result<()> { + let mut clients = self.clients.lock().await; + let mut first_error = None; + let mut index = 0; + while index < clients.len() { + match clients[index].shutdown().await { + Ok(_) => { + clients.remove(index); + } + Err(err) => { + first_error.get_or_insert_with(|| discovery_shutdown_error(&err)); + index += 1; + } + } + } + match first_error { + Some(err) => Err(err), + None => Ok(()), + } + } +} + +fn discovery_shutdown_error(err: &McpClientError) -> io::Error { + io::Error::other(bounded_diagnostic(format!( + "MCP discovery shutdown unconfirmed; client retained for cleanup retry: {err}" + ))) +} + +/// Restore preparation must fail if discovery cannot confirm process shutdown. +/// The caller retains `cleanup` before this future is polled, including during +/// initialization. Ordinary discovery failures remain isolated contributions +/// only after their client has been reaped. +pub async fn discover_stdio_tool_feature_with_cleanup( + config: &McpConfig, + workspace_root: &Path, + cleanup: &McpDiscoveryCleanup, +) -> Result, io::Error> { + if config.stdio_servers.is_empty() { + return Ok(None); + } + let mut feature = McpStdioToolFeature::new(); + for server in &config.stdio_servers { + match resolve_stdio_server(server, workspace_root, None) { + Ok(spec) => { + feature.add_contribution( + discover_server_tools_with_cleanup(spec, cleanup, McpStdioLimits::default()) + .await?, + ); + } + Err(err) => { + feature.add_diagnostic(FeatureDiagnostic::error(bounded_diagnostic(format!( + "failed to resolve MCP stdio server `{}`: {err}", + server.name + )))) + } + } + } + Ok(Some(feature)) +} + /// Discover enabled MCP stdio server tools and return a single feature module /// containing startup contributions for normal ToolRegistry installation. +/// Callers needing retryable startup cleanup must use the cleanup-aware API. pub async fn discover_stdio_tool_feature( config: &McpConfig, workspace_root: &Path, @@ -81,21 +163,66 @@ pub async fn discover_stdio_tool_feature( } async fn discover_server_tools(spec: McpStdioServerSpec) -> ProtocolProviderContribution { + let declaration = provider_declaration(&spec.name, None); + let cleanup = McpDiscoveryCleanup::new(); + match discover_server_tools_with_cleanup(spec, &cleanup, McpStdioLimits::default()).await { + Ok(contribution) => contribution, + Err(err) => { + ProtocolProviderContribution::failed(declaration, bounded_diagnostic(err.to_string())) + } + } +} + +async fn discover_server_tools_with_cleanup( + spec: McpStdioServerSpec, + cleanup: &McpDiscoveryCleanup, + limits: McpStdioLimits, +) -> io::Result { + let declaration = provider_declaration(&spec.name, None); + let mut clients = cleanup.clients.lock().await; + if !clients.is_empty() { + return Err(io::Error::other( + "MCP discovery cleanup is pending; shut down the retained owner before retrying discovery", + )); + } + let client = match McpStdioClient::spawn_uninitialized(spec.clone(), limits) { + Ok(client) => client, + Err(err) => { + // The synchronous spawn API only errors before a process exists; + // pipe/handshake failures retain a client and use the shutdown path. + return Ok(ProtocolProviderContribution::failed( + declaration, + bounded_diagnostic(err.to_string()), + )); + } + }; + // Spawn is synchronous: the owner holds the client before the first await + // that can leave a running process on initialization failure/cancellation. + clients.push(client); + let client = clients.last_mut().expect("retained discovery client"); + let contribution = match client.initialize_with_timeout().await { + Ok(()) => discover_connected_server_tools(spec, client).await, + Err(err) => { + ProtocolProviderContribution::failed(declaration, bounded_diagnostic(err.to_string())) + } + }; + client + .shutdown() + .await + .map_err(|err| discovery_shutdown_error(&err))?; + drop(clients.pop()); + Ok(contribution) +} + +async fn discover_connected_server_tools( + spec: McpStdioServerSpec, + client: &mut McpStdioClient, +) -> ProtocolProviderContribution { let declaration = provider_declaration(&spec.name, None); let mut contribution = ProtocolProviderContribution::ready(declaration.clone()); let server_namespace = sanitize_segment(&spec.name); let execution_spec = spec.clone(); - let mut client = match McpStdioClient::connect(spec, McpStdioLimits::default()).await { - Ok(client) => client, - Err(err) => { - return ProtocolProviderContribution::failed( - declaration, - bounded_diagnostic(err.to_string()), - ); - } - }; - let (server_version, has_tools, has_resources, has_prompts) = client .initialize_result() .map(|result| { @@ -133,18 +260,10 @@ async fn discover_server_tools(spec: McpStdioServerSpec) -> ProtocolProviderCont { Ok(list) => list, Err(err) => { - let mut failed = ProtocolProviderContribution::failed( + return ProtocolProviderContribution::failed( declaration, bounded_diagnostic(err.to_string()), ); - if let Err(shutdown_err) = client.shutdown().await { - failed = failed.with_diagnostic(FeatureDiagnostic::warning( - bounded_diagnostic(format!( - "MCP server shutdown after discovery failure failed: {shutdown_err}" - )), - )); - } - return failed; } }; let list_changes = client.snapshot_list_changes().await; @@ -166,20 +285,12 @@ async fn discover_server_tools(spec: McpStdioServerSpec) -> ProtocolProviderCont { Ok(list) => list, Err(err) => { - let mut failed = ProtocolProviderContribution::failed( + return ProtocolProviderContribution::failed( declaration, bounded_diagnostic(format!( "MCP server sent notifications/tools/list_changed during tool discovery and refresh failed: {err}" )), ); - if let Err(shutdown_err) = client.shutdown().await { - failed = failed.with_diagnostic(FeatureDiagnostic::warning( - bounded_diagnostic(format!( - "MCP server shutdown after discovery refresh failure failed: {shutdown_err}" - )), - )); - } - return failed; } }; let refresh_changes = client.snapshot_list_changes().await; @@ -197,13 +308,6 @@ async fn discover_server_tools(spec: McpStdioServerSpec) -> ProtocolProviderCont } else { None }; - let shutdown_result = client.shutdown().await; - if let Err(err) = shutdown_result { - contribution = - contribution.with_diagnostic(FeatureDiagnostic::warning(bounded_diagnostic(format!( - "MCP server shutdown after capability discovery failed: {err}" - )))); - } if let Some(list) = list { contribution = normalize_listed_tools( @@ -1819,6 +1923,140 @@ read exit_notification || true McpStdioServerSpec::new("shell-mock", "/bin/sh").args(["-c".to_string(), script]) } + #[cfg(unix)] + fn discovery_config(spec: McpStdioServerSpec) -> McpConfig { + McpConfig { + stdio_servers: vec![manifest::McpStdioServerConfig { + name: spec.name, + command: spec.command, + args: spec.args, + cwd: None, + env: Default::default(), + }], + } + } + + #[cfg(unix)] + #[tokio::test] + async fn cancelled_discovery_retains_initializing_client_in_cleanup_clone() { + use std::future::{Future, poll_fn}; + use std::task::Poll; + use std::time::Duration; + + let config = discovery_config( + McpStdioServerSpec::new("initializing", "/bin/sh") + .args(["-c", "while read line; do :; done"]), + ); + let cleanup = McpDiscoveryCleanup::new(); + let retained = cleanup.clone(); + let mut discovery = Box::pin(discover_stdio_tool_feature_with_cleanup( + &config, + Path::new("."), + &cleanup, + )); + poll_fn(|cx| { + assert!(discovery.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + drop(discovery); + drop(cleanup); + assert_eq!(retained.clients.lock().await.len(), 1); + let retry_error = + discover_stdio_tool_feature_with_cleanup(&config, Path::new("."), &retained) + .await + .err() + .expect("discovery must not replace an unreaped client"); + assert!(retry_error.to_string().contains("cleanup is pending")); + assert_eq!(retained.clients.lock().await.len(), 1); + tokio::time::timeout(Duration::from_secs(3), retained.shutdown()) + .await + .expect("cleanup bounded") + .expect("initializing process reaped"); + assert!(retained.clients.lock().await.is_empty()); + retained + .shutdown() + .await + .expect("confirmed cleanup repeatable"); + } + + #[cfg(unix)] + #[tokio::test] + async fn cancelled_owner_shutdown_keeps_client_for_retry() { + use std::future::{Future, poll_fn}; + use std::task::Poll; + use std::time::Duration; + + let cleanup = McpDiscoveryCleanup::new(); + cleanup.clients.lock().await.push( + McpStdioClient::spawn_uninitialized( + McpStdioServerSpec::new("shutdown-pending", "/bin/sh") + .args(["-c", "while :; do :; done"]), + McpStdioLimits { + shutdown_timeout: Duration::from_millis(50), + kill_timeout: Duration::from_millis(100), + ..Default::default() + }, + ) + .expect("local process spawned"), + ); + let mut shutdown = Box::pin(cleanup.shutdown()); + poll_fn(|cx| { + assert!(shutdown.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + drop(shutdown); + assert_eq!(cleanup.clients.lock().await.len(), 1); + tokio::time::timeout(Duration::from_secs(2), cleanup.shutdown()) + .await + .expect("retry bounded") + .expect("retained client shutdown"); + assert!(cleanup.clients.lock().await.is_empty()); + } + + #[cfg(unix)] + #[tokio::test] + async fn cleanup_aware_discovery_reaps_success_and_initialize_failure() { + let cleanup = McpDiscoveryCleanup::new(); + let config = discovery_config(shell_capability_server(r#"{"resources":{}}"#)); + let feature = discover_stdio_tool_feature_with_cleanup(&config, Path::new("."), &cleanup) + .await + .expect("discovery cleanup confirmed") + .expect("configured feature"); + assert!( + feature.contributions[0] + .tools + .iter() + .any(|tool| { tool.name() == "Mcp_shell_mock_resources_read" }) + ); + assert!(cleanup.clients.lock().await.is_empty()); + + let spec = McpStdioServerSpec::new("init-rejected", "/bin/sh").args([ + "-c", + r#"read init || exit 1 +printf '%s\n' '{"jsonrpc":"2.0","id":1,"error":{"code":-32000,"message":"initialization rejected"}}' +while read line; do :; done +"#, + ]); + let config = discovery_config(spec); + let feature = discover_stdio_tool_feature_with_cleanup(&config, Path::new("."), &cleanup) + .await + .expect("failure isolated only after cleanup") + .expect("failed contribution preserved"); + assert!(feature.contributions[0].tools.is_empty()); + assert!( + feature.contributions[0] + .diagnostics + .iter() + .any(|diagnostic| { + diagnostic.severity == FeatureDiagnosticSeverity::Error + && diagnostic.message.contains("initialization rejected") + }) + ); + assert!(cleanup.clients.lock().await.is_empty()); + } + #[test] fn mcp_resource_prompt_operation_tools_are_namespaced_and_untrusted() { let declaration = provider_declaration("demo", Some("1.0.0")); diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index dd45b6f47..f2528b4e0 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -340,6 +340,84 @@ pub struct SpawnedWorkerRegistryLoad { pub reclaimed_unreachable: bool, } +/// Read-only restore preparation using the registry that will be published by the Host. +/// The Host must retain its lifecycle exclusion from preparation through application; +/// the metadata reread detects source drift, but is not an atomic store transaction. +/// Dropping an unapplied plan does not reclaim legacy child state. +pub(crate) struct PreparedSpawnedWorkerRegistryLoad { + pub(crate) registry: Arc, + /// Projected outcome, not evidence that reclamation has already happened. + pub(crate) reclaimed_unreachable: bool, + runtime_dir: Arc, + store: St, + worker_name: String, + parent_scope: Option, + persisted_children: Vec, + valid_records: Vec, + legacy_projection_exists: bool, +} + +impl PreparedSpawnedWorkerRegistryLoad +where + St: WorkerMetadataStore + Clone + Send + Sync + 'static, +{ + /// Apply only after all other fallible tool initialization, before socket publication. + /// Effects retain their original order: projection, canonical reclaim history, then + /// allocation/scope. An error may follow committed effects; no rollback is attempted. + pub(crate) async fn apply(&mut self) -> io::Result { + let metadata = self + .store + .read_by_name(&self.worker_name) + .map_err(store_error_to_io)?; + let reclaimed: Vec<_> = self + .persisted_children + .iter() + .map(reclaimed_child_from_metadata) + .collect(); + // A prior apply may have committed the canonical reclaim and then failed + // in allocation/scope cleanup. Under the host's same lifecycle lock, + // confirm its exact retained facts rather than appending them again. + let already_reclaimed = !reclaimed.is_empty() + && metadata.as_ref().is_some_and(|metadata| { + metadata.spawned_children.is_empty() + && metadata.reclaimed_children.ends_with(&reclaimed) + }); + if !already_reclaimed + && metadata + .as_ref() + .map(|metadata| metadata.spawned_children.as_slice()) + .unwrap_or_default() + != self.persisted_children.as_slice() + { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "spawned-child metadata changed after registry preparation", + )); + } + + // Runtime projection is migration input only; the normal Internal registry is never + // materialized into spawned_workers.json. Only existence was captured, not its contents. + if !self.persisted_children.is_empty() || self.legacy_projection_exists { + self.runtime_dir.write_spawned_workers(&[]).await?; + } + if !reclaimed.is_empty() && !already_reclaimed { + self.store + .reclaim_spawned_children(&self.worker_name, reclaimed) + .map_err(store_error_to_io)?; + } + for record in &self.valid_records { + reclaim_record(&self.worker_name, self.parent_scope.as_ref(), record)?; + } + + let reclaimed_unreachable = self.reclaimed_unreachable; + self.reclaimed_unreachable = false; + Ok(SpawnedWorkerRegistryLoad { + registry: self.registry.clone(), + reclaimed_unreachable, + }) + } +} + impl SpawnedWorkerRegistry { pub(crate) fn bash_output(&self) -> tools::bash_output::BashOutputView { self.bash_output.lock().unwrap().clone() @@ -426,6 +504,23 @@ impl SpawnedWorkerRegistry { worker_name: String, parent_scope: Option, ) -> io::Result + where + St: WorkerMetadataStore + Clone + Send + Sync + 'static, + { + Self::prepare_from_worker_state_with_reclaim(runtime_dir, store, worker_name, parent_scope)? + .apply() + .await + } + + /// Capture legacy child facts without starting children or changing metadata, + /// projection, allocation, or scope. Tool initialization must clone this registry, + /// not construct a replacement, and defer reclamation to the owned plan's `apply`. + pub(crate) fn prepare_from_worker_state_with_reclaim( + runtime_dir: Arc, + store: St, + worker_name: String, + parent_scope: Option, + ) -> io::Result> where St: WorkerMetadataStore + Clone + Send + Sync + 'static, { @@ -442,54 +537,33 @@ impl SpawnedWorkerRegistry { Ok(record) => { warn!( worker = %record.worker_name, - "reclaiming legacy persisted process SubWorker during Internal session restore" + "captured legacy persisted process SubWorker for reclamation during Internal session restore" ); valid_records.push(record); } Err(error) => warn!( error = %error, worker = %child.worker_name, - "clearing corrupt legacy persisted process SubWorker record" + "captured corrupt legacy persisted process SubWorker record for clearing" ), } } - // Runtime projection is migration input only; the normal Internal registry is never - // materialized into spawned_workers.json. let legacy_projection_exists = runtime_dir.path().join("spawned_workers.json").exists(); - if !persisted_children.is_empty() || legacy_projection_exists { - runtime_dir.write_spawned_workers(&[]).await?; - } - if !persisted_children.is_empty() { - let reclaimed = persisted_children - .iter() - .map(reclaimed_child_from_metadata) - .collect(); - store - .reclaim_spawned_children(&worker_name, reclaimed) - .map_err(store_error_to_io)?; - } - for record in &valid_records { - reclaim_record(&worker_name, parent_scope.as_ref(), record)?; - } - - Ok(SpawnedWorkerRegistryLoad { - registry: Arc::new(Self { - internal_records: std::sync::Mutex::new(Vec::new()), - service_records: std::sync::Mutex::new(Vec::new()), - #[cfg(test)] - fail_service_stops: std::sync::Mutex::new(HashMap::new()), - internal_names: std::sync::Mutex::new(HashSet::new()), - internal_shutting_down: AtomicBool::new(false), - pending_internal_spawns: AtomicUsize::new(0), - pending_internal_notify: Notify::new(), - internal_spawn_cleanup_failed: AtomicBool::new(false), - parent_scope, - bash_output: Mutex::new(tools::bash_output::BashOutputView::new()), - parent_protocol: Mutex::new(None), - protocol_publish_lock: Arc::new(Mutex::new(())), - }), + let registry = match &parent_scope { + Some(scope) => Self::new_internal(worker_name.clone(), scope.clone()), + None => Self::new(runtime_dir.clone()), + }; + Ok(PreparedSpawnedWorkerRegistryLoad { + registry, reclaimed_unreachable: !persisted_children.is_empty(), + runtime_dir, + store, + worker_name, + parent_scope, + persisted_children, + valid_records, + legacy_projection_exists, }) } @@ -1140,6 +1214,10 @@ impl Drop for SpawnedWorkerRegistry { .filter(|record| !record.scope_reclaimed.load(Ordering::Acquire)) .flat_map(delegated_write_rules) .collect::>(); + // Empty prepared registries own no delegations, including when abandoned or rejected. + if write_rules.is_empty() { + return; + } let _ = parent_scope.update(|current| current.with_removed_deny_rules(write_rules)); } } @@ -1165,19 +1243,18 @@ fn reclaim_record( parent_scope: Option<&SharedScope>, record: &SpawnedWorkerRecord, ) -> io::Result<()> { - if let Ok(path) = worker_allocation::default_allocation_path() { - if let Ok(mut guard) = worker_allocation::LockFileGuard::open(&path) { - match worker_allocation::reclaim_delegated_scope( - &mut guard, - parent_name, - &record.worker_name, - &record.scope_delegated, - ) { - Ok(()) | Err(worker_allocation::ScopeLockError::UnknownWorker(_)) => {} - Err(error) => return Err(io::Error::other(error)), - } - } + let path = worker_allocation::default_allocation_path().map_err(io::Error::other)?; + let mut guard = worker_allocation::LockFileGuard::open(&path).map_err(io::Error::other)?; + match worker_allocation::reclaim_delegated_scope( + &mut guard, + parent_name, + &record.worker_name, + &record.scope_delegated, + ) { + Ok(()) | Err(worker_allocation::ScopeLockError::UnknownWorker(_)) => {} + Err(error) => return Err(io::Error::other(error)), } + drop(guard); if let Some(parent_scope) = parent_scope { let write_rules = record .scope_delegated @@ -1240,10 +1317,18 @@ fn store_error_to_io(error: WorkerStoreError) -> io::Error { #[cfg(test)] mod tests { + //! Legacy restore cases connect the registry to canonical aggregate storage, + //! runtime projections, and temporary SharedScope rules. Successful apply uses + //! a corrupt legacy record to exercise durable clearing without consulting the + //! process-global legacy allocation file; valid-record scope reclaim is also + //! covered by the existing external restore test. + + use std::path::PathBuf; use std::time::Duration; use manifest::{Scope, ScopeConfig}; - use session_store::LogEntry; + use session_store::{LogEntry, WorkerAggregateStore, WorkerMetadata, WorkerSpawnedScopeRule}; + use tempfile::TempDir; use super::*; use crate::internal_worker::{ @@ -1251,6 +1336,249 @@ mod tests { test_internal_worker_session_with_greeting, }; + struct LegacyRestoreFixture { + runtime_dir: Arc, + store: WorkerAggregateStore, + parent_scope: SharedScope, + scope_root: PathBuf, + metadata_path: PathBuf, + // Keep all storage and scope paths alive until the other fixture fields are dropped. + _root: TempDir, + } + + impl LegacyRestoreFixture { + async fn new(permission: &str) -> Self { + let root = TempDir::new().unwrap(); + let scope_root = root.path().join("scope"); + std::fs::create_dir(&scope_root).unwrap(); + let aggregate = root.path().join("aggregate"); + let store = WorkerAggregateStore::new(&aggregate, "parent").unwrap(); + let runtime_dir = Arc::new( + RuntimeDir::create(&root.path().join("runtime"), "parent") + .await + .unwrap(), + ); + let mut metadata = + WorkerMetadata::new("parent", None).with_workspace_id("restore-test-workspace"); + metadata.spawned_children.push(WorkerSpawnedChild { + worker_name: "legacy-child".into(), + socket_path: root.path().join("legacy.sock"), + callback_address: runtime_dir.socket_path(), + scope_delegated: vec![WorkerSpawnedScopeRule { + target: scope_root.clone(), + permission: permission.into(), + recursive: true, + symlink_policy: Default::default(), + }], + }); + store.write(&metadata).unwrap(); + let write_rule = ScopeRule { + target: scope_root.clone(), + permission: Permission::Write, + recursive: true, + symlink_policy: Default::default(), + }; + let parent_scope = SharedScope::new( + Scope::from_config(&ScopeConfig { + allow: vec![write_rule.clone()], + deny: vec![write_rule], + }) + .unwrap(), + ); + Self { + runtime_dir, + store, + parent_scope, + scope_root, + metadata_path: aggregate.join("metadata.json"), + _root: root, + } + } + + fn prepare(&self) -> PreparedSpawnedWorkerRegistryLoad { + SpawnedWorkerRegistry::prepare_from_worker_state_with_reclaim( + self.runtime_dir.clone(), + self.store.clone(), + "parent".into(), + Some(self.parent_scope.clone()), + ) + .unwrap() + } + + fn projection_path(&self) -> PathBuf { + self.runtime_dir.path().join("spawned_workers.json") + } + } + + #[tokio::test] + async fn registry_preparation_and_abandonment_leave_legacy_state_unchanged() { + for projection_exists in [false, true] { + let fixture = LegacyRestoreFixture::new("write").await; + if projection_exists { + // Projection contents are deliberately not decoded during preparation. + std::fs::write(fixture.projection_path(), b"obsolete projection").unwrap(); + } + let metadata_before = std::fs::read(&fixture.metadata_path).unwrap(); + let projection_before = std::fs::read(fixture.projection_path()).ok(); + let scope_before = fixture.parent_scope.snapshot(); + assert!(!scope_before.is_writable(&fixture.scope_root)); + + let prepared = fixture.prepare(); + + assert!(prepared.reclaimed_unreachable); + assert!(prepared.registry.list_internal().is_empty()); + assert!(prepared.registry.internal_worker_snapshots().is_empty()); + assert_eq!( + std::fs::read(&fixture.metadata_path).unwrap(), + metadata_before + ); + assert_eq!( + std::fs::read(fixture.projection_path()).ok(), + projection_before + ); + assert!(Arc::ptr_eq(&scope_before, &fixture.parent_scope.snapshot())); + assert!(!fixture.runtime_dir.socket_path().exists()); + drop(prepared); + assert_eq!( + std::fs::read(&fixture.metadata_path).unwrap(), + metadata_before + ); + assert_eq!( + std::fs::read(fixture.projection_path()).ok(), + projection_before + ); + assert!(Arc::ptr_eq(&scope_before, &fixture.parent_scope.snapshot())); + } + } + + #[tokio::test] + async fn prepared_apply_clears_canonical_children_and_projection_on_same_registry() { + let fixture = LegacyRestoreFixture::new("unsupported-legacy-permission").await; + std::fs::write(fixture.projection_path(), b"obsolete projection").unwrap(); + let scope_before = fixture.parent_scope.snapshot(); + let mut prepared = fixture.prepare(); + let shared_registry = prepared.registry.clone(); + // Unrelated metadata may advance while tools are initialized. + fixture.store.add_peer("parent", "new-peer").unwrap(); + let mut expected = fixture.store.read_by_name("parent").unwrap().unwrap(); + expected.reclaimed_children.extend( + expected + .spawned_children + .iter() + .map(reclaimed_child_from_metadata), + ); + expected.spawned_children.clear(); + + let loaded = prepared.apply().await.unwrap(); + + assert!(loaded.reclaimed_unreachable); + assert!(Arc::ptr_eq(&shared_registry, &loaded.registry)); + assert!(shared_registry.list_internal().is_empty()); + assert_eq!( + fixture.store.read_by_name("parent").unwrap(), + Some(expected) + ); + let projection: serde_json::Value = + serde_json::from_slice(&std::fs::read(fixture.projection_path()).unwrap()).unwrap(); + assert_eq!(projection, serde_json::json!([])); + // Corrupt records are audited and cleared, not interpreted as scope grants. + assert!(Arc::ptr_eq(&scope_before, &fixture.parent_scope.snapshot())); + assert!( + !fixture + .parent_scope + .snapshot() + .is_writable(&fixture.scope_root) + ); + assert!(!fixture.runtime_dir.socket_path().exists()); + } + + #[tokio::test] + async fn prepared_apply_rejects_child_source_drift_before_any_effects() { + for drift in ["scope", "socket", "removal", "metadata-deletion"] { + let fixture = LegacyRestoreFixture::new("write").await; + std::fs::write(fixture.projection_path(), b"obsolete projection").unwrap(); + let scope_before = fixture.parent_scope.snapshot(); + let mut prepared = fixture.prepare(); + if drift == "metadata-deletion" { + fixture.store.delete_by_name("parent").unwrap(); + } else { + fixture + .store + .update_by_name("parent", |metadata| match drift { + "scope" => { + metadata.spawned_children[0].scope_delegated[0].recursive = false + } + "socket" => { + metadata.spawned_children[0].socket_path = + fixture.runtime_dir.path().join("replacement.sock"); + } + "removal" => metadata.spawned_children.clear(), + _ => unreachable!(), + }) + .unwrap(); + } + let metadata_before = std::fs::read(&fixture.metadata_path).ok(); + let projection_before = std::fs::read(fixture.projection_path()).unwrap(); + + let error = prepared + .apply() + .await + .err() + .expect("source drift must fail"); + + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert_eq!(std::fs::read(&fixture.metadata_path).ok(), metadata_before); + assert_eq!( + std::fs::read(fixture.projection_path()).unwrap(), + projection_before + ); + assert!(Arc::ptr_eq(&scope_before, &fixture.parent_scope.snapshot())); + assert!(!fixture.runtime_dir.socket_path().exists()); + } + } + + #[tokio::test] + async fn prepared_apply_clears_projection_only_without_reporting_child_reclaim() { + let fixture = LegacyRestoreFixture::new("write").await; + fixture + .store + .set_spawned_children("parent", Vec::new()) + .unwrap(); + std::fs::write(fixture.projection_path(), b"obsolete projection").unwrap(); + let metadata_before = std::fs::read(&fixture.metadata_path).unwrap(); + let mut prepared = fixture.prepare(); + assert!(!prepared.reclaimed_unreachable); + + let loaded = prepared.apply().await.unwrap(); + + assert!(!loaded.reclaimed_unreachable); + assert_eq!( + std::fs::read(&fixture.metadata_path).unwrap(), + metadata_before + ); + let projection: serde_json::Value = + serde_json::from_slice(&std::fs::read(fixture.projection_path()).unwrap()).unwrap(); + assert_eq!(projection, serde_json::json!([])); + } + + #[tokio::test] + async fn prepared_apply_projection_failure_keeps_canonical_children_and_scope() { + let fixture = LegacyRestoreFixture::new("write").await; + std::fs::create_dir(fixture.projection_path()).unwrap(); + let metadata_before = std::fs::read(&fixture.metadata_path).unwrap(); + let scope_before = fixture.parent_scope.snapshot(); + let mut prepared = fixture.prepare(); + + assert!(prepared.apply().await.is_err()); + + assert_eq!( + std::fs::read(&fixture.metadata_path).unwrap(), + metadata_before + ); + assert!(fixture.projection_path().is_dir()); + assert!(Arc::ptr_eq(&scope_before, &fixture.parent_scope.snapshot())); + } + fn registry() -> Arc { let scope = Scope::from_config(&ScopeConfig { allow: vec![ScopeRule { diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 59eb439a7..14be7b385 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -3335,6 +3335,7 @@ pub struct WorkerStartupFeatureCleanup { } struct StartupFeatureCleanupInner { + mcp_discovery: crate::feature::mcp::McpDiscoveryCleanup, registries: Mutex>, settlement: tokio::sync::Mutex< Option>>, @@ -3345,12 +3346,17 @@ impl WorkerStartupFeatureCleanup { fn new(registry: FeatureBackgroundTaskRegistry) -> Self { Self { inner: Arc::new(StartupFeatureCleanupInner { + mcp_discovery: crate::feature::mcp::McpDiscoveryCleanup::new(), registries: Mutex::new(vec![registry]), settlement: tokio::sync::Mutex::new(None), }), } } + pub(crate) fn mcp_discovery(&self) -> crate::feature::mcp::McpDiscoveryCleanup { + self.inner.mcp_discovery.clone() + } + fn retain(&self, registry: FeatureBackgroundTaskRegistry) { self.inner .registries @@ -3370,6 +3376,7 @@ impl WorkerStartupFeatureCleanup { .lock() .expect("startup Feature cleanup poisoned") .clone(); + let mcp_discovery = self.inner.mcp_discovery.clone(); Box::pin(async move { let mut failure = None; for registry in registries { @@ -3377,6 +3384,14 @@ impl WorkerStartupFeatureCleanup { failure.get_or_insert(error); } } + if let Err(error) = mcp_discovery.shutdown().await { + failure.get_or_insert_with(|| { + crate::hook::HookError::new( + crate::hook::HookErrorCategory::Dependency, + error.to_string(), + ) + }); + } failure.map_or(Ok(()), Err) }) }); @@ -3414,11 +3429,84 @@ impl RestoreCandidateWorkdirCleanup { pub struct PreparedRestoredWorker { worker: Option, St>>, closing: bool, + activation: Option, +} + +/// Exact existing restore source retained while a Controller is initialized. +/// +/// This carries no new Worker identity or execution authority. The host must keep +/// the per-Worker operation lock and invoke activation only after durable commit. +/// Controller initialization must not change the captured source before then. +pub(crate) struct RestoredWorkerActivation { metadata: Option, entries: serde_json::Value, manifest: serde_json::Value, } +impl RestoredWorkerActivation { + fn revalidate(&self, worker: &Worker, St>) -> Result<(), WorkerError> + where + St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, + { + let changed = |component| WorkerError::RestoreCandidateChanged { + worker_name: worker.manifest.worker.name.clone(), + component, + }; + if worker.store.read_by_name(&worker.manifest.worker.name)? != self.metadata { + return Err(changed("metadata")); + } + if restore_snapshot_value(&worker.manifest)? != self.manifest { + return Err(changed("manifest")); + } + let entries = worker + .store + .read_all(worker.session_id(), worker.segment_id())?; + if restore_snapshot_value(&entries)? != self.entries { + return Err(changed("segment")); + } + Ok(()) + } + + /// Acquire the existing top-level writer allocation after commit, checking + /// exact source both before and after acquisition, then apply restore effects. + /// On failure the Controller still owns the Worker; cleanup handles do not + /// establish rollback and allocation release must wait until it is dropped. + pub(crate) async fn activate( + &self, + worker: &mut Worker, St>, + ) -> Result<(), WorkerError> + where + St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, + { + self.revalidate(worker) + .map_err(|error| worker.construction_failure(error))?; + let allocate = || -> Result { + let socket_path = dir::default_base() + .map_err(ScopeLockError::from)? + .join(&worker.manifest.worker.name) + .join("sock"); + Ok(worker_allocation::install_top_level_with_deny( + worker.manifest.worker.name.clone(), + std::process::id(), + socket_path, + worker.scope.snapshot().allow_rules(), + worker.scope.snapshot().deny_rules(), + worker.segment_id(), + )?) + }; + worker.scope_allocation = + Some(allocate().map_err(|error| worker.construction_failure(error))?); + self.revalidate(worker) + .map_err(|error| worker.construction_failure(error))?; + worker + .finish_restored_construction_in_place(false) + .await + .map_err(|source| WorkerError::RestoreActivationEffects { + source: Box::new(source), + }) + } +} + impl PreparedRestoredWorker where St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, @@ -3460,23 +3548,32 @@ where } fn revalidate(&self, worker: &Worker, St>) -> Result<(), WorkerError> { - let changed = |component| WorkerError::RestoreCandidateChanged { - worker_name: worker.manifest.worker.name.clone(), - component, - }; - if worker.store.read_by_name(&worker.manifest.worker.name)? != self.metadata { - return Err(changed("metadata")); + self.activation + .as_ref() + .expect("restore candidate activation already transferred") + .revalidate(worker) + } + + /// Transfer the read-only reconstruction for Controller initialization, not + /// execution. Validate once without allocation or canonical restore effects; + /// the Controller must retain activation until its durable commit succeeds. + pub(crate) fn into_controller_parts( + mut self, + ) -> Result<(Worker, St>, RestoredWorkerActivation), WorkerError> { + let worker = self + .worker + .take() + .ok_or(WorkerError::RestoreCandidateClosing)?; + if self.closing { + return Err(worker.construction_failure(WorkerError::RestoreCandidateClosing)); } - if restore_snapshot_value(&worker.manifest)? != self.manifest { - return Err(changed("manifest")); - } - let entries = worker - .store - .read_all(worker.session_id(), worker.segment_id())?; - if restore_snapshot_value(&entries)? != self.entries { - return Err(changed("segment")); - } - Ok(()) + self.revalidate(&worker) + .map_err(|error| worker.construction_failure(error))?; + let activation = self + .activation + .take() + .expect("restore candidate activation already transferred"); + Ok((worker, activation)) } /// Acquire writer allocation only after exact source validation, validate @@ -9730,9 +9827,11 @@ where let mut candidate = PreparedRestoredWorker { worker: Some(worker), closing: false, - metadata, - entries: entries_snapshot, - manifest: manifest_snapshot, + activation: Some(RestoredWorkerActivation { + metadata, + entries: entries_snapshot, + manifest: manifest_snapshot, + }), }; if let Err(error) = candidate.revalidate(candidate.worker()) { let worker = candidate.worker.take().expect("candidate just constructed"); @@ -9742,6 +9841,16 @@ where } async fn finish_restored_construction(mut self) -> Result { + self.finish_restored_construction_in_place(true).await?; + Ok(self) + } + + /// Apply the same restore effects while an initialized Controller retains + /// ownership. Errors retain cleanup handles, but do not drop this Worker. + async fn finish_restored_construction_in_place( + &mut self, + reclaim_delegations: bool, + ) -> Result<(), WorkerError> { self.pending_submission_handle() .reconcile_uploaded_file_pins() .map_err(|error| self.construction_failure(error.into()))?; @@ -9752,10 +9861,12 @@ where segment_id: self.segment_id(), }) .map_err(|error| self.construction_failure(error))?; - self.reconcile_restored_delegations() - .await - .map_err(|error| self.construction_failure(error))?; - Ok(self) + if reclaim_delegations { + self.reconcile_restored_delegations() + .await + .map_err(|error| self.construction_failure(error))?; + } + Ok(()) } async fn reconcile_restored_delegations(&mut self) -> Result<(), WorkerError> { @@ -10596,8 +10707,10 @@ fn compact_failure_category(error: &WorkerError) -> CompactFailureCategory { #[derive(Debug, thiserror::Error)] pub enum WorkerError { - /// A constructor failed after acquiring resources. Retain these handles until - /// explicit cleanup succeeds; the Worker has dropped before this is returned. + /// A construction or activation step failed after acquiring resources. Retain + /// these handles until explicit cleanup succeeds. Consuming constructors drop + /// the Worker before returning; in-place activation leaves its owner responsible + /// for dropping it before allocation release. #[error("{source}")] ConstructionCleanup { #[source]