fix: merge worker restore into develop host lifecycle
Integrate integration/worker-restore without reviving removed Worker socket/discovery paths. Keep restore candidate allocation and canonical effects deferred until activation, with registry reclamation and observer publication fences on the in-process controller. Preserve develop BackendJob/provenance, immutable Journal body capture and erasure, and Workdir diagnostics. Adapt obsolete socket fault tests to in-process publication and registry drift boundaries. Validation: root cargo check; ordinary worker, worker-runtime and standalone package tests; Server restore (43) and Journal (49) filters; mcp/server-api/runtime-api/session-store tests; optional Runtime feature check; OpenAPI and six generated TypeScript checks; Web restore unit tests (15); runner tests (21); fmt, naming and diff checks. Full workspace/Server suites, E2E, live/model and fixture generation not run.
This commit is contained in:
@@ -7,6 +7,7 @@ use workdir::WorkdirSessionHandle;
|
||||
|
||||
use crate::PromptCatalogSource;
|
||||
use crate::controller::{ShutdownReceiver, WorkerController, WorkerHandle};
|
||||
pub use crate::worker::{PreparedRestoredWorker, RestoreCandidateWorkdirCleanup};
|
||||
use crate::worker::{Worker, WorkerError, WorkerFilesystemAuthority, WorkerWorkspaceContext};
|
||||
use manifest::WorkerManifest;
|
||||
|
||||
@@ -161,6 +162,38 @@ where
|
||||
Ok(PreparedWorker::new(worker, self.layout))
|
||||
}
|
||||
|
||||
/// Reconstruct a saved materialized Worker without allocation, attribution
|
||||
/// repair, Feature installation, or controller exposure. The caller owns the
|
||||
/// later controller layout and must activate before starting it.
|
||||
pub async fn prepare_restored_candidate(
|
||||
self,
|
||||
worker_name: &str,
|
||||
) -> Result<PreparedRestoredWorker<St>, WorkerBootstrapError> {
|
||||
let mut candidate =
|
||||
Worker::prepare_restored_from_worker_metadata_with_context_and_model_client(
|
||||
worker_name,
|
||||
self.manifest,
|
||||
self.store,
|
||||
self.prompt_catalog,
|
||||
self.workspace_context,
|
||||
self.filesystem_authority,
|
||||
self.model_client,
|
||||
)
|
||||
.await
|
||||
.map_err(WorkerBootstrapError::Worker)?;
|
||||
if let Some(workdir_session) = self.workdir_session {
|
||||
candidate
|
||||
.worker_mut()
|
||||
.map_err(WorkerBootstrapError::Worker)?
|
||||
.bind_single_workdir_session(Some(workdir_session));
|
||||
}
|
||||
candidate
|
||||
.worker_mut()
|
||||
.map_err(WorkerBootstrapError::Worker)?
|
||||
.bind_subjektiv_host(self.subjektiv_host);
|
||||
Ok(candidate)
|
||||
}
|
||||
|
||||
pub async fn prepare_restored(
|
||||
self,
|
||||
worker_name: &str,
|
||||
@@ -243,6 +276,7 @@ where
|
||||
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
|
||||
{
|
||||
let cleanup_sessions = worker.workdir_sessions();
|
||||
let factory_owned_cleanup = worker.startup_feature_cleanup_handle().is_some();
|
||||
let controller = match layout {
|
||||
WorkerBootstrapLayout::Direct { runtime_base } => {
|
||||
WorkerController::spawn_owned(worker, &runtime_base).await
|
||||
@@ -259,7 +293,10 @@ where
|
||||
controller_task,
|
||||
}),
|
||||
Err(source) => {
|
||||
let cleanup_failed = cleanup_sessions.close_all().await.is_err();
|
||||
// Factory-owned task scopes must settle before sessions close.
|
||||
// Cleanup remains unconfirmed here; preserve the startup source.
|
||||
let cleanup_failed =
|
||||
factory_owned_cleanup || cleanup_sessions.close_all().await.is_err();
|
||||
Err(WorkerBootstrapError::Controller {
|
||||
source,
|
||||
cleanup_failed,
|
||||
@@ -267,3 +304,229 @@ 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,
|
||||
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)
|
||||
.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,
|
||||
));
|
||||
}
|
||||
// Preparation's read-only authorization is not a publication grant.
|
||||
// Revalidate current Host authority after commit, before effects/input,
|
||||
// including a publication retry on this same private Controller.
|
||||
initialized
|
||||
.worker
|
||||
.validate_subjektiv_body_publication()
|
||||
.map_err(WorkerBootstrapError::Worker)?;
|
||||
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(())
|
||||
}
|
||||
}
|
||||
|
||||
+388
-45
@@ -21,7 +21,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,
|
||||
@@ -654,7 +654,56 @@ pub type ShutdownReceiver = oneshot::Receiver<()>;
|
||||
|
||||
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,
|
||||
command_events: Option<PreparedWorkdirCommandEvents>,
|
||||
}
|
||||
|
||||
pub(crate) struct PreparedControllerPublication {
|
||||
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
|
||||
/// Controller, input channel, socket, observer, or queue consumer. The caller
|
||||
/// owns the same operation's Feature/session/allocation cleanup on every exit.
|
||||
pub async fn recover_pending_head<C, St>(
|
||||
worker: &mut Worker<C, St>,
|
||||
runtime_base: &Path,
|
||||
) -> std::io::Result<()>
|
||||
where
|
||||
C: LlmClient + Clone + 'static,
|
||||
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
|
||||
{
|
||||
if worker.needs_initial_session_head_materialization() {
|
||||
worker
|
||||
.validate_subjektiv_body_publication()
|
||||
.map_err(std::io::Error::other)?;
|
||||
let registry = SpawnedWorkerRegistry::new_internal(
|
||||
worker.manifest().worker.name.clone(),
|
||||
worker.scope().clone(),
|
||||
);
|
||||
register_worker_tools(worker, runtime_base.to_path_buf(), registry, None, None).await?;
|
||||
worker
|
||||
.materialize_durable_session_head()
|
||||
.await
|
||||
.map_err(std::io::Error::other)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn spawn<C, St>(
|
||||
worker: Worker<C, St>,
|
||||
runtime_base: &Path,
|
||||
@@ -734,8 +783,12 @@ impl WorkerController {
|
||||
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
|
||||
{
|
||||
let sessions = worker.workdir_sessions();
|
||||
let factory_owned_cleanup = worker.startup_feature_cleanup_handle().is_some();
|
||||
let result = Self::spawn_initialized(worker, runtime_base, runtime_run).await;
|
||||
// An opted-in factory must quiesce its captured Feature tasks before
|
||||
// closing their Workdir sessions; this wrapper cannot confirm that order.
|
||||
if result.is_err()
|
||||
&& !factory_owned_cleanup
|
||||
&& let Err(error) = sessions.close_all().await
|
||||
{
|
||||
tracing::warn!(%error, "Workdir sessions close after controller startup failure failed");
|
||||
@@ -744,7 +797,7 @@ impl WorkerController {
|
||||
}
|
||||
|
||||
async fn spawn_initialized<C, St>(
|
||||
mut worker: Worker<C, St>,
|
||||
worker: Worker<C, St>,
|
||||
runtime_base: &Path,
|
||||
runtime_run: Option<&Path>,
|
||||
) -> Result<(WorkerHandle, ShutdownReceiver, tokio::task::JoinHandle<()>), std::io::Error>
|
||||
@@ -752,6 +805,52 @@ impl WorkerController {
|
||||
C: LlmClient + Clone + 'static,
|
||||
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
|
||||
{
|
||||
Self::initialize(worker, runtime_base, runtime_run)
|
||||
.await?
|
||||
.publish()
|
||||
.await
|
||||
}
|
||||
|
||||
async fn initialize<C, St>(
|
||||
worker: Worker<C, St>,
|
||||
runtime_base: &Path,
|
||||
runtime_run: Option<&Path>,
|
||||
) -> 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_run, true).await
|
||||
}
|
||||
|
||||
pub(crate) async fn initialize_restored<St>(
|
||||
worker: Worker<Box<dyn LlmClient>, St>,
|
||||
run_dir: &Path,
|
||||
) -> 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, Some(run_dir), false).await
|
||||
}
|
||||
|
||||
async fn initialize_inner<C, St>(
|
||||
mut worker: Worker<C, St>,
|
||||
runtime_base: &Path,
|
||||
runtime_run: Option<&Path>,
|
||||
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)?;
|
||||
@@ -773,21 +872,13 @@ impl WorkerController {
|
||||
});
|
||||
|
||||
let spawner_name = worker.manifest().worker.name.clone();
|
||||
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
|
||||
@@ -878,7 +969,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()));
|
||||
}
|
||||
|
||||
let artifact_store: Arc<dyn Store> = Arc::new(worker.store().clone());
|
||||
let session_id = worker.session_id();
|
||||
let pending_activations = worker.pending_activation_state();
|
||||
@@ -903,10 +993,71 @@ impl WorkerController {
|
||||
pending_activations,
|
||||
};
|
||||
|
||||
// Start the observer only after fallible head/recovery initialization.
|
||||
let command_observer = fs_for_view
|
||||
let command_events = fs_for_view
|
||||
.as_ref()
|
||||
.and_then(|session| wire_workdir_command_events(session, &in_flight));
|
||||
.map(|session| PreparedWorkdirCommandEvents::new(session, &handle.in_flight));
|
||||
Ok(InitializedWorkerController {
|
||||
worker,
|
||||
handle,
|
||||
method_rx,
|
||||
shutdown_tx,
|
||||
shutdown_rx,
|
||||
registry_reclaim,
|
||||
shutdown_after_idle,
|
||||
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 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 command_observer = self
|
||||
.command_events
|
||||
.take()
|
||||
.and_then(PreparedWorkdirCommandEvents::start);
|
||||
Ok(PreparedControllerPublication { 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 { 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();
|
||||
|
||||
// === 5. controller_loop ===
|
||||
// Clone cancel sender and notification buffer before moving worker
|
||||
// into the controller task so the in-flight turn can be reached
|
||||
@@ -915,6 +1066,7 @@ impl WorkerController {
|
||||
let pause_tx = worker.engine_mut().pause_sender();
|
||||
let notify_buffer = worker.notify_buffer_handle();
|
||||
|
||||
worker.finish_startup_feature_cleanup_capture();
|
||||
let controller_task = tokio::spawn(controller_loop(
|
||||
worker,
|
||||
method_rx,
|
||||
@@ -930,7 +1082,7 @@ impl WorkerController {
|
||||
command_observer,
|
||||
));
|
||||
|
||||
Ok((handle, shutdown_rx, controller_task))
|
||||
(handle, shutdown_rx, controller_task)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -938,30 +1090,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> {
|
||||
@@ -1687,9 +1860,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);
|
||||
}
|
||||
}
|
||||
@@ -3750,6 +3931,159 @@ 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(), None)
|
||||
.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.path().join("sock").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.path().join("sock").exists());
|
||||
assert!(!handle.runtime_dir.path().join("status.json").exists());
|
||||
assert!(!handle.runtime_dir.path().join("manifest.toml").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.path().join("sock").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(), None).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,
|
||||
@@ -4077,10 +4411,19 @@ 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 reclaim = startup
|
||||
.find("self.registry_reclaim.apply().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 && observer < loop_start);
|
||||
assert!(recover < observer_preparation && observer_preparation < reclaim);
|
||||
assert!(reclaim < observer_start && observer_start < loop_start);
|
||||
assert!(
|
||||
startup[materialize_gate..materialize]
|
||||
.contains("needs_initial_session_head_materialization")
|
||||
|
||||
@@ -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"));
|
||||
|
||||
@@ -19,6 +19,9 @@ pub mod model_client;
|
||||
mod paste_artifact_tool;
|
||||
pub mod profile_manifest;
|
||||
pub mod prompt;
|
||||
#[cfg(test)]
|
||||
#[path = "../../../tests/support/restore_environment.rs"]
|
||||
mod restore_test_environment;
|
||||
pub mod runtime;
|
||||
pub mod segment_log_sink;
|
||||
mod session_capture;
|
||||
@@ -70,9 +73,10 @@ pub use session_history::{
|
||||
pub use shared_state::WorkerSharedState;
|
||||
pub use worker::{
|
||||
LocalWorkingDirectory, SubjektivSessionAttributionLifecycle, Worker, WorkerError,
|
||||
WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, WorkspaceBinaryRequest,
|
||||
WorkspaceBinaryResponse, WorkspaceClient, WorkspaceClientError, WorkspaceId, WorkspaceIdError,
|
||||
WorkspacePromptCatalogResolution, WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse,
|
||||
WorkspaceServerOperation, WorkspaceWorkerDiscoveryRequest, apply_worker_manifest,
|
||||
marker_workspace_client, unavailable_workspace_client,
|
||||
WorkerFilesystemAuthority, WorkerRunResult, WorkerStartupFeatureCleanup,
|
||||
WorkerWorkspaceContext, WorkspaceBinaryRequest, WorkspaceBinaryResponse, WorkspaceClient,
|
||||
WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspacePromptCatalogResolution,
|
||||
WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse, WorkspaceServerOperation,
|
||||
WorkspaceWorkerDiscoveryRequest, apply_worker_manifest, marker_workspace_client,
|
||||
unavailable_workspace_client,
|
||||
};
|
||||
|
||||
@@ -21,8 +21,9 @@ pub use error::ScopeLockError;
|
||||
#[cfg(test)]
|
||||
pub(crate) use lifecycle::install_top_level_at_for_test;
|
||||
pub use lifecycle::{
|
||||
ScopeAllocationGuard, SegmentActivationCommitError, SegmentActivationGuard, SegmentLockInfo,
|
||||
adopt_allocation, install_top_level, install_top_level_with_deny, lookup_segment,
|
||||
ScopeAllocationCleanup, ScopeAllocationGuard, SegmentActivationCommitError,
|
||||
SegmentActivationGuard, SegmentLockInfo, adopt_allocation, install_top_level,
|
||||
install_top_level_with_deny, lookup_segment,
|
||||
};
|
||||
pub use mutate::{
|
||||
delegate_scope, reclaim_delegated_scope, reclaim_stale, reclaim_stale_with, register_worker,
|
||||
|
||||
@@ -11,6 +11,14 @@ use session_store::SegmentId;
|
||||
pub enum ScopeLockError {
|
||||
#[error("I/O error on workers.json: {0}")]
|
||||
Io(#[from] io::Error),
|
||||
#[error("worker `{worker_name}` allocation still has a live guard")]
|
||||
AllocationGuardAttached { worker_name: String },
|
||||
#[error("allocation acquisition was published but its save failed: {source}")]
|
||||
AcquisitionSaveFailed {
|
||||
#[source]
|
||||
source: io::Error,
|
||||
cleanup: super::lifecycle::ScopeAllocationCleanup,
|
||||
},
|
||||
#[error("worker `{0}` is already allocated")]
|
||||
DuplicateWorkerName(String),
|
||||
#[error("requested scope `{}` conflicts with worker allocation `{competitor}` rule `{}`", .rule.target.display(), .competitor_rule.target.display())]
|
||||
@@ -46,3 +54,15 @@ pub enum ScopeLockError {
|
||||
actual: Option<SegmentId>,
|
||||
},
|
||||
}
|
||||
|
||||
impl ScopeLockError {
|
||||
/// Retain the release authority for an owned acquisition whose table rename
|
||||
/// succeeded before save returned an error. Refusals and unpublished saves
|
||||
/// do not confer ownership of any preexisting allocation.
|
||||
pub fn allocation_cleanup_handle(&self) -> Option<super::lifecycle::ScopeAllocationCleanup> {
|
||||
match self {
|
||||
Self::AcquisitionSaveFailed { cleanup, .. } => Some(cleanup.clone()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
//! up on drop.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use manifest::ScopeRule;
|
||||
use session_store::SegmentId;
|
||||
@@ -11,13 +12,82 @@ use super::error::ScopeLockError;
|
||||
use super::mutate::release_worker;
|
||||
use super::table::{LockFileGuard, default_allocation_path};
|
||||
|
||||
/// Owned allocation: on drop, opens the lock file and releases this
|
||||
/// Worker's entry. The guard keeps only the name + lock-file path; it
|
||||
/// does not hold the `flock` for the Worker's lifetime.
|
||||
/// Owned allocation shared with any retained cleanup handles. Only the final
|
||||
/// owner's drop attempts best-effort release; no `flock` is held for its lifetime.
|
||||
#[derive(Debug)]
|
||||
pub struct ScopeAllocationGuard {
|
||||
owner: Arc<ScopeAllocationOwner>,
|
||||
}
|
||||
|
||||
/// Retains an allocation's release authority even after its Worker is consumed
|
||||
/// or dropped. A successful controller transfer should drop this handle without
|
||||
/// explicitly releasing the live Worker's allocation.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ScopeAllocationCleanup {
|
||||
owner: Arc<ScopeAllocationOwner>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct ScopeAllocationOwner {
|
||||
worker_name: String,
|
||||
lock_path: PathBuf,
|
||||
state: Mutex<ScopeAllocationOwnerState>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct ScopeAllocationOwnerState {
|
||||
guard_attached: bool,
|
||||
released: bool,
|
||||
}
|
||||
|
||||
impl ScopeAllocationOwner {
|
||||
fn release(&self) -> Result<(), ScopeLockError> {
|
||||
let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner());
|
||||
if state.guard_attached {
|
||||
return Err(ScopeLockError::AllocationGuardAttached {
|
||||
worker_name: self.worker_name.clone(),
|
||||
});
|
||||
}
|
||||
if state.released {
|
||||
return Ok(());
|
||||
}
|
||||
let mut guard = LockFileGuard::open(&self.lock_path)?;
|
||||
if guard.data().find(&self.worker_name).is_some() {
|
||||
release_worker(&mut guard, &self.worker_name)?;
|
||||
} else {
|
||||
// A previous save may have removed the entry but failed after rename.
|
||||
// Complete the durable save before reporting release as confirmed.
|
||||
guard.save()?;
|
||||
}
|
||||
state.released = true;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ScopeAllocationOwner {
|
||||
fn drop(&mut self) {
|
||||
let _ = self.release();
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ScopeAllocationGuard {
|
||||
fn drop(&mut self) {
|
||||
self.owner
|
||||
.state
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner())
|
||||
.guard_attached = false;
|
||||
}
|
||||
}
|
||||
|
||||
impl ScopeAllocationCleanup {
|
||||
/// Release only after the factory has returned an error and no live
|
||||
/// Controller can still use this allocation. An attached guard rejects
|
||||
/// release. Success is idempotent across all clones; failure retains the
|
||||
/// authority for an explicit retry.
|
||||
pub fn release(&self) -> Result<(), ScopeLockError> {
|
||||
self.owner.release()
|
||||
}
|
||||
}
|
||||
|
||||
/// Exclusive authority to activate a replacement Segment for one allocated Worker.
|
||||
@@ -56,12 +126,58 @@ impl SegmentActivationCommitError {
|
||||
}
|
||||
|
||||
impl ScopeAllocationGuard {
|
||||
fn new(worker_name: String, lock_path: PathBuf) -> Self {
|
||||
Self {
|
||||
owner: Arc::new(ScopeAllocationOwner {
|
||||
worker_name,
|
||||
lock_path,
|
||||
state: Mutex::new(ScopeAllocationOwnerState {
|
||||
guard_attached: true,
|
||||
released: false,
|
||||
}),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn commit_acquisition(self, table: &mut LockFileGuard) -> Result<Self, ScopeLockError> {
|
||||
let mut published = false;
|
||||
match table.save_with_publication(&mut published) {
|
||||
Ok(()) => Ok(self),
|
||||
Err(source) => {
|
||||
let cleanup = if published {
|
||||
Some(self.cleanup_handle())
|
||||
} else {
|
||||
// A staged entry is not ours until rename. Disarm before
|
||||
// detaching so Drop cannot remove a preexisting allocation
|
||||
// (or reopen the table lock still held by this caller).
|
||||
self.owner
|
||||
.state
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner())
|
||||
.released = true;
|
||||
None
|
||||
};
|
||||
drop(self);
|
||||
match cleanup {
|
||||
Some(cleanup) => Err(ScopeLockError::AcquisitionSaveFailed { source, cleanup }),
|
||||
None => Err(ScopeLockError::Io(source)),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn worker_name(&self) -> &str {
|
||||
&self.worker_name
|
||||
&self.owner.worker_name
|
||||
}
|
||||
|
||||
pub fn lock_path(&self) -> &Path {
|
||||
&self.lock_path
|
||||
&self.owner.lock_path
|
||||
}
|
||||
|
||||
pub fn cleanup_handle(&self) -> ScopeAllocationCleanup {
|
||||
ScopeAllocationCleanup {
|
||||
owner: Arc::clone(&self.owner),
|
||||
}
|
||||
}
|
||||
|
||||
/// Lock the machine-wide table and validate a Segment replacement before
|
||||
@@ -71,24 +187,25 @@ impl ScopeAllocationGuard {
|
||||
expected: SegmentId,
|
||||
replacement: SegmentId,
|
||||
) -> Result<SegmentActivationGuard, ScopeLockError> {
|
||||
let guard = LockFileGuard::open(&self.lock_path)?;
|
||||
let guard = LockFileGuard::open(self.lock_path())?;
|
||||
let actual = guard
|
||||
.data()
|
||||
.allocations
|
||||
.iter()
|
||||
.find(|allocation| allocation.worker_name == self.worker_name)
|
||||
.find(|allocation| allocation.worker_name == self.worker_name())
|
||||
.and_then(|allocation| allocation.segment_id);
|
||||
|
||||
if actual != Some(expected) {
|
||||
return Err(ScopeLockError::SegmentChanged {
|
||||
worker_name: self.worker_name.clone(),
|
||||
worker_name: self.worker_name().to_owned(),
|
||||
expected,
|
||||
actual,
|
||||
});
|
||||
}
|
||||
|
||||
if let Some(existing) = guard.data().allocations.iter().find(|allocation| {
|
||||
allocation.segment_id == Some(replacement) && allocation.worker_name != self.worker_name
|
||||
allocation.segment_id == Some(replacement)
|
||||
&& allocation.worker_name != self.worker_name()
|
||||
}) {
|
||||
return Err(ScopeLockError::SegmentConflict {
|
||||
segment_id: replacement,
|
||||
@@ -99,7 +216,7 @@ impl ScopeAllocationGuard {
|
||||
|
||||
Ok(SegmentActivationGuard {
|
||||
guard,
|
||||
worker_name: self.worker_name.clone(),
|
||||
worker_name: self.worker_name().to_owned(),
|
||||
expected,
|
||||
replacement,
|
||||
})
|
||||
@@ -145,14 +262,6 @@ impl SegmentActivationGuard {
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ScopeAllocationGuard {
|
||||
fn drop(&mut self) {
|
||||
if let Ok(mut guard) = LockFileGuard::open(&self.lock_path) {
|
||||
let _ = release_worker(&mut guard, &self.worker_name);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Open the default lock file, register a top-level Worker, and return a
|
||||
/// guard that will release the allocation on drop.
|
||||
pub fn install_top_level(
|
||||
@@ -185,7 +294,7 @@ pub fn install_top_level_with_deny(
|
||||
) -> Result<ScopeAllocationGuard, ScopeLockError> {
|
||||
let lock_path = default_allocation_path()?;
|
||||
let mut guard = LockFileGuard::open(&lock_path)?;
|
||||
super::mutate::register_worker_with_deny(
|
||||
super::mutate::prepare_worker_with_deny(
|
||||
&mut guard,
|
||||
worker_name.clone(),
|
||||
pid,
|
||||
@@ -194,10 +303,7 @@ pub fn install_top_level_with_deny(
|
||||
scope_deny,
|
||||
segment_id,
|
||||
)?;
|
||||
Ok(ScopeAllocationGuard {
|
||||
worker_name,
|
||||
lock_path,
|
||||
})
|
||||
ScopeAllocationGuard::new(worker_name, lock_path).commit_acquisition(&mut guard)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -210,7 +316,7 @@ pub(crate) fn install_top_level_at_for_test(
|
||||
segment_id: SegmentId,
|
||||
) -> Result<ScopeAllocationGuard, ScopeLockError> {
|
||||
let mut guard = LockFileGuard::open(&lock_path)?;
|
||||
super::mutate::register_worker_with_deny(
|
||||
super::mutate::prepare_worker_with_deny(
|
||||
&mut guard,
|
||||
worker_name.clone(),
|
||||
pid,
|
||||
@@ -219,10 +325,7 @@ pub(crate) fn install_top_level_at_for_test(
|
||||
Vec::new(),
|
||||
segment_id,
|
||||
)?;
|
||||
Ok(ScopeAllocationGuard {
|
||||
worker_name,
|
||||
lock_path,
|
||||
})
|
||||
ScopeAllocationGuard::new(worker_name, lock_path).commit_acquisition(&mut guard)
|
||||
}
|
||||
|
||||
/// Take ownership of an existing allocation that was pre-registered by
|
||||
@@ -240,6 +343,15 @@ pub fn adopt_allocation(
|
||||
segment_id: SegmentId,
|
||||
) -> Result<ScopeAllocationGuard, ScopeLockError> {
|
||||
let lock_path = default_allocation_path()?;
|
||||
adopt_allocation_at(lock_path, worker_name, new_pid, segment_id)
|
||||
}
|
||||
|
||||
fn adopt_allocation_at(
|
||||
lock_path: PathBuf,
|
||||
worker_name: String,
|
||||
new_pid: u32,
|
||||
segment_id: SegmentId,
|
||||
) -> Result<ScopeAllocationGuard, ScopeLockError> {
|
||||
let mut guard = LockFileGuard::open(&lock_path)?;
|
||||
let alloc = guard
|
||||
.data_mut()
|
||||
@@ -247,11 +359,7 @@ pub fn adopt_allocation(
|
||||
.ok_or_else(|| ScopeLockError::UnknownWorker(worker_name.clone()))?;
|
||||
alloc.pid = new_pid;
|
||||
alloc.segment_id = Some(segment_id);
|
||||
guard.save()?;
|
||||
Ok(ScopeAllocationGuard {
|
||||
worker_name,
|
||||
lock_path,
|
||||
})
|
||||
ScopeAllocationGuard::new(worker_name, lock_path).commit_acquisition(&mut guard)
|
||||
}
|
||||
|
||||
/// Information about a Worker that currently holds an allocation for a
|
||||
@@ -329,6 +437,388 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn retained_cleanup_handle_defers_release_until_final_owner_drop() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let lock_path = dir.path().join("workers.json");
|
||||
let guard = install_top_level_at_for_test(
|
||||
lock_path.clone(),
|
||||
"retained".into(),
|
||||
std::process::id(),
|
||||
sock("retained"),
|
||||
Vec::new(),
|
||||
sid(),
|
||||
)
|
||||
.unwrap();
|
||||
let cleanup = guard.cleanup_handle();
|
||||
let retained = cleanup.clone();
|
||||
|
||||
drop(guard);
|
||||
drop(cleanup);
|
||||
{
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert!(table.data().find("retained").is_some());
|
||||
}
|
||||
drop(retained);
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert!(table.data().find("retained").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dropping_cleanup_handle_after_transfer_preserves_live_guard() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let lock_path = dir.path().join("workers.json");
|
||||
let segment = sid();
|
||||
let guard = install_top_level_at_for_test(
|
||||
lock_path.clone(),
|
||||
"transferred".into(),
|
||||
std::process::id(),
|
||||
sock("transferred"),
|
||||
Vec::new(),
|
||||
segment,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
drop(guard.cleanup_handle());
|
||||
// A successful controller transfer leaves the guard usable, including
|
||||
// the existing activation protocol, without retaining a cleanup handle.
|
||||
let replacement = sid();
|
||||
guard
|
||||
.begin_segment_activation(segment, replacement)
|
||||
.unwrap()
|
||||
.commit()
|
||||
.unwrap();
|
||||
{
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert_eq!(
|
||||
table.data().find("transferred").unwrap().segment_id,
|
||||
Some(replacement)
|
||||
);
|
||||
}
|
||||
drop(guard);
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert!(table.data().find("transferred").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cleanup_rejects_live_guard_then_releases_only_detached_allocation() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let lock_path = dir.path().join("workers.json");
|
||||
let guard = install_top_level_at_for_test(
|
||||
lock_path.clone(),
|
||||
"live".into(),
|
||||
std::process::id(),
|
||||
sock("live"),
|
||||
Vec::new(),
|
||||
sid(),
|
||||
)
|
||||
.unwrap();
|
||||
let _competitor = install_top_level_at_for_test(
|
||||
lock_path.clone(),
|
||||
"competitor".into(),
|
||||
std::process::id(),
|
||||
sock("competitor"),
|
||||
Vec::new(),
|
||||
sid(),
|
||||
)
|
||||
.unwrap();
|
||||
let cleanup = guard.cleanup_handle();
|
||||
assert!(matches!(
|
||||
cleanup.release(),
|
||||
Err(ScopeLockError::AllocationGuardAttached { worker_name }) if worker_name == "live"
|
||||
));
|
||||
drop(cleanup.clone());
|
||||
{
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert!(table.data().find("live").is_some());
|
||||
assert!(table.data().find("competitor").is_some());
|
||||
}
|
||||
|
||||
drop(guard);
|
||||
cleanup.release().unwrap();
|
||||
cleanup.clone().release().unwrap();
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert!(table.data().find("live").is_none());
|
||||
assert!(table.data().find("competitor").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn published_install_failure_retains_cleanup_across_error_drop_and_retry() {
|
||||
use super::super::table::{SaveFaultPoint, fail_next_save_at_for_test};
|
||||
|
||||
let dir = TempDir::new().unwrap();
|
||||
let lock_path = dir.path().join("workers.json");
|
||||
let _competitor = install_top_level_at_for_test(
|
||||
lock_path.clone(),
|
||||
"competitor".into(),
|
||||
std::process::id(),
|
||||
sock("competitor"),
|
||||
Vec::new(),
|
||||
sid(),
|
||||
)
|
||||
.unwrap();
|
||||
let segment = sid();
|
||||
fail_next_save_at_for_test(&lock_path, SaveFaultPoint::AfterRename);
|
||||
let error = install_top_level_at_for_test(
|
||||
lock_path.clone(),
|
||||
"failed".into(),
|
||||
std::process::id(),
|
||||
sock("failed"),
|
||||
Vec::new(),
|
||||
segment,
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(matches!(
|
||||
&error,
|
||||
ScopeLockError::AcquisitionSaveFailed { .. }
|
||||
));
|
||||
let cleanup = error.allocation_cleanup_handle().unwrap();
|
||||
let retry = error.allocation_cleanup_handle().unwrap();
|
||||
drop(error);
|
||||
{
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert_eq!(
|
||||
table.data().find("failed").unwrap().segment_id,
|
||||
Some(segment)
|
||||
);
|
||||
assert!(table.data().find("competitor").is_some());
|
||||
}
|
||||
|
||||
fail_next_save_at_for_test(&lock_path, SaveFaultPoint::BeforeRename);
|
||||
assert!(matches!(cleanup.release(), Err(ScopeLockError::Io(_))));
|
||||
drop(cleanup);
|
||||
{
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert!(table.data().find("failed").is_some());
|
||||
}
|
||||
retry.release().unwrap();
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert!(table.data().find("failed").is_none());
|
||||
assert!(table.data().find("competitor").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refused_install_has_no_cleanup_authority_over_existing_allocation() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let lock_path = dir.path().join("workers.json");
|
||||
let segment = sid();
|
||||
let _competitor = install_top_level_at_for_test(
|
||||
lock_path.clone(),
|
||||
"competitor".into(),
|
||||
std::process::id(),
|
||||
sock("competitor"),
|
||||
vec![write_rule("/shared", true)],
|
||||
segment,
|
||||
)
|
||||
.unwrap();
|
||||
let cases = [
|
||||
("competitor", sid(), Vec::new()),
|
||||
("other", segment, Vec::new()),
|
||||
("other", sid(), vec![write_rule("/shared", true)]),
|
||||
];
|
||||
for (name, requested_segment, scope) in cases {
|
||||
let error = install_top_level_at_for_test(
|
||||
lock_path.clone(),
|
||||
name.into(),
|
||||
std::process::id(),
|
||||
sock("refused"),
|
||||
scope,
|
||||
requested_segment,
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(matches!(
|
||||
&error,
|
||||
ScopeLockError::DuplicateWorkerName(_)
|
||||
| ScopeLockError::SegmentConflict { .. }
|
||||
| ScopeLockError::WriteConflict { .. }
|
||||
));
|
||||
assert!(error.allocation_cleanup_handle().is_none());
|
||||
drop(error);
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert_eq!(table.data().allocations.len(), 1);
|
||||
assert_eq!(
|
||||
table.data().find("competitor").unwrap().socket,
|
||||
sock("competitor")
|
||||
);
|
||||
assert_eq!(
|
||||
table.data().find("competitor").unwrap().segment_id,
|
||||
Some(segment)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unpublished_install_failure_has_no_cleanup_authority() {
|
||||
use super::super::table::{SaveFaultPoint, fail_next_save_at_for_test};
|
||||
|
||||
let dir = TempDir::new().unwrap();
|
||||
let lock_path = dir.path().join("workers.json");
|
||||
let _competitor = install_top_level_at_for_test(
|
||||
lock_path.clone(),
|
||||
"competitor".into(),
|
||||
std::process::id(),
|
||||
sock("competitor"),
|
||||
Vec::new(),
|
||||
sid(),
|
||||
)
|
||||
.unwrap();
|
||||
fail_next_save_at_for_test(&lock_path, SaveFaultPoint::BeforeRename);
|
||||
let error = install_top_level_at_for_test(
|
||||
lock_path.clone(),
|
||||
"failed".into(),
|
||||
std::process::id(),
|
||||
sock("failed"),
|
||||
Vec::new(),
|
||||
sid(),
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(matches!(&error, ScopeLockError::Io(_)));
|
||||
assert!(error.allocation_cleanup_handle().is_none());
|
||||
drop(error);
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert!(table.data().find("failed").is_none());
|
||||
assert!(table.data().find("competitor").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn published_adoption_failure_retains_cleanup_for_retry() {
|
||||
use super::super::table::{SaveFaultPoint, fail_next_save_at_for_test};
|
||||
|
||||
let dir = TempDir::new().unwrap();
|
||||
let lock_path = dir.path().join("workers.json");
|
||||
{
|
||||
let mut table = LockFileGuard::open(&lock_path).unwrap();
|
||||
delegate_placeholder(&mut table, "child", std::process::id());
|
||||
delegate_placeholder(&mut table, "competitor", std::process::id());
|
||||
}
|
||||
let segment = sid();
|
||||
fail_next_save_at_for_test(&lock_path, SaveFaultPoint::AfterRename);
|
||||
let error = adopt_allocation_at(
|
||||
lock_path.clone(),
|
||||
"child".into(),
|
||||
std::process::id(),
|
||||
segment,
|
||||
)
|
||||
.unwrap_err();
|
||||
let cleanup = error.allocation_cleanup_handle().unwrap();
|
||||
drop(error);
|
||||
{
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert_eq!(
|
||||
table.data().find("child").unwrap().segment_id,
|
||||
Some(segment)
|
||||
);
|
||||
}
|
||||
fail_next_save_at_for_test(&lock_path, SaveFaultPoint::AfterRename);
|
||||
assert!(matches!(cleanup.release(), Err(ScopeLockError::Io(_))));
|
||||
{
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert!(table.data().find("child").is_none());
|
||||
assert!(table.data().find("competitor").is_some());
|
||||
}
|
||||
cleanup.release().unwrap();
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert!(table.data().find("child").is_none());
|
||||
assert!(table.data().find("competitor").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unpublished_adoption_failure_preserves_placeholder_without_cleanup_authority() {
|
||||
use super::super::table::{SaveFaultPoint, fail_next_save_at_for_test};
|
||||
|
||||
let dir = TempDir::new().unwrap();
|
||||
let lock_path = dir.path().join("workers.json");
|
||||
let placeholder_pid = std::process::id();
|
||||
{
|
||||
let mut table = LockFileGuard::open(&lock_path).unwrap();
|
||||
delegate_placeholder(&mut table, "child", placeholder_pid);
|
||||
}
|
||||
fail_next_save_at_for_test(&lock_path, SaveFaultPoint::BeforeRename);
|
||||
let error = adopt_allocation_at(
|
||||
lock_path.clone(),
|
||||
"child".into(),
|
||||
placeholder_pid.wrapping_add(1),
|
||||
sid(),
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(matches!(&error, ScopeLockError::Io(_)));
|
||||
assert!(error.allocation_cleanup_handle().is_none());
|
||||
drop(error);
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
let allocation = table.data().find("child").unwrap();
|
||||
assert_eq!(allocation.pid, placeholder_pid);
|
||||
assert_eq!(allocation.segment_id, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failed_cleanup_release_retains_authority_for_retry() {
|
||||
use super::super::table::{SaveFaultPoint, fail_next_save_at_for_test};
|
||||
|
||||
for point in [SaveFaultPoint::BeforeRename, SaveFaultPoint::AfterRename] {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let lock_path = dir.path().join("workers.json");
|
||||
let guard = install_top_level_at_for_test(
|
||||
lock_path.clone(),
|
||||
"failed".into(),
|
||||
std::process::id(),
|
||||
sock("failed"),
|
||||
vec![read_rule("/shared", true)],
|
||||
sid(),
|
||||
)
|
||||
.unwrap();
|
||||
let _other = install_top_level_at_for_test(
|
||||
lock_path.clone(),
|
||||
"other".into(),
|
||||
std::process::id(),
|
||||
sock("other"),
|
||||
vec![read_rule("/shared", true)],
|
||||
sid(),
|
||||
)
|
||||
.unwrap();
|
||||
let cleanup = guard.cleanup_handle();
|
||||
let retry = cleanup.clone();
|
||||
drop(guard);
|
||||
fail_next_save_at_for_test(&lock_path, point);
|
||||
|
||||
assert!(matches!(cleanup.release(), Err(ScopeLockError::Io(_))));
|
||||
drop(cleanup);
|
||||
{
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert_eq!(
|
||||
table.data().find("failed").is_some(),
|
||||
point == SaveFaultPoint::BeforeRename
|
||||
);
|
||||
assert!(table.data().find("other").is_some());
|
||||
}
|
||||
retry.release().unwrap();
|
||||
{
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert!(table.data().find("failed").is_none());
|
||||
assert!(table.data().find("other").is_some());
|
||||
}
|
||||
|
||||
// Success disarms every clone and final Drop, even if the name is
|
||||
// subsequently allocated to a different Worker.
|
||||
let _replacement = install_top_level_at_for_test(
|
||||
lock_path.clone(),
|
||||
"failed".into(),
|
||||
std::process::id(),
|
||||
sock("replacement"),
|
||||
vec![read_rule("/shared", true)],
|
||||
sid(),
|
||||
)
|
||||
.unwrap();
|
||||
retry.clone().release().unwrap();
|
||||
drop(retry);
|
||||
let table = LockFileGuard::open(&lock_path).unwrap();
|
||||
assert_eq!(
|
||||
table.data().find("failed").unwrap().socket,
|
||||
sock("replacement")
|
||||
);
|
||||
assert!(table.data().find("other").is_some());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn adopt_allocation_rewrites_pid_and_releases_on_drop() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
@@ -358,6 +848,7 @@ mod tests {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let _sandbox = RuntimeDirSandbox::new(dir.path());
|
||||
let err = adopt_allocation("ghost".into(), 42, sid()).unwrap_err();
|
||||
assert!(err.allocation_cleanup_handle().is_none());
|
||||
assert!(matches!(err, ScopeLockError::UnknownWorker(ref n) if n == "ghost"));
|
||||
}
|
||||
|
||||
|
||||
@@ -57,6 +57,30 @@ pub fn register_worker_with_deny(
|
||||
scope_allow: Vec<ScopeRule>,
|
||||
scope_deny: Vec<ScopeRule>,
|
||||
segment_id: SegmentId,
|
||||
) -> Result<(), ScopeLockError> {
|
||||
prepare_worker_with_deny(
|
||||
guard,
|
||||
worker_name,
|
||||
pid,
|
||||
socket,
|
||||
scope_allow,
|
||||
scope_deny,
|
||||
segment_id,
|
||||
)?;
|
||||
guard.save()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Admit and stage a registration without publishing it. Owned installation
|
||||
/// constructs its cleanup authority after admission and before committing.
|
||||
pub(super) fn prepare_worker_with_deny(
|
||||
guard: &mut LockFileGuard,
|
||||
worker_name: String,
|
||||
pid: u32,
|
||||
socket: PathBuf,
|
||||
scope_allow: Vec<ScopeRule>,
|
||||
scope_deny: Vec<ScopeRule>,
|
||||
segment_id: SegmentId,
|
||||
) -> Result<(), ScopeLockError> {
|
||||
reclaim_stale(guard);
|
||||
if guard.data().find(&worker_name).is_some() {
|
||||
@@ -101,7 +125,6 @@ pub fn register_worker_with_deny(
|
||||
delegated_from: None,
|
||||
segment_id: Some(segment_id),
|
||||
});
|
||||
guard.save()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -246,6 +246,14 @@ impl LockFileGuard {
|
||||
/// Persist with atomic replacement while the separate allocation lock stays
|
||||
/// held. A crash exposes either the old complete table or the new one.
|
||||
pub fn save(&mut self) -> io::Result<()> {
|
||||
let mut published = false;
|
||||
self.save_with_publication(&mut published)
|
||||
}
|
||||
|
||||
/// Report this save's rename while the caller still holds the table lock.
|
||||
/// In-memory mutation alone is not evidence that an acquisition took place.
|
||||
pub(super) fn save_with_publication(&mut self, published: &mut bool) -> io::Result<()> {
|
||||
*published = false;
|
||||
let json = serde_json::to_vec_pretty(&self.data).map_err(io::Error::other)?;
|
||||
let mut temp_os = self.data_path.as_os_str().to_owned();
|
||||
temp_os.push(".tmp");
|
||||
@@ -265,6 +273,7 @@ impl LockFileGuard {
|
||||
));
|
||||
}
|
||||
fs::rename(&temp_path, &self.data_path)?;
|
||||
*published = true;
|
||||
#[cfg(test)]
|
||||
if take_save_fault(&self.data_path, SaveFaultPoint::AfterRename) {
|
||||
return Err(io::Error::other(
|
||||
|
||||
@@ -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.path().join("sock"),
|
||||
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.path().join("sock").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.path().join("sock").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.path().join("sock").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 {
|
||||
|
||||
+1863
-51
File diff suppressed because it is too large
Load Diff
@@ -1,3 +1,4 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::pin::Pin;
|
||||
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
@@ -150,6 +151,19 @@ impl LlmClient for MockClient {
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[path = "../../../tests/support/restore_environment.rs"]
|
||||
mod restore_test_environment;
|
||||
use restore_test_environment::require_restore_test_environment;
|
||||
|
||||
fn owned_relative_runtime_base() -> (tempfile::TempDir, PathBuf) {
|
||||
require_restore_test_environment();
|
||||
let temp = tempfile::tempdir_in(".").unwrap();
|
||||
// The relative reference and the TempDir name the same owned tree, without
|
||||
// putting the checkout prefix into sockaddr_un or changing process cwd.
|
||||
let runtime_base = Path::new(temp.path().file_name().unwrap()).join("runtime");
|
||||
(temp, runtime_base)
|
||||
}
|
||||
|
||||
fn simple_text_events() -> Vec<LlmEvent> {
|
||||
vec![
|
||||
LlmEvent::text_block_start(0),
|
||||
@@ -295,6 +309,7 @@ async fn make_worker_with_manifest_and_workspace_context(
|
||||
manifest: WorkerManifest,
|
||||
workspace_context: WorkerWorkspaceContext,
|
||||
) -> (Worker<MockClient, TestStore>, std::path::PathBuf) {
|
||||
require_restore_test_environment();
|
||||
let store_tmp = tempfile::tempdir().unwrap();
|
||||
let store = CombinedStore::new(
|
||||
FsStore::new(store_tmp.path()).unwrap(),
|
||||
@@ -320,18 +335,21 @@ async fn make_worker_with_manifest_and_workspace_context(
|
||||
}
|
||||
|
||||
async fn spawn_controller(worker: Worker<MockClient, TestStore>) -> WorkerHandle {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let runtime_base = tmp.path().to_owned();
|
||||
std::mem::forget(tmp);
|
||||
let (handle, _shutdown_rx) = WorkerController::spawn(worker, &runtime_base)
|
||||
let (tmp, runtime_base) = owned_relative_runtime_base();
|
||||
let (handle, shutdown_rx) = WorkerController::spawn(worker, &runtime_base)
|
||||
.await
|
||||
.unwrap();
|
||||
// Retain ownership until controller shutdown (or test-runtime teardown).
|
||||
tokio::spawn(async move {
|
||||
let _ = shutdown_rx.await;
|
||||
drop(tmp);
|
||||
});
|
||||
handle
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn canonical_worker_handle_uploads_and_submits_attachment() {
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let (root, runtime_base) = owned_relative_runtime_base();
|
||||
let manifest = WorkerManifest::from_toml(MANIFEST_TOML).unwrap();
|
||||
let session = session_store::WorkerSessionStore::new(root.path().join("session")).unwrap();
|
||||
let store = CombinedStore::new(
|
||||
@@ -353,7 +371,7 @@ async fn canonical_worker_handle_uploads_and_submits_attachment() {
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let (handle, shutdown_rx) = WorkerController::spawn(worker, &root.path().join("runtime"))
|
||||
let (handle, shutdown_rx) = WorkerController::spawn(worker, &runtime_base)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut rx = handle.subscribe();
|
||||
@@ -762,6 +780,227 @@ async fn controller_startup_failure_closes_bound_workdir_session() {
|
||||
));
|
||||
}
|
||||
|
||||
mod registry_startup_failure {
|
||||
//! Connect real controller preparation, registry source fences and session cleanup.
|
||||
//! The Workdir telemetry boundary substitutes a concurrent metadata change at subscription;
|
||||
//! the real registry must reject drift before starting its prepared observer.
|
||||
use super::*;
|
||||
use session_store::WorkerMetadataStore;
|
||||
use tokio::sync::broadcast;
|
||||
use workdir::WorkdirSession;
|
||||
|
||||
struct ObservedSession {
|
||||
local: LocalWorkdirSession,
|
||||
metadata_store: TestStore,
|
||||
worker_name: String,
|
||||
events: broadcast::Sender<workdir::CommandEvent>,
|
||||
subscriptions: AtomicUsize,
|
||||
receivers_at_close: Mutex<Vec<usize>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for ObservedSession {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("ObservedSession")
|
||||
.field("worker_name", &self.worker_name)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WorkdirSession for ObservedSession {
|
||||
fn workdir(&self) -> &Workdir {
|
||||
self.local.workdir()
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> WorkdirSessionCapabilities {
|
||||
self.local.capabilities()
|
||||
}
|
||||
|
||||
async fn stat(
|
||||
&self,
|
||||
request: workdir::StatRequest,
|
||||
) -> Result<workdir::StatResult, WorkdirError> {
|
||||
self.local.stat(request).await
|
||||
}
|
||||
|
||||
async fn read(
|
||||
&self,
|
||||
request: workdir::ReadRequest,
|
||||
) -> Result<workdir::ReadResult, WorkdirError> {
|
||||
self.local.read(request).await
|
||||
}
|
||||
|
||||
async fn write(
|
||||
&self,
|
||||
request: workdir::WriteRequest,
|
||||
) -> Result<workdir::WriteResult, WorkdirError> {
|
||||
self.local.write(request).await
|
||||
}
|
||||
|
||||
async fn edit(
|
||||
&self,
|
||||
request: workdir::EditRequest,
|
||||
) -> Result<workdir::EditResult, WorkdirError> {
|
||||
self.local.edit(request).await
|
||||
}
|
||||
|
||||
async fn list(
|
||||
&self,
|
||||
request: workdir::ListRequest,
|
||||
) -> Result<workdir::ListResult, WorkdirError> {
|
||||
self.local.list(request).await
|
||||
}
|
||||
|
||||
async fn glob(
|
||||
&self,
|
||||
request: workdir::GlobRequest,
|
||||
) -> Result<workdir::GlobResult, WorkdirError> {
|
||||
self.local.glob(request).await
|
||||
}
|
||||
|
||||
async fn grep(
|
||||
&self,
|
||||
request: workdir::GrepRequest,
|
||||
) -> Result<workdir::GrepResult, WorkdirError> {
|
||||
self.local.grep(request).await
|
||||
}
|
||||
|
||||
async fn start_command(
|
||||
&self,
|
||||
request: CommandRequest,
|
||||
) -> Result<workdir::CommandHandle, WorkdirError> {
|
||||
self.local.start_command(request).await
|
||||
}
|
||||
|
||||
async fn command_status(
|
||||
&self,
|
||||
handle: workdir::CommandHandle,
|
||||
) -> Result<workdir::CommandStatus, WorkdirError> {
|
||||
self.local.command_status(handle).await
|
||||
}
|
||||
|
||||
async fn command_output(
|
||||
&self,
|
||||
request: CommandOutputRequest,
|
||||
) -> Result<workdir::CommandOutput, WorkdirError> {
|
||||
self.local.command_output(request).await
|
||||
}
|
||||
|
||||
async fn cancel_command(&self, handle: workdir::CommandHandle) -> Result<(), WorkdirError> {
|
||||
self.local.cancel_command(handle).await
|
||||
}
|
||||
|
||||
fn subscribe_command_events(&self) -> Option<broadcast::Receiver<workdir::CommandEvent>> {
|
||||
// Interleave a canonical change after read-only registry preparation.
|
||||
let mut metadata = self
|
||||
.metadata_store
|
||||
.read_by_name(&self.worker_name)
|
||||
.unwrap()
|
||||
.unwrap_or_else(|| session_store::WorkerMetadata::new(&self.worker_name, None));
|
||||
metadata
|
||||
.spawned_children
|
||||
.push(session_store::WorkerSpawnedChild {
|
||||
worker_name: "legacy-drift".into(),
|
||||
socket_path: PathBuf::new(),
|
||||
callback_address: PathBuf::new(),
|
||||
scope_delegated: vec![],
|
||||
});
|
||||
self.metadata_store.write(&metadata).unwrap();
|
||||
self.subscriptions.fetch_add(1, Ordering::SeqCst);
|
||||
Some(self.events.subscribe())
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), WorkdirError> {
|
||||
// Record before any await: aborting without joining must not pass by getting a
|
||||
// scheduler turn during the provider's asynchronous close. The sender stays alive
|
||||
// even after close, so channel closure cannot accidentally dispose of the observer.
|
||||
self.receivers_at_close
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(self.events.receiver_count());
|
||||
self.local.close().await
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn registry_drift_drops_prepared_observer_before_closing_session() {
|
||||
let root = tempfile::tempdir().unwrap();
|
||||
let run_dir = root.path().join("runs/1");
|
||||
let client = MockClient::new(simple_text_events());
|
||||
let manifest = WorkerManifest::from_toml(MANIFEST_TOML).unwrap();
|
||||
let store = CombinedStore::new(
|
||||
FsStore::new(root.path().join("sessions")).unwrap(),
|
||||
FsWorkerStore::new(root.path().join("workers")).unwrap(),
|
||||
);
|
||||
let engine =
|
||||
Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(
|
||||
client.clone(),
|
||||
);
|
||||
let mut worker = Worker::new(
|
||||
manifest,
|
||||
engine,
|
||||
store.clone(),
|
||||
WorkerWorkspaceContext::local_filesystem(None),
|
||||
WorkerFilesystemAuthority::local(root.path().into(), root.path().into()),
|
||||
manifest::Scope::writable(root.path()).unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let (events, unused_receiver) = broadcast::channel(8);
|
||||
drop(unused_receiver);
|
||||
let session = Arc::new(ObservedSession {
|
||||
local: LocalWorkdirSession::materialized_bound(
|
||||
Workdir::new("registry-startup-failure"),
|
||||
root.path().into(),
|
||||
root.path().into(),
|
||||
worker.scope().clone(),
|
||||
WorkdirSessionCapabilities::ALL,
|
||||
),
|
||||
metadata_store: store,
|
||||
worker_name: worker.manifest().worker.name.clone(),
|
||||
events,
|
||||
subscriptions: AtomicUsize::new(0),
|
||||
receivers_at_close: Mutex::new(Vec::new()),
|
||||
});
|
||||
worker.bind_single_workdir_session(Some(session.clone()));
|
||||
let sessions = worker.workdir_sessions();
|
||||
|
||||
let result = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(5),
|
||||
WorkerController::spawn_runtime_managed_run(worker, &run_dir),
|
||||
)
|
||||
.await
|
||||
.expect("registry rejection and cleanup must settle");
|
||||
let error = match result {
|
||||
Err(error) => error,
|
||||
Ok(_) => panic!("registry source drift must reject controller startup"),
|
||||
};
|
||||
assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
|
||||
assert_eq!(session.subscriptions.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(*session.receivers_at_close.lock().unwrap(), vec![0]);
|
||||
assert_eq!(session.events.receiver_count(), 0);
|
||||
assert!(sessions.aliases().is_empty());
|
||||
// No controller retains the Worker/Engine after the failed startup.
|
||||
assert_eq!(Arc::strong_count(&client.call_count), 1);
|
||||
assert_eq!(client.call_count.load(Ordering::SeqCst), 0);
|
||||
assert!(!run_dir.join("worker.sock").exists());
|
||||
assert!(!run_dir.join("status.json").exists());
|
||||
assert!(!run_dir.join("manifest.toml").exists());
|
||||
assert!(matches!(
|
||||
session
|
||||
.start_command(CommandRequest {
|
||||
command: "printf unreachable".to_owned(),
|
||||
timeout_secs: 5,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
tool_call_id: None,
|
||||
})
|
||||
.await,
|
||||
Err(ref error) if matches!(error.classification_source(), WorkdirError::Unavailable(_))
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_for_status(handle: &WorkerHandle, status: WorkerStatus) {
|
||||
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
|
||||
loop {
|
||||
@@ -1561,11 +1800,12 @@ permission = "write"
|
||||
"#;
|
||||
let client = MockClient::new(simple_text_events());
|
||||
let worker = make_worker_with_pwd_and_manifest(client, manifest).await.0;
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let result = WorkerController::spawn(worker, tmp.path()).await;
|
||||
let (_tmp, runtime_base) = owned_relative_runtime_base();
|
||||
let result = WorkerController::spawn(worker, &runtime_base).await;
|
||||
assert!(
|
||||
result.is_ok(),
|
||||
"feature exposure must not imply delegation authority"
|
||||
"feature exposure must not imply delegation authority: {:?}",
|
||||
result.as_ref().err()
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user