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.
This commit is contained in:
@@ -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<crate::BackendJobExecutionBinding>,
|
||||
pub workdir_sessions: std::sync::Arc<workdir::WorkdirSessionRouter>,
|
||||
pub observation_provider:
|
||||
Option<std::sync::Arc<dyn crate::feature::builtin::WorkerObservationProvider>>,
|
||||
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<St: Store> {
|
||||
initialized: Option<crate::controller::InitializedWorkerController<Box<dyn LlmClient>, St>>,
|
||||
activation: Option<crate::worker::RestoredWorkerActivation>,
|
||||
streams: Option<crate::ipc::protocol_session::WorkerProtocolSessionStreams>,
|
||||
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<St> PreparedRestoredWorker<St>
|
||||
where
|
||||
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
|
||||
{
|
||||
pub fn workspace_client_handle(
|
||||
&mut self,
|
||||
) -> Result<std::sync::Arc<dyn crate::WorkspaceClient>, WorkerError> {
|
||||
Ok(self.worker_mut()?.workspace_client_handle())
|
||||
}
|
||||
|
||||
pub fn retain_startup_feature_cleanup(
|
||||
&mut self,
|
||||
) -> Result<crate::WorkerStartupFeatureCleanup, WorkerError> {
|
||||
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<UnpublishedRestoredController<St>, 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<St> UnpublishedRestoredController<St>
|
||||
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<crate::runtime::worker_allocation::ScopeAllocationCleanup> {
|
||||
self.initialized
|
||||
.as_ref()
|
||||
.and_then(|initialized| initialized.worker.restore_allocation_cleanup_handle())
|
||||
}
|
||||
|
||||
pub async fn activate(&mut self) -> Result<ActivatedRestoredController, WorkerBootstrapError> {
|
||||
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(())
|
||||
}
|
||||
}
|
||||
|
||||
+404
-69
@@ -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<C: LlmClient, St: Store> {
|
||||
pub(crate) worker: Worker<C, St>,
|
||||
pub(crate) handle: WorkerHandle,
|
||||
method_rx: mpsc::Receiver<Method>,
|
||||
shutdown_tx: oneshot::Sender<()>,
|
||||
shutdown_rx: ShutdownReceiver,
|
||||
registry_reclaim: PreparedSpawnedWorkerRegistryLoad<St>,
|
||||
shutdown_after_idle: ShutdownAfterIdleRequest,
|
||||
transport: WorkerControllerTransport,
|
||||
command_events: Option<PreparedWorkdirCommandEvents>,
|
||||
}
|
||||
|
||||
pub(crate) struct PreparedControllerPublication {
|
||||
socket_server: Option<SocketServer>,
|
||||
command_observer: Option<tokio::task::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
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<C, St>(
|
||||
mut worker: Worker<C, St>,
|
||||
worker: Worker<C, St>,
|
||||
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<C, St>(
|
||||
worker: Worker<C, St>,
|
||||
runtime_base: &Path,
|
||||
runtime_managed: bool,
|
||||
runtime_run: Option<&Path>,
|
||||
transport: WorkerControllerTransport,
|
||||
) -> Result<InitializedWorkerController<C, St>, 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<St>(
|
||||
worker: Worker<Box<dyn LlmClient>, St>,
|
||||
run_dir: &Path,
|
||||
transport: WorkerControllerTransport,
|
||||
) -> std::io::Result<InitializedWorkerController<Box<dyn LlmClient>, 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<C, St>(
|
||||
mut worker: Worker<C, St>,
|
||||
runtime_base: &Path,
|
||||
runtime_managed: bool,
|
||||
runtime_run: Option<&Path>,
|
||||
transport: WorkerControllerTransport,
|
||||
canonical_startup: bool,
|
||||
) -> Result<InitializedWorkerController<C, St>, 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<dyn Store> = 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<C, St> InitializedWorkerController<C, St>
|
||||
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<PreparedControllerPublication> {
|
||||
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<dyn WorkdirSession>,
|
||||
in_flight: &InFlightEvents,
|
||||
) -> Option<tokio::task::JoinHandle<()>> {
|
||||
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<dyn WorkdirSession>,
|
||||
in_flight: InFlightEvents,
|
||||
events: Option<broadcast::Receiver<WorkdirCommandEvent>>,
|
||||
}
|
||||
|
||||
impl PreparedWorkdirCommandEvents {
|
||||
fn new(session: &Arc<dyn WorkdirSession>, 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<tokio::task::JoinHandle<()>> {
|
||||
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<ProtocolCommandSnapshot> {
|
||||
@@ -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]
|
||||
|
||||
@@ -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<Mutex<Vec<McpStdioClient>>>,
|
||||
}
|
||||
|
||||
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<Option<McpStdioToolFeature>, 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<ProtocolProviderContribution> {
|
||||
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"));
|
||||
|
||||
@@ -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<St> {
|
||||
pub(crate) registry: Arc<SpawnedWorkerRegistry>,
|
||||
/// Projected outcome, not evidence that reclamation has already happened.
|
||||
pub(crate) reclaimed_unreachable: bool,
|
||||
runtime_dir: Arc<RuntimeDir>,
|
||||
store: St,
|
||||
worker_name: String,
|
||||
parent_scope: Option<SharedScope>,
|
||||
persisted_children: Vec<WorkerSpawnedChild>,
|
||||
valid_records: Vec<SpawnedWorkerRecord>,
|
||||
legacy_projection_exists: bool,
|
||||
}
|
||||
|
||||
impl<St> PreparedSpawnedWorkerRegistryLoad<St>
|
||||
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<SpawnedWorkerRegistryLoad> {
|
||||
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<SharedScope>,
|
||||
) -> io::Result<SpawnedWorkerRegistryLoad>
|
||||
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<St>(
|
||||
runtime_dir: Arc<RuntimeDir>,
|
||||
store: St,
|
||||
worker_name: String,
|
||||
parent_scope: Option<SharedScope>,
|
||||
) -> io::Result<PreparedSpawnedWorkerRegistryLoad<St>>
|
||||
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::<Vec<_>>();
|
||||
// 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<RuntimeDir>,
|
||||
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<WorkerAggregateStore> {
|
||||
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<SpawnedWorkerRegistry> {
|
||||
let scope = Scope::from_config(&ScopeConfig {
|
||||
allow: vec![ScopeRule {
|
||||
|
||||
+138
-25
@@ -3335,6 +3335,7 @@ pub struct WorkerStartupFeatureCleanup {
|
||||
}
|
||||
|
||||
struct StartupFeatureCleanupInner {
|
||||
mcp_discovery: crate::feature::mcp::McpDiscoveryCleanup,
|
||||
registries: Mutex<Vec<FeatureBackgroundTaskRegistry>>,
|
||||
settlement: tokio::sync::Mutex<
|
||||
Option<futures_util::future::BoxFuture<'static, Result<(), crate::hook::HookError>>>,
|
||||
@@ -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<St: Store> {
|
||||
worker: Option<Worker<Box<dyn LlmClient>, St>>,
|
||||
closing: bool,
|
||||
activation: Option<RestoredWorkerActivation>,
|
||||
}
|
||||
|
||||
/// 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<WorkerMetadata>,
|
||||
entries: serde_json::Value,
|
||||
manifest: serde_json::Value,
|
||||
}
|
||||
|
||||
impl RestoredWorkerActivation {
|
||||
fn revalidate<St>(&self, worker: &Worker<Box<dyn LlmClient>, 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<St>(
|
||||
&self,
|
||||
worker: &mut Worker<Box<dyn LlmClient>, St>,
|
||||
) -> Result<(), WorkerError>
|
||||
where
|
||||
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
|
||||
{
|
||||
self.revalidate(worker)
|
||||
.map_err(|error| worker.construction_failure(error))?;
|
||||
let allocate = || -> Result<ScopeAllocationGuard, WorkerError> {
|
||||
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<St> PreparedRestoredWorker<St>
|
||||
where
|
||||
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
|
||||
@@ -3460,23 +3548,32 @@ where
|
||||
}
|
||||
|
||||
fn revalidate(&self, worker: &Worker<Box<dyn LlmClient>, 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<Box<dyn LlmClient>, 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, WorkerError> {
|
||||
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]
|
||||
|
||||
Reference in New Issue
Block a user