diff --git a/crates/worker-runtime/src/execution.rs b/crates/worker-runtime/src/execution.rs index e169e6dfe..86af999e1 100644 --- a/crates/worker-runtime/src/execution.rs +++ b/crates/worker-runtime/src/execution.rs @@ -207,34 +207,6 @@ impl WorkerExecutionResult { } } -/// Opaque per-Worker execution handle returned by a backend. -/// -/// The handle is a typed token for routing calls back into the same backend. It -/// intentionally contains no socket path, process id, credential, manifest path, -/// or session path. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct WorkerExecutionHandle { - worker_ref: WorkerRef, - backend_id: String, -} - -impl WorkerExecutionHandle { - pub fn new(worker_ref: WorkerRef, backend_id: impl Into) -> Self { - Self { - worker_ref, - backend_id: backend_id.into(), - } - } - - pub fn worker_ref(&self) -> &WorkerRef { - &self.worker_ref - } - - pub fn backend_id(&self) -> &str { - &self.backend_id - } -} - /// Runtime hooks available to an execution backend for one Worker. #[cfg(feature = "ws-server")] type WorkerObservationPublisher = Arc< @@ -348,7 +320,17 @@ impl WorkerExecutionContext { let Some(event) = event else { return Ok(()); }; - (self.observation_publisher)(self.worker_ref.clone(), event)?; + if let Err(error) = (self.observation_publisher)(self.worker_ref.clone(), event.clone()) + { + let mut state = self + .publication_state + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + if let WorkerExecutionPublicationState::Activating(events) = &mut *state { + events.push_front(event); + } + return Err(error); + } } } @@ -443,7 +425,6 @@ pub struct WorkerExecutionRestoreRequest { pub struct WorkerExecutionStopRequest { pub operation_id: WorkerLifecycleOperationId, pub worker_ref: WorkerRef, - pub handle: Option, } /// Runtime-side request to refresh the latest Workspace Config before Worker creation. @@ -466,7 +447,6 @@ pub enum WorkspaceConfigFetchResult { #[derive(Clone, Debug)] pub enum WorkerExecutionSpawnResult { Connected { - handle: WorkerExecutionHandle, worker_state: protocol::WorkerStateSnapshot, workdir_attachments: Vec, }, @@ -475,11 +455,10 @@ pub enum WorkerExecutionSpawnResult { /// Live work started, failed, and all operation-owned resources were joined/released. RolledBack(WorkerExecutionResult), /// The operation crossed the live side-effect boundary and could not prove - /// either commit or complete cleanup. Optional handle/state fields preserve + /// either commit or complete cleanup. Optional state fields preserve /// concrete execution evidence for retry reconciliation. ReconciliationRequired { result: WorkerExecutionResult, - handle: Option, worker_state: Option, workdir_attachments: Vec, }, @@ -490,12 +469,10 @@ pub enum WorkerExecutionSpawnResult { impl WorkerExecutionSpawnResult { pub fn connected( - handle: WorkerExecutionHandle, worker_state: protocol::WorkerStateSnapshot, workdir_attachments: Vec, ) -> Self { Self::Connected { - handle, worker_state, workdir_attachments, } @@ -521,6 +498,57 @@ pub struct WorkerSessionAttachmentRequest { pub attachment_id: String, } +/// One real protocol connection, bound to the execution that was attached. +/// This is not a routing token and is never persisted as Worker authority. +#[cfg(feature = "ws-server")] +pub struct WorkerProtocolTransport { + pub(crate) worker_ref: WorkerRef, + pub snapshot: protocol::Event, + pub events: tokio::sync::mpsc::Receiver, + dispatch: + Arc Result, WorkerExecutionResult> + Send + Sync>, + validate: Arc Result<(), WorkerExecutionResult> + Send + Sync>, +} + +#[cfg(feature = "ws-server")] +impl WorkerProtocolTransport { + pub fn new( + worker_ref: WorkerRef, + snapshot: protocol::Event, + events: tokio::sync::mpsc::Receiver, + dispatch: Arc< + dyn Fn(Method) -> Result, WorkerExecutionResult> + Send + Sync, + >, + validate: Arc Result<(), WorkerExecutionResult> + Send + Sync>, + ) -> Self { + Self { + worker_ref, + snapshot, + events, + dispatch, + validate, + } + } + + pub fn dispatch(&self, method: Method) -> Result, WorkerExecutionResult> { + (self.dispatch)(method) + } + + pub fn validate(&self) -> Result<(), WorkerExecutionResult> { + (self.validate)() + } +} + +#[cfg(feature = "ws-server")] +impl fmt::Debug for WorkerProtocolTransport { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("WorkerProtocolTransport") + .finish_non_exhaustive() + } +} + +/// The connection selects this backend; WorkerRef identifies an execution, +/// never a capability or authorization token. pub trait WorkerExecutionBackend: Send + Sync + 'static { fn backend_id(&self) -> &str; @@ -580,8 +608,9 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static { fn activate_restored_worker( &self, _operation_id: WorkerLifecycleOperationId, - _handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, ) -> Result<(), String> { + let _ = worker_ref; Ok(()) } @@ -673,20 +702,17 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static { Ok(()) } - fn dispatch_input( - &self, - handle: &WorkerExecutionHandle, - input: WorkerInput, - ) -> WorkerExecutionResult; + fn dispatch_input(&self, worker_ref: &WorkerRef, input: WorkerInput) -> WorkerExecutionResult; fn upload_file( &self, - _handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, _file_name: &str, _media_type: &str, _content: &[u8], _context: Option<&session_store::UploadedFileUploadContext>, ) -> Result { + let _ = worker_ref; Err(WorkerExecutionResult::unsupported( WorkerExecutionOperation::UploadFile, "execution backend does not support file upload", @@ -695,20 +721,18 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static { fn delete_uploaded_file( &self, - _handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, _artifact_id: &str, ) -> WorkerExecutionResult { + let _ = worker_ref; WorkerExecutionResult::unsupported( WorkerExecutionOperation::DeleteUploadedFile, "execution backend does not support uploaded-file deletion", ) } - fn dispatch_method( - &self, - _handle: &WorkerExecutionHandle, - _method: Method, - ) -> WorkerExecutionResult { + fn dispatch_method(&self, worker_ref: &WorkerRef, _method: Method) -> WorkerExecutionResult { + let _ = worker_ref; WorkerExecutionResult::unsupported( WorkerExecutionOperation::ProtocolMethod, "execution backend does not support direct Worker protocol methods", @@ -717,32 +741,29 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static { fn worker_completions( &self, - _handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, _kind: protocol::CompletionKind, _prefix: &str, _context: Option<&protocol::CompletionContext>, ) -> Vec { + let _ = worker_ref; Vec::new() } fn stop_worker_operation(&self, request: WorkerExecutionStopRequest) -> WorkerExecutionResult { - let Some(handle) = request.handle.as_ref() else { - return WorkerExecutionResult::unsupported( - WorkerExecutionOperation::Stop, - "execution backend cannot reconcile a stopped Worker without a live handle", - ); - }; - self.stop_worker(handle) + self.stop_worker(&request.worker_ref) } - fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { + fn stop_worker(&self, worker_ref: &WorkerRef) -> WorkerExecutionResult { + let _ = worker_ref; WorkerExecutionResult::unsupported( WorkerExecutionOperation::Stop, "execution backend does not support stopping workers", ) } - fn cancel_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { + fn cancel_worker(&self, worker_ref: &WorkerRef) -> WorkerExecutionResult { + let _ = worker_ref; WorkerExecutionResult::unsupported( WorkerExecutionOperation::Cancel, "execution backend does not support cancelling workers", @@ -750,7 +771,20 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static { } #[cfg(feature = "ws-server")] - fn worker_snapshot(&self, _handle: &WorkerExecutionHandle) -> Option { + fn attach_worker_protocol( + self: Arc, + worker_ref: &WorkerRef, + ) -> Result { + let _ = worker_ref; + Err(WorkerExecutionResult::unsupported( + WorkerExecutionOperation::ProtocolMethod, + "execution backend does not support bound Worker protocol connections", + )) + } + + #[cfg(feature = "ws-server")] + fn worker_snapshot(&self, worker_ref: &WorkerRef) -> Option { + let _ = worker_ref; None } } @@ -828,9 +862,10 @@ impl WorkerExecutionBackendRef { pub(crate) fn activate_restored_worker( &self, operation_id: WorkerLifecycleOperationId, - handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, ) -> Result<(), String> { - self.backend.activate_restored_worker(operation_id, handle) + self.backend + .activate_restored_worker(operation_id, worker_ref) } pub(crate) fn restore_worker( @@ -896,57 +931,62 @@ impl WorkerExecutionBackendRef { pub(crate) fn dispatch_input( &self, - handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, input: WorkerInput, ) -> WorkerExecutionResult { - self.backend.dispatch_input(handle, input) + self.backend.dispatch_input(worker_ref, input) } pub(crate) fn upload_file( &self, - handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, file_name: &str, media_type: &str, content: &[u8], context: Option<&session_store::UploadedFileUploadContext>, ) -> Result { self.backend - .upload_file(handle, file_name, media_type, content, context) + .upload_file(worker_ref, file_name, media_type, content, context) } pub(crate) fn delete_uploaded_file( &self, - handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, artifact_id: &str, ) -> WorkerExecutionResult { - self.backend.delete_uploaded_file(handle, artifact_id) + self.backend.delete_uploaded_file(worker_ref, artifact_id) } pub(crate) fn dispatch_method( &self, - handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, method: Method, ) -> WorkerExecutionResult { - self.backend.dispatch_method(handle, method) + self.backend.dispatch_method(worker_ref, method) } #[cfg(feature = "ws-server")] - pub(crate) fn worker_snapshot( + pub(crate) fn attach_worker_protocol( &self, - handle: &WorkerExecutionHandle, - ) -> Option { - self.backend.worker_snapshot(handle) + worker_ref: &WorkerRef, + ) -> Result { + Arc::clone(&self.backend).attach_worker_protocol(worker_ref) + } + + #[cfg(feature = "ws-server")] + pub(crate) fn worker_snapshot(&self, worker_ref: &WorkerRef) -> Option { + self.backend.worker_snapshot(worker_ref) } pub(crate) fn worker_completions( &self, - handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, kind: protocol::CompletionKind, prefix: &str, context: Option<&protocol::CompletionContext>, ) -> Vec { self.backend - .worker_completions(handle, kind, prefix, context) + .worker_completions(worker_ref, kind, prefix, context) } pub(crate) fn stop_worker_operation( @@ -956,12 +996,12 @@ impl WorkerExecutionBackendRef { self.backend.stop_worker_operation(request) } - pub(crate) fn stop_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult { - self.backend.stop_worker(handle) + pub(crate) fn stop_worker(&self, worker_ref: &WorkerRef) -> WorkerExecutionResult { + self.backend.stop_worker(worker_ref) } - pub(crate) fn cancel_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult { - self.backend.cancel_worker(handle) + pub(crate) fn cancel_worker(&self, worker_ref: &WorkerRef) -> WorkerExecutionResult { + self.backend.cancel_worker(worker_ref) } } diff --git a/crates/worker-runtime/src/fs_store.rs b/crates/worker-runtime/src/fs_store.rs index c2fb52bf7..b5a9e54fe 100644 --- a/crates/worker-runtime/src/fs_store.rs +++ b/crates/worker-runtime/src/fs_store.rs @@ -493,7 +493,12 @@ pub(crate) struct PersistedWorkerRestoreOperation { #[serde(deny_unknown_fields)] pub(crate) struct PersistedWorkerStopOperation { pub(crate) operation_id: WorkerLifecycleOperationId, - pub(crate) request: CreateWorkerRequest, + /// Restore resources remain operation-owned until stop confirms cleanup. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) pending_restore: Option, + /// Stop can reconcile an identity whose execution metadata is unavailable. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) request: Option, pub(crate) binding: Option, pub(crate) restore_intent: WorkerRestoreIntent, pub(crate) last_settled_status: WorkerStatus, @@ -1942,7 +1947,7 @@ impl WorkerAggregateRecord { }); } ( - &restore.request, + Some(&restore.request), &restore.binding, if restore.last_settled_status.is_active() { WorkerRestoreIntent::Automatic @@ -1953,7 +1958,7 @@ impl WorkerAggregateRecord { ) } PersistedWorkerLifecycleOperation::Stop(stop) => { - if !stop.last_settled_status.is_active() { + if !stop.last_settled_status.is_active() && stop.pending_restore.is_none() { return Err(RuntimeError::StoreCorrupt { operation: "read Worker stop reconciliation", path: path.to_path_buf(), @@ -1961,8 +1966,18 @@ impl WorkerAggregateRecord { .to_string(), }); } + if stop.request.is_none() + && (stop.binding.is_some() || stop.pending_restore.is_some()) + { + return Err(RuntimeError::StoreCorrupt { + operation: "read Worker stop reconciliation", + path: path.to_path_buf(), + message: "metadata-free stop cannot carry restore authority" + .to_string(), + }); + } ( - &stop.request, + stop.request.as_ref(), &stop.binding, stop.restore_intent, stop.last_settled_status, @@ -1978,13 +1993,15 @@ impl WorkerAggregateRecord { .to_string(), }); } - WorkerExecutionRecord { - schema_version: SCHEMA_VERSION, - request: request.clone(), - binding: binding.clone(), - restore_intent, + if let Some(request) = request { + WorkerExecutionRecord { + schema_version: SCHEMA_VERSION, + request: request.clone(), + binding: binding.clone(), + restore_intent, + } + .validate_for_schema(&identity, path, SCHEMA_VERSION)?; } - .validate_for_schema(&identity, path, SCHEMA_VERSION)?; PersistedWorkerExecutionState::ReconciliationRequired(operation) } WorkerExecutionStateRecord::Unavailable => PersistedWorkerExecutionState::Unavailable, diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index ed1ac467a..a2ae8d744 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -1589,78 +1589,30 @@ async fn worker_protocol_ws_session( input_source: Option, mut socket: WebSocket, ) { - let mut cursor = match query.cursor.as_deref() { - Some(raw) => match WorkerObservationCursor::decode(raw) { - Some(cursor) => cursor, - None => { - let event = - protocol_error_event(format!("malformed worker observation cursor: {raw}")); - let _ = send_protocol_event(&mut socket, &event).await; - return; - } - }, - None => match runtime.worker_observation_cursor_now(&worker_ref) { - Ok(cursor) => cursor, - Err(error) => { - let event = protocol_error_event(error.to_string()); - let _ = send_protocol_event(&mut socket, &event).await; - return; - } - }, - }; - // Observation cursors are process-local. After a Runtime restart (or - // bounded backlog expiry), the current Worker snapshot is authoritative - // and replay resumes from the current in-memory tail. - if runtime - .read_worker_observation_events(&worker_ref, cursor) - .is_err() + if let Some(raw) = query.cursor.as_deref() + && WorkerObservationCursor::decode(raw).is_none() { - cursor = match runtime.worker_observation_cursor_now(&worker_ref) { - Ok(cursor) => cursor, - Err(error) => { - let event = protocol_error_event(error.to_string()); - let _ = send_protocol_event(&mut socket, &event).await; - return; - } - }; - } - - let mut receiver = match runtime.subscribe_worker_observation() { - Ok(receiver) => receiver, - Err(error) => { - let event = - protocol_error_event(format!("runtime observation bus unavailable: {error}")); - let _ = send_protocol_event(&mut socket, &event).await; - return; - } - }; - - let snapshot = match runtime.worker_observation_snapshot(&worker_ref) { - Ok(snapshot) => snapshot, - Err(error) => { - let event = protocol_error_event(error.to_string()); - let _ = send_protocol_event(&mut socket, &event).await; - return; - } - }; - if !send_protocol_event(&mut socket, &snapshot).await { + let event = protocol_error_event("malformed worker observation cursor"); + let _ = send_protocol_event(&mut socket, &event).await; return; } - - match runtime.read_worker_observation_events(&worker_ref, cursor) { - Ok(backlog) => { - for event in backlog { - cursor = WorkerObservationCursor::new(event.sequence); - if !send_protocol_event(&mut socket, &event.payload).await { - return; - } - } - } + // The complete Controller snapshot reconstitutes this protocol connection. + // A WorkerRef-wide catalog cursor cannot attribute historical operational + // events to this captured execution; never append that bus after its snapshot. + let attached = match scope.as_ref() { + Some(scope) => runtime.attach_worker_protocol_scoped(scope, &worker_ref), + None => runtime.attach_worker_protocol(&worker_ref), + }; + let mut transport = match attached { + Ok(transport) => transport, Err(error) => { let event = protocol_error_event(error.to_string()); let _ = send_protocol_event(&mut socket, &event).await; return; } + }; + if !send_protocol_event(&mut socket, &transport.snapshot).await { + return; } loop { @@ -1673,9 +1625,13 @@ async fn worker_protocol_ws_session( authorize_runtime_protocol_method(method, input_source.as_ref()); let result = match scope.as_ref() { Some(scope) => { - runtime.send_protocol_method_scoped(scope, &worker_ref, method) + runtime.send_connected_protocol_method_scoped( + scope, &worker_ref, &transport, method, + ) } - None => runtime.send_protocol_method(&worker_ref, method), + None => runtime.send_connected_protocol_method( + &worker_ref, &transport, method, + ), }; match result { Ok(events) => { @@ -1690,6 +1646,10 @@ async fn worker_protocol_ws_session( if !send_protocol_event(&mut socket, &event).await { return; } + if transport.validate().is_err() { + let _ = socket.send(WsMessage::Close(None)).await; + return; + } } } }, @@ -1716,23 +1676,20 @@ async fn worker_protocol_ws_session( } } } - event = receiver.recv() => { + event = transport.events.recv() => { match event { - Ok(event) if event.worker_ref == worker_ref && event.sequence > cursor.sequence => { - cursor = WorkerObservationCursor::new(event.sequence); - if !send_protocol_event(&mut socket, &event.payload).await { + Some(event) => { + let terminal = matches!(event, protocol::Event::Shutdown); + if !send_protocol_event(&mut socket, &event).await { + return; + } + if terminal { + let _ = socket.send(WsMessage::Close(None)).await; return; } } - Ok(_) => {} - Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => { - let event = protocol_error_event("runtime observation backlog was overrun"); - let _ = send_protocol_event(&mut socket, &event).await; - return; - } - Err(tokio::sync::broadcast::error::RecvError::Closed) => { - let event = protocol_error_event("runtime observation bus closed"); - let _ = send_protocol_event(&mut socket, &event).await; + None => { + let _ = socket.send(WsMessage::Close(None)).await; return; } } @@ -2641,9 +2598,8 @@ mod tests { ConfigBundle, ConfigBundleMetadata, ConfigBundleProvenance, ConfigProfileDescriptor, }; use crate::execution::{ - WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, - WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionSpawnRequest, - WorkerExecutionSpawnResult, + WorkerExecutionBackend, WorkerExecutionOperation, WorkerExecutionRestoreRequest, + WorkerExecutionResult, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, }; use crate::management::RuntimeOptions; use crate::retention::{DiagnosticsDisposition, SessionDisposition}; @@ -3018,6 +2974,7 @@ mod tests { assert_eq!(error.error.code, "worker_not_found"); } + #[cfg(feature = "ws-server")] #[test] fn runtime_protocol_replaces_serialized_tracked_source() { let wire = serde_json::to_string(&protocol::Method::SubmitTracked { @@ -3045,6 +3002,7 @@ mod tests { )); } + #[cfg(feature = "ws-server")] #[test] fn runtime_protocol_uses_transport_authenticated_account_source() { let mut headers = HeaderMap::new(); @@ -3398,12 +3356,9 @@ mod tests { } fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { - WorkerExecutionSpawnResult::Connected { - handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), - worker_state: protocol::WorkerStateSnapshot { - ..protocol::WorkerStatus::Idle.into() - }, - workdir_attachments: request + WorkerExecutionSpawnResult::connected( + protocol::WorkerStatus::Idle.into(), + request .workdir_attachments .iter() .map( @@ -3413,25 +3368,22 @@ mod tests { }, ) .collect(), - } + ) } fn restore_worker( &self, request: WorkerExecutionRestoreRequest, ) -> WorkerExecutionSpawnResult { - WorkerExecutionSpawnResult::Connected { - handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), - worker_state: protocol::WorkerStateSnapshot { - ..protocol::WorkerStatus::Idle.into() - }, - workdir_attachments: request.previous_workdir_attachments, - } + WorkerExecutionSpawnResult::connected( + protocol::WorkerStatus::Idle.into(), + request.previous_workdir_attachments, + ) } fn dispatch_input( &self, - _handle: &WorkerExecutionHandle, + _worker_ref: &WorkerRef, input: WorkerInput, ) -> WorkerExecutionResult { if let Some(submission_id) = input.submission_request_id { @@ -3446,7 +3398,57 @@ mod tests { } } - fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { + fn stop_worker(&self, _worker_ref: &WorkerRef) -> WorkerExecutionResult { + WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop) + } + } + + #[cfg(feature = "fs-store")] + #[derive(Default)] + struct RejectingRestoreBackend { + restores: Mutex>, + inputs: Mutex>, + stops: Mutex>, + } + + #[cfg(feature = "fs-store")] + impl WorkerExecutionBackend for RejectingRestoreBackend { + fn backend_id(&self) -> &str { + "http-test" + } + + fn spawn_worker( + &self, + _request: WorkerExecutionSpawnRequest, + ) -> WorkerExecutionSpawnResult { + panic!("a saved Worker must be restored, not spawned again"); + } + + fn restore_worker( + &self, + request: WorkerExecutionRestoreRequest, + ) -> WorkerExecutionSpawnResult { + self.restores.lock().unwrap().push(request.worker_ref); + WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + "repository access must be reacquired", + )) + } + + fn dispatch_input( + &self, + worker_ref: &WorkerRef, + _input: WorkerInput, + ) -> WorkerExecutionResult { + self.inputs.lock().unwrap().push(worker_ref.clone()); + WorkerExecutionResult::rejected( + WorkerExecutionOperation::Input, + "Worker execution is unavailable after restore rejection", + ) + } + + fn stop_worker(&self, worker_ref: &WorkerRef) -> WorkerExecutionResult { + self.stops.lock().unwrap().push(worker_ref.clone()); WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop) } } @@ -3831,6 +3833,192 @@ mod tests { assert_eq!(summary.runtime.stopped_worker_count, 1); } + #[cfg(feature = "fs-store")] + #[tokio::test] + async fn rest_saved_active_worker_restore_rejection_is_observable_and_can_be_stopped_removed() { + let temp = tempfile::tempdir().unwrap(); + let options = crate::fs_store::FsRuntimeStoreOptions::new(temp.path().join("runtime")); + let runtime = Runtime::with_fs_store_and_execution_backend( + options.clone(), + Arc::new(AcceptingBackend), + ) + .unwrap(); + runtime + .store_config_bundle(test_bundle(ProfileSelector::Builtin( + "builtin:coder".to_string(), + ))) + .unwrap(); + let mut request = task_request("saved active Worker"); + request.workspace_api = Some(WorkspaceApiRef { + workspace_id: "local".to_string(), + base_url: "http://127.0.0.1:8787".to_string(), + }); + let created = runtime + .create_worker_scoped(&RuntimeWorkspaceScope::new("local", "local-token"), request) + .unwrap(); + assert_eq!(created.status, crate::catalog::WorkerStatus::Idle); + assert!(created.worker_state.is_some()); + drop(runtime); + + let backend = Arc::new(RejectingRestoreBackend::default()); + let runtime = + Runtime::with_fs_store_and_execution_backend(options, backend.clone()).unwrap(); + // Startup attempted automatic restore without turning rejection into a + // fabricated live Worker or silently changing its saved lifecycle status. + assert_eq!( + *backend.restores.lock().unwrap(), + vec![created.worker_ref.clone()], + "startup diagnostics: {:?}; Worker: {:?}", + runtime.diagnostics().unwrap(), + runtime.worker_detail(&created.worker_ref) + ); + let token = "local-token"; + let app = runtime_http_router(runtime.clone(), token.to_string()); + let worker_uri = format!("/v1/workers/{}", created.worker_id); + + let response = authed_empty_request(app.clone(), Method::GET, &worker_uri, token).await; + assert_eq!(response.status(), StatusCode::OK); + let detail: RuntimeHttpWorkerResponse = read_json(response).await; + assert_eq!(detail.worker.worker_ref, created.worker_ref); + assert_eq!(detail.worker.status, crate::catalog::WorkerStatus::Idle); + assert!(detail.worker.execution_metadata_available); + assert!(detail.worker.worker_state.is_none()); + + let response = authed_empty_request(app.clone(), Method::GET, "/v1/workers", token).await; + assert_eq!(response.status(), StatusCode::OK); + let listed: RuntimeHttpWorkersResponse = read_json(response).await; + assert_eq!(listed.workers.len(), 1); + assert_eq!(listed.workers[0].worker_ref, created.worker_ref); + assert_eq!(listed.workers[0].status, crate::catalog::WorkerStatus::Idle); + assert!(listed.workers[0].worker_state.is_none()); + + let response = authed_empty_request( + app.clone(), + Method::POST, + &format!("{worker_uri}/restore"), + token, + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + let rejected: runtime_api::WorkerRestoreResponse = read_json(response).await; + assert_eq!(rejected.state, WorkerRestoreState::Rejected); + assert_eq!( + rejected.reason_code.as_deref(), + Some("worker_restore_rejected") + ); + assert!(rejected.worker.is_none()); + assert_eq!( + *backend.restores.lock().unwrap(), + vec![created.worker_ref.clone(), created.worker_ref.clone()] + ); + + let response = authed_json_request( + app.clone(), + Method::POST, + &format!("{worker_uri}/input"), + token, + &WorkerInput::user("must not be accepted by a fabricated execution"), + ) + .await; + assert_eq!(response.status(), StatusCode::CONFLICT); + let error: RuntimeHttpErrorResponse = read_json(response).await; + assert_eq!(error.error.code, "worker_execution_rejected"); + assert_eq!( + *backend.inputs.lock().unwrap(), + vec![created.worker_ref.clone()] + ); + + #[cfg(feature = "ws-server")] + { + use tokio_tungstenite::tungstenite::Message; + use tokio_tungstenite::tungstenite::client::IntoClientRequest; + + assert!(matches!( + runtime.worker_observation_snapshot(&created.worker_ref), + Err(RuntimeError::WorkerExecutionUnavailable { worker_id, .. }) + if worker_id == created.worker_id + )); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server_app = app.clone(); + let server = tokio::spawn(async move { + axum::serve(listener, server_app).await.unwrap(); + }); + let mut request = format!("ws://{address}{worker_uri}/protocol/ws") + .into_client_request() + .unwrap(); + request.headers_mut().insert( + tokio_tungstenite::tungstenite::http::header::AUTHORIZATION, + format!("Bearer {token}").parse().unwrap(), + ); + let (mut stream, _) = tokio_tungstenite::connect_async(request).await.unwrap(); + let frame = tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()) + .await + .expect("unavailable snapshot must produce a protocol error promptly") + .unwrap() + .unwrap(); + let Message::Text(text) = frame else { + panic!("expected protocol error frame, not a fabricated snapshot"); + }; + let event: protocol::Event = serde_json::from_str(&text).unwrap(); + assert!(matches!(event, protocol::Event::Error { .. })); + drop(stream); + server.abort(); + let _ = server.await; + } + + // Saved active status still requires a normal stop before removal, + // even though automatic restore never connected a live execution. + let response = authed_empty_request(app.clone(), Method::DELETE, &worker_uri, token).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let error: RuntimeHttpErrorResponse = read_json(response).await; + assert_eq!(error.error.code, "invalid_request"); + assert!(backend.stops.lock().unwrap().is_empty()); + + let response = authed_empty_request( + app.clone(), + Method::POST, + &format!("{worker_uri}/stop"), + token, + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + let stopped: RuntimeHttpWorkerLifecycleResponse = read_json(response).await; + assert_eq!(stopped.ack.worker_ref, created.worker_ref); + assert_eq!(stopped.ack.status, crate::catalog::WorkerStatus::Stopped); + assert_eq!( + *backend.stops.lock().unwrap(), + vec![created.worker_ref.clone()] + ); + + let response = authed_empty_request(app.clone(), Method::GET, &worker_uri, token).await; + assert_eq!(response.status(), StatusCode::OK); + let detail: RuntimeHttpWorkerResponse = read_json(response).await; + assert_eq!(detail.worker.status, crate::catalog::WorkerStatus::Stopped); + + let response = authed_empty_request(app.clone(), Method::DELETE, &worker_uri, token).await; + assert_eq!(response.status(), StatusCode::OK); + let removed: RuntimeHttpWorkerDeleteResponse = read_json(response).await; + assert_eq!(removed.worker.worker_id, created.worker_id); + assert!(removed.worker.deleted); + // Removal asks the backend to reconfirm shutdown even for a saved + // Stopped Worker; it must use the same WorkerRef as the original stop. + assert_eq!( + *backend.stops.lock().unwrap(), + vec![created.worker_ref.clone(), created.worker_ref.clone()] + ); + assert!( + !temp + .path() + .join("runtime/workers") + .join(created.worker_id.to_string()) + .exists() + ); + + let response = authed_empty_request(app, Method::GET, &worker_uri, token).await; + assert_eq!(response.status(), StatusCode::NOT_FOUND); + } + #[tokio::test] async fn local_token_placeholder_rejects_missing_bearer_token() { let app = runtime_http_router(Runtime::new_memory(), "local-token".to_string()); @@ -4040,8 +4228,9 @@ mod ws_tests { ConfigBundle, ConfigBundleMetadata, ConfigBundleProvenance, ConfigProfileDescriptor, }; use crate::execution::{ - WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, - WorkerExecutionResult, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, + WorkerExecutionBackend, WorkerExecutionContext, WorkerExecutionOperation, + WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionSpawnRequest, + WorkerExecutionSpawnResult, WorkerProtocolTransport, }; use crate::management::RuntimeOptions; use futures::{SinkExt, StreamExt}; @@ -4051,7 +4240,133 @@ mod ws_tests { use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::header as ws_header; - struct WsBackend; + struct WsController { + commands: tokio::sync::mpsc::UnboundedReceiver, + connections: Vec>, + } + + struct WsExecution { + snapshot: protocol::Event, + context: WorkerExecutionContext, + methods: tokio::sync::mpsc::UnboundedSender, + controller: Mutex, + } + + impl WsExecution { + fn validate(&self) -> Result<(), WorkerExecutionResult> { + if self.methods.is_closed() { + Err(WorkerExecutionResult::rejected( + WorkerExecutionOperation::ProtocolMethod, + "mock Controller method channel is closed", + )) + } else { + Ok(()) + } + } + + fn dispatch( + &self, + method: protocol::Method, + ) -> Result, WorkerExecutionResult> { + let mut controller = self.controller.lock().unwrap(); + self.validate()?; + match method { + protocol::Method::ListCompletions { + kind, + prefix, + request_id, + context, + } => Ok(vec![protocol::Event::Completions { + kind, + prefix, + request_id, + context, + entries: Vec::new(), + }]), + protocol::Method::Shutdown { .. } => { + controller.commands.close(); + controller.connections.clear(); + Ok(Vec::new()) + } + method => { + self.methods.send(method).map_err(|error| { + WorkerExecutionResult::rejected( + WorkerExecutionOperation::ProtocolMethod, + error.to_string(), + ) + })?; + Ok(Vec::new()) + } + } + } + + fn close(&self) { + let mut controller = self.controller.lock().unwrap(); + controller.commands.close(); + controller.connections.clear(); + } + } + + #[derive(Default)] + struct WsBackend { + executions: Mutex>>, + } + + impl WsBackend { + fn insert_execution(&self, worker_ref: WorkerRef, context: WorkerExecutionContext) { + let (methods, commands) = tokio::sync::mpsc::unbounded_channel(); + let execution = Arc::new(WsExecution { + snapshot: ws_snapshot(&worker_ref), + context, + methods, + controller: Mutex::new(WsController { + commands, + connections: Vec::new(), + }), + }); + if let Some(previous) = self + .executions + .lock() + .unwrap() + .insert(worker_ref, execution) + { + previous.close(); + } + } + + fn publish( + &self, + worker_ref: &WorkerRef, + event: protocol::Event, + ) -> crate::observation::WorkerObservationEvent { + let execution = self + .executions + .lock() + .unwrap() + .get(worker_ref) + .unwrap() + .clone(); + let observation = execution + .context + .publish_protocol_event(event.clone()) + .unwrap(); + // The fixture publishes separately to Runtime history and to this + // Controller's attached streams, just like the backend event relay. + execution + .controller + .lock() + .unwrap() + .connections + .retain(|sender| match sender.try_send(event.clone()) { + Ok(()) => true, + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => false, + Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { + panic!("mock protocol event receiver lagged") + } + }); + observation + } + } impl WorkerExecutionBackend for WsBackend { fn backend_id(&self) -> &str { @@ -4059,12 +4374,10 @@ mod ws_tests { } fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { - WorkerExecutionSpawnResult::Connected { - handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), - worker_state: protocol::WorkerStateSnapshot { - ..protocol::WorkerStatus::Idle.into() - }, - workdir_attachments: request + self.insert_execution(request.worker_ref, request.context); + WorkerExecutionSpawnResult::connected( + protocol::WorkerStatus::Idle.into(), + request .workdir_attachments .iter() .map( @@ -4074,12 +4387,12 @@ mod ws_tests { }, ) .collect(), - } + ) } fn dispatch_input( &self, - _handle: &WorkerExecutionHandle, + _worker_ref: &WorkerRef, input: WorkerInput, ) -> WorkerExecutionResult { if let Some(submission_id) = input.submission_request_id { @@ -4094,40 +4407,91 @@ mod ws_tests { } } - fn worker_snapshot(&self, handle: &WorkerExecutionHandle) -> Option { - Some(protocol::Event::Snapshot { - session: protocol::SessionSnapshot { - pending_submissions: protocol::PendingSubmissionsSnapshot::default(), - entries: Vec::new(), - }, - greeting: protocol::Greeting { - worker_name: handle.worker_ref().worker_id.to_string(), - cwd: String::new(), - provider: "ws-test".to_string(), - model: "ws-test".to_string(), - scope_summary: "WebSocket test execution snapshot".to_string(), - tools: Vec::new(), - context_window: 0, - context_tokens: 0, - reasoning: None, - context_usage: None, - }, - state: protocol::WorkerStateSnapshot::initial(), - in_flight: protocol::InFlightSnapshot { - blocks: Vec::new(), - commands: Vec::new(), - compaction: None, - }, - internal_workers: Vec::new(), - }) + fn worker_snapshot(&self, worker_ref: &WorkerRef) -> Option { + self.executions + .lock() + .unwrap() + .get(worker_ref) + .map(|execution| execution.snapshot.clone()) } - fn dispatch_method( + fn restore_worker( &self, - _handle: &WorkerExecutionHandle, - _method: protocol::Method, - ) -> WorkerExecutionResult { - WorkerExecutionResult::accepted(WorkerExecutionOperation::ProtocolMethod) + request: WorkerExecutionRestoreRequest, + ) -> WorkerExecutionSpawnResult { + self.insert_execution(request.worker_ref, request.context); + WorkerExecutionSpawnResult::connected( + protocol::WorkerStatus::Idle.into(), + request.previous_workdir_attachments, + ) + } + + fn stop_worker(&self, worker_ref: &WorkerRef) -> WorkerExecutionResult { + if let Some(execution) = self.executions.lock().unwrap().remove(worker_ref) { + execution.close(); + } + WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop) + } + + fn attach_worker_protocol( + self: Arc, + worker_ref: &WorkerRef, + ) -> Result { + let execution = self + .executions + .lock() + .unwrap() + .get(worker_ref) + .cloned() + .ok_or_else(|| { + WorkerExecutionResult::rejected( + WorkerExecutionOperation::ProtocolMethod, + "mock Controller is unavailable", + ) + })?; + let (sender, events) = tokio::sync::mpsc::channel(32); + { + let mut controller = execution.controller.lock().unwrap(); + execution.validate()?; + controller.connections.push(sender); + } + let dispatcher = execution.clone(); + let validator = execution.clone(); + Ok(WorkerProtocolTransport::new( + worker_ref.clone(), + execution.snapshot.clone(), + events, + Arc::new(move |method| dispatcher.dispatch(method)), + Arc::new(move || validator.validate()), + )) + } + } + + fn ws_snapshot(worker_ref: &WorkerRef) -> protocol::Event { + protocol::Event::Snapshot { + session: protocol::SessionSnapshot { + pending_submissions: protocol::PendingSubmissionsSnapshot::default(), + entries: Vec::new(), + }, + greeting: protocol::Greeting { + worker_name: worker_ref.worker_id.to_string(), + cwd: String::new(), + provider: "ws-test".to_string(), + model: "ws-test".to_string(), + scope_summary: "WebSocket test execution snapshot".to_string(), + tools: Vec::new(), + context_window: 0, + context_tokens: 0, + reasoning: None, + context_usage: None, + }, + state: protocol::WorkerStateSnapshot::initial(), + in_flight: protocol::InFlightSnapshot { + blocks: Vec::new(), + commands: Vec::new(), + compaction: None, + }, + internal_workers: Vec::new(), } } @@ -4199,10 +4563,10 @@ mod ws_tests { } } - async fn spawn_runtime_server() -> (Runtime, WorkerRef, String) { + async fn spawn_runtime_server() -> (Runtime, Arc, WorkerRef, String) { + let backend = Arc::new(WsBackend::default()); let runtime = - Runtime::with_execution_backend(RuntimeOptions::default(), Arc::new(WsBackend)) - .unwrap(); + Runtime::with_execution_backend(RuntimeOptions::default(), backend.clone()).unwrap(); runtime .store_config_bundle(ws_test_bundle(ProfileSelector::Builtin( "builtin:companion".to_string(), @@ -4226,6 +4590,7 @@ mod ws_tests { }); ( runtime, + backend, worker.worker_ref.clone(), format!( "ws://{addr}/v1/workers/{}/protocol/ws", @@ -4248,7 +4613,11 @@ mod ws_tests { tokio_tungstenite::MaybeTlsStream, >, ) -> protocol::Event { - let message = stream.next().await.unwrap().unwrap(); + let message = tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()) + .await + .expect("Worker protocol frame must arrive promptly") + .unwrap() + .unwrap(); let Message::Text(text) = message else { panic!("expected text frame"); }; @@ -4260,7 +4629,11 @@ mod ws_tests { tokio_tungstenite::MaybeTlsStream, >, ) -> SubscriptionFrame { - let message = stream.next().await.unwrap().unwrap(); + let message = tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()) + .await + .expect("Worker protocol frame must arrive promptly") + .unwrap() + .unwrap(); let Message::Text(text) = message else { panic!("expected text frame"); }; @@ -4290,7 +4663,7 @@ mod ws_tests { #[tokio::test] async fn runtime_protocol_ws_subscribes_and_filters_worker_lifecycle() { - let (runtime, worker_ref, worker_url) = spawn_runtime_server().await; + let (runtime, backend, worker_ref, worker_url) = spawn_runtime_server().await; let other = runtime .create_worker_scoped( &RuntimeWorkspaceScope::new("local", "local-token"), @@ -4347,22 +4720,18 @@ mod ws_tests { )); snapshot }; - runtime - .observe_worker_event( - &other.worker_ref, - protocol::Event::WorkerState { - snapshot: running_snapshot(&other.worker_ref), - }, - ) - .unwrap(); - runtime - .observe_worker_event( - &worker_ref, - protocol::Event::WorkerState { - snapshot: running_snapshot(&worker_ref), - }, - ) - .unwrap(); + backend.publish( + &other.worker_ref, + protocol::Event::WorkerState { + snapshot: running_snapshot(&other.worker_ref), + }, + ); + backend.publish( + &worker_ref, + protocol::Event::WorkerState { + snapshot: running_snapshot(&worker_ref), + }, + ); let event = next_subscription_frame(&mut stream).await; assert!(matches!( event.payload, @@ -4409,7 +4778,7 @@ mod ws_tests { #[tokio::test] async fn protocol_ws_connect_sends_snapshot_and_live_worker_events() { - let (runtime, worker_ref, url) = spawn_runtime_server().await; + let (_runtime, backend, worker_ref, url) = spawn_runtime_server().await; let (mut stream, _) = connect_async(authed_ws_request(&url)).await.unwrap(); assert!(matches!( @@ -4417,40 +4786,129 @@ mod ws_tests { protocol::Event::Snapshot { .. } )); - runtime - .observe_worker_event( - &worker_ref, - protocol::Event::TextDelta { - text: "started".into(), - }, - ) - .unwrap(); + backend.publish( + &worker_ref, + protocol::Event::TextDelta { + text: "started".into(), + }, + ); assert!(matches!( next_frame(&mut stream).await, protocol::Event::TextDelta { .. } )); } + #[tokio::test] + async fn protocol_ws_connected_dispatch_preserves_transport_source_and_completion_reply() { + let (_runtime, backend, worker_ref, url) = spawn_runtime_server().await; + let mut request = authed_ws_request(&url); + request.headers_mut().insert( + protocol::AUTHENTICATED_ACCOUNT_ID_HEADER, + "account-1".parse().unwrap(), + ); + let (mut stream, _) = connect_async(request).await.unwrap(); + let _ = next_frame(&mut stream).await; + let notification = protocol::Method::NotifyTracked { + notification_request_id: "notification-1".into(), + message: "hello".into(), + source: protocol::AuthenticatedInputSource::Account { + account_id: "forged".into(), + }, + }; + let completions = protocol::Method::ListCompletions { + kind: protocol::CompletionKind::Feature, + prefix: "builtin:".into(), + request_id: Some("completion-1".into()), + context: None, + }; + for method in [notification, completions] { + stream + .send(Message::Text( + serde_json::to_string(&method).unwrap().into(), + )) + .await + .unwrap(); + } + assert!(matches!( + next_frame(&mut stream).await, + protocol::Event::Completions { request_id: Some(request_id), prefix, .. } + if request_id == "completion-1" && prefix == "builtin:" + )); + // The reply orders this assertion after dispatch of the preceding frame. + let execution = backend + .executions + .lock() + .unwrap() + .get(&worker_ref) + .unwrap() + .clone(); + let method = execution + .controller + .lock() + .unwrap() + .commands + .try_recv() + .unwrap(); + assert!(matches!( + method, + protocol::Method::NotifyTracked { + source: protocol::AuthenticatedInputSource::Account { account_id }, + .. + } if account_id == "account-1" + )); + } + + #[tokio::test] + async fn protocol_ws_stop_closes_old_stream_and_restore_requires_new_connection() { + let (runtime, backend, worker_ref, url) = spawn_runtime_server().await; + let (mut old_stream, _) = connect_async(authed_ws_request(&url)).await.unwrap(); + let _ = next_frame(&mut old_stream).await; + runtime.stop_worker(&worker_ref, None).unwrap(); + runtime.restore_worker(&worker_ref).unwrap(); + backend.publish( + &worker_ref, + protocol::Event::TextDelta { + text: "new execution".into(), + }, + ); + let closed = tokio::time::timeout(std::time::Duration::from_secs(5), old_stream.next()) + .await + .expect("the old protocol connection must close on stop"); + assert!(matches!(closed, Some(Ok(Message::Close(_))) | None)); + + let (mut new_stream, _) = connect_async(authed_ws_request(&url)).await.unwrap(); + assert!(matches!( + next_frame(&mut new_stream).await, + protocol::Event::Snapshot { .. } + )); + backend.publish( + &worker_ref, + protocol::Event::TextDone { + text: "restored execution".into(), + }, + ); + assert!(matches!( + next_frame(&mut new_stream).await, + protocol::Event::TextDone { text } if text == "restored execution" + )); + } + #[tokio::test] async fn protocol_ws_cursor_resume_is_duplicate_safe_and_filters_workers() { - let (runtime, worker_ref, url) = spawn_runtime_server().await; + let (runtime, backend, worker_ref, url) = spawn_runtime_server().await; let other = runtime.create_worker(ws_create_request()).unwrap(); - let first = runtime - .observe_worker_event( - &worker_ref, - protocol::Event::TextDelta { - text: "started".into(), - }, - ) - .unwrap(); - runtime - .observe_worker_event( - &other.worker_ref, - protocol::Event::TextDelta { - text: "other".into(), - }, - ) - .unwrap(); + let first = backend.publish( + &worker_ref, + protocol::Event::TextDelta { + text: "started".into(), + }, + ); + backend.publish( + &other.worker_ref, + protocol::Event::TextDelta { + text: "other".into(), + }, + ); let resume_url = format!("{url}?cursor={}", first.cursor); let (mut stream, _) = connect_async(authed_ws_request(&resume_url)).await.unwrap(); @@ -4459,23 +4917,56 @@ mod ws_tests { protocol::Event::Snapshot { .. } )); - runtime - .observe_worker_event( - &worker_ref, - protocol::Event::TextDone { - text: "done".into(), - }, - ) - .unwrap(); + backend.publish( + &worker_ref, + protocol::Event::TextDone { + text: "done".into(), + }, + ); assert!(matches!( next_frame(&mut stream).await, protocol::Event::TextDone { .. } )); } + #[tokio::test] + async fn protocol_ws_old_cursor_never_replays_stopped_execution_after_new_snapshot() { + let (runtime, backend, worker_ref, url) = spawn_runtime_server().await; + let old = backend.publish( + &worker_ref, + protocol::Event::TextDelta { + text: "old cursor".into(), + }, + ); + backend.publish( + &worker_ref, + protocol::Event::TextDone { + text: "old output".into(), + }, + ); + backend.publish(&worker_ref, protocol::Event::Shutdown); + runtime.stop_worker(&worker_ref, None).unwrap(); + runtime.restore_worker(&worker_ref).unwrap(); + let resume_url = format!("{url}?cursor={}", old.cursor); + let (mut stream, _) = connect_async(authed_ws_request(&resume_url)).await.unwrap(); + assert!(matches!( + next_frame(&mut stream).await, + protocol::Event::Snapshot { .. } + )); + backend.publish( + &worker_ref, + protocol::Event::TextDone { + text: "new output".into(), + }, + ); + assert!( + matches!(next_frame(&mut stream).await, protocol::Event::TextDone { text } if text == "new output") + ); + } + #[tokio::test] async fn protocol_ws_reports_malformed_cursor_and_method_frame() { - let (_runtime, _worker_ref, url) = spawn_runtime_server().await; + let (_runtime, _backend, _worker_ref, url) = spawn_runtime_server().await; let malformed_url = format!("{url}?cursor=bad"); let (mut malformed, _) = connect_async(authed_ws_request(&malformed_url)) .await diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 93e82971f..3d37c59d6 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -13,11 +13,10 @@ use crate::diagnostics::{DiagnosticSeverity, RuntimeDiagnostic}; use crate::error::RuntimeError; use crate::execution::WorkerExecutionRestoreRequest; use crate::execution::{ - WorkerExecutionBackend, WorkerExecutionBackendRef, WorkerExecutionHandle, - WorkerExecutionOperation, WorkerExecutionResult, WorkerExecutionSpawnRequest, - WorkerExecutionSpawnResult, WorkerExecutionStopRequest, WorkerLifecycleOperationId, - WorkerSessionHistoryRequest, WorkerSessionObservationRequest, WorkspaceConfigFetchRequest, - WorkspaceConfigFetchResult, + WorkerExecutionBackend, WorkerExecutionBackendRef, WorkerExecutionOperation, + WorkerExecutionResult, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, + WorkerExecutionStopRequest, WorkerLifecycleOperationId, WorkerSessionHistoryRequest, + WorkerSessionObservationRequest, WorkspaceConfigFetchRequest, WorkspaceConfigFetchResult, }; #[cfg(feature = "fs-store")] use crate::fs_store::{ @@ -164,12 +163,16 @@ impl RuntimeWorkerRestoreResult { } } - fn failed(state: WorkerRestoreState, reason_code: &'static str, message: &'static str) -> Self { + fn failed( + state: WorkerRestoreState, + reason_code: &'static str, + message: impl Into, + ) -> Self { Self { state, worker: None, reason_code: Some(reason_code.to_string()), - message: Some(message.to_string()), + message: Some(message.into()), } } @@ -1118,7 +1121,6 @@ impl Runtime { restore_candidate_context: None, workdir_attachments: Vec::new(), logical_workdir_attachments: Vec::new(), - execution_handle: None, internal_workers: InternalWorkerActivityProjection::default(), }; state.workers.insert(worker_id, record); @@ -1136,12 +1138,11 @@ impl Runtime { }; let spawn_result = backend.spawn_worker(spawn_request); - let (handle, initial_worker_state, workdir_attachments) = match spawn_result { + let (initial_worker_state, workdir_attachments) = match spawn_result { WorkerExecutionSpawnResult::Connected { - handle, worker_state, workdir_attachments, - } => (handle, worker_state, workdir_attachments), + } => (worker_state, workdir_attachments), WorkerExecutionSpawnResult::Rejected(result) | WorkerExecutionSpawnResult::RolledBack(result) | WorkerExecutionSpawnResult::Errored(result) => { @@ -1156,21 +1157,18 @@ impl Runtime { } WorkerExecutionSpawnResult::ReconciliationRequired { result, - handle, worker_state, workdir_attachments, } => { - let retain_result = - if let (Some(handle), Some(worker_state)) = (handle, worker_state) { - self.retain_create_execution_evidence( - &worker_ref, - handle, - worker_state, - workdir_attachments, - ) - } else { - self.retain_create_reconciliation_pending(&worker_ref, workdir_attachments) - }; + let retain_result = if let Some(worker_state) = worker_state { + self.retain_create_execution_evidence( + &worker_ref, + worker_state, + workdir_attachments, + ) + } else { + self.retain_create_reconciliation_pending(&worker_ref, workdir_attachments) + }; if let Err(error) = retain_result { tracing::error!( worker_id = %worker_ref.worker_id, @@ -1202,9 +1200,9 @@ impl Runtime { .filter(|request_id| !request_id.trim().is_empty()) .unwrap_or_else(|| Uuid::now_v7().to_string()); initial_input.submission_request_id = Some(expected_submission_id.clone()); - let dispatch_result = backend.dispatch_input(&handle, initial_input.clone()); + let dispatch_result = backend.dispatch_input(&worker_ref, initial_input.clone()); if !dispatch_result.is_accepted() { - self.cleanup_connected_failed_create(&backend, &worker_ref, &handle)?; + self.cleanup_connected_failed_create(&backend, &worker_ref)?; return Err(RuntimeError::WorkerExecutionRejected { worker_id: worker_ref.worker_id.clone(), operation: dispatch_result.operation, @@ -1218,7 +1216,7 @@ impl Runtime { .as_ref() .is_some_and(|ack| ack.submission_request_id == expected_submission_id); if !has_durable_acceptance { - self.cleanup_connected_failed_create(&backend, &worker_ref, &handle)?; + self.cleanup_connected_failed_create(&backend, &worker_ref)?; let result = WorkerExecutionResult::rejected( WorkerExecutionOperation::Input, "execution backend accepted initial input without a durable submission acknowledgement", @@ -1233,33 +1231,31 @@ impl Runtime { } let detail = match self.commit_created_worker( &worker_ref, - handle.clone(), initial_worker_state.clone(), workdir_attachments, dispatch_result, ) { Ok(detail) => detail, Err(error) => { - self.cleanup_connected_failed_create(&backend, &worker_ref, &handle)?; + self.cleanup_connected_failed_create(&backend, &worker_ref)?; return Err(error); } }; if let Err(error) = self.record_input_observation(&worker_ref, initial_input) { - self.cleanup_connected_failed_create(&backend, &worker_ref, &handle)?; + self.cleanup_connected_failed_create(&backend, &worker_ref)?; return Err(error); } Ok(detail) } else { match self.commit_created_worker( &worker_ref, - handle.clone(), initial_worker_state, workdir_attachments, WorkerExecutionResult::accepted(WorkerExecutionOperation::Spawn), ) { Ok(detail) => Ok(detail), Err(error) => { - self.cleanup_connected_failed_create(&backend, &worker_ref, &handle)?; + self.cleanup_connected_failed_create(&backend, &worker_ref)?; Err(error) } } @@ -1516,22 +1512,33 @@ impl Runtime { worker_ref: &WorkerRef, attachments: Vec, ) -> Result { - self.ensure_worker_in_workspace(scope, worker_ref)?; - self.replace_worker_workdir_attachments(worker_ref, attachments) + self.replace_worker_workdir_attachments_with_scope(worker_ref, attachments, Some(scope)) } /// Persist the Workspace-authoritative logical attachment set used by live /// mutation and future restore. Runtime-local bindings remain a separate, /// optional implementation detail and never authorize logical attachment. pub fn replace_worker_workdir_attachments( + &self, + worker_ref: &WorkerRef, + attachments: Vec, + ) -> Result { + self.replace_worker_workdir_attachments_with_scope(worker_ref, attachments, None) + } + + fn replace_worker_workdir_attachments_with_scope( &self, worker_ref: &WorkerRef, mut attachments: Vec, + scope: Option<&RuntimeWorkspaceScope>, ) -> Result { let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; let _operation_guard = operation_lock .lock() .map_err(|_| RuntimeError::StatePoisoned)?; + if let Some(scope) = scope { + self.ensure_worker_in_workspace(scope, worker_ref)?; + } validate_logical_workdir_attachments(&attachments)?; attachments.sort_by(|left, right| left.alias.cmp(&right.alias)); @@ -1626,10 +1633,7 @@ impl Runtime { self.ensure_worker_in_workspace(scope, worker_ref)?; let backend = { let state = self.lock()?; - let worker = state.worker(worker_ref)?; - if worker.execution_handle.is_some() && worker.status.is_active() { - return Ok(runtime_api::WorkerSessionAvailability::LiveProtocol); - } + state.worker(worker_ref)?; state.execution_backend.clone() }; Ok(match backend { @@ -1782,33 +1786,6 @@ impl Runtime { .map(|pending| pending.operation_id) .unwrap_or(fresh_operation_id); let operation_mode = pending_restore.map(|pending| pending.mode).unwrap_or(mode); - if !reconcile && worker.execution_handle.is_some() { - if worker.status.is_active() { - return Ok(RuntimeWorkerRestoreResult::accepted(worker.detail())); - } - return Ok(RuntimeWorkerRestoreResult::failed( - WorkerRestoreState::ReconciliationRequired, - "worker_restore_cleanup_pending", - "Worker restore cleanup or reconciliation is still pending", - )); - } - // An active persisted Worker can still have no live handle when its - // startup restore was rejected before side effects (for example, - // before Workspace Backend renewed operation-scoped Repository - // access). Preserve its automatic intent so an explicit restore can - // serve as the authorized retry for that state. - if !reconcile - && mode == WorkerRestoreMode::Explicit - && worker.execution_metadata_available - && worker.status != WorkerStatus::Stopped - && worker.restore_intent != WorkerRestoreIntent::Automatic - { - return Ok(RuntimeWorkerRestoreResult::failed( - WorkerRestoreState::Rejected, - "worker_restore_not_stopped", - "Worker restore requires a stopped Worker", - )); - } if !reconcile && mode == WorkerRestoreMode::Automatic && (!worker.status.is_active() @@ -1861,11 +1838,39 @@ impl Runtime { ) }; - if !reconcile && let Err(_result) = backend.preflight_restore(&request) { + if !reconcile && let Err(result) = backend.preflight_restore(&request) { + // A preflight may attest an already-connected Controller snapshot. + // Repeated restore is read-only; do not replace its operation owner. + if result.outcome == crate::execution::WorkerExecutionOutcome::Busy + && let Some(snapshot) = result.worker_state.as_ref() + { + let unfinished_context = { + self.lock()? + .worker(worker_ref)? + .restore_candidate_context + .clone() + }; + if let Some(context) = unfinished_context { + // The prior aggregate commit succeeded, but observation + // activation can still be unfinished after a transient error. + if let Err(error) = + self.activate_restore_candidate_context(worker_ref, &context) + { + return Ok(RuntimeWorkerRestoreResult::failed( + WorkerRestoreState::ReconciliationRequired, + "worker_restore_observation_activation_pending", + sanitize_worker_create_failure_message(&error.to_string()), + )); + } + } + let mut detail = self.lock()?.worker(worker_ref)?.detail(); + detail.worker_state = Some(snapshot.clone()); + return Ok(RuntimeWorkerRestoreResult::accepted(detail)); + } return Ok(RuntimeWorkerRestoreResult::failed( WorkerRestoreState::Rejected, "worker_restore_preflight_rejected", - "Worker restore preflight rejected the current request", + sanitize_worker_create_failure_message(&result.message_or_default()), )); } @@ -1894,20 +1899,19 @@ impl Runtime { }; match restore_result { WorkerExecutionSpawnResult::Connected { - handle, worker_state, workdir_attachments, } => { let commit = self.commit_restored_worker_execution( worker_ref, - handle.clone(), worker_state.clone(), WorkerStatus::Idle, workdir_attachments.clone(), ); match commit { Ok(worker) => { - if let Err(error) = backend.activate_restored_worker(operation_id, &handle) + if let Err(error) = + backend.activate_restored_worker(operation_id, worker_ref) { tracing::error!( worker_id = %worker_ref.worker_id, @@ -1915,12 +1919,19 @@ impl Runtime { "Worker restore committed but Backend observation activation failed" ); } - if let Err(error) = candidate_context.activate_candidate() { + if let Err(error) = + self.activate_restore_candidate_context(worker_ref, &candidate_context) + { tracing::error!( worker_id = %worker_ref.worker_id, error = %error, "Worker restore committed but buffered protocol observations could not be activated" ); + return Ok(RuntimeWorkerRestoreResult::failed( + WorkerRestoreState::ReconciliationRequired, + "worker_restore_observation_activation_pending", + sanitize_worker_create_failure_message(&error.to_string()), + )); } Ok(RuntimeWorkerRestoreResult::accepted(worker)) } @@ -1932,7 +1943,6 @@ impl Runtime { ); if let Err(retain_error) = self.retain_restore_execution_evidence( worker_ref, - handle, worker_state, workdir_attachments, candidate_context.clone(), @@ -1949,73 +1959,69 @@ impl Runtime { "Worker restore commit outcome is uncertain; reread or retry the Worker restore", )) } - Err(error) => { - match self.cleanup_failed_restore(&backend, worker_ref, &handle) { - Ok(()) => { - candidate_context.discard_candidate(); - tracing::warn!( - worker_id = %worker_ref.worker_id, - error = %error, - "Worker restore commit failed and was rolled back" - ); - if let Err(settle_error) = self.settle_restore_rollback(worker_ref) - { - tracing::error!( - worker_id = %worker_ref.worker_id, - error = %settle_error, - "Worker restore cleanup completed but its durable operation could not be settled" - ); - return Ok(RuntimeWorkerRestoreResult::failed( - WorkerRestoreState::ReconciliationRequired, - "worker_restore_rollback_commit_failed", - "Worker restore cleanup completed but durable reconciliation is still required", - )); - } - Ok(RuntimeWorkerRestoreResult::failed( - WorkerRestoreState::RolledBack, - "worker_restore_commit_rolled_back", - "Worker restore commit failed and cleanup completed", - )) - } - Err(cleanup_error) => { + Err(error) => match self.cleanup_failed_restore(&backend, worker_ref) { + Ok(()) => { + self.discard_restore_candidate_context(worker_ref, &candidate_context)?; + tracing::warn!( + worker_id = %worker_ref.worker_id, + error = %error, + "Worker restore commit failed and was rolled back" + ); + if let Err(settle_error) = self.settle_restore_rollback(worker_ref) { tracing::error!( worker_id = %worker_ref.worker_id, - error = %error, - cleanup_error = %cleanup_error, - "Worker restore requires reconciliation after commit and cleanup failures" + error = %settle_error, + "Worker restore cleanup completed but its durable operation could not be settled" ); - if let Err(retain_error) = self.retain_restore_execution_evidence( - worker_ref, - handle, - worker_state, - workdir_attachments, - candidate_context.clone(), - ) { - tracing::error!( - worker_id = %worker_ref.worker_id, - error = %retain_error, - "failed to persist restore execution evidence; in-process reconciliation remains required" - ); - } - Ok(RuntimeWorkerRestoreResult::failed( + return Ok(RuntimeWorkerRestoreResult::failed( WorkerRestoreState::ReconciliationRequired, - "worker_restore_reconciliation_required", - "Worker restore result is uncertain; reread or retry the Worker restore", - )) + "worker_restore_rollback_commit_failed", + "Worker restore cleanup completed but durable reconciliation is still required", + )); } + Ok(RuntimeWorkerRestoreResult::failed( + WorkerRestoreState::RolledBack, + "worker_restore_commit_rolled_back", + "Worker restore commit failed and cleanup completed", + )) } - } + Err(cleanup_error) => { + tracing::error!( + worker_id = %worker_ref.worker_id, + error = %error, + cleanup_error = %cleanup_error, + "Worker restore requires reconciliation after commit and cleanup failures" + ); + if let Err(retain_error) = self.retain_restore_execution_evidence( + worker_ref, + worker_state, + workdir_attachments, + candidate_context.clone(), + ) { + tracing::error!( + worker_id = %worker_ref.worker_id, + error = %retain_error, + "failed to persist restore execution evidence; in-process reconciliation remains required" + ); + } + Ok(RuntimeWorkerRestoreResult::failed( + WorkerRestoreState::ReconciliationRequired, + "worker_restore_reconciliation_required", + "Worker restore result is uncertain; reread or retry the Worker restore", + )) + } + }, } } - WorkerExecutionSpawnResult::Rejected(_result) => { + WorkerExecutionSpawnResult::Rejected(result) => { if reconcile { return Ok(RuntimeWorkerRestoreResult::failed( WorkerRestoreState::ReconciliationRequired, "worker_restore_reconciliation_rejected", - "Worker restore reconciliation was not completed", + sanitize_worker_create_failure_message(&result.message_or_default()), )); } - candidate_context.discard_candidate(); + self.discard_restore_candidate_context(worker_ref, &candidate_context)?; if let Err(error) = self.settle_restore_rollback(worker_ref) { tracing::error!( worker_id = %worker_ref.worker_id, @@ -2031,13 +2037,13 @@ impl Runtime { Ok(RuntimeWorkerRestoreResult::failed( WorkerRestoreState::Rejected, "worker_restore_rejected", - "Worker restore was rejected before live work started", + sanitize_worker_create_failure_message(&result.message_or_default()), )) } WorkerExecutionSpawnResult::RolledBack(_result) => { - candidate_context.discard_candidate(); + self.discard_restore_candidate_context(worker_ref, &candidate_context)?; #[cfg(feature = "fs-store")] - if let Err(error) = self.lock()?.record_restore_failure(worker_ref, _result) { + if let Err(error) = self.lock()?.record_restore_failure(worker_ref, &_result) { tracing::error!( worker_id = %worker_ref.worker_id, error = %error, @@ -2054,19 +2060,17 @@ impl Runtime { Ok(RuntimeWorkerRestoreResult::failed( WorkerRestoreState::RolledBack, "worker_restore_rolled_back", - "Worker restore failed after live work started; cleanup completed", + sanitize_worker_create_failure_message(&_result.message_or_default()), )) } WorkerExecutionSpawnResult::ReconciliationRequired { result: _result, - handle, worker_state, workdir_attachments, } => { - if let (Some(handle), Some(worker_state)) = (handle, worker_state) { + if let Some(worker_state) = worker_state { if let Err(retain_error) = self.retain_restore_execution_evidence( worker_ref, - handle, worker_state, workdir_attachments, candidate_context.clone(), @@ -2105,13 +2109,39 @@ impl Runtime { Ok(RuntimeWorkerRestoreResult::failed( WorkerRestoreState::ReconciliationRequired, "worker_restore_backend_outcome_unknown", - "Worker restore backend outcome is uncertain; reread or retry the Worker restore", + sanitize_worker_create_failure_message(&_result.message_or_default()), )) } } } - fn ensure_worker_execution(&self, worker_ref: &WorkerRef) -> Result<(), RuntimeError> { + fn activate_restore_candidate_context( + &self, + worker_ref: &WorkerRef, + context: &crate::execution::WorkerExecutionContext, + ) -> Result<(), RuntimeError> { + context.activate_candidate()?; + self.lock()? + .worker_mut(worker_ref)? + .restore_candidate_context = None; + Ok(()) + } + + fn discard_restore_candidate_context( + &self, + worker_ref: &WorkerRef, + context: &crate::execution::WorkerExecutionContext, + ) -> Result<(), RuntimeError> { + // Cleanup/rejection proved no live candidate remains. Detach its + // discarded publication owner even if journal settlement later fails. + self.lock()? + .worker_mut(worker_ref)? + .restore_candidate_context = None; + context.discard_candidate(); + Ok(()) + } + + fn ensure_worker_operation_settled(&self, worker_ref: &WorkerRef) -> Result<(), RuntimeError> { let state = self.lock()?; let worker = state.worker(worker_ref)?; if let Some(operation) = worker.pending_lifecycle_operation_name() { @@ -2120,18 +2150,7 @@ impl Runtime { message: format!("worker {operation} reconciliation is pending"), }); } - if worker.execution_handle.is_some() { - return Ok(()); - } - let message = if worker.status == WorkerStatus::Stopped { - "stopped worker requires an explicit restore" - } else { - "worker has no live execution handle" - }; - Err(RuntimeError::WorkerExecutionUnavailable { - worker_id: worker_ref.worker_id, - message: message.to_string(), - }) + Ok(()) } /// Accept input into a Worker through a workspace-scoped Runtime authorization context. @@ -2141,16 +2160,31 @@ impl Runtime { worker_ref: &WorkerRef, input: WorkerInput, ) -> Result { - self.ensure_worker_in_workspace(scope, worker_ref)?; - self.send_input(worker_ref, input) + self.send_input_with_scope(worker_ref, input, Some(scope)) } /// Accept input into a Worker. pub fn send_input( &self, worker_ref: &WorkerRef, - mut input: WorkerInput, + input: WorkerInput, ) -> Result { + self.send_input_with_scope(worker_ref, input, None) + } + + fn send_input_with_scope( + &self, + worker_ref: &WorkerRef, + mut input: WorkerInput, + scope: Option<&RuntimeWorkspaceScope>, + ) -> Result { + let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; + let _operation_guard = operation_lock + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + if let Some(scope) = scope { + self.ensure_worker_in_workspace(scope, worker_ref)?; + } validate_worker_input(&input)?; if matches!( input.kind, @@ -2190,32 +2224,21 @@ impl Runtime { )); } } - self.ensure_worker_execution(worker_ref)?; - let (backend, handle) = { + self.ensure_worker_operation_settled(worker_ref)?; + let backend = { let state = self.lock()?; state.ensure_running()?; state.ensure_worker_ref(worker_ref)?; - let worker = state.worker(worker_ref)?; - if !worker.status.is_active() { - return Err(RuntimeError::InvalidRequest(format!( - "worker {} is not running", - worker_ref.worker_id - ))); - } - let backend = state.execution_backend.clone(); - let handle = worker.execution_handle.clone(); - match (backend, handle) { - (Some(backend), Some(handle)) => (backend, handle), - _ => { - return Err(RuntimeError::WorkerExecutionUnavailable { - worker_id: worker_ref.worker_id.clone(), - message: "worker has no live execution handle".to_string(), - }); + state.worker(worker_ref)?; + state.execution_backend.clone().ok_or_else(|| { + RuntimeError::WorkerExecutionUnavailable { + worker_id: worker_ref.worker_id, + message: "runtime has no execution backend".to_string(), } - } + })? }; - let dispatch_result = backend.dispatch_input(&handle, input.clone()); + let dispatch_result = backend.dispatch_input(worker_ref, input.clone()); if !dispatch_result.is_accepted() { self.record_execution_result(worker_ref, dispatch_result.clone())?; return Err(RuntimeError::WorkerExecutionRejected { @@ -2300,8 +2323,14 @@ impl Runtime { media_type: &str, content: &[u8], ) -> Result { - self.ensure_worker_in_workspace(scope, worker_ref)?; - self.upload_worker_file(worker_ref, file_name, media_type, content) + self.upload_worker_file_inner( + worker_ref, + file_name, + media_type, + content, + None, + Some(scope), + ) } pub fn upload_worker_file( @@ -2311,7 +2340,7 @@ impl Runtime { media_type: &str, content: &[u8], ) -> Result { - self.upload_worker_file_inner(worker_ref, file_name, media_type, content, None) + self.upload_worker_file_inner(worker_ref, file_name, media_type, content, None, None) } pub fn upload_worker_file_with_context_scoped( @@ -2323,8 +2352,14 @@ impl Runtime { content: &[u8], context: &session_store::UploadedFileUploadContext, ) -> Result { - self.ensure_worker_in_workspace(scope, worker_ref)?; - self.upload_worker_file_inner(worker_ref, file_name, media_type, content, Some(context)) + self.upload_worker_file_inner( + worker_ref, + file_name, + media_type, + content, + Some(context), + Some(scope), + ) } pub fn upload_worker_file_with_context( @@ -2335,7 +2370,14 @@ impl Runtime { content: &[u8], context: &session_store::UploadedFileUploadContext, ) -> Result { - self.upload_worker_file_inner(worker_ref, file_name, media_type, content, Some(context)) + self.upload_worker_file_inner( + worker_ref, + file_name, + media_type, + content, + Some(context), + None, + ) } fn upload_worker_file_inner( @@ -2345,27 +2387,30 @@ impl Runtime { media_type: &str, content: &[u8], context: Option<&session_store::UploadedFileUploadContext>, + scope: Option<&RuntimeWorkspaceScope>, ) -> Result { - let (backend, handle) = { + let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; + let _operation_guard = operation_lock + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + if let Some(scope) = scope { + self.ensure_worker_in_workspace(scope, worker_ref)?; + } + self.ensure_worker_operation_settled(worker_ref)?; + let backend = { let state = self.lock()?; state.ensure_running()?; state.ensure_worker_ref(worker_ref)?; - let worker = state.worker(worker_ref)?; - match ( - state.execution_backend.clone(), - worker.execution_handle.clone(), - ) { - (Some(backend), Some(handle)) => (backend, handle), - _ => { - return Err(RuntimeError::WorkerExecutionUnavailable { - worker_id: worker_ref.worker_id.clone(), - message: "worker has no live execution handle".to_string(), - }); + state.worker(worker_ref)?; + state.execution_backend.clone().ok_or_else(|| { + RuntimeError::WorkerExecutionUnavailable { + worker_id: worker_ref.worker_id, + message: "runtime has no execution backend".to_string(), } - } + })? }; backend - .upload_file(&handle, file_name, media_type, content, context) + .upload_file(worker_ref, file_name, media_type, content, context) .map_err(|result| RuntimeError::WorkerExecutionRejected { worker_id: worker_ref.worker_id.clone(), operation: result.operation, @@ -2382,8 +2427,7 @@ impl Runtime { worker_ref: &WorkerRef, artifact_id: &str, ) -> Result<(), RuntimeError> { - self.ensure_worker_in_workspace(scope, worker_ref)?; - self.delete_worker_uploaded_file(worker_ref, artifact_id) + self.delete_worker_uploaded_file_with_scope(worker_ref, artifact_id, Some(scope)) } pub fn delete_worker_uploaded_file( @@ -2391,25 +2435,36 @@ impl Runtime { worker_ref: &WorkerRef, artifact_id: &str, ) -> Result<(), RuntimeError> { - let (backend, handle) = { + self.delete_worker_uploaded_file_with_scope(worker_ref, artifact_id, None) + } + + fn delete_worker_uploaded_file_with_scope( + &self, + worker_ref: &WorkerRef, + artifact_id: &str, + scope: Option<&RuntimeWorkspaceScope>, + ) -> Result<(), RuntimeError> { + let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; + let _operation_guard = operation_lock + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + if let Some(scope) = scope { + self.ensure_worker_in_workspace(scope, worker_ref)?; + } + self.ensure_worker_operation_settled(worker_ref)?; + let backend = { let state = self.lock()?; state.ensure_running()?; state.ensure_worker_ref(worker_ref)?; - let worker = state.worker(worker_ref)?; - match ( - state.execution_backend.clone(), - worker.execution_handle.clone(), - ) { - (Some(backend), Some(handle)) => (backend, handle), - _ => { - return Err(RuntimeError::WorkerExecutionUnavailable { - worker_id: worker_ref.worker_id.clone(), - message: "worker has no live execution handle".to_string(), - }); + state.worker(worker_ref)?; + state.execution_backend.clone().ok_or_else(|| { + RuntimeError::WorkerExecutionUnavailable { + worker_id: worker_ref.worker_id, + message: "runtime has no execution backend".to_string(), } - } + })? }; - let result = backend.delete_uploaded_file(&handle, artifact_id); + let result = backend.delete_uploaded_file(worker_ref, artifact_id); if result.is_accepted() { Ok(()) } else { @@ -2432,8 +2487,7 @@ impl Runtime { prefix: &str, context: Option<&protocol::CompletionContext>, ) -> Result, RuntimeError> { - self.ensure_worker_in_workspace(scope, worker_ref)?; - self.worker_completions(worker_ref, kind, prefix, context) + self.worker_completions_with_scope(worker_ref, kind, prefix, context, Some(scope)) } /// Return live completion entries for the Worker composer. @@ -2444,22 +2498,37 @@ impl Runtime { prefix: &str, context: Option<&protocol::CompletionContext>, ) -> Result, RuntimeError> { - let (backend, handle) = { + self.worker_completions_with_scope(worker_ref, kind, prefix, context, None) + } + + fn worker_completions_with_scope( + &self, + worker_ref: &WorkerRef, + kind: protocol::CompletionKind, + prefix: &str, + context: Option<&protocol::CompletionContext>, + scope: Option<&RuntimeWorkspaceScope>, + ) -> Result, RuntimeError> { + let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; + let _operation_guard = operation_lock + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + if let Some(scope) = scope { + self.ensure_worker_in_workspace(scope, worker_ref)?; + } + let backend = { let state = self.lock()?; state.ensure_worker_ref(worker_ref)?; let worker = state.worker(worker_ref)?; if worker.has_pending_lifecycle_operation() { return Ok(Vec::new()); } - ( - state.execution_backend.clone(), - worker.execution_handle.clone(), - ) + state.execution_backend.clone() }; - let Some((backend, handle)) = backend.zip(handle) else { + let Some(backend) = backend else { return Ok(Vec::new()); }; - Ok(backend.worker_completions(&handle, kind, prefix, context)) + Ok(backend.worker_completions(worker_ref, kind, prefix, context)) } /// Accept a protocol method through a workspace-scoped Runtime authorization context. @@ -2469,8 +2538,7 @@ impl Runtime { worker_ref: &WorkerRef, method: Method, ) -> Result, RuntimeError> { - self.ensure_worker_in_workspace(scope, worker_ref)?; - self.send_protocol_method(worker_ref, method) + self.send_protocol_method_with_scope(worker_ref, method, Some(scope)) } /// Accept a protocol method for a Worker through a Backend/runtime transport. @@ -2483,6 +2551,15 @@ impl Runtime { &self, worker_ref: &WorkerRef, method: Method, + ) -> Result, RuntimeError> { + self.send_protocol_method_with_scope(worker_ref, method, None) + } + + fn send_protocol_method_with_scope( + &self, + worker_ref: &WorkerRef, + method: Method, + scope: Option<&RuntimeWorkspaceScope>, ) -> Result, RuntimeError> { if let Method::ListCompletions { kind, @@ -2491,7 +2568,13 @@ impl Runtime { context, } = method { - let entries = self.worker_completions(worker_ref, kind, &prefix, context.as_ref())?; + let entries = self.worker_completions_with_scope( + worker_ref, + kind, + &prefix, + context.as_ref(), + scope, + )?; return Ok(vec![Event::Completions { kind, prefix, @@ -2501,36 +2584,36 @@ impl Runtime { }]); } if matches!(&method, Method::Shutdown { .. }) { - self.stop_worker(worker_ref, Some("worker protocol shutdown".to_string()))?; + self.stop_worker_with_scope( + worker_ref, + Some("worker protocol shutdown".to_string()), + scope, + )?; return Ok(Vec::new()); } - self.ensure_worker_execution(worker_ref)?; + let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; + let _operation_guard = operation_lock + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + if let Some(scope) = scope { + self.ensure_worker_in_workspace(scope, worker_ref)?; + } + self.ensure_worker_operation_settled(worker_ref)?; - let (backend, handle) = { + let backend = { let state = self.lock()?; state.ensure_running()?; state.ensure_worker_ref(worker_ref)?; - let worker = state.worker(worker_ref)?; - if !worker.status.is_active() { - return Err(RuntimeError::InvalidRequest(format!( - "worker {} is not running", - worker_ref.worker_id - ))); - } - let backend = state.execution_backend.clone(); - let handle = worker.execution_handle.clone(); - match (backend, handle) { - (Some(backend), Some(handle)) => (backend, handle), - _ => { - return Err(RuntimeError::WorkerExecutionUnavailable { - worker_id: worker_ref.worker_id.clone(), - message: "worker has no live execution handle".to_string(), - }); + state.worker(worker_ref)?; + state.execution_backend.clone().ok_or_else(|| { + RuntimeError::WorkerExecutionUnavailable { + worker_id: worker_ref.worker_id, + message: "runtime has no execution backend".to_string(), } - } + })? }; - let dispatch_result = backend.dispatch_method(&handle, method); + let dispatch_result = backend.dispatch_method(worker_ref, method); if !dispatch_result.is_accepted() { self.record_execution_result(worker_ref, dispatch_result.clone())?; return Err(RuntimeError::WorkerExecutionRejected { @@ -2546,10 +2629,138 @@ impl Runtime { Ok(Vec::new()) } + /// Attach one protocol transport under the same admission/authorization lock + /// as lifecycle operations. The backend owns the real Controller endpoint; + /// this session is never a Worker catalog field or a routing token. + #[cfg(feature = "ws-server")] + pub fn attach_worker_protocol( + &self, + worker_ref: &WorkerRef, + ) -> Result { + self.attach_worker_protocol_with_scope(worker_ref, None) + } + + #[cfg(feature = "ws-server")] + fn attach_worker_protocol_with_scope( + &self, + worker_ref: &WorkerRef, + scope: Option<&RuntimeWorkspaceScope>, + ) -> Result { + let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; + let _operation_guard = operation_lock + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + if let Some(scope) = scope { + self.ensure_worker_in_workspace(scope, worker_ref)?; + } + self.ensure_worker_operation_settled(worker_ref)?; + let backend = { + let state = self.lock()?; + state.ensure_running()?; + state.worker(worker_ref)?; + state.execution_backend.clone().ok_or_else(|| { + RuntimeError::WorkerExecutionUnavailable { + worker_id: worker_ref.worker_id, + message: "runtime has no execution backend".to_string(), + } + })? + }; + let transport = backend + .attach_worker_protocol(worker_ref) + .map_err(|result| RuntimeError::WorkerExecutionRejected { + worker_id: worker_ref.worker_id, + operation: result.operation, + outcome: result.outcome, + message: result.message_or_default(), + result, + })?; + if &transport.worker_ref != worker_ref { + return Err(RuntimeError::WorkerExecutionUnavailable { + worker_id: worker_ref.worker_id, + message: "protocol endpoint target does not match the authorized Worker".into(), + }); + } + Ok(transport) + } + + #[cfg(feature = "ws-server")] + pub fn attach_worker_protocol_scoped( + &self, + scope: &RuntimeWorkspaceScope, + worker_ref: &WorkerRef, + ) -> Result { + self.attach_worker_protocol_with_scope(worker_ref, Some(scope)) + } + + #[cfg(feature = "ws-server")] + pub fn send_connected_protocol_method_scoped( + &self, + scope: &RuntimeWorkspaceScope, + worker_ref: &WorkerRef, + transport: &crate::execution::WorkerProtocolTransport, + method: Method, + ) -> Result, RuntimeError> { + self.send_connected_protocol_method_with_scope(worker_ref, transport, method, Some(scope)) + } + + #[cfg(feature = "ws-server")] + pub fn send_connected_protocol_method( + &self, + worker_ref: &WorkerRef, + transport: &crate::execution::WorkerProtocolTransport, + method: Method, + ) -> Result, RuntimeError> { + self.send_connected_protocol_method_with_scope(worker_ref, transport, method, None) + } + + #[cfg(feature = "ws-server")] + fn send_connected_protocol_method_with_scope( + &self, + worker_ref: &WorkerRef, + transport: &crate::execution::WorkerProtocolTransport, + method: Method, + scope: Option<&RuntimeWorkspaceScope>, + ) -> Result, RuntimeError> { + // Authorization must address the endpoint's actual target, not a + // different WorkerRef supplied beside a transport captured elsewhere. + if &transport.worker_ref != worker_ref { + return Err(RuntimeError::WorkerNotFound { + worker_id: worker_ref.worker_id, + }); + } + let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; + let _operation_guard = operation_lock + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + if let Some(scope) = scope { + self.ensure_worker_in_workspace(scope, worker_ref)?; + } + self.lock()?.ensure_running()?; + self.ensure_worker_operation_settled(worker_ref)?; + let rejection = |result: WorkerExecutionResult| RuntimeError::WorkerExecutionRejected { + worker_id: worker_ref.worker_id, + operation: result.operation, + outcome: result.outcome, + message: result.message_or_default(), + result, + }; + // Validate the captured endpoint before ANY durable mutation. Even if + // this socket processes terminal closure late, it cannot stop or send + // to the next execution under the same WorkerRef. + transport.validate().map_err(rejection)?; + if matches!(&method, Method::Shutdown { .. }) { + let operation_id = WorkerLifecycleOperationId::new(); + self.begin_stop_operation(worker_ref, operation_id)?; + let events = transport.dispatch(method).map_err(rejection)?; + self.commit_stopped_worker(worker_ref, operation_id)?; + return Ok(events); + } + transport.dispatch(method).map_err(rejection) + } + fn commit_created_worker( &self, worker_ref: &WorkerRef, - handle: WorkerExecutionHandle, initial_worker_state: protocol::WorkerStateSnapshot, workdir_attachments: Vec, result: WorkerExecutionResult, @@ -2557,7 +2768,7 @@ impl Runtime { let mut state = self.lock()?; let detail = { let worker = state.worker_mut(worker_ref)?; - worker.execution_handle = Some(handle); + worker.execution_metadata_available = true; worker.execution_bound = true; worker.status = WorkerStatus::Idle; @@ -2586,15 +2797,14 @@ impl Runtime { &self, backend: &WorkerExecutionBackendRef, worker_ref: &WorkerRef, - handle: &WorkerExecutionHandle, ) -> Result<(), RuntimeError> { - let stop_result = backend.stop_worker(handle); + let stop_result = backend.stop_worker(worker_ref); if stop_result.is_accepted() { return self.rollback_failed_create(worker_ref); } let mut state = self.lock()?; let record = state.worker_mut(worker_ref)?; - record.execution_handle = Some(handle.clone()); + record.worker_state = stop_result.worker_state.clone(); state.persist_runtime_snapshot()?; state.persist_worker(&worker_ref.worker_id)?; @@ -2629,7 +2839,7 @@ impl Runtime { result: WorkerExecutionResult, ) -> Result<(), RuntimeError> { // Accepted dispatch without a state snapshot is transport evidence only; - // the revisioned protocol stream remains live authority. Test/detached + // the Controller protocol stream remains live authority. Test/detached // backends may return an exact full snapshot as their acknowledgement. if !result.is_accepted() { return Ok(()); @@ -2651,27 +2861,24 @@ impl Runtime { worker_ref: &WorkerRef, operation: WorkerExecutionOperation, ) -> Result<(), RuntimeError> { - let Some((backend, handle)) = ({ + let backend = { let state = self.lock()?; state.ensure_worker_ref(worker_ref)?; let worker = state.worker(worker_ref)?; if !worker.status.is_active() { return Ok(()); } - match ( - state.execution_backend.clone(), - worker.execution_handle.clone(), - ) { - (Some(backend), Some(handle)) => Some((backend, handle)), - _ => None, - } - }) else { - return Ok(()); + state.execution_backend.clone().ok_or_else(|| { + RuntimeError::WorkerExecutionUnavailable { + worker_id: worker_ref.worker_id, + message: "runtime has no execution backend".to_string(), + } + })? }; let result = match operation { - WorkerExecutionOperation::Stop => backend.stop_worker(&handle), - WorkerExecutionOperation::Cancel => backend.cancel_worker(&handle), + WorkerExecutionOperation::Stop => backend.stop_worker(worker_ref), + WorkerExecutionOperation::Cancel => backend.cancel_worker(worker_ref), WorkerExecutionOperation::Spawn | WorkerExecutionOperation::Restore | WorkerExecutionOperation::Input @@ -2726,24 +2933,18 @@ impl Runtime { self.ensure_worker_in_workspace(scope, worker_ref)?; } - let (backend, handle, pending_stop, status) = { + let (backend, pending_stop, pending_restore, status) = { let state = self.lock()?; state.ensure_running()?; let worker = state.worker(worker_ref)?; - if worker.pending_restore.is_some() { - return Err(RuntimeError::InvalidRequest(format!( - "worker {} has pending restore reconciliation", - worker_ref.worker_id - ))); - } ( state.execution_backend.clone(), - worker.execution_handle.clone(), worker.pending_stop, + worker.pending_restore, worker.status, ) }; - if status == WorkerStatus::Stopped && pending_stop.is_none() { + if status == WorkerStatus::Stopped && pending_stop.is_none() && pending_restore.is_none() { let state = self.lock()?; let worker = state.worker(worker_ref)?; return Ok(WorkerLifecycleAck { @@ -2756,12 +2957,6 @@ impl Runtime { worker_id: worker_ref.worker_id, message: "runtime has no execution backend".to_string(), })?; - if handle.is_none() && pending_stop.is_none() { - return Err(RuntimeError::WorkerExecutionUnavailable { - worker_id: worker_ref.worker_id, - message: "worker has no execution handle".to_string(), - }); - } let operation_id = match pending_stop { Some(pending) => pending.operation_id, None => { @@ -2773,7 +2968,6 @@ impl Runtime { let result = backend.stop_worker_operation(WorkerExecutionStopRequest { operation_id, worker_ref: worker_ref.clone(), - handle, }); if !result.is_accepted() { return Err(RuntimeError::WorkerExecutionRejected { @@ -2796,8 +2990,7 @@ impl Runtime { worker_ref: &WorkerRef, reason: Option, ) -> Result { - self.ensure_worker_in_workspace(scope, worker_ref)?; - self.cancel_worker(worker_ref, reason) + self.cancel_worker_with_scope(worker_ref, reason, Some(scope)) } /// Cancel the current run while keeping the Worker session available. @@ -2806,6 +2999,22 @@ impl Runtime { worker_ref: &WorkerRef, reason: Option, ) -> Result { + self.cancel_worker_with_scope(worker_ref, reason, None) + } + + fn cancel_worker_with_scope( + &self, + worker_ref: &WorkerRef, + reason: Option, + scope: Option<&RuntimeWorkspaceScope>, + ) -> Result { + let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; + let _operation_guard = operation_lock + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + if let Some(scope) = scope { + self.ensure_worker_in_workspace(scope, worker_ref)?; + } { let state = self.lock()?; state.ensure_running()?; @@ -2864,7 +3073,7 @@ impl Runtime { if let Some(scope) = scope { self.ensure_worker_in_workspace(scope, worker_ref)?; } - let (backend, execution_handle) = { + let backend = { let state = self.lock()?; state.ensure_running()?; state.ensure_worker_ref(worker_ref)?; @@ -2875,27 +3084,16 @@ impl Runtime { worker_ref.worker_id ))); } - if worker.status.is_active() - && (worker.execution_handle.is_some() - || (!worker.execution_metadata_available - && !worker.workdir_attachments.is_empty())) - { + if worker.status.is_active() { return Err(RuntimeError::InvalidRequest(format!( "worker {} is active or requires reconciliation and must be stopped before deletion", worker_ref.worker_id ))); } - ( - state.execution_backend.clone(), - worker.execution_handle.clone(), - ) + state.execution_backend.clone() }; - if let Some(handle) = execution_handle { - let backend = backend.ok_or_else(|| RuntimeError::ExecutionBackendUnavailable { - message: "Worker deletion requires its execution backend to confirm shutdown" - .to_string(), - })?; - let result = backend.stop_worker(&handle); + if let Some(backend) = backend { + let result = backend.stop_worker(worker_ref); if !result.is_accepted() { return Err(RuntimeError::WorkerExecutionRejected { worker_id: worker_ref.worker_id, @@ -2916,10 +3114,7 @@ impl Runtime { worker_ref.worker_id ))); } - if worker.status.is_active() - && (worker.execution_handle.is_some() - || (!worker.execution_metadata_available && !worker.workdir_attachments.is_empty())) - { + if worker.status.is_active() { return Err(RuntimeError::InvalidRequest(format!( "worker {} became active or reconciliation-required before deletion", worker_ref.worker_id @@ -2978,16 +3173,17 @@ impl Runtime { &self, worker_ref: &WorkerRef, ) -> Result { - let (backend, handle) = { + let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; + let _operation_guard = operation_lock + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + let backend = { let state = self.lock()?; - let worker = state.worker(worker_ref)?; - ( - state.execution_backend.clone(), - worker.execution_handle.clone(), - ) + state.worker(worker_ref)?; + state.execution_backend.clone() }; - if let (Some(backend), Some(handle)) = (backend, handle) { - if let Some(snapshot) = backend.worker_snapshot(&handle) { + if let Some(backend) = backend { + if let Some(snapshot) = backend.worker_snapshot(worker_ref) { return Ok(snapshot); } } @@ -3152,6 +3348,25 @@ impl Runtime { #[cfg(feature = "fs-store")] fn restore_persisted_worker_executions(&self) -> Result<(), RuntimeError> { + let candidates = { + let state = self.lock()?; + if state.execution_backend.is_none() { + return Ok(()); + } + state + .workers + .values() + .filter(|worker| { + worker.pending_restore.is_none() + && worker.pending_stop.is_none() + && worker.execution_bound + && worker.status.is_active() + && worker.restore_intent == WorkerRestoreIntent::Automatic + }) + .map(|worker| worker.worker_ref.clone()) + .collect::>() + }; + let pending_stops = { let state = self.lock()?; if state.execution_backend.is_none() { @@ -3183,6 +3398,7 @@ impl Runtime { state .workers .values() + .filter(|worker| worker.pending_stop.is_none()) .filter_map(|worker| { worker .pending_restore @@ -3195,57 +3411,59 @@ impl Runtime { let _operation_guard = operation_lock .lock() .map_err(|_| RuntimeError::StatePoisoned)?; - if let Err(error) = self.restore_worker_under_lock(&worker_ref, mode) { - tracing::warn!( - worker_id = %worker_ref.worker_id, - error = %error, - "Worker restore reconciliation remains pending after Runtime startup" - ); - } + let result = self.restore_worker_under_lock(&worker_ref, mode)?; + self.record_startup_restore_result(&worker_ref, &result)?; } - let candidates = { - let state = self.lock()?; - if state.execution_backend.is_none() { - return Ok(()); - } - state - .workers - .values() - .filter(|worker| { - worker.execution_handle.is_none() - && worker.pending_restore.is_none() - && worker.pending_stop.is_none() - && worker.execution_bound - && worker.status.is_active() - && worker.restore_intent == WorkerRestoreIntent::Automatic - }) - .map(|worker| worker.worker_ref.clone()) - .collect::>() - }; - for worker_ref in candidates { let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; let _operation_guard = operation_lock .lock() .map_err(|_| RuntimeError::StatePoisoned)?; - match self.restore_worker_under_lock(&worker_ref, WorkerRestoreMode::Automatic) { - Ok(_) => {} - Err(RuntimeError::WorkerExecutionRejected { .. }) => { - // The failed restore and fail-closed terminal status were - // persisted by restore_worker_under_lock. - } - Err(error) => return Err(error), - } + let result = + self.restore_worker_under_lock(&worker_ref, WorkerRestoreMode::Automatic)?; + self.record_startup_restore_result(&worker_ref, &result)?; } Ok(()) } + #[cfg(feature = "fs-store")] + fn record_startup_restore_result( + &self, + worker_ref: &WorkerRef, + result: &RuntimeWorkerRestoreResult, + ) -> Result<(), RuntimeError> { + match result.state { + WorkerRestoreState::Accepted => return Ok(()), + WorkerRestoreState::Rejected + | WorkerRestoreState::RolledBack + | WorkerRestoreState::ReconciliationRequired => {} + } + let mut state = self.lock()?; + let diagnostic_id = state.next_diagnostic_id; + state.next_diagnostic_id = state.next_diagnostic_id.saturating_add(1); + let message = sanitize_worker_create_failure_message( + result + .message + .as_deref() + .unwrap_or("Worker restore did not commit"), + ); + tracing::warn!(worker_id = %worker_ref.worker_id, restore_state = ?result.state, reason = ?result.reason_code, %message, "Worker startup restore did not commit"); + state.diagnostics.push(RuntimeDiagnostic { + id: diagnostic_id, + severity: DiagnosticSeverity::Warning, + code: "worker_execution_restore_failed".to_string(), + message, + worker_ref: Some(worker_ref.clone()), + }); + state.persist_runtime_snapshot() + } + fn reconcile_pending_stop_under_lock( &self, worker_ref: &WorkerRef, ) -> Result<(), RuntimeError> { - let (backend, handle, operation_id) = { + let (backend, operation_id) = { let state = self.lock()?; let worker = state.worker(worker_ref)?; let pending = worker.pending_stop.ok_or_else(|| { @@ -3260,16 +3478,11 @@ impl Runtime { message: "runtime has no execution backend".to_string(), } })?; - ( - backend, - worker.execution_handle.clone(), - pending.operation_id, - ) + (backend, pending.operation_id) }; let result = backend.stop_worker_operation(WorkerExecutionStopRequest { operation_id, worker_ref: worker_ref.clone(), - handle, }); if !result.is_accepted() { return Err(RuntimeError::WorkerExecutionRejected { @@ -3291,13 +3504,13 @@ impl Runtime { ) -> Result<(), RuntimeError> { let mut state = self.lock()?; let mut candidate = state.worker(worker_ref)?.clone(); - if candidate.has_pending_lifecycle_operation() { + if candidate.pending_stop.is_some() { return Err(RuntimeError::InvalidRequest(format!( - "worker {} already has pending lifecycle reconciliation", + "worker {} already has pending stop reconciliation", worker_ref.worker_id ))); } - if !candidate.status.is_active() || candidate.request.is_none() { + if !candidate.status.is_active() && candidate.pending_restore.is_none() { return Err(RuntimeError::InvalidRequest(format!( "worker {} cannot begin a stop operation from its current state", worker_ref.worker_id @@ -3337,10 +3550,14 @@ impl Runtime { ))); } candidate.pending_stop = None; + candidate.pending_restore = None; + if let Some(context) = candidate.restore_candidate_context.take() { + context.discard_candidate(); + } candidate.status = WorkerStatus::Stopped; candidate.worker_state = None; candidate.restore_intent = WorkerRestoreIntent::Explicit; - candidate.execution_handle = None; + candidate.internal_workers.clear(); state.persist_worker_record(&candidate)?; let ack = WorkerLifecycleAck { @@ -3414,7 +3631,7 @@ impl Runtime { fn commit_restored_worker_execution( &self, worker_ref: &WorkerRef, - handle: WorkerExecutionHandle, + worker_state: protocol::WorkerStateSnapshot, status: WorkerStatus, workdir_attachments: Vec, @@ -3422,18 +3639,11 @@ impl Runtime { let mut state = self.lock()?; state.ensure_worker_ref(worker_ref)?; let mut candidate = state.worker(worker_ref)?.clone(); - if candidate.execution_handle.is_some() { - return Err(RuntimeError::InvalidRequest(format!( - "worker {} already has a current execution", - worker_ref.worker_id - ))); - } - candidate.execution_handle = Some(handle); candidate.execution_metadata_available = true; candidate.execution_bound = true; candidate.pending_restore = None; candidate.pending_stop = None; - candidate.restore_candidate_context = None; + // Keep the actual publication owner until buffered observations activate. candidate.status = status; let mut worker_state = worker_state; if worker_state.last_finished_submission_request_id.is_none() { @@ -3474,13 +3684,13 @@ impl Runtime { fn retain_create_execution_evidence( &self, worker_ref: &WorkerRef, - handle: WorkerExecutionHandle, + worker_state: protocol::WorkerStateSnapshot, workdir_attachments: Vec, ) -> Result<(), RuntimeError> { let mut state = self.lock()?; let mut candidate = state.worker(worker_ref)?.clone(); - candidate.execution_handle = Some(handle); + candidate.execution_metadata_available = true; candidate.execution_bound = true; candidate.status = WorkerStatus::Idle; @@ -3496,7 +3706,7 @@ impl Runtime { fn retain_restore_execution_evidence( &self, worker_ref: &WorkerRef, - handle: WorkerExecutionHandle, + _worker_state: protocol::WorkerStateSnapshot, workdir_attachments: Vec, candidate_context: crate::execution::WorkerExecutionContext, @@ -3509,7 +3719,7 @@ impl Runtime { worker_ref.worker_id ))); } - candidate.execution_handle = Some(handle); + candidate.execution_metadata_available = true; candidate.execution_bound = true; candidate.restore_candidate_context = Some(candidate_context); @@ -3526,7 +3736,7 @@ impl Runtime { ) -> Result<(), RuntimeError> { let mut state = self.lock()?; let mut candidate = state.worker(worker_ref)?.clone(); - candidate.execution_handle = None; + candidate.execution_metadata_available = false; candidate.execution_bound = true; candidate.status = WorkerStatus::Idle; @@ -3567,9 +3777,8 @@ impl Runtime { &self, backend: &WorkerExecutionBackendRef, worker_ref: &WorkerRef, - handle: &WorkerExecutionHandle, ) -> Result<(), RuntimeError> { - let result = backend.stop_worker(handle); + let result = backend.stop_worker(worker_ref); if !result.is_accepted() { return Err(RuntimeError::WorkerExecutionRejected { worker_id: worker_ref.worker_id, @@ -4019,11 +4228,17 @@ impl RuntimeState { ) } PersistedWorkerLifecycleOperation::Stop(operation) => ( - Some(operation.request), - true, + operation.request.clone(), + operation.request.is_some(), operation.binding.is_some(), operation.restore_intent, - None, + operation + .pending_restore + .map(|restore| PendingWorkerRestore { + operation_id: restore.operation_id, + mode: restore.mode.into(), + last_settled_status: restore.last_settled_status, + }), Some(PendingWorkerStop { operation_id: operation.operation_id, last_settled_status: operation.last_settled_status, @@ -4072,7 +4287,6 @@ impl RuntimeState { restore_candidate_context: None, workdir_attachments: worker.workdir_attachments, logical_workdir_attachments, - execution_handle: None, internal_workers: InternalWorkerActivityProjection::default(), }, ); @@ -4552,12 +4766,9 @@ impl RuntimeState { fn record_restore_failure( &mut self, worker_ref: &WorkerRef, - result: WorkerExecutionResult, + result: &WorkerExecutionResult, ) -> Result<(), RuntimeError> { - let message = result - .message - .clone() - .unwrap_or_else(|| "worker execution restore failed".to_string()); + let message = sanitize_worker_create_failure_message(&result.message_or_default()); let diagnostic_id = self.next_diagnostic_id; let workspace_id = self .workers @@ -4565,7 +4776,7 @@ impl RuntimeState { .and_then(|worker| worker.workspace_id.clone()) .unwrap_or_else(|| "".to_string()); let mut candidate = self.worker(worker_ref)?.clone(); - candidate.execution_handle = None; + candidate.execution_bound = false; candidate.pending_restore = None; candidate.pending_stop = None; @@ -4971,7 +5182,6 @@ struct WorkerRecord { restore_candidate_context: Option, workdir_attachments: Vec, logical_workdir_attachments: Vec, - execution_handle: Option, internal_workers: InternalWorkerActivityProjection, } @@ -5038,16 +5248,26 @@ impl WorkerRecord { #[cfg(feature = "fs-store")] fn persisted_record(&self) -> PersistedWorkerRecord { - debug_assert!(self.pending_restore.is_none() || self.pending_stop.is_none()); - let execution_state = if let (Some(pending), Some(request)) = - (self.pending_stop, self.request.clone()) - { + let execution_state = if let Some(pending) = self.pending_stop { PersistedWorkerExecutionState::ReconciliationRequired( PersistedWorkerLifecycleOperation::Stop(PersistedWorkerStopOperation { operation_id: pending.operation_id, - request, - binding: self - .execution_bound + pending_restore: self.pending_restore.map(|restore| { + PersistedWorkerRestoreOperation { + operation_id: restore.operation_id, + mode: restore.mode.into(), + request: self + .request + .clone() + .expect("pending restore retains its operation specification"), + binding: self + .execution_bound + .then_some(PersistedWorkerExecutionBinding {}), + last_settled_status: restore.last_settled_status, + } + }), + request: self.request.clone(), + binding: (self.execution_bound && self.request.is_some()) .then_some(PersistedWorkerExecutionBinding {}), restore_intent: self.restore_intent, last_settled_status: pending.last_settled_status, @@ -5595,8 +5815,7 @@ mod tests { ConfigDeclarationKind, ConfigProfileDescriptor, }; use crate::execution::{ - WorkerExecutionBackend, WorkerExecutionContext, WorkerExecutionHandle, - WorkerExecutionRestoreRequest, + WorkerExecutionBackend, WorkerExecutionContext, WorkerExecutionRestoreRequest, }; use crate::working_directory::WorkingDirectoryDiagnostic; use async_trait::async_trait; @@ -6647,7 +6866,6 @@ mod tests { .unwrap() .insert(request.worker_ref.worker_id.clone(), request.context); WorkerExecutionSpawnResult::Connected { - handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), worker_state: protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, @@ -6664,10 +6882,41 @@ mod tests { } } + fn worker_session( + &self, + request: WorkerSessionObservationRequest, + ) -> runtime_api::WorkerSessionAvailability { + if self + .contexts + .lock() + .unwrap() + .contains_key(&request.worker_ref.worker_id) + { + runtime_api::WorkerSessionAvailability::LiveProtocol + } else { + runtime_api::WorkerSessionAvailability::Unavailable { + reason: runtime_api::WorkerSessionUnavailableReason::StorageUnavailable, + message: "test backend has no retained session".to_string(), + } + } + } + fn preflight_restore( &self, - _request: &WorkerExecutionRestoreRequest, + request: &WorkerExecutionRestoreRequest, ) -> Result<(), WorkerExecutionResult> { + if self + .contexts + .lock() + .unwrap() + .contains_key(&request.worker_ref.worker_id) + { + return Err(WorkerExecutionResult::busy( + WorkerExecutionOperation::Restore, + "Worker is already connected", + ) + .with_worker_state(protocol::WorkerStatus::Idle.into())); + } match self.preflight_restore_result.lock().unwrap().clone() { Some(result) => Err(result), None => Ok(()), @@ -6703,7 +6952,6 @@ mod tests { .unwrap() .insert(request.worker_ref.worker_id.clone(), request.context); WorkerExecutionSpawnResult::Connected { - handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), worker_state: protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, @@ -6722,7 +6970,7 @@ mod tests { fn worker_completions( &self, - _handle: &WorkerExecutionHandle, + _worker_ref: &WorkerRef, kind: protocol::CompletionKind, prefix: &str, context: Option<&protocol::CompletionContext>, @@ -6739,9 +6987,20 @@ mod tests { fn dispatch_input( &self, - _handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, input: WorkerInput, ) -> WorkerExecutionResult { + if !self + .contexts + .lock() + .unwrap() + .contains_key(&worker_ref.worker_id) + { + return WorkerExecutionResult::rejected( + WorkerExecutionOperation::Input, + "WorkerRef does not reference a live Worker execution", + ); + } let request_id = input.submission_request_id.clone(); let input_kind = input.kind.clone(); self.dispatched_inputs.lock().unwrap().push(input); @@ -6785,35 +7044,35 @@ mod tests { .lock() .unwrap() .push(request.operation_id); - let handle = request.handle.unwrap_or_else(|| { - WorkerExecutionHandle::new(request.worker_ref, self.backend_id()) - }); - self.stop_worker(&handle) + self.stop_worker(&request.worker_ref) } - fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { + fn stop_worker(&self, worker_ref: &WorkerRef) -> WorkerExecutionResult { *self.stop_count.lock().unwrap() += 1; let stop_gate = self.stop_gate.lock().unwrap().clone(); if let Some(gate) = stop_gate { gate.enter_and_wait(); } - self.stop_result - .lock() - .unwrap() - .take() - .unwrap_or_else(|| WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop)) + let result = + self.stop_result.lock().unwrap().take().unwrap_or_else(|| { + WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop) + }); + if result.is_accepted() { + self.contexts.lock().unwrap().remove(&worker_ref.worker_id); + } + result } - fn cancel_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { + fn cancel_worker(&self, _worker_ref: &WorkerRef) -> WorkerExecutionResult { WorkerExecutionResult::accepted(WorkerExecutionOperation::Cancel) } #[cfg(feature = "ws-server")] - fn worker_snapshot(&self, handle: &WorkerExecutionHandle) -> Option { + fn worker_snapshot(&self, worker_ref: &WorkerRef) -> Option { self.snapshots .lock() .unwrap() - .get(&handle.worker_ref().worker_id) + .get(&worker_ref.worker_id) .cloned() } } @@ -6977,6 +7236,7 @@ mod tests { assert!(backend.dispatched_inputs.lock().unwrap().is_empty()); } + #[cfg(feature = "ws-server")] #[test] fn runtime_worker_subscription_has_gap_free_snapshot_and_live_updates() { let runtime = runtime_with_backend(); @@ -7918,6 +8178,62 @@ mod tests { ); } + #[cfg(feature = "fs-store")] + #[test] + fn fs_store_metadata_free_stop_retains_operation_across_unknown_commit_restart() { + let dir = tempfile::tempdir().unwrap(); + let options = crate::fs_store::FsRuntimeStoreOptions { + root: dir.path().join("runtime"), + runtime_id: "test-runtime".into(), + display_name: None, + }; + let runtime = Runtime::with_fs_store_and_execution_backend( + options.clone(), + Arc::new(TestExecutionBackend::default()), + ) + .unwrap(); + runtime.store_config_bundle(test_bundle()).unwrap(); + let worker = runtime + .create_worker(task_request("metadata-free stop")) + .unwrap(); + { + let mut state = runtime.lock().unwrap(); + let record = state.worker_mut(&worker.worker_ref).unwrap(); + record.request = None; + record.execution_metadata_available = false; + record.execution_bound = false; + state.persist_worker(&worker.worker_id).unwrap(); + } + runtime_store(&runtime) + .fail_next_worker_write(crate::fs_store::AtomicWriteFault::AfterRename); + assert!(matches!( + runtime.stop_worker(&worker.worker_ref, None), + Err(RuntimeError::StoreCommitOutcomeUnknown { .. }) + )); + let operation_id = runtime + .lock() + .unwrap() + .worker(&worker.worker_ref) + .unwrap() + .pending_stop + .unwrap() + .operation_id; + drop(runtime); + let backend = Arc::new(TestExecutionBackend::default()); + let reopened = + Runtime::with_fs_store_and_execution_backend(options, backend.clone()).unwrap(); + let stopped = reopened.worker_detail(&worker.worker_ref).unwrap(); + assert_eq!(stopped.status, WorkerStatus::Stopped); + assert!(!stopped.execution_metadata_available); + assert!(stopped.worker_state.is_none()); + assert_eq!(*backend.restore_count.lock().unwrap(), 0); + assert_eq!( + backend.stop_operation_ids.lock().unwrap().as_slice(), + &[operation_id] + ); + assert!(reopened.delete_worker(&worker.worker_ref).unwrap().deleted); + } + #[cfg(feature = "fs-store")] #[test] fn fs_store_reconciles_pending_stop_before_automatic_restore() { @@ -8059,6 +8375,95 @@ mod tests { ); } + #[cfg(feature = "fs-store")] + #[test] + fn stop_preserves_pending_restore_owner_across_unknown_commit_and_restart() { + for mode in [WorkerRestoreMode::Automatic, WorkerRestoreMode::Explicit] { + let dir = tempfile::tempdir().unwrap(); + let options = crate::fs_store::FsRuntimeStoreOptions { + root: dir.path().join("runtime"), + runtime_id: "stop-pending-restore".to_string(), + display_name: None, + }; + let backend = Arc::new(TestExecutionBackend::default()); + let runtime = + Runtime::with_fs_store_and_execution_backend(options.clone(), backend.clone()) + .unwrap(); + runtime.store_config_bundle(test_bundle()).unwrap(); + let worker = runtime + .create_worker(task_request("stop pending restore")) + .unwrap(); + runtime.stop_worker(&worker.worker_ref, None).unwrap(); + if mode == WorkerRestoreMode::Automatic { + let mut state = runtime.lock().unwrap(); + let record = state.worker_mut(&worker.worker_ref).unwrap(); + record.status = WorkerStatus::Idle; + record.restore_intent = WorkerRestoreIntent::Automatic; + state.persist_worker(&worker.worker_id).unwrap(); + } + *backend.restore_result.lock().unwrap() = + Some(WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + "cleanup owner is pending", + ), + worker_state: None, + workdir_attachments: Vec::new(), + }); + let restored = runtime + .restore_worker_under_lock(&worker.worker_ref, mode) + .unwrap(); + assert_eq!(restored.state, WorkerRestoreState::ReconciliationRequired); + let restore_id = runtime + .lock() + .unwrap() + .worker(&worker.worker_ref) + .unwrap() + .pending_restore + .unwrap() + .operation_id; + runtime_store(&runtime) + .fail_next_worker_write(crate::fs_store::AtomicWriteFault::AfterRename); + assert!(matches!( + runtime.stop_worker(&worker.worker_ref, None), + Err(RuntimeError::StoreCommitOutcomeUnknown { .. }) + )); + let state = runtime.lock().unwrap(); + let record = state.worker(&worker.worker_ref).unwrap(); + let stop_id = record.pending_stop.unwrap().operation_id; + assert_eq!(record.pending_restore.unwrap().operation_id, restore_id); + drop(state); + drop(runtime); + let retry_backend = Arc::new(TestExecutionBackend::default()); + let retry = + Runtime::with_fs_store_and_execution_backend(options, retry_backend.clone()) + .unwrap(); + assert_eq!( + *retry_backend.restore_count.lock().unwrap(), + 0, + "Stop must not launch a Controller to clean up restore" + ); + assert_eq!( + retry_backend.stop_operation_ids.lock().unwrap().as_slice(), + &[stop_id] + ); + assert_eq!( + retry.worker_detail(&worker.worker_ref).unwrap().status, + WorkerStatus::Stopped + ); + assert!( + retry + .lock() + .unwrap() + .worker(&worker.worker_ref) + .unwrap() + .pending_restore + .is_none() + ); + assert!(retry.delete_worker(&worker.worker_ref).unwrap().deleted); + } + } + #[test] fn concurrent_exact_restores_share_one_execution() { let backend = Arc::new(TestExecutionBackend::default()); @@ -8115,6 +8520,60 @@ mod tests { assert!(runtime.worker_operations.lock().unwrap().is_empty()); } + #[test] + fn scoped_attachment_replacement_revalidates_workspace_after_worker_lock() { + let runtime = Arc::new(runtime_with_backend()); + runtime.store_config_bundle(test_bundle()).unwrap(); + let workspace_scope = scope("workspace-a", "server-a"); + let worker = runtime + .create_worker_scoped( + &workspace_scope, + scoped_task_request("attachment scope race", "workspace-a"), + ) + .unwrap(); + let lease = runtime.worker_operation_lock(worker.worker_id).unwrap(); + let guard = lease.lock().unwrap(); + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let replacing_runtime = runtime.clone(); + let worker_ref = worker.worker_ref.clone(); + let replacing = std::thread::spawn(move || { + started_tx.send(()).unwrap(); + replacing_runtime.replace_worker_workdir_attachments_scoped( + &workspace_scope, + &worker_ref, + Vec::new(), + ) + }); + started_rx.recv().unwrap(); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2); + while Arc::strong_count(&lease.operation_lock) < 3 { + assert!( + std::time::Instant::now() < deadline, + "attachment replacement did not reach its Worker lock" + ); + std::thread::yield_now(); + } + { + let mut state = runtime.lock().unwrap(); + state.worker_mut(&worker.worker_ref).unwrap().workspace_id = Some("workspace-b".into()); + } + drop(guard); + drop(lease); + assert!(matches!( + replacing.join().unwrap(), + Err(RuntimeError::WorkerNotFound { .. }) + )); + assert!( + runtime + .lock() + .unwrap() + .worker(&worker.worker_ref) + .unwrap() + .logical_workdir_attachments + .is_empty() + ); + } + #[test] fn scoped_stop_revalidates_workspace_after_waiting_for_worker_lock() { let (runtime, backend) = runtime_and_backend(); @@ -8140,6 +8599,14 @@ mod tests { stopping_runtime.stop_worker_scoped(&stopping_scope, &stopping_ref, None) }); started_rx.recv().unwrap(); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2); + while Arc::strong_count(&lease.operation_lock) < 3 { + assert!( + std::time::Instant::now() < deadline, + "scoped operation did not reach its Worker lock" + ); + std::thread::yield_now(); + } { let mut state = runtime.lock().unwrap(); state.worker_mut(&worker.worker_ref).unwrap().workspace_id = @@ -8154,6 +8621,199 @@ mod tests { assert!(runtime.worker_operations.lock().unwrap().is_empty()); } + #[cfg(all(feature = "fs-store", feature = "ws-server"))] + #[test] + fn rejected_restore_settlement_failure_detaches_discarded_publication_owner() { + for rolled_back in [false, true] { + let dir = tempfile::tempdir().unwrap(); + let backend = Arc::new(TestExecutionBackend::default()); + let runtime = Arc::new( + Runtime::with_fs_store_and_execution_backend( + crate::fs_store::FsRuntimeStoreOptions { + root: dir.path().join("runtime"), + runtime_id: "test-runtime".into(), + display_name: None, + }, + backend.clone(), + ) + .unwrap(), + ); + runtime.store_config_bundle(test_bundle()).unwrap(); + let worker = runtime + .create_worker(task_request("discarded context settlement")) + .unwrap(); + runtime.stop_worker(&worker.worker_ref, None).unwrap(); + let failure = WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + "injected restore rejection", + ); + backend + .restore_result + .lock() + .unwrap() + .replace(if rolled_back { + WorkerExecutionSpawnResult::RolledBack(failure) + } else { + WorkerExecutionSpawnResult::Rejected(failure) + }); + let gate = Arc::new(RestoreGate::default()); + backend.restore_gate.lock().unwrap().replace(gate.clone()); + let restoring_runtime = runtime.clone(); + let worker_ref = worker.worker_ref.clone(); + let restoring = + std::thread::spawn(move || restoring_runtime.restore_worker_operation(&worker_ref)); + assert!(gate.wait_for_entered(1, std::time::Duration::from_secs(2))); + let operation_id = runtime + .lock() + .unwrap() + .worker(&worker.worker_ref) + .unwrap() + .pending_restore + .unwrap() + .operation_id; + runtime_store(&runtime) + .fail_next_worker_write(crate::fs_store::AtomicWriteFault::BeforeRename); + gate.release(); + assert_eq!( + restoring.join().unwrap().unwrap().state, + WorkerRestoreState::ReconciliationRequired + ); + { + let state = runtime.lock().unwrap(); + let record = state.worker(&worker.worker_ref).unwrap(); + assert_eq!(record.pending_restore.unwrap().operation_id, operation_id); + assert!(record.restore_candidate_context.is_none()); + } + backend.restore_result.lock().unwrap().take(); + backend.restore_gate.lock().unwrap().take(); + backend + .restore_event + .lock() + .unwrap() + .replace(protocol::Event::TextDelta { + text: "new candidate".into(), + }); + let cursor = runtime + .worker_observation_cursor_now(&worker.worker_ref) + .unwrap(); + assert_eq!( + runtime + .restore_worker_operation(&worker.worker_ref) + .unwrap() + .state, + WorkerRestoreState::Accepted + ); + backend + .publish_text_delta(&worker.worker_ref, "live candidate") + .unwrap(); + let events = runtime + .read_worker_observation_events(&worker.worker_ref, cursor) + .unwrap(); + assert_eq!(events.len(), 2); + assert_eq!( + backend.restore_operation_ids.lock().unwrap().as_slice(), + &[operation_id, operation_id] + ); + } + } + + #[cfg(feature = "ws-server")] + #[test] + fn committed_restore_retries_retained_observation_activation() { + let (runtime, backend) = runtime_and_backend(); + runtime.store_config_bundle(test_bundle()).unwrap(); + let worker = runtime + .create_worker(task_request("activation retry")) + .unwrap(); + runtime.stop_worker(&worker.worker_ref, None).unwrap(); + let cursor = runtime + .worker_observation_cursor_now(&worker.worker_ref) + .unwrap(); + let actual = runtime.execution_context(worker.worker_ref.clone()); + let fail_once = Arc::new(AtomicBool::new(true)); + let failure = fail_once.clone(); + let context = WorkerExecutionContext::candidate( + worker.worker_ref.clone(), + Arc::new(move |_, event| { + if failure.swap(false, Ordering::SeqCst) { + return Err(RuntimeError::InvalidRequest( + "injected publication failure".into(), + )); + } + actual.publish_observation(event) + }), + ); + let operation_id = WorkerLifecycleOperationId::new(); + runtime + .begin_restore_operation( + &worker.worker_ref, + operation_id, + WorkerRestoreMode::Explicit, + context, + ) + .unwrap(); + backend + .restore_event + .lock() + .unwrap() + .replace(protocol::Event::TextDelta { + text: "first".into(), + }); + let first = runtime + .restore_worker_operation(&worker.worker_ref) + .unwrap(); + assert_eq!(first.state, WorkerRestoreState::ReconciliationRequired); + assert!( + runtime + .lock() + .unwrap() + .worker(&worker.worker_ref) + .unwrap() + .restore_candidate_context + .is_some() + ); + assert!( + runtime + .read_worker_observation_events(&worker.worker_ref, cursor) + .unwrap() + .is_empty() + ); + backend + .publish_text_delta(&worker.worker_ref, "second") + .unwrap(); + let retry = runtime + .restore_worker_operation(&worker.worker_ref) + .unwrap(); + assert_eq!(retry.state, WorkerRestoreState::Accepted); + assert_eq!( + *backend.restore_count.lock().unwrap(), + 1, + "activation retry must not replace the execution" + ); + assert!( + runtime + .lock() + .unwrap() + .worker(&worker.worker_ref) + .unwrap() + .restore_candidate_context + .is_none() + ); + backend + .publish_text_delta(&worker.worker_ref, "third") + .unwrap(); + let texts = runtime + .read_worker_observation_events(&worker.worker_ref, cursor) + .unwrap() + .into_iter() + .filter_map(|event| match event.payload { + protocol::Event::TextDelta { text } => Some(text), + _ => None, + }) + .collect::>(); + assert_eq!(texts, ["first", "second", "third"]); + } + #[cfg(feature = "ws-server")] #[test] fn restore_candidate_observation_is_hidden_until_aggregate_commit() { @@ -8337,7 +8997,6 @@ mod tests { WorkerExecutionOperation::Spawn, "Workdir cleanup could not be proven", ), - handle: None, worker_state: None, workdir_attachments: vec![WorkingDirectoryAttachmentStatus { alias: workdir::WorkdirAttachmentAlias::new("checkout").unwrap(), @@ -8581,7 +9240,6 @@ mod tests { fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { WorkerExecutionSpawnResult::Connected { - handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), worker_state: protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, @@ -8600,7 +9258,7 @@ mod tests { fn dispatch_input( &self, - _handle: &WorkerExecutionHandle, + _worker_ref: &WorkerRef, input: WorkerInput, ) -> WorkerExecutionResult { WorkerExecutionResult::accepted_submission( @@ -8638,6 +9296,109 @@ mod tests { ); } + #[cfg(feature = "ws-server")] + #[test] + fn connected_protocol_cannot_authorize_a_different_transport_target() { + let (runtime, _) = runtime_and_backend(); + runtime.store_config_bundle(test_bundle()).unwrap(); + let a_scope = scope("workspace-a", "server-a"); + let b_scope = scope("workspace-b", "server-b"); + let a = runtime + .create_worker_scoped(&a_scope, scoped_task_request("a", "workspace-a")) + .unwrap(); + let b = runtime + .create_worker_scoped(&b_scope, scoped_task_request("b", "workspace-b")) + .unwrap(); + let (_, events) = tokio::sync::mpsc::channel(1); + let transport = crate::execution::WorkerProtocolTransport::new( + a.worker_ref.clone(), + protocol::Event::TextDelta { + text: "endpoint a".into(), + }, + events, + Arc::new(|_| panic!("scope for B must not authorize endpoint A")), + Arc::new(|| Ok(())), + ); + let method = Method::Shutdown { + command: protocol::WorkerCommandEnvelope { command_id: 1 }, + }; + assert!(matches!( + runtime.send_connected_protocol_method_scoped( + &b_scope, + &b.worker_ref, + &transport, + method + ), + Err(RuntimeError::WorkerNotFound { .. }) + )); + for worker_ref in [&a.worker_ref, &b.worker_ref] { + let state = runtime.lock().unwrap(); + let record = state.worker(worker_ref).unwrap(); + assert!(record.pending_stop.is_none()); + assert_eq!(record.status, WorkerStatus::Idle); + } + } + + #[cfg(feature = "ws-server")] + #[test] + fn closed_protocol_transport_cannot_journal_shutdown_or_dispatch_into_restored_worker() { + let (runtime, backend) = runtime_and_backend(); + let worker = runtime + .create_worker(task_request("old protocol endpoint")) + .unwrap(); + let (_, events) = tokio::sync::mpsc::channel(1); + let transport = crate::execution::WorkerProtocolTransport::new( + worker.worker_ref.clone(), + protocol::Event::TextDelta { + text: "old transport".to_string(), + }, + events, + Arc::new(|_| panic!("closed transport must be rejected before dispatch")), + Arc::new(|| { + Err(WorkerExecutionResult::rejected( + WorkerExecutionOperation::ProtocolMethod, + "original Controller transport is closed", + )) + }), + ); + runtime.stop_worker(&worker.worker_ref, None).unwrap(); + runtime.restore_worker(&worker.worker_ref).unwrap(); + let stop_count = *backend.stop_count.lock().unwrap(); + let detail = runtime.worker_detail(&worker.worker_ref).unwrap(); + for method in [ + Method::Submit { + submission_request_id: "old-request".to_string(), + input: Vec::new(), + }, + Method::ListCompletions { + kind: protocol::CompletionKind::File, + prefix: "".to_string(), + request_id: None, + context: None, + }, + Method::Shutdown { + command: test_command(), + }, + ] { + assert!(matches!( + runtime.send_connected_protocol_method(&worker.worker_ref, &transport, method), + Err(RuntimeError::WorkerExecutionRejected { .. }) + )); + } + assert_eq!(*backend.stop_count.lock().unwrap(), stop_count); + assert!(backend.dispatched_inputs.lock().unwrap().is_empty()); + let state = runtime.lock().unwrap(); + assert!( + state + .worker(&worker.worker_ref) + .unwrap() + .pending_stop + .is_none() + ); + assert_eq!(runtime.worker_operations.lock().unwrap().len(), 0); + assert_eq!(state.worker(&worker.worker_ref).unwrap().detail(), detail); + } + #[test] fn protocol_shutdown_transitions_worker_to_stopped() { let runtime = runtime_with_backend(); @@ -8711,7 +9472,10 @@ mod tests { assert!(matches!( runtime.send_input(&detail.worker_ref, WorkerInput::user("do not wake")), - Err(RuntimeError::WorkerExecutionUnavailable { .. }) + Err(RuntimeError::WorkerExecutionRejected { + outcome: crate::execution::WorkerExecutionOutcome::Rejected, + .. + }) )); assert_eq!(*backend.restore_count.lock().unwrap(), 0); @@ -9365,7 +10129,7 @@ mod tests { #[cfg(feature = "fs-store")] #[test] - fn fs_store_restores_active_worker_execution_handles() { + fn fs_store_restores_active_worker_executions() { let root = fs_store_root("execution-restore"); let runtime = Runtime::with_fs_store_and_execution_backend( crate::fs_store::FsRuntimeStoreOptions { @@ -9488,6 +10252,96 @@ mod tests { let _ = std::fs::remove_dir_all(root); } + #[cfg(feature = "fs-store")] + #[test] + fn active_startup_preflight_rejection_is_observable_bounded_and_stop_reconciles_absence() { + let dir = tempfile::tempdir().unwrap(); + let options = crate::fs_store::FsRuntimeStoreOptions { + root: dir.path().join("runtime"), + runtime_id: "active-preflight-rejection".to_string(), + display_name: None, + }; + let runtime = Runtime::with_fs_store_and_execution_backend( + options.clone(), + Arc::new(TestExecutionBackend::default()), + ) + .unwrap(); + runtime.store_config_bundle(test_bundle()).unwrap(); + let worker = runtime + .create_worker(task_request("saved active preflight rejection")) + .unwrap(); + drop(runtime); + let backend = Arc::new(TestExecutionBackend::default()); + let reason = format!( + "saved Manifest attachment mismatch token credential-value /private/path {}", + "details ".repeat(100) + ); + *backend.preflight_restore_result.lock().unwrap() = Some(WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + reason, + )); + let runtime = + Runtime::with_fs_store_and_execution_backend(options, backend.clone()).unwrap(); + let detail = runtime.worker_detail(&worker.worker_ref).unwrap(); + assert_eq!(detail.status, WorkerStatus::Idle); + assert!(detail.worker_state.is_none()); + assert_eq!(*backend.restore_count.lock().unwrap(), 0); + let diagnostics = runtime.diagnostics().unwrap(); + let diagnostic = diagnostics + .iter() + .find(|entry| entry.worker_ref.as_ref() == Some(&worker.worker_ref)) + .unwrap(); + assert!( + diagnostic + .message + .starts_with("saved Manifest attachment mismatch") + ); + assert!(!diagnostic.message.contains("credential-value")); + assert!(!diagnostic.message.contains("/private/path")); + assert!(diagnostic.message.chars().count() <= WORKER_CREATE_FAILURE_MESSAGE_LIMIT + 1); + let before = std::fs::read( + dir.path() + .join("runtime/workers") + .join(worker.worker_id.to_string()) + .join("worker.json"), + ) + .unwrap(); + assert_eq!( + runtime + .restore_worker_operation(&worker.worker_ref) + .unwrap() + .state, + WorkerRestoreState::Rejected + ); + assert_eq!( + std::fs::read( + dir.path() + .join("runtime/workers") + .join(worker.worker_id.to_string()) + .join("worker.json") + ) + .unwrap(), + before, + "explicit preflight rejection is read-only" + ); + assert!(matches!( + runtime.send_input( + &worker.worker_ref, + WorkerInput::notify("no implicit restore") + ), + Err(RuntimeError::WorkerExecutionRejected { + outcome: crate::execution::WorkerExecutionOutcome::Rejected, + .. + }) + )); + let stopped = runtime.stop_worker(&worker.worker_ref, None).unwrap(); + assert_eq!(stopped.status, WorkerStatus::Stopped); + assert!(stopped.worker_state.is_none()); + assert_eq!(*backend.restore_count.lock().unwrap(), 0); + assert_eq!(backend.stop_operation_ids.lock().unwrap().len(), 1); + assert!(runtime.delete_worker(&worker.worker_ref).unwrap().deleted); + } + #[cfg(feature = "fs-store")] #[test] fn fs_store_explicitly_retries_rejected_automatic_restore() { @@ -9635,7 +10489,6 @@ mod tests { WorkerExecutionOperation::Restore, "controller handshake outcome was not observable", ), - handle: None, worker_state: None, workdir_attachments: Vec::new(), }, @@ -9904,7 +10757,10 @@ mod tests { ); assert!(matches!( restored.send_input(&worker.worker_ref, WorkerInput::user("implicit restore")), - Err(RuntimeError::WorkerExecutionUnavailable { .. }) + Err(RuntimeError::WorkerExecutionRejected { + outcome: crate::execution::WorkerExecutionOutcome::Rejected, + .. + }) )); assert_eq!(*backend.restore_count.lock().unwrap(), 0); diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 698388853..e47eadb8f 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -2,7 +2,7 @@ //! `worker` crate controller/run lifecycle. //! //! The adapter intentionally owns real `WorkerHandle`s internally and exposes -//! only the opaque `worker-runtime` execution handle to callers. Browser/API +//! operations addressed by `WorkerRef` to callers. Browser/API //! projections therefore keep the existing runtime redaction boundary: no raw //! socket paths, session paths, manifests, credentials, or handles leave this //! module. @@ -22,11 +22,13 @@ use crate::catalog::{ WorkingDirectoryStatusKind, }; use crate::config_bundle::{ConfigBundle, workspace_config_etag}; +#[cfg(feature = "ws-server")] +use crate::execution::WorkerProtocolTransport; use crate::execution::{ - WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, - WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionSpawnRequest, - WorkerExecutionSpawnResult, WorkerExecutionStopRequest, WorkerSessionObservationRequest, - WorkspaceConfigFetchRequest, WorkspaceConfigFetchResult, + WorkerExecutionBackend, WorkerExecutionOperation, WorkerExecutionRestoreRequest, + WorkerExecutionResult, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, + WorkerExecutionStopRequest, WorkerSessionObservationRequest, WorkspaceConfigFetchRequest, + WorkspaceConfigFetchResult, }; use crate::identity::WorkerRef; use crate::interaction::{WorkerInput, WorkerInputKind}; @@ -83,7 +85,6 @@ fn rollback_materialized_spawn_workdirs( cleanup_failures.join("; ") ), ), - handle: None, worker_state: None, workdir_attachments: uncertain_attachments, }) @@ -96,6 +97,12 @@ fn next_internal_command( .read() .map_err(|_| "worker state lock is poisoned".to_string())? .clone(); + Ok(next_internal_command_for_snapshot(&snapshot)) +} + +fn next_internal_command_for_snapshot( + snapshot: &protocol::WorkerStateSnapshot, +) -> WorkerCommandEnvelope { let floor = snapshot.last_command_id.saturating_add(1); let command_id = NEXT_INTERNAL_COMMAND_ID .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { @@ -103,7 +110,7 @@ fn next_internal_command( }) .unwrap_or(floor) .max(floor); - Ok(WorkerCommandEnvelope::new(command_id)) + WorkerCommandEnvelope::new(command_id) } use session_store::{ CombinedStore, Store, WorkerAggregateStore, WorkerMetadataStore, WorkerSessionStore, @@ -214,6 +221,23 @@ pub trait RuntimeWorkerFactory: Send + Sync + 'static { request: WorkerExecutionRestoreRequest, ) -> Result; + /// Prove cleanup of a restore future that completed without returning a + /// controller. An error is not absence evidence; unsupported factories keep + /// the operation-owned resource record for retry. + async fn cleanup_failed_restore( + &self, + _request: &WorkerExecutionRestoreRequest, + ) -> Result<(), String> { + Err("Runtime Worker factory cannot prove failed restore cleanup".to_string()) + } + + /// Verify/release factory-owned resources when the backend registry has no + /// Controller. Factories retaining resources outside returned controllers + /// must override this; it must never start a new Controller. + async fn reconcile_stopped_worker(&self, _worker_ref: &WorkerRef) -> Result<(), String> { + Ok(()) + } + fn activate_restored_controller( &self, _worker_ref: &WorkerRef, @@ -418,6 +442,17 @@ pub struct ProfileRuntimeWorkerFactory { workspace_request_clients: Arc>, embedded_worker_mutation_dispatcher: Option>, controller_transport: WorkerControllerTransport, + failed_restore_sessions: Arc< + Mutex< + HashMap< + WorkerRef, + ( + crate::execution::WorkerLifecycleOperationId, + Arc, + ), + >, + >, + >, } impl ProfileRuntimeWorkerFactory { @@ -434,6 +469,7 @@ impl ProfileRuntimeWorkerFactory { workspace_request_clients: Arc::new(HashMap::new()), embedded_worker_mutation_dispatcher: None, controller_transport: WorkerControllerTransport::InProcess, + failed_restore_sessions: Arc::new(Mutex::new(HashMap::new())), } } @@ -516,6 +552,87 @@ impl ProfileRuntimeWorkerFactory { .join(format!("restore-{operation_id}"))) } + // Called only after the backend has excluded actual and pending executions + // under its per-Worker operation lock. Persisted run artifacts survive a + // process restart even when the in-memory failed-session map does not. + fn stopped_restore_run_dirs(&self, worker_ref: &WorkerRef) -> Result, String> { + let aggregate = self.worker_aggregate_dir(worker_ref)?; + let root = aggregate + .parent() + .expect("Worker aggregate has configured parent"); + let runs = aggregate.join("runs"); + for path in [root, aggregate.as_path(), runs.as_path()] { + match std::fs::symlink_metadata(path) { + Ok(metadata) if metadata.is_dir() && !metadata.file_type().is_symlink() => {} + Ok(_) => { + return Err(format!( + "unsafe restore artifact directory: {}", + path.display() + )); + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(error) => { + return Err(format!( + "inspect restore artifact directory {}: {error}", + path.display() + )); + } + } + } + let mut candidates = Vec::new(); + for entry in std::fs::read_dir(&runs) + .map_err(|error| format!("read restore artifact directory: {error}"))? + { + let entry = entry.map_err(|error| format!("read restore artifact entry: {error}"))?; + let name = entry.file_name(); + let name = name + .to_str() + .ok_or_else(|| "invalid restore artifact run name".to_string())?; + let Some(operation) = name.strip_prefix("restore-") else { + continue; + }; + let id = uuid::Uuid::parse_str(operation) + .map_err(|_| format!("invalid restore artifact operation name: {name}"))?; + if id.to_string() != operation { + return Err(format!( + "noncanonical restore artifact operation name: {name}" + )); + } + let path = entry.path(); + Self::validate_restore_artifact_tree(&path, true)?; + candidates.push(path); + } + candidates.sort(); + Ok(candidates) + } + + fn validate_restore_artifact_tree(path: &Path, require_directory: bool) -> Result<(), String> { + let metadata = std::fs::symlink_metadata(path) + .map_err(|error| format!("inspect restore artifact {}: {error}", path.display()))?; + if metadata.file_type().is_symlink() || (require_directory && !metadata.is_dir()) { + return Err(format!("unsafe restore artifact: {}", path.display())); + } + if metadata.is_dir() { + for entry in std::fs::read_dir(path) + .map_err(|error| format!("read restore artifact {}: {error}", path.display()))? + { + let entry = entry.map_err(|error| format!("read restore artifact: {error}"))?; + Self::validate_restore_artifact_tree(&entry.path(), false)?; + } + } + Ok(()) + } + + fn remove_restore_run_artifacts(run_dir: &Path) -> Result<(), String> { + match std::fs::symlink_metadata(run_dir) { + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()), + Err(error) => return Err(format!("inspect restore run artifacts: {error}")), + Ok(_) => Self::validate_restore_artifact_tree(run_dir, true)?, + } + std::fs::remove_dir_all(run_dir) + .map_err(|error| format!("failed to clean restore run artifacts: {error}")) + } + fn runtime_worker_name_for_ref(worker_ref: &crate::identity::WorkerRef) -> String { format!("worker-runtime-{}", worker_ref.worker_id) } @@ -558,6 +675,25 @@ impl ProfileRuntimeWorkerFactory { .map_err(|err| format!("failed to build restore fallback manifest: {err}"))?; Ok((manifest, PromptCatalogSource::builtins_only())) } + fn manifest_for_restore( + metadata: &session_store::WorkerMetadata, + request: &CreateWorkerRequest, + ) -> Result<(manifest::WorkerManifest, PromptCatalogSource), String> { + let (fallback, loader) = Self::restore_fallback_manifest(&metadata.worker_name)?; + let manifest = match metadata.resolved_manifest_snapshot.clone() { + Some(snapshot) => manifest::read_persisted_worker_manifest_snapshot(snapshot) + .map_err(|error| format!("failed to read saved Worker Manifest: {error}"))?, + None if request.workspace_api.is_some() || request.subjektiv_attached => { + return Err("Workspace Worker metadata has no saved Manifest; replacement Worker is required".to_string()); + } + None => fallback, + }; + // Saved policy is execution authority. Rebinding current settings here + // would hide revision and attachment mismatches. + validate_worker_memory_settings(&manifest, request)?; + Ok((manifest, loader)) + } + fn observe_bundle_prompt_projection( &self, bundle: &crate::config_bundle::ConfigBundle, @@ -1008,14 +1144,18 @@ fn validate_worker_memory_settings( request: &CreateWorkerRequest, ) -> Result<(), String> { let Some(expected) = request.memory_settings.as_ref() else { - if request.subjektiv_attached { + if request.subjektiv_attached || request.workspace_api.is_some() { return Err( - "subject-attached Worker restore is missing its trusted settings snapshot" - .to_string(), + "Workspace Worker restore is missing its trusted settings snapshot".to_string(), ); } return Ok(()); }; + if let Some(workspace_api) = request.workspace_api.as_ref() + && expected.workspace_id != workspace_api.workspace_id + { + return Err("trusted Memory settings do not match Workspace API scope".to_string()); + } manifest .feature .memory @@ -1368,8 +1508,6 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { request: &WorkerExecutionRestoreRequest, ) -> Result<(), String> { let worker_name = Self::runtime_worker_name_for_ref(&request.worker_ref); - let (mut manifest, _) = Self::restore_fallback_manifest(&worker_name)?; - bind_workspace_memory_settings(&mut manifest, &request.request)?; let worker_aggregate_dir = self.worker_aggregate_dir(&request.worker_ref)?; if !worker_aggregate_dir.is_dir() { return Err("Persisted Worker aggregate metadata is unavailable".to_string()); @@ -1383,6 +1521,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { .read_by_name(&worker_name) .map_err(|error| format!("failed to read Worker metadata: {error}"))? .ok_or_else(|| "Persisted Worker metadata is unavailable".to_string())?; + let (manifest, loader) = Self::manifest_for_restore(&metadata, &request.request)?; if let Some(active) = metadata.active.as_ref() && let Some(segment_id) = active.segment_id { @@ -1398,17 +1537,35 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { .read_all(active.session_id, segment_id) .map_err(|error| format!("failed to read Worker Session segment: {error}"))?; } - if request.request.workspace_api.is_some() + if let Some(api) = request.request.workspace_api.as_ref() && metadata .active .as_ref() .is_some_and(|active| active.segment_id.is_none()) - && request.config_bundle.is_none() { - return Err( - "Pending Workspace Worker restore requires current profile launch authority" - .to_string(), - ); + let bundle = request.config_bundle.as_ref().ok_or_else(|| { + "Pending Workspace Worker restore requires operation-owned launch material" + .to_string() + })?; + crate::config_bundle::validate_config_bundle(bundle) + .map_err(|error| format!("invalid pending Worker launch material: {error}"))?; + if bundle.metadata.workspace_id != api.workspace_id { + return Err(format!( + "Workspace Prompt projection scope mismatch: expected {}, got {}", + api.workspace_id, bundle.metadata.workspace_id + )); + } + let prompt_catalog = bundle.prompt_catalog.as_ref().ok_or_else(|| { + "pending Workspace Worker restore requires a saved Workspace Prompt projection" + .to_string() + })?; + // Validate the exact launch catalog without publishing it into the + // current projection cache or allocating live Worker resources. + worker::SystemPromptTemplate::parse( + &manifest.engine.instruction, + loader.with_effective_catalog(prompt_catalog.clone()), + ) + .map_err(|error| format!("invalid pending Worker launch Prompt: {error}"))?; } Ok(()) } @@ -1417,6 +1574,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { &self, request: WorkerExecutionRestoreRequest, ) -> Result { + self.preflight_restore(&request).await?; let worker_name = Self::runtime_worker_name_for_ref(&request.worker_ref); let only_binding = (request.workdir_attachments.len() == 1) .then(|| request.workdir_attachments.values().next()) @@ -1454,8 +1612,6 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { self.embedded_worker_mutation_dispatcher.as_ref(), Some(self.prompt_projection_cache.clone()), ); - let (mut manifest, loader) = Self::restore_fallback_manifest(&worker_name)?; - bind_workspace_memory_settings(&mut manifest, &request.request)?; let worker_aggregate_dir = self.worker_aggregate_dir(&request.worker_ref)?; let session_dir = worker_aggregate_dir.join("session"); @@ -1474,6 +1630,11 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { ) }, )?; + let metadata = worker_metadata_store + .read_by_name(&worker_name) + .map_err(|error| format!("failed to read Worker metadata: {error}"))? + .ok_or_else(|| "Persisted Worker metadata is unavailable".to_string())?; + let (manifest, loader) = Self::manifest_for_restore(&metadata, &request.request)?; let store = CombinedStore::new(session_store, worker_metadata_store); let mut worker = match Worker::restore_from_worker_metadata_with_context( @@ -1573,6 +1734,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { let workspace_client = worker.workspace_client_handle(); let run_dir = self.worker_restore_run_dir(&request.worker_ref, request.operation_id)?; let bash_output_dir = bash_output_dir_for_worker_id(&request.worker_ref.worker_id); + let cleanup_sessions = worker.workdir_sessions(); let started = PreparedWorker::new( worker, WorkerBootstrapLayout::RuntimeManagedRun { @@ -1587,10 +1749,24 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { WorkerBootstrapError::Worker(source) => { format!("failed to prepare restored Worker: {source}") } - WorkerBootstrapError::Controller { source, .. } => format!( - "failed to spawn restored Worker controller in {}: {source}", - run_dir.display() - ), + WorkerBootstrapError::Controller { + source, + cleanup_failed, + } => { + if cleanup_failed { + self.failed_restore_sessions + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert( + request.worker_ref.clone(), + (request.operation_id, cleanup_sessions), + ); + } + format!( + "failed to spawn restored Worker controller in {}: {source}", + run_dir.display() + ) + } })?; let (handle, shutdown_rx, controller_task) = (started.handle, started.shutdown, started.controller_task); @@ -1605,6 +1781,66 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { }) } + async fn cleanup_failed_restore( + &self, + request: &WorkerExecutionRestoreRequest, + ) -> Result<(), String> { + let retained = self + .failed_restore_sessions + .lock() + .map_err(|_| "failed restore sessions lock is poisoned".to_string())? + .get(&request.worker_ref) + .cloned(); + if let Some((operation_id, sessions)) = retained { + if operation_id != request.operation_id { + return Err("failed restore sessions belong to another operation".to_string()); + } + sessions + .close_all() + .await + .map_err(|error| format!("failed restore session cleanup: {error}"))?; + } + // restore_controller starts the Controller only as its final fallible + // step. A completed error leaves no live Controller; remove only this + // operation's disposable run artifacts, never Session/metadata authority. + self.stopped_restore_run_dirs(&request.worker_ref)?; + let run_dir = self.worker_restore_run_dir(&request.worker_ref, request.operation_id)?; + Self::remove_restore_run_artifacts(&run_dir)?; + self.failed_restore_sessions + .lock() + .map_err(|_| "failed restore sessions lock is poisoned".to_string())? + .remove(&request.worker_ref); + Ok(()) + } + + async fn reconcile_stopped_worker(&self, worker_ref: &WorkerRef) -> Result<(), String> { + let retained = self + .failed_restore_sessions + .lock() + .map_err(|_| "failed restore sessions lock is poisoned".to_string())? + .get(worker_ref) + .cloned(); + if let Some((_, sessions)) = retained.as_ref() { + sessions + .close_all() + .await + .map_err(|error| format!("failed restore session cleanup: {error}"))?; + } + // A restart loses the map, not the operation-owned filesystem artifacts. + // Validate every candidate before deleting any, and retain map/journal + // retry authority whenever validation or cleanup remains unproven. + for run_dir in self.stopped_restore_run_dirs(worker_ref)? { + Self::remove_restore_run_artifacts(&run_dir)?; + } + if retained.is_some() { + self.failed_restore_sessions + .lock() + .map_err(|_| "failed restore sessions lock is poisoned".to_string())? + .remove(worker_ref); + } + Ok(()) + } + fn activate_restored_controller( &self, worker_ref: &WorkerRef, @@ -1618,23 +1854,167 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { #[derive(Clone)] struct RuntimeExecutionTaskScope { - tasks: Arc>>, + tasks: Arc>>>, + shutdown_admission: Arc>>, + // Actual unconnected Controllers whose cleanup failed, owned by this + // execution's structured scope. No WorkerRef index or presence mirror. + cleanup_candidates: Arc>>>, +} + +// Evidence for one exact command on this actual Controller. Keep the send task +// and receiver across bounded waits: unknown admission never authorizes resend. +struct RuntimeShutdownAdmission { + command_id: u64, + send: Option>>, + events: tokio::sync::broadcast::Receiver, + acknowledgement: Option, + uncertainty: Option, +} + +impl RuntimeShutdownAdmission { + async fn wait(&mut self, timeout: Duration) -> Result<(), String> { + let result = tokio::time::timeout(timeout, async { + if let Some(send) = self.send.as_mut() { + match send.await { + Ok(Ok(())) => { self.send.take(); } + Ok(Err(message)) => { + self.send.take(); + self.uncertainty = Some(message.clone()); + return Err(message); + } + Err(error) => { + self.send.take(); + let message = format!("Shutdown send task failed: {error}; admission is unknown"); + self.uncertainty = Some(message.clone()); + return Err(message); + } + } + } + if self.acknowledgement.is_some() { + return Ok(()); + } + loop { + match self.events.recv().await { + Ok(Event::CommandAcknowledged { acknowledgement }) + if acknowledgement.command_id == self.command_id + && acknowledgement.command == protocol::WorkerCommandKind::Shutdown => { + self.acknowledgement = Some(acknowledgement); + self.uncertainty = None; + return Ok(()); + } + Ok(_) => {} + Err(error) => { + let message = format!("Shutdown command {} acknowledgement unavailable: {error}; admission is unknown", self.command_id); + self.uncertainty = Some(message.clone()); + return Err(message); + } + } + } + }).await; + match result { + Ok(result) => result, + Err(_) => { + let message = format!( + "Shutdown command {} acknowledgement timed out; admission is unknown and waiter is retained", + self.command_id + ); + self.uncertainty = Some(message.clone()); + Err(message) + } + } + } } struct RuntimeExecutionTask { name: &'static str, - task: tokio::task::JoinHandle<()>, + completion: tokio::sync::Mutex, abort_before_join: bool, } +// The scope retains the actual handle while awaiting it and retains its abnormal +// completion after it has been consumed. Neither cancellation of a cleanup +// future nor another Stop attempt can turn an unknown cleanup into an empty scope. +struct RuntimeExecutionTaskCompletion { + task: Option>, + failure: Option, +} + impl RuntimeExecutionTaskScope { fn new(controller_task: tokio::task::JoinHandle<()>) -> Self { - Self { - tasks: Arc::new(Mutex::new(vec![RuntimeExecutionTask { - name: "controller", - task: controller_task, - abort_before_join: false, - }])), + let scope = Self { + tasks: Arc::new(Mutex::new(Vec::new())), + shutdown_admission: Arc::new(tokio::sync::Mutex::new(None)), + cleanup_candidates: Arc::new(Mutex::new(Vec::new())), + }; + scope.push("controller", controller_task, false); + scope + } + + fn has_cleanup_candidates(&self) -> Result { + self.cleanup_candidates + .lock() + .map(|candidates| !candidates.is_empty()) + .map_err(|_| "execution cleanup candidate lock is poisoned".to_string()) + } + + async fn request_shutdown( + &self, + handle: &WorkerHandle, + shutdown_requested: &AtomicBool, + timeout: Duration, + ) -> Result<(), String> { + let mut retained = self.shutdown_admission.lock().await; + if handle.protocol_is_closed() { + // Endpoint closure alone is not cleanup proof. Drain our send task + // too, then let the existing completion + all-task join barrier prove it. + if let Some(pending) = retained.as_mut() { + if let Some(send) = pending.send.as_mut() { + let _ = tokio::time::timeout(timeout, send).await.map_err(|_| { + "Shutdown send task still pending on closed endpoint".to_string() + })?; + pending.send.take(); + } + } + shutdown_requested.store(true, Ordering::Release); + return Ok(()); + } + if retained.is_none() { + if shutdown_requested.load(Ordering::Acquire) { + return Err("Shutdown admission evidence missing; cleanup is unproven".to_string()); + } + // The bridge cache can lag admission of already queued commands. + // Rejection retries use the real Controller snapshot, never that cache. + let command = next_internal_command_for_snapshot(&handle.shared_state.snapshot()); + let events = handle.subscribe(); + let worker = handle.clone(); + *retained = Some(RuntimeShutdownAdmission { + command_id: command.command_id, + send: Some(tokio::spawn(async move { + worker + .send(Method::Shutdown { command }) + .await + .map_err(|error| format!("failed to send Shutdown: {error}")) + })), + events, + acknowledgement: None, + uncertainty: None, + }); + shutdown_requested.store(true, Ordering::Release); + } + let pending = retained.as_mut().unwrap(); + pending.wait(timeout).await?; + let acknowledgement = pending.acknowledgement.as_ref().unwrap(); + match acknowledgement.disposition { + protocol::WorkerCommandDisposition::Accepted => Ok(()), + disposition => { + let message = format!( + "Shutdown command {} rejected: {disposition:?}; stop remains retryable", + acknowledgement.command_id + ); + retained.take(); + shutdown_requested.store(false, Ordering::Release); + Err(message) + } } } @@ -1643,70 +2023,128 @@ impl RuntimeExecutionTaskScope { Ok(tasks) => tasks, Err(poisoned) => poisoned.into_inner(), }; - tasks.push(RuntimeExecutionTask { + tasks.push(Arc::new(RuntimeExecutionTask { name, - task, + completion: tokio::sync::Mutex::new(RuntimeExecutionTaskCompletion { + task: Some(task), + failure: None, + }), abort_before_join, - }); + })); } - fn abort_all(&self) { - let mut tasks = match self.tasks.lock() { - Ok(tasks) => tasks, - Err(poisoned) => poisoned.into_inner(), + fn retain_completion_failure(&self, message: String) -> Result<(), String> { + self.tasks + .lock() + .map_err(|_| "execution task registry lock is poisoned".to_string())? + .push(Arc::new(RuntimeExecutionTask { + name: "controller shutdown completion", + completion: tokio::sync::Mutex::new(RuntimeExecutionTaskCompletion { + task: None, + failure: Some(message), + }), + abort_before_join: false, + })); + Ok(()) + } + + async fn confirm_shutdown_and_join( + &self, + shutdown: &Arc>>, + ) -> Result<(), String> { + // Do not consume a completion receiver if its failure cannot be retained. + drop( + self.tasks + .lock() + .map_err(|_| "execution task registry lock is poisoned".to_string())?, + ); + let confirmation = { + let mut guard = shutdown.lock().await; + if let Some(receiver) = guard.as_mut() { + match tokio::time::timeout(Duration::from_secs(5), receiver).await { + Ok(Ok(())) => { + guard.take(); + Ok(()) + } + Ok(Err(_)) => { + let message = + "Worker shutdown completion channel closed; cleanup is unproven" + .to_string(); + self.retain_completion_failure(message.clone())?; + guard.take(); + Err(message) + } + Err(_) => Err( + "Worker shutdown confirmation timed out; stop remains retryable" + .to_string(), + ), + } + } else { + Ok(()) + } }; - for task in tasks.drain(..) { - task.task.abort(); - } + // Drain bridges/relays even after an abnormal completion, while keeping + // the failed completion and all timed-out handles in their owning scope. + self.join().await?; + confirmation } async fn join(&self) -> Result<(), String> { + let tasks = self + .tasks + .lock() + .map_err(|_| "execution task registry lock is poisoned".to_string())? + .clone(); let mut first_failure = None; - let mut pending = Vec::new(); - loop { - let next = { - let mut tasks = self - .tasks - .lock() - .map_err(|_| "execution task registry lock is poisoned".to_string())?; - if tasks.is_empty() { - None - } else { - Some(tasks.remove(0)) - } - }; - let Some(mut task) = next else { - break; - }; - let name = task.name; - if task.abort_before_join { - task.task.abort(); + for task in tasks { + let mut completion = task.completion.lock().await; + if let Some(message) = completion.failure.as_ref() { + first_failure.get_or_insert_with(|| message.clone()); + continue; } - match tokio::time::timeout(Duration::from_secs(5), &mut task.task).await { - Ok(Ok(())) => {} - Ok(Err(error)) if task.abort_before_join && error.is_cancelled() => {} - Ok(Err(error)) => { - first_failure.get_or_insert_with(|| { - format!("{name} task failed while stopping Worker: {error}") - }); + if let Some(handle) = completion.task.as_mut() { + if task.abort_before_join { + handle.abort(); } - Err(_) => { - first_failure.get_or_insert_with(|| { - format!("{name} task did not stop before timeout; stop remains retryable") - }); - pending.push(task); + match tokio::time::timeout(Duration::from_secs(5), handle).await { + Ok(Ok(())) => { + completion.task.take(); + } + Ok(Err(error)) if task.abort_before_join && error.is_cancelled() => { + completion.task.take(); + } + Ok(Err(error)) => { + let message = format!( + "{} task failed while stopping Worker: {error}; cleanup is unproven", + task.name + ); + completion.task.take(); + completion.failure = Some(message.clone()); + first_failure.get_or_insert(message); + continue; + } + Err(_) => { + first_failure.get_or_insert_with(|| { + format!( + "{} task did not stop before timeout; stop remains retryable", + task.name + ) + }); + continue; + } } } - } - if !pending.is_empty() { - let mut tasks = match self.tasks.lock() { - Ok(tasks) => tasks, - Err(poisoned) => poisoned.into_inner(), - }; - tasks.extend(pending); + drop(completion); + self.tasks + .lock() + .map_err(|_| "execution task registry lock is poisoned".to_string())? + .retain(|retained| !Arc::ptr_eq(retained, &task)); } match first_failure { Some(message) => Err(message), + None if self.has_cleanup_candidates()? => Err( + "unconnected Controller cleanup remains unproven; scope is retained".to_string(), + ), None => Ok(()), } } @@ -1725,14 +2163,59 @@ struct RuntimeWorkerExecution { workdir_attachments: Vec, } -/// `worker-runtime` execution backend backed by real `worker` crate Workers. +// Actual operation-owned resources, not a second Worker catalog or presence +// marker. Timeout never drops the factory future or a late Controller result. +struct PendingRuntimeRestore { + request: WorkerExecutionRestoreRequest, + task: tokio::sync::Mutex>>, + result: Mutex>>, + failure: Mutex>, +} + +impl PendingRuntimeRestore { + async fn wait(&self, timeout: Duration) -> Result<(), String> { + let mut task = self.task.lock().await; + if let Some(message) = self + .failure + .lock() + .map_err(|_| "restore factory failure lock is poisoned".to_string())? + .as_ref() + { + return Err(message.clone()); + } + if let Some(handle) = task.as_mut() { + match tokio::time::timeout(timeout, handle).await { + Ok(Ok(())) => { + task.take(); + } + Ok(Err(error)) => { + task.take(); + let message = + format!("restore factory task failed: {error}; cleanup is unproven"); + *self + .failure + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(message.clone()); + return Err(message); + } + Err(_) => { + return Err("restore factory is still pending; cleanup is unproven".to_string()); + } + } + } + Ok(()) + } +} + +/// Worker execution backend backed by real worker crate Workers. pub struct WorkerRuntimeExecutionBackend { backend_id: String, factory: Arc, working_directory_materializer: Option>, runtime: Mutex>, workers: Mutex>, - restore_lock: Mutex<()>, + pending_restores: Mutex>>, + worker_locks: Mutex>>>, spawn_restore_timeout: Duration, } @@ -1752,7 +2235,8 @@ where working_directory_materializer: None, runtime: Mutex::new(Some(runtime)), workers: Mutex::new(HashMap::new()), - restore_lock: Mutex::new(()), + pending_restores: Mutex::new(HashMap::new()), + worker_locks: Mutex::new(HashMap::new()), spawn_restore_timeout: SPAWN_RESTORE_TASK_TIMEOUT, }) } @@ -1776,6 +2260,17 @@ where self } + fn worker_lock(&self, worker_ref: &WorkerRef) -> Result>, String> { + let mut locks = self + .worker_locks + .lock() + .map_err(|_| "worker operation lock registry is poisoned".to_string())?; + Ok(locks + .entry(worker_ref.clone()) + .or_insert_with(|| Arc::new(Mutex::new(()))) + .clone()) + } + fn wait_for_runtime_task(receiver: mpsc::Receiver>) -> Result { receiver .recv_timeout(RUNTIME_TASK_TIMEOUT) @@ -1816,6 +2311,26 @@ where Self::wait_for_runtime_task(rx) } + // Cleanup contains its own bounded waits. Do not return while an outer + // waiter still owns removed task handles: a retry must see complete scope + // ownership, rather than mistake an in-flight join for an empty scope. + fn run_joined_on_adapter_runtime(&self, task: Fut) -> Result + where + T: Send + 'static, + Fut: Future> + Send + 'static, + { + let (tx, rx) = mpsc::sync_channel(1); + self.spawn_on_adapter_runtime(async move { + let result = match tokio::spawn(task).await { + Ok(result) => result, + Err(error) => Err(format!("worker adapter cleanup task failed: {error}")), + }; + let _ = tx.send(result); + })?; + rx.recv() + .map_err(|error| format!("worker adapter cleanup did not complete: {error}"))? + } + fn run_cancellable_on_adapter_runtime( &self, timeout: Duration, @@ -1852,9 +2367,38 @@ where .map_err(|err| format!("worker adapter task did not complete: {err}"))? } + fn pending_restore( + &self, + worker_ref: &WorkerRef, + ) -> Result>, String> { + self.pending_restores + .lock() + .map(|pending| pending.get(worker_ref).cloned()) + .map_err(|_| "pending restore resources lock is poisoned".to_string()) + } + + fn wait_pending_restore( + &self, + pending: &Arc, + timeout: Duration, + ) -> Result<(), String> { + let pending = Arc::clone(pending); + self.run_cancellable_on_adapter_runtime(timeout + Duration::from_secs(1), async move { + pending.wait(timeout).await + }) + } + + fn remove_pending_restore(&self, worker_ref: &WorkerRef) -> Result<(), String> { + self.pending_restores + .lock() + .map_err(|_| "pending restore resources lock is poisoned".to_string())? + .remove(worker_ref); + Ok(()) + } + fn get_execution( &self, - handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, ) -> Result< ( WorkerHandle, @@ -1863,16 +2407,6 @@ where ), WorkerExecutionResult, > { - if handle.backend_id() != self.backend_id() { - return Err(WorkerExecutionResult::rejected( - WorkerExecutionOperation::Input, - format!( - "execution handle belongs to backend {}, not {}", - handle.backend_id(), - self.backend_id() - ), - )); - } let workers = self.workers.lock().map_err(|_| { WorkerExecutionResult::errored( WorkerExecutionOperation::Input, @@ -1880,7 +2414,7 @@ where ) })?; workers - .get(handle.worker_ref()) + .get(worker_ref) .map(|execution| { ( execution.handle.clone(), @@ -1891,11 +2425,115 @@ where .ok_or_else(|| { WorkerExecutionResult::rejected( WorkerExecutionOperation::Input, - "execution handle does not reference a live Worker", + "WorkerRef does not reference a live Worker execution", ) }) } + #[cfg(feature = "ws-server")] + fn validate_protocol_execution_under_lock( + &self, + worker_ref: &WorkerRef, + execution: &RuntimeWorkerExecution, + ) -> Result<(), WorkerExecutionResult> { + if execution.shutdown_requested.load(Ordering::Acquire) + || execution.handle.protocol_is_closed() + { + return Err(WorkerExecutionResult::rejected( + WorkerExecutionOperation::ProtocolMethod, + "attached Worker protocol endpoint is closed", + )); + } + if self + .pending_restore(worker_ref) + .map_err(|message| { + WorkerExecutionResult::errored(WorkerExecutionOperation::ProtocolMethod, message) + })? + .is_some() + { + return Err(WorkerExecutionResult::busy( + WorkerExecutionOperation::ProtocolMethod, + "Worker restore still owns pending resources", + )); + } + let workers = self.workers.lock().map_err(|_| { + WorkerExecutionResult::errored( + WorkerExecutionOperation::ProtocolMethod, + "worker adapter registry lock is poisoned", + ) + })?; + if workers + .get(worker_ref) + .is_none_or(|current| !current.handle.same_controller(&execution.handle)) + { + return Err(WorkerExecutionResult::rejected( + WorkerExecutionOperation::ProtocolMethod, + "attached Worker protocol endpoint is no longer current", + )); + } + Ok(()) + } + + #[cfg(feature = "ws-server")] + fn dispatch_protocol_execution_under_lock( + &self, + worker_ref: &WorkerRef, + execution: &RuntimeWorkerExecution, + method: Method, + ) -> Result, WorkerExecutionResult> { + self.validate_protocol_execution_under_lock(worker_ref, execution)?; + match method { + Method::ListCompletions { + request_id, + kind, + prefix, + context, + } => { + let handle = execution.handle.clone(); + self.run_cancellable_on_adapter_runtime(USER_INPUT_TASK_TIMEOUT, async move { + let entries = handle + .completion_entries(kind, &prefix, context.as_ref()) + .await; + Ok(vec![Event::Completions { + request_id, + kind, + prefix, + context, + entries, + }]) + }) + .map_err(|message| { + WorkerExecutionResult::errored( + WorkerExecutionOperation::ProtocolMethod, + message, + ) + }) + } + Method::Shutdown { .. } => { + // Identity was checked against the captured Controller while + // this same operation lock is held. Never unlock/rebind here. + let result = self.stop_worker_under_lock(worker_ref); + if result.is_accepted() { + Ok(Vec::new()) + } else { + Err(result) + } + } + method => { + let result = self.send_method( + WorkerExecutionOperation::ProtocolMethod, + execution.handle.clone(), + method, + ); + if result.is_accepted() { + Ok(Vec::new()) + } else { + Err(result) + } + } + } + } + fn send_method( &self, operation: WorkerExecutionOperation, @@ -2058,40 +2696,45 @@ where worker_ref: &crate::identity::WorkerRef, handle: WorkerHandle, shutdown: Arc>>, + shutdown_requested: Arc, tasks: RuntimeExecutionTaskScope, worker_state: Arc>, workspace_client: Option>, restore_operation_id: Option, workdir_attachments: Vec, - ) -> Result { + ) { + // Retention must not return the candidate as an error and drop it. Even + // a poisoned registry keeps its existing actual resource ownership. let mut workers = self .workers .lock() - .map_err(|_| "worker adapter registry lock is poisoned".to_string())?; - if workers.contains_key(worker_ref) { - return Err("Worker already has a retained execution handle".to_string()); - } + .unwrap_or_else(std::sync::PoisonError::into_inner); let workspace_id = workspace_client .as_ref() .and_then(|client| client.workspace_id().map(str::to_string)); - workers.insert( - worker_ref.clone(), - RuntimeWorkerExecution { - handle, - shutdown, - shutdown_requested: Arc::new(AtomicBool::new(false)), - tasks, - worker_state, - workspace_client, - workspace_id, - restore_operation_id, - workdir_attachments, - }, - ); - Ok(WorkerExecutionHandle::new( - worker_ref.clone(), - self.backend_id(), - )) + let candidate = RuntimeWorkerExecution { + handle, + shutdown, + shutdown_requested, + tasks, + worker_state, + workspace_client, + workspace_id, + restore_operation_id, + workdir_attachments, + }; + if let Some(existing) = workers.get(worker_ref) { + // Keep the actual endpoint, admission waiter, shutdown receiver and + // all task handles without replacing the registered Controller. + existing + .tasks + .cleanup_candidates + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(Arc::new(candidate)); + } else { + workers.insert(worker_ref.clone(), candidate); + } } fn cleanup_unconnected_controller( @@ -2099,36 +2742,19 @@ where handle: &WorkerHandle, shutdown: &Arc>>, tasks: &RuntimeExecutionTaskScope, - worker_state: &Arc>, + shutdown_requested: &Arc, + _worker_state: &Arc>, ) -> Result<(), String> { - let command = next_internal_command(worker_state)?; let handle = handle.clone(); + let shutdown_requested = Arc::clone(shutdown_requested); let shutdown = shutdown.clone(); let tasks_for_join = tasks.clone(); - let cleanup = self.run_on_adapter_runtime(async move { - handle - .send(Method::Shutdown { command }) - .await - .map_err(|error| format!("failed to request controller cleanup: {error}"))?; - let mut guard = shutdown.lock().await; - if let Some(mut receiver) = guard.take() { - match tokio::time::timeout(Duration::from_secs(5), &mut receiver).await { - Ok(Ok(())) => {} - Ok(Err(_)) => { - return Err("controller cleanup completion channel closed".to_string()); - } - Err(_) => { - *guard = Some(receiver); - return Err("controller cleanup confirmation timed out".to_string()); - } - } - } - drop(guard); - tasks_for_join.join().await + let cleanup = self.run_joined_on_adapter_runtime(async move { + tasks_for_join + .request_shutdown(&handle, &shutdown_requested, Duration::from_secs(5)) + .await?; + tasks_for_join.confirm_shutdown_and_join(&shutdown).await }); - if cleanup.is_err() { - tasks.abort_all(); - } cleanup } @@ -2143,6 +2769,36 @@ where restore_operation_id: Option, workdir_attachments: BTreeMap, workspace_client: Option>, + ) -> WorkerExecutionSpawnResult { + self.connect_controller_scope( + operation, + worker_ref, + bridge_context, + handle, + shutdown, + RuntimeExecutionTaskScope::new(controller_task), + Arc::new(AtomicBool::new(false)), + restore_operation_id, + workdir_attachments, + workspace_client, + ) + } + + // Own the actual Controller scope before doing any fallible connection work. + // The split also permits testing duplicate cleanup with an already-pending + // real admission waiter, rather than relying on a scheduler race. + fn connect_controller_scope( + &self, + operation: WorkerExecutionOperation, + worker_ref: WorkerRef, + bridge_context: crate::execution::WorkerExecutionContext, + handle: WorkerHandle, + shutdown: Arc>>, + tasks: RuntimeExecutionTaskScope, + shutdown_requested: Arc, + restore_operation_id: Option, + workdir_attachments: BTreeMap, + workspace_client: Option>, ) -> WorkerExecutionSpawnResult { #[cfg(feature = "ws-server")] let streams = subscribe_worker_protocol_session(&handle); @@ -2154,7 +2810,6 @@ where #[cfg(not(feature = "ws-server"))] let worker_state_snapshot = handle.shared_state.snapshot(); let worker_state = Arc::new(RwLock::new(worker_state_snapshot)); - let tasks = RuntimeExecutionTaskScope::new(controller_task); let workdir_attachment_statuses = workdir_attachments .iter() .map( @@ -2164,6 +2819,50 @@ where }, ) .collect::>(); + let mut workers = self + .workers + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if workers.contains_key(&worker_ref) { + drop(workers); + let cleanup = self.cleanup_unconnected_controller( + &handle, + &shutdown, + &tasks, + &shutdown_requested, + &worker_state, + ); + let result = WorkerExecutionResult::busy( + operation, + match &cleanup { + Ok(()) => "Worker is already connected to execution backend".to_string(), + Err(cleanup) => format!( + "Worker is already connected to execution backend; controller cleanup failed: {cleanup}" + ), + }, + ); + return match cleanup { + Ok(()) => WorkerExecutionSpawnResult::RolledBack(result), + Err(_) => { + self.retain_uncertain_unconnected_controller( + &worker_ref, + handle, + shutdown, + shutdown_requested, + tasks, + worker_state, + workspace_client, + restore_operation_id, + workdir_attachment_statuses.clone(), + ); + WorkerExecutionSpawnResult::ReconciliationRequired { + result, + worker_state: None, + workdir_attachments: workdir_attachment_statuses, + } + } + }; + } #[cfg(feature = "ws-server")] { let mut events = streams.events; @@ -2209,10 +2908,12 @@ where }) { Ok(task) => task, Err(message) => { + drop(workers); let cleanup = self.cleanup_unconnected_controller( &handle, &shutdown, &tasks, + &shutdown_requested, &worker_state, ); let result = WorkerExecutionResult::errored( @@ -2224,36 +2925,27 @@ where } }, ); - return if operation == WorkerExecutionOperation::Restore { - match cleanup { - Ok(()) => WorkerExecutionSpawnResult::RolledBack(result), - Err(_) => { - let retained_state = - worker_state.read().ok().map(|state| state.clone()); - let retained_workdir_attachments = - workdir_attachment_statuses.clone(); - let retained_handle = self - .retain_uncertain_unconnected_controller( - &worker_ref, - handle, - shutdown, - tasks, - worker_state, - workspace_client, - restore_operation_id, - retained_workdir_attachments.clone(), - ) - .ok(); - WorkerExecutionSpawnResult::ReconciliationRequired { - result, - handle: retained_handle, - worker_state: retained_state, - workdir_attachments: retained_workdir_attachments, - } + return match cleanup { + Ok(()) => WorkerExecutionSpawnResult::RolledBack(result), + Err(_) => { + let retained_state = worker_state.read().ok().map(|state| state.clone()); + self.retain_uncertain_unconnected_controller( + &worker_ref, + handle, + shutdown, + shutdown_requested, + tasks, + worker_state, + workspace_client, + restore_operation_id, + workdir_attachment_statuses.clone(), + ); + WorkerExecutionSpawnResult::ReconciliationRequired { + result, + worker_state: retained_state, + workdir_attachments: workdir_attachment_statuses, } } - } else { - WorkerExecutionSpawnResult::Errored(result) }; } }; @@ -2264,41 +2956,6 @@ where let _ = bridge_context; } - let mut workers = self - .workers - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if workers.contains_key(&worker_ref) { - let existing_worker_state = workers - .get(&worker_ref) - .and_then(|worker| worker.worker_state.read().ok().map(|state| state.clone())); - let existing_handle = WorkerExecutionHandle::new(worker_ref.clone(), self.backend_id()); - drop(workers); - let cleanup = - self.cleanup_unconnected_controller(&handle, &shutdown, &tasks, &worker_state); - let result = WorkerExecutionResult::busy( - operation, - match &cleanup { - Ok(()) => "Worker is already connected to execution backend".to_string(), - Err(cleanup) => format!( - "Worker is already connected to execution backend; controller cleanup failed: {cleanup}" - ), - }, - ); - return if operation == WorkerExecutionOperation::Restore { - match cleanup { - Ok(()) => WorkerExecutionSpawnResult::RolledBack(result), - Err(_) => WorkerExecutionSpawnResult::ReconciliationRequired { - result, - handle: Some(existing_handle), - worker_state: existing_worker_state, - workdir_attachments: Vec::new(), - }, - } - } else { - WorkerExecutionSpawnResult::Rejected(result) - }; - } let connected_worker_state = worker_state .read() .unwrap_or_else(std::sync::PoisonError::into_inner) @@ -2311,7 +2968,7 @@ where RuntimeWorkerExecution { handle, shutdown, - shutdown_requested: Arc::new(AtomicBool::new(false)), + shutdown_requested, tasks, worker_state, workspace_client, @@ -2322,11 +2979,176 @@ where ); WorkerExecutionSpawnResult::Connected { - handle: WorkerExecutionHandle::new(worker_ref, self.backend_id()), worker_state: connected_worker_state, workdir_attachments: workdir_attachment_statuses, } } + + // Called under the same per-Worker operation lock as connect/Stop. Each + // candidate remains in its actual owning scope until its own cleanup proves + // success. A failed child does not prevent joining the other Controllers. + fn cleanup_execution_resources( + &self, + execution: &RuntimeWorkerExecution, + ) -> Result<(), String> { + let candidates = execution + .tasks + .cleanup_candidates + .lock() + .map_err(|_| "execution cleanup candidate lock is poisoned".to_string())? + .clone(); + let mut first_failure = None; + for candidate in candidates { + match self.cleanup_execution_resources(&candidate) { + Ok(()) => { + execution + .tasks + .cleanup_candidates + .lock() + .map_err(|_| "execution cleanup candidate lock is poisoned".to_string())? + .retain(|retained| !Arc::ptr_eq(retained, &candidate)); + } + Err(message) => { + first_failure.get_or_insert(message); + } + } + } + let handle = execution.handle.clone(); + let shutdown_requested = execution.shutdown_requested.clone(); + let shutdown = execution.shutdown.clone(); + let tasks = execution.tasks.clone(); + let cleanup = self.run_joined_on_adapter_runtime(async move { + tasks + .request_shutdown(&handle, &shutdown_requested, Duration::from_secs(5)) + .await?; + tasks.confirm_shutdown_and_join(&shutdown).await?; + handle.delete_uncommitted_uploaded_files().map_err(|error| { + format!("uploaded_file_cleanup_failed: {error}; stop remains retryable") + }) + }); + if let Err(message) = cleanup { + first_failure.get_or_insert(message); + } + match first_failure { + Some(message) => Err(message), + None => Ok(()), + } + } + + // Caller owns the per-Worker operation lock for this entire cleanup. + fn stop_worker_under_lock(&self, worker_ref: &WorkerRef) -> WorkerExecutionResult { + let pending = match self.pending_restore(worker_ref) { + Ok(pending) => pending, + Err(message) => { + return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message); + } + }; + if let Some(pending) = pending { + if let Err(message) = self.wait_pending_restore(&pending, Duration::from_secs(5)) { + return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message); + } + if self.workers.lock().is_err() { + return WorkerExecutionResult::errored( + WorkerExecutionOperation::Stop, + "worker adapter registry lock is poisoned", + ); + } + let mut retained = match pending.result.lock() { + Ok(retained) => retained, + Err(_) => { + return WorkerExecutionResult::errored( + WorkerExecutionOperation::Stop, + "restore result lock is poisoned", + ); + } + }; + match retained.as_ref() { + Some(Ok(_)) => { + let controller = match retained.take().unwrap() { + Ok(controller) => controller, + Err(_) => unreachable!(), + }; + let state = Arc::new(RwLock::new(controller.handle.shared_state.snapshot())); + self.retain_uncertain_unconnected_controller( + worker_ref, + controller.handle, + controller.shutdown, + Arc::new(AtomicBool::new(false)), + RuntimeExecutionTaskScope::new(controller.controller_task), + state, + Some(controller.workspace_client), + Some(pending.request.operation_id), + Vec::new(), + ); + } + Some(Err(_)) => { + let factory = Arc::clone(&self.factory); + let request = pending.request.clone(); + if let Err(message) = self + .run_cancellable_on_adapter_runtime(RUNTIME_TASK_TIMEOUT, async move { + factory.cleanup_failed_restore(&request).await + }) + { + return WorkerExecutionResult::errored( + WorkerExecutionOperation::Stop, + message, + ); + } + } + None => { + return WorkerExecutionResult::errored( + WorkerExecutionOperation::Stop, + "restore factory has no completion evidence; cleanup is unproven", + ); + } + } + drop(retained); + if let Err(message) = self.remove_pending_restore(worker_ref) { + return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message); + } + } + let execution = match self.workers.lock() { + Ok(workers) => workers.get(worker_ref).cloned(), + Err(_) => { + return WorkerExecutionResult::errored( + WorkerExecutionOperation::Stop, + "worker adapter registry lock is poisoned", + ); + } + }; + let Some(execution) = execution else { + // The execution backend cleanup may have committed before the + // Runtime catalog commit failed. Treat the retry as converged so + // the Runtime can durably finish its Stopped transition. + let factory = Arc::clone(&self.factory); + let worker_ref = worker_ref.clone(); + return match self.run_cancellable_on_adapter_runtime(RUNTIME_TASK_TIMEOUT, async move { + factory.reconcile_stopped_worker(&worker_ref).await + }) { + Ok(()) => WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop), + Err(message) => { + WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message) + } + }; + }; + + if let Err(message) = self.cleanup_execution_resources(&execution) { + return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message); + } + + match self.workers.lock() { + Ok(mut workers) => { + workers.remove(worker_ref); + } + Err(_) => { + return WorkerExecutionResult::errored( + WorkerExecutionOperation::Stop, + "worker adapter registry lock is poisoned; cleanup commit remains retryable", + ); + } + } + WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop) + } } impl Drop for WorkerRuntimeExecutionBackend { @@ -2376,6 +3198,27 @@ where &self, request: WorkerSessionObservationRequest, ) -> runtime_api::WorkerSessionAvailability { + let unavailable = |message: String| runtime_api::WorkerSessionAvailability::Unavailable { + reason: runtime_api::WorkerSessionUnavailableReason::StorageUnavailable, + message, + }; + let operation_lock = match self.worker_lock(&request.worker_ref) { + Ok(lock) => lock, + Err(message) => return unavailable(message), + }; + let _operation_guard = match operation_lock.lock() { + Ok(guard) => guard, + Err(_) => return unavailable("worker operation lock is poisoned".to_string()), + }; + let workers = match self.workers.lock() { + Ok(workers) => workers, + Err(_) => return unavailable("worker adapter registry lock is poisoned".to_string()), + }; + if workers.contains_key(&request.worker_ref) { + return runtime_api::WorkerSessionAvailability::LiveProtocol; + } + drop(workers); + match self.factory.retained_session_snapshot(&request.worker_ref) { Ok(retained) => runtime_api::WorkerSessionAvailability::RetainedSnapshot { identity: runtime_api::RetainedSessionIdentity { @@ -2461,6 +3304,12 @@ where request: crate::execution::WorkerSessionAttachmentRequest, ) -> Result { + let operation_lock = self + .worker_lock(&request.worker_ref) + .map_err(|_| session_store::RetainedAttachmentReadError::StorageUnavailable)?; + let _operation_guard = operation_lock + .lock() + .map_err(|_| session_store::RetainedAttachmentReadError::StorageUnavailable)?; let live = self .workers .lock() @@ -2596,21 +3445,56 @@ where } fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { - let _start_guard = self - .restore_lock - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if self - .workers - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .contains_key(&request.worker_ref) - { + let operation_lock = match self.worker_lock(&request.worker_ref) { + Ok(lock) => lock, + Err(message) => { + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::errored( + WorkerExecutionOperation::Spawn, + message, + )); + } + }; + let _operation_guard = match operation_lock.lock() { + Ok(guard) => guard, + Err(_) => { + let message = "worker operation lock is poisoned".to_string(); + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::errored( + WorkerExecutionOperation::Spawn, + message, + )); + } + }; + match self.pending_restore(&request.worker_ref) { + Ok(None) => {} + Ok(Some(_)) => { + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::busy( + WorkerExecutionOperation::Spawn, + "Worker restore still owns pending resources", + )); + } + Err(message) => { + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::errored( + WorkerExecutionOperation::Spawn, + message, + )); + } + } + let workers = match self.workers.lock() { + Ok(workers) => workers, + Err(_) => { + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::errored( + WorkerExecutionOperation::Spawn, + "worker adapter registry lock is poisoned", + )); + } + }; + if workers.contains_key(&request.worker_ref) { return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::busy( WorkerExecutionOperation::Spawn, "Worker is already connected to execution backend", )); } + drop(workers); let mut request = request; let Some(materializer) = self.working_directory_materializer.as_ref().or_else(|| { @@ -2786,16 +3670,69 @@ where &self, request: &WorkerExecutionRestoreRequest, ) -> Result<(), WorkerExecutionResult> { - let workers = self - .workers - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if workers.contains_key(&request.worker_ref) { + let operation_lock = match self.worker_lock(&request.worker_ref) { + Ok(lock) => lock, + Err(message) => { + return Err(WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + message, + )); + } + }; + let _operation_guard = match operation_lock.lock() { + Ok(guard) => guard, + Err(_) => { + let message = "worker operation lock is poisoned".to_string(); + return Err(WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + message, + )); + } + }; + if self + .pending_restore(&request.worker_ref) + .map_err(|message| { + WorkerExecutionResult::errored(WorkerExecutionOperation::Restore, message) + })? + .is_some() + { return Err(WorkerExecutionResult::busy( WorkerExecutionOperation::Restore, - "Worker is already connected to execution backend", + "Worker restore still owns pending resources", )); } + let workers = self.workers.lock().map_err(|_| { + WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + "worker adapter registry lock is poisoned", + ) + })?; + if let Some(existing) = workers.get(&request.worker_ref) { + let result = WorkerExecutionResult::busy( + WorkerExecutionOperation::Restore, + "Worker is already connected to execution backend", + ); + if existing.shutdown_requested.load(Ordering::Acquire) + || existing.tasks.has_cleanup_candidates().map_err(|message| { + WorkerExecutionResult::errored(WorkerExecutionOperation::Restore, message) + })? + { + return Err(result); + } + let snapshot = existing + .worker_state + .read() + .map_err(|_| { + WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + "worker state lock is poisoned", + ) + })? + .clone(); + // Read-only AlreadyConnected evidence. The Runtime may accept this + // exact restore without journaling or replacing the current owner. + return Err(result.with_worker_state(snapshot)); + } drop(workers); if !request.previous_workdir_attachments.is_empty() { @@ -2878,18 +3815,51 @@ where &self, mut request: WorkerExecutionRestoreRequest, ) -> WorkerExecutionSpawnResult { - // Serialize backend restore starts so duplicate detection is completed - // before any controller or bridge side effect. Runtime also serializes - // per Worker; this guard protects direct/concurrent backend callers. - let _restore_guard = self - .restore_lock - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let workers = self - .workers - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let operation_lock = match self.worker_lock(&request.worker_ref) { + Ok(lock) => lock, + Err(message) => { + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + message, + )); + } + }; + let _operation_guard = match operation_lock.lock() { + Ok(guard) => guard, + Err(_) => { + let message = "worker operation lock is poisoned".to_string(); + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + message, + )); + } + }; + let workers = self.workers.lock().map_err(|_| { + WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + "worker adapter registry lock is poisoned", + ) + }); + let workers = match workers { + Ok(workers) => workers, + Err(result) => return WorkerExecutionSpawnResult::Rejected(result), + }; if let Some(existing) = workers.get(&request.worker_ref) { + let pending_cleanup = existing.tasks.has_cleanup_candidates(); + if existing.shutdown_requested.load(Ordering::Acquire) + || !matches!(pending_cleanup.as_ref(), Ok(false)) + { + return WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + pending_cleanup.err().unwrap_or_else(|| { + "Worker still owns unproven Controller cleanup".to_string() + }), + ), + worker_state: None, + workdir_attachments: existing.workdir_attachments.clone(), + }; + } if existing.restore_operation_id == Some(request.operation_id) { let worker_state = existing .worker_state @@ -2897,10 +3867,6 @@ where .unwrap_or_else(std::sync::PoisonError::into_inner) .clone(); return WorkerExecutionSpawnResult::Connected { - handle: WorkerExecutionHandle::new( - request.worker_ref.clone(), - self.backend_id(), - ), worker_state, workdir_attachments: existing.workdir_attachments.clone(), }; @@ -2995,85 +3961,212 @@ where } request.workdir_attachments = workdir_attachments.clone(); - let factory = self.factory.clone(); - let bridge_context = request.context.clone(); let worker_ref = request.worker_ref.clone(); let operation_id = request.operation_id; - let restore_result = self - .run_cancellable_on_adapter_runtime(self.spawn_restore_timeout, async move { - factory.restore_controller(request).await - }); - - let controller = match restore_result { - Ok(controller) => controller, + let pending = match self.pending_restore(&worker_ref) { + Ok(Some(pending)) if pending.request.operation_id == operation_id => pending, + Ok(Some(_)) => { + return WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::busy( + WorkerExecutionOperation::Restore, + "another restore operation still owns pending resources", + ), + worker_state: None, + workdir_attachments: Vec::new(), + }; + } + Ok(None) => { + let factory = Arc::clone(&self.factory); + let preflight_request = request.clone(); + if let Err(message) = self.run_on_adapter_runtime(async move { + factory.preflight_restore(&preflight_request).await + }) { + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + message, + )); + } + let pending = Arc::new(PendingRuntimeRestore { + request: request.clone(), + task: tokio::sync::Mutex::new(None), + result: Mutex::new(None), + failure: Mutex::new(None), + }); + let mut resources = match self.pending_restores.lock() { + Ok(resources) => resources, + Err(_) => { + return WorkerExecutionSpawnResult::Rejected( + WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + "pending restore resources lock is poisoned", + ), + ); + } + }; + resources.insert(worker_ref.clone(), Arc::clone(&pending)); + let factory = Arc::clone(&self.factory); + let result_owner = Arc::clone(&pending); + let task = self.spawn_on_adapter_runtime(async move { + let result = factory.restore_controller(request).await; + let mut retained = result_owner + .result + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + *retained = Some(result); + }); + match task { + Ok(task) => { + *pending.task.try_lock().expect("new restore task lock") = Some(task) + } + Err(message) => { + resources.remove(&worker_ref); + return WorkerExecutionSpawnResult::Rejected( + WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + message, + ), + ); + } + } + drop(resources); + pending + } Err(message) => { return WorkerExecutionSpawnResult::ReconciliationRequired { result: WorkerExecutionResult::errored( WorkerExecutionOperation::Restore, message, ), - handle: None, worker_state: None, workdir_attachments: Vec::new(), }; } }; + if let Err(message) = self.wait_pending_restore(&pending, self.spawn_restore_timeout) { + return WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::errored(WorkerExecutionOperation::Restore, message), + worker_state: None, + workdir_attachments: Vec::new(), + }; + } + let controller = { + let mut retained = match pending.result.lock() { + Ok(retained) => retained, + Err(_) => { + return WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + "restore result lock is poisoned", + ), + worker_state: None, + workdir_attachments: Vec::new(), + }; + } + }; + match retained.as_ref() { + Some(Ok(_)) => match retained.take().unwrap() { + Ok(controller) => controller, + Err(_) => unreachable!(), + }, + Some(Err(message)) => { + return WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + message.clone(), + ), + worker_state: None, + workdir_attachments: Vec::new(), + }; + } + None => { + return WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + "restore factory completion has no result; cleanup is unproven", + ), + worker_state: None, + workdir_attachments: Vec::new(), + }; + } + } + }; - self.connect_handle( + let result = self.connect_handle( WorkerExecutionOperation::Restore, - worker_ref, - bridge_context, + worker_ref.clone(), + pending.request.context.clone(), controller.handle, controller.shutdown, controller.controller_task, Some(operation_id), - workdir_attachments, + pending.request.workdir_attachments.clone(), Some(controller.workspace_client), - ) + ); + if let Err(message) = self.remove_pending_restore(&worker_ref) { + return WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::errored(WorkerExecutionOperation::Restore, message), + worker_state: None, + workdir_attachments: Vec::new(), + }; + } + result } fn activate_restored_worker( &self, operation_id: crate::execution::WorkerLifecycleOperationId, - handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, ) -> Result<(), String> { - if handle.backend_id() != self.backend_id() { - return Err(format!( - "execution handle belongs to backend {}, not {}", - handle.backend_id(), - self.backend_id() - )); - } + let operation_lock = match self.worker_lock(worker_ref) { + Ok(lock) => lock, + Err(message) => return Err(message), + }; + let _operation_guard = match operation_lock.lock() { + Ok(guard) => guard, + Err(_) => { + let message = "worker operation lock is poisoned".to_string(); + return Err(message); + } + }; let workers = self .workers .lock() .map_err(|_| "worker adapter registry lock is poisoned".to_string())?; let execution = workers - .get(handle.worker_ref()) + .get(worker_ref) .ok_or_else(|| "restored Worker execution is not registered".to_string())?; if execution.restore_operation_id != Some(operation_id) { return Err("restored Worker operation identity does not match".to_string()); } self.factory.activate_restored_controller( - handle.worker_ref(), + worker_ref, execution.workspace_id.as_deref(), &execution.handle, ); Ok(()) } - fn dispatch_input( - &self, - handle: &WorkerExecutionHandle, - input: WorkerInput, - ) -> WorkerExecutionResult { + fn dispatch_input(&self, worker_ref: &WorkerRef, input: WorkerInput) -> WorkerExecutionResult { + let operation_lock = match self.worker_lock(worker_ref) { + Ok(lock) => lock, + Err(message) => { + return WorkerExecutionResult::errored(WorkerExecutionOperation::Input, message); + } + }; + let _operation_guard = match operation_lock.lock() { + Ok(guard) => guard, + Err(_) => { + let message = "worker operation lock is poisoned".to_string(); + return WorkerExecutionResult::errored(WorkerExecutionOperation::Input, message); + } + }; if let Err(error) = crate::runtime::validate_worker_input(&input) { return WorkerExecutionResult::rejected( WorkerExecutionOperation::Input, error.to_string(), ); } - let (worker, worker_state, _workspace_client) = match self.get_execution(handle) { + let (worker, worker_state, _workspace_client) = match self.get_execution(worker_ref) { Ok(execution) => execution, Err(mut result) => { result.operation = WorkerExecutionOperation::Input; @@ -3173,13 +4266,32 @@ where fn upload_file( &self, - handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, file_name: &str, media_type: &str, content: &[u8], context: Option<&session_store::UploadedFileUploadContext>, ) -> Result { - let (worker, _, _) = self.get_execution(handle).map_err(|mut result| { + let operation_lock = match self.worker_lock(worker_ref) { + Ok(lock) => lock, + Err(message) => { + return Err(WorkerExecutionResult::errored( + WorkerExecutionOperation::UploadFile, + message, + )); + } + }; + let _operation_guard = match operation_lock.lock() { + Ok(guard) => guard, + Err(_) => { + let message = "worker operation lock is poisoned".to_string(); + return Err(WorkerExecutionResult::errored( + WorkerExecutionOperation::UploadFile, + message, + )); + } + }; + let (worker, _, _) = self.get_execution(worker_ref).map_err(|mut result| { result.operation = WorkerExecutionOperation::UploadFile; result })?; @@ -3199,10 +4311,29 @@ where fn delete_uploaded_file( &self, - handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, artifact_id: &str, ) -> WorkerExecutionResult { - let (worker, _, _) = match self.get_execution(handle) { + let operation_lock = match self.worker_lock(worker_ref) { + Ok(lock) => lock, + Err(message) => { + return WorkerExecutionResult::errored( + WorkerExecutionOperation::DeleteUploadedFile, + message, + ); + } + }; + let _operation_guard = match operation_lock.lock() { + Ok(guard) => guard, + Err(_) => { + let message = "worker operation lock is poisoned".to_string(); + return WorkerExecutionResult::errored( + WorkerExecutionOperation::DeleteUploadedFile, + message, + ); + } + }; + let (worker, _, _) = match self.get_execution(worker_ref) { Ok(execution) => execution, Err(mut result) => { result.operation = WorkerExecutionOperation::DeleteUploadedFile; @@ -3218,12 +4349,27 @@ where } } - fn dispatch_method( - &self, - handle: &WorkerExecutionHandle, - method: Method, - ) -> WorkerExecutionResult { - let (worker, _worker_state, _workspace_client) = match self.get_execution(handle) { + fn dispatch_method(&self, worker_ref: &WorkerRef, method: Method) -> WorkerExecutionResult { + let operation_lock = match self.worker_lock(worker_ref) { + Ok(lock) => lock, + Err(message) => { + return WorkerExecutionResult::errored( + WorkerExecutionOperation::ProtocolMethod, + message, + ); + } + }; + let _operation_guard = match operation_lock.lock() { + Ok(guard) => guard, + Err(_) => { + let message = "worker operation lock is poisoned".to_string(); + return WorkerExecutionResult::errored( + WorkerExecutionOperation::ProtocolMethod, + message, + ); + } + }; + let (worker, _worker_state, _workspace_client) = match self.get_execution(worker_ref) { Ok(execution) => execution, Err(mut result) => { result.operation = WorkerExecutionOperation::ProtocolMethod; @@ -3235,109 +4381,41 @@ where } fn stop_worker_operation(&self, request: WorkerExecutionStopRequest) -> WorkerExecutionResult { - let handle = request - .handle - .unwrap_or_else(|| WorkerExecutionHandle::new(request.worker_ref, self.backend_id())); - self.stop_worker(&handle) + self.stop_worker(&request.worker_ref) } - fn stop_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult { - if handle.backend_id() != self.backend_id() { - return WorkerExecutionResult::rejected( - WorkerExecutionOperation::Stop, - format!( - "execution handle belongs to backend {}, not {}", - handle.backend_id(), - self.backend_id() - ), - ); - } - let execution = match self.workers.lock() { - Ok(workers) => workers.get(handle.worker_ref()).cloned(), + fn stop_worker(&self, worker_ref: &WorkerRef) -> WorkerExecutionResult { + let operation_lock = match self.worker_lock(worker_ref) { + Ok(lock) => lock, + Err(message) => { + return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message); + } + }; + let _operation_guard = match operation_lock.lock() { + Ok(guard) => guard, Err(_) => { - return WorkerExecutionResult::errored( - WorkerExecutionOperation::Stop, - "worker adapter registry lock is poisoned", - ); + let message = "worker operation lock is poisoned".to_string(); + return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message); } }; - let Some(execution) = execution else { - // The execution backend cleanup may have committed before the - // Runtime catalog commit failed. Treat the retry as converged so - // the Runtime can durably finish its Stopped transition. - return WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop); - }; - - let first_request = !execution.shutdown_requested.swap(true, Ordering::AcqRel); - let result = if first_request { - let command = match next_internal_command(&execution.worker_state) { - Ok(command) => command, - Err(error) => { - execution.shutdown_requested.store(false, Ordering::Release); - return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, error); - } - }; - let result = self.send_method( - WorkerExecutionOperation::Stop, - execution.handle.clone(), - Method::Shutdown { command }, - ); - if result.outcome != crate::execution::WorkerExecutionOutcome::Accepted { - execution.shutdown_requested.store(false, Ordering::Release); - return result; - } - result - } else { - WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop) - }; - - let shutdown = execution.shutdown.clone(); - let tasks = execution.tasks.clone(); - let shutdown_wait = self.run_on_adapter_runtime(async move { - { - let mut guard = shutdown.lock().await; - if let Some(mut receiver) = guard.take() { - match tokio::time::timeout(Duration::from_secs(5), &mut receiver).await { - Ok(Ok(())) => {} - Ok(Err(_)) => { - return Err("Worker shutdown completion channel closed".to_string()); - } - Err(_) => { - *guard = Some(receiver); - return Err( - "Worker shutdown confirmation timed out; stop remains retryable" - .to_string(), - ); - } - } - } - } - tasks.join().await - }); - if let Err(message) = shutdown_wait { - return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message); - } - - if let Err(error) = execution.handle.delete_uncommitted_uploaded_files() { - return WorkerExecutionResult::errored( - WorkerExecutionOperation::Stop, - format!("uploaded_file_cleanup_failed: {error}; stop remains retryable"), - ); - } - - match self.workers.lock() { - Ok(mut workers) => { - workers.remove(handle.worker_ref()); - } - Err(poisoned) => { - poisoned.into_inner().remove(handle.worker_ref()); - } - } - result + self.stop_worker_under_lock(worker_ref) } - fn cancel_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult { - let (worker, worker_state, _workspace_client) = match self.get_execution(handle) { + fn cancel_worker(&self, worker_ref: &WorkerRef) -> WorkerExecutionResult { + let operation_lock = match self.worker_lock(worker_ref) { + Ok(lock) => lock, + Err(message) => { + return WorkerExecutionResult::errored(WorkerExecutionOperation::Cancel, message); + } + }; + let _operation_guard = match operation_lock.lock() { + Ok(guard) => guard, + Err(_) => { + let message = "worker operation lock is poisoned".to_string(); + return WorkerExecutionResult::errored(WorkerExecutionOperation::Cancel, message); + } + }; + let (worker, worker_state, _workspace_client) = match self.get_execution(worker_ref) { Ok(execution) => execution, Err(mut result) => { result.operation = WorkerExecutionOperation::Cancel; @@ -3358,31 +4436,157 @@ where } #[cfg(feature = "ws-server")] - fn worker_snapshot(&self, handle: &WorkerExecutionHandle) -> Option { - if handle.backend_id() != self.backend_id() { - return None; - } + fn attach_worker_protocol( + self: Arc, + worker_ref: &WorkerRef, + ) -> Result { + let operation_lock = self.worker_lock(worker_ref).map_err(|message| { + WorkerExecutionResult::errored(WorkerExecutionOperation::ProtocolMethod, message) + })?; + let _operation_guard = operation_lock.lock().map_err(|_| { + WorkerExecutionResult::errored( + WorkerExecutionOperation::ProtocolMethod, + "worker operation lock is poisoned", + ) + })?; + let execution = self + .workers + .lock() + .map_err(|_| { + WorkerExecutionResult::errored( + WorkerExecutionOperation::ProtocolMethod, + "worker adapter registry lock is poisoned", + ) + })? + .get(worker_ref) + .cloned() + .ok_or_else(|| { + WorkerExecutionResult::rejected( + WorkerExecutionOperation::ProtocolMethod, + "Worker has no live protocol endpoint", + ) + })?; + self.validate_protocol_execution_under_lock(worker_ref, &execution)?; + let streams = subscribe_worker_protocol_session(&execution.handle); + let snapshot = streams.snapshot_event; + let (tx, events) = tokio::sync::mpsc::channel(128); + let relay_handle = execution.handle.clone(); + let task = self + .spawn_on_adapter_runtime(async move { + for alert in streams.alert_snapshot { + tokio::select! { + biased; + _ = relay_handle.protocol_closed() => return, + sent = tx.send(Event::Alert(alert)) => if sent.is_err() { return; }, + } + } + let mut protocol_events = streams.events; + let mut log_entries = streams.log_entries; + loop { + let event = tokio::select! { + biased; + _ = relay_handle.protocol_closed() => break, + _ = tx.closed() => break, + event = protocol_events.recv() => match event { + Ok(event) => Some(event), + Err(_) => break, + }, + entry = log_entries.recv() => match entry { + Ok(entry) => live_log_entry_event(entry), + Err(_) => break, + }, + }; + if let Some(event) = event { + tokio::select! { + biased; + _ = relay_handle.protocol_closed() => break, + sent = tx.send(event) => if sent.is_err() { break; }, + } + } + } + }) + .map_err(|message| { + WorkerExecutionResult::errored(WorkerExecutionOperation::ProtocolMethod, message) + })?; + // Stop owns the Controller and every attached relay in this same scope. + // Its abort/join barrier closes the bounded stream before cleanup succeeds. + execution.tasks.push("protocol transport", task, true); + + let dispatcher_backend = Arc::clone(&self); + let dispatcher_lock = Arc::clone(&operation_lock); + let dispatcher_execution = execution.clone(); + let dispatcher_worker_ref = worker_ref.clone(); + let dispatch = Arc::new(move |method| { + let _guard = dispatcher_lock.lock().map_err(|_| { + WorkerExecutionResult::errored( + WorkerExecutionOperation::ProtocolMethod, + "worker operation lock is poisoned", + ) + })?; + dispatcher_backend.dispatch_protocol_execution_under_lock( + &dispatcher_worker_ref, + &dispatcher_execution, + method, + ) + }); + let validator_backend = Arc::clone(&self); + let validator_lock = Arc::clone(&operation_lock); + let validator_worker_ref = worker_ref.clone(); + let validate = Arc::new(move || { + let _guard = validator_lock.lock().map_err(|_| { + WorkerExecutionResult::errored( + WorkerExecutionOperation::ProtocolMethod, + "worker operation lock is poisoned", + ) + })?; + validator_backend + .validate_protocol_execution_under_lock(&validator_worker_ref, &execution) + }); + Ok(WorkerProtocolTransport::new( + worker_ref.clone(), + snapshot, + events, + dispatch, + validate, + )) + } + + #[cfg(feature = "ws-server")] + fn worker_snapshot(&self, worker_ref: &WorkerRef) -> Option { + let operation_lock = match self.worker_lock(worker_ref) { + Ok(lock) => lock, + Err(_) => return None, + }; + let _operation_guard = match operation_lock.lock() { + Ok(guard) => guard, + Err(_) => return None, + }; let workers = self.workers.lock().ok()?; workers - .get(handle.worker_ref()) + .get(worker_ref) .map(|execution| execution.handle.snapshot_event()) } fn worker_completions( &self, - handle: &WorkerExecutionHandle, + worker_ref: &WorkerRef, kind: protocol::CompletionKind, prefix: &str, context: Option<&protocol::CompletionContext>, ) -> Vec { - if handle.backend_id() != self.backend_id() { - return Vec::new(); - } + let operation_lock = match self.worker_lock(worker_ref) { + Ok(lock) => lock, + Err(_) => return Vec::new(), + }; + let _operation_guard = match operation_lock.lock() { + Ok(guard) => guard, + Err(_) => return Vec::new(), + }; let Ok(workers) = self.workers.lock() else { return Vec::new(); }; workers - .get(handle.worker_ref()) + .get(worker_ref) .map(|execution| { futures::executor::block_on( execution.handle.completion_entries(kind, prefix, context), @@ -3412,6 +4616,7 @@ mod tests { use crate::identity::WorkerId; use crate::identity::WorkerRef; use crate::management::RuntimeOptions; + #[cfg(feature = "ws-server")] use crate::observation::WorkerObservationCursor; use crate::working_directory::RuntimeGitMaterializer; use agen::Engine; @@ -3512,6 +4717,1419 @@ mod tests { (session_id, segment_id, expected_entry_id) } + const SAVED_SUBJEKTIV_SYSTEM_PROMPT: &str = "Committed saved Subject Worker system prompt"; + + async fn saved_subjektiv_restore_fixture( + root: &Path, + materialize_head: bool, + ) -> WorkerExecutionRestoreRequest { + let worker_ref = WorkerRef::new(WorkerId::now_v7()); + let worker_name = ProfileRuntimeWorkerFactory::runtime_worker_name_for_ref(&worker_ref); + let aggregate_dir = root.join("workers").join(worker_ref.worker_id.to_string()); + let session_store = WorkerSessionStore::new(aggregate_dir.join("session")).unwrap(); + let store = WorkerAggregateStore::new(&aggregate_dir, &worker_name).unwrap(); + let mut saved = WorkerManifest::from_toml(&format!( + r#" + [worker] + name = "{worker_name}" + [model] + scheme = "anthropic" + model_id = "saved-test-model" + auth = {{ kind = "none" }} + [engine] + instruction = "default" + max_tokens = 100 + [[scope.allow]] + target = "{}" + permission = "read" + recursive = true + symlink_policy = "resolved" + [feature.memory.profile] + enabled = false + [feature.subjektiv.profile] + enabled = true + [feature.subjektiv.profile.extraction] + enabled = false + "#, + root.display() + )) + .unwrap(); + Scope::from_config(&saved.scope) + .expect("saved fixture must declare a valid resolved scope"); + let mut request = create_request("saved subject policy"); + request.worker_id = worker_ref.worker_id; + request.workspace_api = Some(WorkspaceApiRef { + workspace_id: "workspace-saved-subjektiv".to_string(), + base_url: "http://workspace.invalid".to_string(), + }); + request.subjektiv_attached = true; + request.memory_settings = Some(manifest::WorkspaceMemorySettingsSnapshot { + workspace_id: "workspace-saved-subjektiv".to_string(), + settings_revision: 17, + language: "English".to_string(), + }); + bind_workspace_memory_settings(&mut saved, &request).unwrap(); + let client = MockClient::new(Vec::new()); + let mut engine = + Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated( + client.clone(), + ); + engine.set_system_prompt(SAVED_SUBJEKTIV_SYSTEM_PROMPT); + let mut worker = Worker::new( + saved.clone(), + engine, + CombinedStore::new(session_store.clone(), store.clone()), + WorkerWorkspaceContext::unavailable( + Some(WorkspaceId::new("workspace-saved-subjektiv").unwrap()), + "saved session fixture has no network client", + ), + WorkerFilesystemAuthority::None, + Scope::empty(), + ) + .await + .unwrap(); + worker.enable_worker_metadata_write_through().unwrap(); + let session_id = worker.session_id(); + let segment_id = worker.segment_id(); + if materialize_head { + // Use the Worker's normal typed startup writer, not the retained + // public-history fixture (which intentionally has no system prompt). + worker.materialize_durable_session_head().await.unwrap(); + let entries = session_store.read_all(session_id, segment_id).unwrap(); + assert!( + matches!(entries.as_slice(), [LogEntry::AnnotatedSegmentStart { + session_id: committed_session_id, + system_prompt: Some(prompt), + history, + forked_from: None, + compacted_from: None, + .. + }] if *committed_session_id == session_id + && prompt == SAVED_SUBJEKTIV_SYSTEM_PROMPT + && history.is_empty()) + ); + assert_eq!( + session_store::collect_state(&entries) + .system_prompt + .as_deref(), + Some(SAVED_SUBJEKTIV_SYSTEM_PROMPT) + ); + } else { + assert!(matches!( + session_store.exists(session_id, segment_id), + Err(session_store::StoreError::Corrupt { message, .. }) + if message.contains("no materialized Session") + )); + } + assert_eq!(client.call_count.load(Ordering::SeqCst), 0); + drop(worker); + let mut metadata = store.read_by_name(&worker_name).unwrap().unwrap(); + let (manifest, _) = ProfileRuntimeWorkerFactory::manifest_for_restore(&metadata, &request) + .expect("fixture must persist valid saved Manifest authority"); + assert_eq!( + manifest::write_persisted_worker_manifest_snapshot(&manifest).unwrap(), + manifest::write_persisted_worker_manifest_snapshot(&saved).unwrap() + ); + assert_eq!( + metadata.active.as_ref().unwrap().session_id, + session_id, + "metadata write-through must preserve the real Worker Session" + ); + assert_eq!( + metadata.active.as_ref().unwrap().segment_id, + materialize_head.then_some(segment_id) + ); + assert_eq!( + metadata.workspace_id.as_deref(), + Some("workspace-saved-subjektiv") + ); + // Attribution already committed before restart; the real restore needs + // no network fixture or replacement current Profile. + metadata.subjektiv_session_attribution = + Some(session_store::SubjektivSessionAttributionState::Confirmed { + session_id, + subject_id: "saved-subject".to_string(), + }); + store.write(&metadata).unwrap(); + WorkerExecutionRestoreRequest { + operation_id: crate::execution::WorkerLifecycleOperationId::new(), + worker_ref: worker_ref.clone(), + request, + workspace_scope: Some(crate::runtime::RuntimeWorkspaceScope::new( + "workspace-saved-subjektiv", + "server-main", + )), + context: test_execution_context(worker_ref), + previous_workdir_attachments: Vec::new(), + logical_workdir_attachments: Vec::new(), + workdir_attachments: BTreeMap::new(), + config_bundle: None, + } + } + + async fn shutdown_profile_controller(controller: RuntimeWorkerController) { + let state = Arc::new(RwLock::new(controller.handle.shared_state.snapshot())); + controller + .handle + .send(Method::Shutdown { + command: next_internal_command(&state).unwrap(), + }) + .await + .unwrap(); + if let Some(receiver) = controller.shutdown.lock().await.take() { + tokio::time::timeout(Duration::from_secs(5), receiver) + .await + .unwrap() + .unwrap(); + } + tokio::time::timeout(Duration::from_secs(5), controller.controller_task) + .await + .unwrap() + .unwrap(); + } + + #[test] + #[serial_test::serial(worker_allocation)] + fn profile_backend_existing_restore_attestation_is_read_only_under_concurrency() { + let root = tempfile::tempdir().unwrap(); + let runtime_store = root.path().join("runtime"); + let backend = Arc::new( + WorkerRuntimeExecutionBackend::new( + ProfileRuntimeWorkerFactory::new(root.path()) + .with_runtime_store_dir(&runtime_store) + .with_runtime_id("runtime-subjektiv-test"), + ) + .unwrap(), + ); + let fixture_root = runtime_store.clone(); + let request = backend + .run_on_adapter_runtime(async move { + Ok(saved_subjektiv_restore_fixture(&fixture_root, true).await) + }) + .unwrap(); + backend.preflight_restore(&request).unwrap(); + let expected = match backend.restore_worker(request.clone()) { + WorkerExecutionSpawnResult::Connected { worker_state, .. } => worker_state, + other => panic!("real Profile backend restore failed: {other:?}"), + }; + assert!(matches!( + backend.worker_session(WorkerSessionObservationRequest { + worker_ref: request.worker_ref.clone(), + }), + runtime_api::WorkerSessionAvailability::LiveProtocol + )); + let current = backend + .workers + .lock() + .unwrap() + .get(&request.worker_ref) + .unwrap() + .clone(); + let authority_files = || { + let mut files = persisted_files(&runtime_store); + files.retain(|path, _| { + !path + .components() + .any(|component| component.as_os_str() == "runs") + }); + files + }; + let files_before = authority_files(); + let runs = backend + .factory + .worker_aggregate_dir(&request.worker_ref) + .unwrap() + .join("runs"); + let run_directories = || { + fs::read_dir(&runs) + .unwrap() + .map(|entry| entry.unwrap().path()) + .collect::>() + }; + let runs_before = run_directories(); + std::thread::scope(|threads| { + for _ in 0..8 { + let backend = Arc::clone(&backend); + let request = request.clone(); + let expected = expected.clone(); + threads.spawn(move || { + for _ in 0..2 { + let connected = backend.preflight_restore(&request).unwrap_err(); + assert_eq!( + connected.outcome, + crate::execution::WorkerExecutionOutcome::Busy + ); + assert_eq!(connected.worker_state, Some(expected.clone())); + match backend.reconcile_restore(request.clone()) { + WorkerExecutionSpawnResult::Connected { worker_state, .. } => { + assert_eq!(worker_state, expected) + } + other => { + panic!("repeated exact operation was not connected: {other:?}") + } + } + } + }); + } + }); + let mut repeated_request = request.clone(); + repeated_request.operation_id = crate::execution::WorkerLifecycleOperationId::new(); + let connected = backend.preflight_restore(&repeated_request).unwrap_err(); + assert_eq!( + connected.outcome, + crate::execution::WorkerExecutionOutcome::Busy + ); + assert_eq!(connected.worker_state, Some(expected)); + let workers = backend.workers.lock().unwrap(); + let unchanged = workers.get(&request.worker_ref).unwrap(); + assert!(Arc::ptr_eq( + &unchanged.handle.shared_state, + ¤t.handle.shared_state + )); + assert!(Arc::ptr_eq(&unchanged.worker_state, ¤t.worker_state)); + assert_eq!(unchanged.restore_operation_id, Some(request.operation_id)); + assert_eq!(workers.len(), 1); + drop(workers); + assert_eq!(authority_files(), files_before); + assert_eq!(run_directories(), runs_before); + assert!( + backend + .pending_restore(&request.worker_ref) + .unwrap() + .is_none() + ); + current.shutdown_requested.store(true, Ordering::Release); + let pending_cleanup = backend.preflight_restore(&request).unwrap_err(); + assert_eq!( + pending_cleanup.outcome, + crate::execution::WorkerExecutionOutcome::Busy + ); + assert!( + pending_cleanup.worker_state.is_none(), + "pending cleanup is not AlreadyConnected evidence" + ); + current.shutdown_requested.store(false, Ordering::Release); + assert!(backend.stop_worker(&request.worker_ref).is_accepted()); + assert!(matches!( + backend.worker_session(WorkerSessionObservationRequest { + worker_ref: request.worker_ref, + }), + runtime_api::WorkerSessionAvailability::RetainedSnapshot { .. } + )); + } + + #[cfg(feature = "ws-server")] + #[test] + #[serial_test::serial(worker_allocation)] + fn protocol_transport_stays_bound_to_stopped_execution_after_restore() { + let root = tempfile::tempdir().unwrap(); + let runtime_store = root.path().join("runtime"); + let backend = Arc::new( + WorkerRuntimeExecutionBackend::new( + ProfileRuntimeWorkerFactory::new(root.path()) + .with_runtime_store_dir(&runtime_store) + .with_runtime_id("runtime-subjektiv-test"), + ) + .unwrap(), + ); + let fixture_root = runtime_store.clone(); + let mut request = backend + .run_on_adapter_runtime(async move { + Ok(saved_subjektiv_restore_fixture(&fixture_root, true).await) + }) + .unwrap(); + assert!(matches!( + backend.restore_worker(request.clone()), + WorkerExecutionSpawnResult::Connected { .. } + )); + let mut transport_a = Arc::clone(&backend) + .attach_worker_protocol(&request.worker_ref) + .unwrap(); + assert!(matches!(&transport_a.snapshot, Event::Snapshot { .. })); + transport_a.validate().unwrap(); + let completion = |id: &str| Method::ListCompletions { + request_id: Some(id.to_string()), + kind: protocol::CompletionKind::File, + prefix: String::new(), + context: None, + }; + assert!( + matches!(transport_a.dispatch(completion("a-completion")).unwrap().as_slice(), + [Event::Completions { request_id: Some(id), .. }] if id == "a-completion") + ); + assert!( + transport_a + .dispatch(Method::Shutdown { + command: WorkerCommandEnvelope::new(1), + }) + .unwrap() + .is_empty() + ); + assert!( + transport_a.events.is_closed(), + "stop must join the connection relay before success" + ); + while transport_a.events.try_recv().is_ok() {} + assert!(matches!( + transport_a.events.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) + )); + request.operation_id = crate::execution::WorkerLifecycleOperationId::new(); + assert!(matches!( + backend.restore_worker(request.clone()), + WorkerExecutionSpawnResult::Connected { .. } + )); + let current_b = backend + .workers + .lock() + .unwrap() + .get(&request.worker_ref) + .unwrap() + .clone(); + let snapshot_b = current_b.handle.shared_state.snapshot(); + assert!(transport_a.validate().is_err()); + for delayed in [ + Method::Cancel { + command: WorkerCommandEnvelope::new(900), + }, + completion("delayed-a-completion"), + Method::Shutdown { + command: WorkerCommandEnvelope::new(901), + }, + ] { + let rejected = transport_a.dispatch(delayed).unwrap_err(); + assert_eq!( + rejected.outcome, + crate::execution::WorkerExecutionOutcome::Rejected + ); + assert!(rejected.worker_state.is_none()); + } + assert_eq!(current_b.handle.shared_state.snapshot(), snapshot_b); + assert_eq!( + backend + .workers + .lock() + .unwrap() + .get(&request.worker_ref) + .unwrap() + .restore_operation_id, + Some(request.operation_id) + ); + let transport_b = Arc::clone(&backend) + .attach_worker_protocol(&request.worker_ref) + .unwrap(); + transport_b.validate().unwrap(); + assert!( + matches!(transport_b.dispatch(completion("b-completion")).unwrap().as_slice(), + [Event::Completions { request_id: Some(id), .. }] if id == "b-completion") + ); + let command_id = snapshot_b.last_command_id.saturating_add(1); + transport_b + .dispatch(Method::Cancel { + command: WorkerCommandEnvelope::new(command_id), + }) + .unwrap(); + let mut b_events = transport_b.events; + backend.run_on_adapter_runtime(async move { + tokio::time::timeout(Duration::from_secs(5), async { + while let Some(event) = b_events.recv().await { + if matches!(event, Event::CommandAcknowledged { acknowledgement } if acknowledgement.command_id == command_id) { + return Ok(()); + } + } + Err("fresh B protocol stream closed before command acknowledgement".to_string()) + }).await.map_err(|_| "fresh B protocol command acknowledgement timed out".to_string())? + }).unwrap(); + assert!( + matches!( + transport_a.events.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) + ), + "A must receive no restored B protocol events" + ); + assert!(backend.stop_worker(&request.worker_ref).is_accepted()); + } + + fn profile_stop_fixture( + root: &Path, + ) -> (WorkerRuntimeExecutionBackend, WorkerExecutionRestoreRequest) { + let runtime_store = root.join("runtime"); + let backend = WorkerRuntimeExecutionBackend::new( + ProfileRuntimeWorkerFactory::new(root) + .with_runtime_store_dir(&runtime_store) + .with_runtime_id("runtime-subjektiv-test"), + ) + .unwrap(); + let request = backend + .run_on_adapter_runtime(async move { + Ok(saved_subjektiv_restore_fixture(&runtime_store, true).await) + }) + .unwrap(); + assert!(matches!( + backend.restore_worker(request.clone()), + WorkerExecutionSpawnResult::Connected { .. } + )); + (backend, request) + } + + // Queue a real Cancel first so the exact generated Shutdown ID is either + // already used by another kind (Conflict) or below the Controller floor. + // Install the same owned evidence as request_shutdown before enqueueing. + async fn prepare_rejected_shutdown( + execution: &RuntimeWorkerExecution, + stale: bool, + release: Option>, + ) -> u64 { + let command = next_internal_command_for_snapshot(&execution.handle.shared_state.snapshot()); + let cancel_id = command.command_id + u64::from(stale); + let mut cancel_events = execution.handle.subscribe(); + execution + .handle + .send(Method::Cancel { + command: WorkerCommandEnvelope::new(cancel_id), + }) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if matches!(cancel_events.recv().await.unwrap(), Event::CommandAcknowledged { acknowledgement } + if acknowledgement.command_id == cancel_id + && acknowledgement.command == protocol::WorkerCommandKind::Cancel) { + break; + } + } + }).await.unwrap(); + let events = execution.handle.subscribe(); + let handle = execution.handle.clone(); + *execution.tasks.shutdown_admission.lock().await = Some(RuntimeShutdownAdmission { + command_id: command.command_id, + send: Some(tokio::spawn(async move { + if let Some(release) = release { + release.await.map_err(|error| error.to_string())?; + } + handle + .send(Method::Shutdown { command }) + .await + .map_err(|error| error.to_string()) + })), + events, + acknowledgement: None, + uncertainty: None, + }); + execution.shutdown_requested.store(true, Ordering::Release); + command.command_id + } + + #[test] + #[serial_test::serial(worker_allocation)] + fn stop_rejected_shutdown_resets_request_and_retries_from_actual_snapshot() { + for stale in [false, true] { + let root = tempfile::tempdir().unwrap(); + let (backend, request) = profile_stop_fixture(root.path()); + let execution = backend + .workers + .lock() + .unwrap() + .get(&request.worker_ref) + .unwrap() + .clone(); + let preparing = execution.clone(); + let rejected_id = backend + .run_on_adapter_runtime(async move { + Ok(prepare_rejected_shutdown(&preparing, stale, None).await) + }) + .unwrap(); + // Explicitly stale bridge cache must not determine the retry ID. + execution.worker_state.write().unwrap().last_command_id = 0; + let result = backend.stop_worker(&request.worker_ref); + assert!(!result.is_accepted()); + let message = result.message.unwrap(); + assert!( + message.contains(if stale { "StaleCommandId" } else { "Conflict" }), + "{message}" + ); + assert!(!execution.shutdown_requested.load(Ordering::Acquire)); + assert!( + backend + .workers + .lock() + .unwrap() + .contains_key(&request.worker_ref) + ); + assert!(backend.stop_worker(&request.worker_ref).is_accepted()); + assert!(execution.handle.shared_state.snapshot().last_command_id > rejected_id); + assert!(backend.workers.lock().unwrap().is_empty()); + assert!(execution.tasks.tasks.lock().unwrap().is_empty()); + } + } + + #[test] + #[serial_test::serial(worker_allocation)] + fn stop_late_shutdown_rejection_after_timeout_retains_waiter_without_double_enqueue() { + let root = tempfile::tempdir().unwrap(); + let (backend, request) = profile_stop_fixture(root.path()); + let execution = backend + .workers + .lock() + .unwrap() + .get(&request.worker_ref) + .unwrap() + .clone(); + let checking = execution.clone(); + backend + .run_on_adapter_runtime(async move { + let (release, gate) = tokio::sync::oneshot::channel(); + let command_id = prepare_rejected_shutdown(&checking, false, Some(gate)).await; + for _ in 0..2 { + let error = checking + .tasks + .request_shutdown( + &checking.handle, + &checking.shutdown_requested, + Duration::from_millis(10), + ) + .await + .unwrap_err(); + assert!(error.contains("timed out"), "{error}"); + assert!(checking.shutdown_requested.load(Ordering::Acquire)); + let retained = checking.tasks.shutdown_admission.lock().await; + let pending = retained.as_ref().unwrap(); + assert_eq!(pending.command_id, command_id); + assert!( + pending + .send + .as_ref() + .is_some_and(|task| !task.is_finished()) + ); + assert!(pending.acknowledgement.is_none()); + } + // A mismatched ID or command kind is not this Shutdown's verdict. + for (id, kind) in [ + (command_id + 1, protocol::WorkerCommandKind::Shutdown), + (command_id, protocol::WorkerCommandKind::Cancel), + ] { + checking + .handle + .send_event(Event::CommandAcknowledged { + acknowledgement: protocol::WorkerCommandAcknowledgement { + command_id: id, + command: kind, + disposition: protocol::WorkerCommandDisposition::Accepted, + state: checking.handle.shared_state.snapshot(), + }, + }) + .unwrap(); + } + release.send(()).unwrap(); + let error = checking + .tasks + .request_shutdown( + &checking.handle, + &checking.shutdown_requested, + Duration::from_secs(5), + ) + .await + .unwrap_err(); + assert!(error.contains("Conflict"), "{error}"); + assert!(!checking.shutdown_requested.load(Ordering::Acquire)); + assert!(checking.tasks.shutdown_admission.lock().await.is_none()); + Ok(()) + }) + .unwrap(); + assert!(backend.stop_worker(&request.worker_ref).is_accepted()); + assert!(execution.tasks.tasks.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn shutdown_ack_timeout_then_late_rejection_keeps_exact_receiver() { + let (events, receiver) = tokio::sync::broadcast::channel(4); + let mut pending = RuntimeShutdownAdmission { + command_id: 37, + send: None, + events: receiver, + acknowledgement: None, + uncertainty: None, + }; + assert!( + pending + .wait(Duration::from_millis(10)) + .await + .unwrap_err() + .contains("timed out") + ); + assert!(pending.uncertainty.as_ref().unwrap().contains("timed out")); + let acknowledgement = protocol::WorkerCommandAcknowledgement { + command_id: 37, + command: protocol::WorkerCommandKind::Shutdown, + disposition: protocol::WorkerCommandDisposition::StaleCommandId, + state: protocol::WorkerStateSnapshot::initial(), + }; + events + .send(Event::CommandAcknowledged { + acknowledgement: acknowledgement.clone(), + }) + .unwrap(); + pending.wait(Duration::from_secs(1)).await.unwrap(); + assert_eq!(pending.acknowledgement, Some(acknowledgement)); + } + + #[tokio::test] + async fn shutdown_ack_lag_and_closed_retain_correlated_unknown_evidence() { + let (events, receiver) = tokio::sync::broadcast::channel(1); + let mut pending = RuntimeShutdownAdmission { + command_id: 38, + send: None, + events: receiver, + acknowledgement: None, + uncertainty: None, + }; + for _ in 0..2 { + events.send(Event::Shutdown).unwrap(); + } + assert!(pending.wait(Duration::from_secs(1)).await.is_err()); + assert!(pending.uncertainty.is_some()); + assert!(pending.acknowledgement.is_none()); + drop(events); + assert!(pending.wait(Duration::from_secs(1)).await.is_err()); + assert_eq!(pending.command_id, 38); + assert!(pending.uncertainty.is_some()); + assert!(pending.acknowledgement.is_none()); + } + + #[test] + #[serial_test::serial(worker_allocation)] + fn unconnected_cleanup_rejected_shutdown_retains_scope_for_retry() { + let root = tempfile::tempdir().unwrap(); + let (backend, request) = profile_stop_fixture(root.path()); + let execution = backend + .workers + .lock() + .unwrap() + .get(&request.worker_ref) + .unwrap() + .clone(); + let preparing = execution.clone(); + backend + .run_on_adapter_runtime(async move { + prepare_rejected_shutdown(&preparing, true, None).await; + Ok(()) + }) + .unwrap(); + let cleanup = || { + backend.cleanup_unconnected_controller( + &execution.handle, + &execution.shutdown, + &execution.tasks, + &execution.shutdown_requested, + &execution.worker_state, + ) + }; + assert!(cleanup().unwrap_err().contains("StaleCommandId")); + assert!(!execution.shutdown_requested.load(Ordering::Acquire)); + assert!(!execution.tasks.tasks.lock().unwrap().is_empty()); + assert!(cleanup().is_ok()); + assert!(execution.tasks.tasks.lock().unwrap().is_empty()); + assert!(backend.stop_worker(&request.worker_ref).is_accepted()); + } + + #[test] + #[serial_test::serial(worker_allocation)] + fn duplicate_connect_rejected_shutdown_retains_actual_candidate_until_stop_retry() { + for operation in [ + WorkerExecutionOperation::Spawn, + WorkerExecutionOperation::Restore, + ] { + let root = tempfile::tempdir().unwrap(); + let candidate_root = tempfile::tempdir().unwrap(); + let (backend, request) = profile_stop_fixture(root.path()); + let primary = backend + .workers + .lock() + .unwrap() + .get(&request.worker_ref) + .unwrap() + .clone(); + let (candidate_backend, candidate_request) = + profile_stop_fixture(candidate_root.path()); + let candidate = candidate_backend + .workers + .lock() + .unwrap() + .remove(&candidate_request.worker_ref) + .unwrap(); + assert!(!primary.handle.same_controller(&candidate.handle)); + let preparing = candidate.clone(); + candidate_backend + .run_on_adapter_runtime(async move { + prepare_rejected_shutdown(&preparing, true, None).await; + Ok(()) + }) + .unwrap(); + let result = { + let lock = backend.worker_lock(&request.worker_ref).unwrap(); + let _guard = lock.lock().unwrap(); + backend.connect_controller_scope( + operation, + request.worker_ref.clone(), + request.context.clone(), + candidate.handle.clone(), + candidate.shutdown.clone(), + candidate.tasks.clone(), + candidate.shutdown_requested.clone(), + Some(request.operation_id), + BTreeMap::new(), + candidate.workspace_client.clone(), + ) + }; + match result { + WorkerExecutionSpawnResult::ReconciliationRequired { + result, + worker_state, + .. + } => { + assert!(result.message.unwrap().contains("StaleCommandId")); + assert!( + worker_state.is_none(), + "existing snapshot is not candidate cleanup proof" + ); + } + other => panic!( + "duplicate cleanup failure was classified as side-effect-free: {other:?}" + ), + } + let retained = primary.tasks.cleanup_candidates.lock().unwrap()[0].clone(); + assert!(retained.handle.same_controller(&candidate.handle)); + assert!(Arc::ptr_eq(&retained.shutdown, &candidate.shutdown)); + assert!(Arc::ptr_eq(&retained.tasks.tasks, &candidate.tasks.tasks)); + assert!(Arc::ptr_eq( + &retained.tasks.shutdown_admission, + &candidate.tasks.shutdown_admission + )); + let registered = backend + .workers + .lock() + .unwrap() + .get(&request.worker_ref) + .unwrap() + .clone(); + assert!(registered.handle.same_controller(&primary.handle)); + assert_eq!(backend.workers.lock().unwrap().len(), 1); + assert!( + backend + .preflight_restore(&request) + .unwrap_err() + .worker_state + .is_none() + ); + assert!(matches!( + backend.restore_worker(request.clone()), + WorkerExecutionSpawnResult::ReconciliationRequired { .. } + )); + + // Fail candidate cleanup once more during Stop. The main Controller + // must still join, but that cannot turn the failed child into Stopped. + let preparing = candidate.clone(); + candidate_backend + .run_on_adapter_runtime(async move { + prepare_rejected_shutdown(&preparing, false, None).await; + Ok(()) + }) + .unwrap(); + let stopped = backend.stop_worker(&request.worker_ref); + assert!(!stopped.is_accepted()); + assert!(stopped.message.unwrap().contains("Conflict")); + assert!(primary.handle.protocol_is_closed()); + assert!(primary.tasks.tasks.lock().unwrap().is_empty()); + let checking_scope = primary.tasks.clone(); + let incomplete_scope = + backend.run_on_adapter_runtime(async move { checking_scope.join().await }); + assert!( + incomplete_scope + .unwrap_err() + .contains("unconnected Controller cleanup remains unproven") + ); + assert!( + backend + .workers + .lock() + .unwrap() + .contains_key(&request.worker_ref) + ); + assert_eq!(primary.tasks.cleanup_candidates.lock().unwrap().len(), 1); + assert!(!candidate.tasks.tasks.lock().unwrap().is_empty()); + assert!(backend.stop_worker(&request.worker_ref).is_accepted()); + assert!(candidate.tasks.tasks.lock().unwrap().is_empty()); + assert!(primary.tasks.cleanup_candidates.lock().unwrap().is_empty()); + assert!(backend.workers.lock().unwrap().is_empty()); + } + } + + #[test] + #[serial_test::serial(worker_allocation)] + fn duplicate_connect_timeout_retains_candidate_waiter_and_no_false_stop() { + let root = tempfile::tempdir().unwrap(); + let candidate_root = tempfile::tempdir().unwrap(); + let (backend, request) = profile_stop_fixture(root.path()); + let primary = backend + .workers + .lock() + .unwrap() + .get(&request.worker_ref) + .unwrap() + .clone(); + let (candidate_backend, candidate_request) = profile_stop_fixture(candidate_root.path()); + let candidate = candidate_backend + .workers + .lock() + .unwrap() + .remove(&candidate_request.worker_ref) + .unwrap(); + let (release, gate) = tokio::sync::oneshot::channel(); + let preparing = candidate.clone(); + let command_id = candidate_backend + .run_on_adapter_runtime(async move { + Ok(prepare_rejected_shutdown(&preparing, false, Some(gate)).await) + }) + .unwrap(); + let result = { + let lock = backend.worker_lock(&request.worker_ref).unwrap(); + let _guard = lock.lock().unwrap(); + backend.connect_controller_scope( + WorkerExecutionOperation::Restore, + request.worker_ref.clone(), + request.context.clone(), + candidate.handle.clone(), + candidate.shutdown.clone(), + candidate.tasks.clone(), + candidate.shutdown_requested.clone(), + Some(request.operation_id), + BTreeMap::new(), + candidate.workspace_client.clone(), + ) + }; + assert!(matches!( + result, + WorkerExecutionSpawnResult::ReconciliationRequired { .. } + )); + let stopped = backend.stop_worker(&request.worker_ref); + assert!(!stopped.is_accepted()); + assert!(stopped.message.unwrap().contains("timed out")); + assert!(primary.handle.protocol_is_closed()); + let retained = primary.tasks.cleanup_candidates.lock().unwrap()[0].clone(); + assert!(retained.handle.same_controller(&candidate.handle)); + assert!(Arc::ptr_eq( + &retained.tasks.shutdown_admission, + &candidate.tasks.shutdown_admission + )); + let checking = retained.clone(); + backend + .run_on_adapter_runtime(async move { + let admission = checking.tasks.shutdown_admission.lock().await; + let pending = admission.as_ref().unwrap(); + assert_eq!( + pending.command_id, command_id, + "unknown admission must not enqueue another command" + ); + assert!( + pending + .send + .as_ref() + .is_some_and(|task| !task.is_finished()) + ); + assert!(pending.uncertainty.as_ref().unwrap().contains("timed out")); + assert!(pending.acknowledgement.is_none()); + Ok(()) + }) + .unwrap(); + assert!( + backend + .workers + .lock() + .unwrap() + .contains_key(&request.worker_ref) + ); + release.send(()).unwrap(); + let rejected = backend.stop_worker(&request.worker_ref); + assert!(!rejected.is_accepted()); + assert!(rejected.message.unwrap().contains("Conflict")); + assert!(!candidate.shutdown_requested.load(Ordering::Acquire)); + assert_eq!(primary.tasks.cleanup_candidates.lock().unwrap().len(), 1); + assert!(backend.stop_worker(&request.worker_ref).is_accepted()); + assert!(candidate.tasks.tasks.lock().unwrap().is_empty()); + assert!(primary.tasks.tasks.lock().unwrap().is_empty()); + assert!(primary.tasks.cleanup_candidates.lock().unwrap().is_empty()); + assert!(backend.workers.lock().unwrap().is_empty()); + } + + #[test] + #[serial_test::serial(worker_allocation)] + fn pending_restore_join_failure_never_becomes_absence_on_stop_retry() { + let root = tempfile::tempdir().unwrap(); + let runtime_store = root.path().join("runtime"); + let backend = WorkerRuntimeExecutionBackend::new( + ProfileRuntimeWorkerFactory::new(root.path()) + .with_runtime_store_dir(&runtime_store) + .with_runtime_id("runtime-subjektiv-test"), + ) + .unwrap(); + let request = backend + .run_on_adapter_runtime(async move { + Ok(saved_subjektiv_restore_fixture(&runtime_store, true).await) + }) + .unwrap(); + let failed = backend + .spawn_on_adapter_runtime(async { + panic!("factory scope failed before publishing result") + }) + .unwrap(); + let pending = Arc::new(PendingRuntimeRestore { + request: request.clone(), + task: tokio::sync::Mutex::new(Some(failed)), + result: Mutex::new(None), + failure: Mutex::new(None), + }); + backend + .pending_restores + .lock() + .unwrap() + .insert(request.worker_ref.clone(), pending.clone()); + for _ in 0..3 { + let stopped = backend.stop_worker(&request.worker_ref); + assert!(!stopped.is_accepted()); + assert!( + stopped + .message + .unwrap() + .contains("restore factory task failed") + ); + assert!(pending.failure.lock().unwrap().is_some()); + assert!(pending.result.lock().unwrap().is_none()); + assert!(Arc::ptr_eq( + &backend + .pending_restore(&request.worker_ref) + .unwrap() + .unwrap(), + &pending + )); + assert!(backend.workers.lock().unwrap().is_empty()); + } + assert!(matches!( + backend.restore_worker(request), + WorkerExecutionSpawnResult::ReconciliationRequired { .. } + )); + } + + #[test] + #[serial_test::serial(worker_allocation)] + fn stop_joins_healthy_self_terminated_controller_with_closed_endpoint() { + let root = tempfile::tempdir().unwrap(); + let (backend, request) = profile_stop_fixture(root.path()); + let execution = backend + .workers + .lock() + .unwrap() + .get(&request.worker_ref) + .unwrap() + .clone(); + let handle = execution.handle.clone(); + let controller = execution.tasks.tasks.lock().unwrap()[0].clone(); + backend + .run_on_adapter_runtime(async move { + handle + .send(Method::Shutdown { + command: WorkerCommandEnvelope::new(1), + }) + .await + .map_err(|error| error.to_string())?; + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if handle.protocol_is_closed() + && controller + .completion + .lock() + .await + .task + .as_ref() + .unwrap() + .is_finished() + { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .map_err(|error| error.to_string())?; + Ok(()) + }) + .unwrap(); + assert!(!execution.shutdown_requested.load(Ordering::Acquire)); + assert!(backend.stop_worker(&request.worker_ref).is_accepted()); + assert!(backend.workers.lock().unwrap().is_empty()); + assert!(execution.tasks.tasks.lock().unwrap().is_empty()); + assert!(backend.stop_worker(&request.worker_ref).is_accepted()); + } + + #[test] + #[serial_test::serial(worker_allocation)] + fn stop_and_unconnected_cleanup_retain_abnormal_controller_completion_across_retries() { + let root = tempfile::tempdir().unwrap(); + let (backend, request) = profile_stop_fixture(root.path()); + let execution = backend + .workers + .lock() + .unwrap() + .get(&request.worker_ref) + .unwrap() + .clone(); + let controller = execution.tasks.tasks.lock().unwrap()[0].clone(); + let aborted_controller = Arc::clone(&controller); + backend + .run_on_adapter_runtime(async move { + aborted_controller + .completion + .lock() + .await + .task + .as_ref() + .unwrap() + .abort(); + tokio::time::timeout(Duration::from_secs(5), async { + while !aborted_controller + .completion + .lock() + .await + .task + .as_ref() + .unwrap() + .is_finished() + { + tokio::task::yield_now().await; + } + }) + .await + .map_err(|error| error.to_string())?; + Ok(()) + }) + .unwrap(); + for _ in 0..3 { + let cleanup = backend + .cleanup_unconnected_controller( + &execution.handle, + &execution.shutdown, + &execution.tasks, + &execution.shutdown_requested, + &execution.worker_state, + ) + .unwrap_err(); + assert!(cleanup.contains("controller task failed"), "{cleanup}"); + let stopped = backend.stop_worker(&request.worker_ref); + assert_eq!( + stopped.outcome, + crate::execution::WorkerExecutionOutcome::Errored + ); + assert!(stopped.message.unwrap().contains("controller task failed")); + let workers = backend.workers.lock().unwrap(); + let retained = workers + .get(&request.worker_ref) + .expect("unknown cleanup must retain owner"); + assert!(retained.handle.same_controller(&execution.handle)); + assert!(Arc::ptr_eq( + &retained.tasks.tasks.lock().unwrap()[0], + &controller + )); + } + let retained = execution.tasks.tasks.lock().unwrap().clone(); + assert_eq!( + retained.len(), + 2, + "failed task and failed shutdown proof remain; relays drain" + ); + } + + #[tokio::test] + #[serial_test::serial(worker_allocation)] + async fn profile_factory_saved_subjektiv_manifest_survives_restart_and_explicit_restore() { + let root = tempfile::tempdir().unwrap(); + let runtime_store = root.path().join("runtime"); + let mut request = saved_subjektiv_restore_fixture(&runtime_store, true).await; + let worker_name = + ProfileRuntimeWorkerFactory::runtime_worker_name_for_ref(&request.worker_ref); + let aggregate_dir = runtime_store + .join("workers") + .join(request.worker_ref.worker_id.to_string()); + let metadata_store = WorkerAggregateStore::new(&aggregate_dir, &worker_name).unwrap(); + let saved_metadata = metadata_store.read_by_name(&worker_name).unwrap().unwrap(); + let saved_active = saved_metadata.active.as_ref().unwrap(); + let session_store = WorkerSessionStore::new(aggregate_dir.join("session")).unwrap(); + for _ in 0..2 { + // Fresh factory: exercise the same saved authority after both a + // restart and another explicit restore, not a mock Controller path. + let factory = + ProfileRuntimeWorkerFactory::new(root.path().join("nonexistent-current-profiles")) + .with_runtime_store_dir(&runtime_store) + .with_runtime_id("runtime-subjektiv-test"); + let before = persisted_files(&runtime_store); + factory.preflight_restore(&request).await.unwrap(); + assert_eq!( + persisted_files(&runtime_store), + before, + "preflight must be read-only" + ); + assert!(factory.observation_hub.workers.lock().unwrap().is_empty()); + let controller = factory + .restore_controller(request.clone()) + .await + .unwrap_or_else(|error| panic!("saved Subjektiv restore failed: {error}")); + assert_eq!( + controller.handle.shared_state.catalog_status(), + WorkerStatus::Idle + ); + shutdown_profile_controller(controller).await; + let restored_metadata = metadata_store.read_by_name(&worker_name).unwrap().unwrap(); + assert_eq!(restored_metadata.active, saved_metadata.active); + assert_eq!( + restored_metadata.resolved_manifest_snapshot, + saved_metadata.resolved_manifest_snapshot + ); + assert_eq!( + restored_metadata.subjektiv_session_attribution, + saved_metadata.subjektiv_session_attribution + ); + let entries = session_store + .read_all(saved_active.session_id, saved_active.segment_id.unwrap()) + .unwrap(); + assert!( + matches!(entries.first(), Some(LogEntry::AnnotatedSegmentStart { + system_prompt: Some(prompt), + .. + }) if prompt == SAVED_SUBJEKTIV_SYSTEM_PROMPT) + ); + assert_eq!( + session_store::collect_state(&entries) + .system_prompt + .as_deref(), + Some(SAVED_SUBJEKTIV_SYSTEM_PROMPT), + "actual restore must replay the committed prompt rather than re-render it" + ); + request.operation_id = crate::execution::WorkerLifecycleOperationId::new(); + } + assert!(!root.path().join("nonexistent-current-profiles").exists()); + } + + #[tokio::test] + #[serial_test::serial(worker_allocation)] + async fn profile_factory_pending_subjektiv_restore_requires_operation_owned_launch_material() { + let root = tempfile::tempdir().unwrap(); + let runtime_store = root.path().join("runtime"); + let mut request = saved_subjektiv_restore_fixture(&runtime_store, false).await; + let factory = ProfileRuntimeWorkerFactory::new(root.path()) + .with_runtime_store_dir(&runtime_store) + .with_runtime_id("runtime-subjektiv-test"); + let worker_name = + ProfileRuntimeWorkerFactory::runtime_worker_name_for_ref(&request.worker_ref); + let store = WorkerAggregateStore::new( + factory.worker_aggregate_dir(&request.worker_ref).unwrap(), + &worker_name, + ) + .unwrap(); + let metadata = store.read_by_name(&worker_name).unwrap().unwrap(); + let pending_session_id = metadata.active.as_ref().unwrap().session_id; + assert!(metadata.active.as_ref().unwrap().segment_id.is_none()); + let before = persisted_files(&runtime_store); + let error = factory.preflight_restore(&request).await.unwrap_err(); + assert!( + error.contains("Pending Workspace Worker restore requires"), + "{error}" + ); + assert_eq!(persisted_files(&runtime_store), before); + let mut bundle = test_bundle(); + bundle.metadata.workspace_id = "workspace-saved-subjektiv".to_string(); + let builtins = worker::PromptCatalog::builtins_only().unwrap(); + let projection = builtins.projection(); + let mut templates = projection.templates.clone(); + templates.insert( + "default".to_string(), + "OPERATION-OWNED-PENDING-SYSTEM-PROMPT".to_string(), + ); + templates.insert( + "internal.notify_wrapper".to_string(), + "SAVED {{ message }}".to_string(), + ); + let mut catalog = worker::EffectivePromptCatalog::new( + templates, + 17, + projection.schema_fingerprint.clone(), + projection.toolchain_fingerprint.clone(), + ) + .unwrap(); + catalog.source_digest = "saved-source".to_string(); + bundle.prompt_catalog = Some(catalog); + let bundle = bundle.with_computed_digest(); + let mut missing_projection = request.clone(); + let mut missing_projection_bundle = bundle.clone(); + missing_projection_bundle.prompt_catalog = None; + missing_projection.config_bundle = Some(missing_projection_bundle.with_computed_digest()); + let mut wrong_workspace = request.clone(); + let mut wrong_workspace_bundle = bundle.clone(); + wrong_workspace_bundle.metadata.workspace_id = "another-workspace".to_string(); + wrong_workspace.config_bundle = Some(wrong_workspace_bundle.with_computed_digest()); + let mut missing_instruction = request.clone(); + let mut missing_instruction_bundle = bundle.clone(); + let mut templates = projection.templates.clone(); + templates.remove("default"); + missing_instruction_bundle.prompt_catalog = Some( + worker::EffectivePromptCatalog::new( + templates, + 17, + projection.schema_fingerprint.clone(), + projection.toolchain_fingerprint.clone(), + ) + .unwrap(), + ); + missing_instruction.config_bundle = Some(missing_instruction_bundle.with_computed_digest()); + for (invalid, expected) in [ + (request.clone(), "requires operation-owned launch material"), + ( + missing_projection, + "requires a saved Workspace Prompt projection", + ), + ( + wrong_workspace, + "Workspace Prompt projection scope mismatch", + ), + (missing_instruction, "invalid pending Worker launch Prompt"), + ] { + let error = factory.preflight_restore(&invalid).await.unwrap_err(); + assert!(error.contains(expected), "{error}"); + let error = match factory.restore_controller(invalid).await { + Ok(controller) => { + shutdown_profile_controller(controller).await; + panic!("invalid pending launch authority started a Controller"); + } + Err(error) => error, + }; + assert!(error.contains(expected), "{error}"); + assert_eq!(persisted_files(&runtime_store), before); + assert!( + factory + .prompt_projection_cache + .active("workspace-saved-subjektiv") + .unwrap() + .is_none() + ); + assert!(factory.observation_hub.workers.lock().unwrap().is_empty()); + } + request.config_bundle = Some(bundle); + factory.preflight_restore(&request).await.unwrap(); + assert_eq!(persisted_files(&runtime_store), before); + assert!( + factory + .prompt_projection_cache + .active("workspace-saved-subjektiv") + .unwrap() + .is_none() + ); + let controller = factory + .restore_controller(request.clone()) + .await + .unwrap_or_else(|error| panic!("pending saved Subjektiv restore failed: {error}")); + assert_eq!( + controller.handle.shared_state.catalog_status(), + WorkerStatus::Idle + ); + shutdown_profile_controller(controller).await; + let restored = store.read_by_name(&worker_name).unwrap().unwrap(); + assert_eq!( + restored.resolved_manifest_snapshot, metadata.resolved_manifest_snapshot, + "pending launch material must not replace the saved Manifest" + ); + assert_eq!( + restored.subjektiv_session_attribution, metadata.subjektiv_session_attribution, + "pending restore must retain committed Subject attribution" + ); + let active = restored.active.unwrap(); + assert_eq!(active.session_id, pending_session_id); + let segment_id = active + .segment_id + .expect("operation-owned launch must commit the initial Segment before exposure"); + let entries = WorkerSessionStore::new( + factory + .worker_aggregate_dir(&request.worker_ref) + .unwrap() + .join("session"), + ) + .unwrap() + .read_all(active.session_id, segment_id) + .unwrap(); + assert!( + matches!(entries.first(), Some(LogEntry::AnnotatedSegmentStart { + session_id, + system_prompt: Some(prompt), + history, + forked_from: None, + compacted_from: None, + .. + }) if *session_id == pending_session_id + && prompt.contains("OPERATION-OWNED-PENDING-SYSTEM-PROMPT") + && history.is_empty()) + ); + assert!( + session_store::collect_state(&entries) + .system_prompt + .is_some() + ); + } + + #[tokio::test] + #[serial_test::serial(worker_allocation)] + async fn profile_factory_saved_manifest_mismatch_rejected_before_live_work() { + let root = tempfile::tempdir().unwrap(); + let runtime_store = root.path().join("runtime"); + let request = saved_subjektiv_restore_fixture(&runtime_store, true).await; + let factory = ProfileRuntimeWorkerFactory::new(root.path()) + .with_runtime_store_dir(&runtime_store) + .with_runtime_id("runtime-subjektiv-test"); + let before = persisted_files(&runtime_store); + let mut wrong_revision = request.clone(); + wrong_revision + .request + .memory_settings + .as_mut() + .unwrap() + .settings_revision = 18; + let mut detached = request.clone(); + detached.request.subjektiv_attached = false; + let mut missing_settings = request.clone(); + missing_settings.request.memory_settings = None; + for (invalid, expected) in [ + (wrong_revision, "subjektiv settings snapshot mismatch"), + (detached, "unauthorized subjektiv attachment"), + (missing_settings, "missing its trusted settings snapshot"), + ] { + let error = factory.preflight_restore(&invalid).await.unwrap_err(); + assert!(error.contains(expected), "{error}"); + let error = match factory.restore_controller(invalid).await { + Ok(controller) => { + shutdown_profile_controller(controller).await; + panic!("invalid restore started a Controller") + } + Err(error) => error, + }; + assert!(error.contains(expected), "{error}"); + assert_eq!(persisted_files(&runtime_store), before); + assert!(factory.observation_hub.workers.lock().unwrap().is_empty()); + } + let run_dir = factory + .worker_aggregate_dir(&request.worker_ref) + .unwrap() + .join("runs"); + assert!( + !run_dir.exists(), + "mismatch must reject before creating any run resources" + ); + } + #[test] fn profile_factory_reads_retained_session_from_canonical_worker_aggregate() { let root = tempfile::tempdir().unwrap(); @@ -3870,7 +6488,7 @@ mod tests { } #[tokio::test] - async fn execution_scope_drains_remaining_tasks_after_join_failure_and_allows_retry() { + async fn execution_scope_drains_remaining_tasks_and_retains_abnormal_completion() { let failed = tokio::spawn(async { panic!("injected controller failure") }); let scope = RuntimeExecutionTaskScope::new(failed); scope.push( @@ -3881,8 +6499,81 @@ mod tests { let error = scope.join().await.unwrap_err(); assert!(error.contains("controller task failed")); - assert!(scope.tasks.lock().unwrap().is_empty()); + let retained = scope.tasks.lock().unwrap().clone(); + assert_eq!(retained.len(), 1, "only the failed completion must remain"); + assert_eq!(retained[0].name, "controller"); + assert!(retained[0].completion.lock().await.failure.is_some()); + for _ in 0..3 { + assert_eq!(scope.join().await.unwrap_err(), error); + } + } + + #[tokio::test] + async fn execution_scope_retains_closed_shutdown_completion_after_other_tasks_drain() { + let scope = RuntimeExecutionTaskScope::new(tokio::spawn(async {})); + scope.push( + "protocol bridge", + tokio::spawn(std::future::pending()), + true, + ); + let (sender, receiver) = tokio::sync::oneshot::channel(); + drop(sender); + let shutdown = Arc::new(tokio::sync::Mutex::new(Some(receiver))); + for _ in 0..3 { + let error = scope + .confirm_shutdown_and_join(&shutdown) + .await + .unwrap_err(); + assert!(error.contains("completion channel closed"), "{error}"); + assert_eq!(scope.tasks.lock().unwrap().len(), 1); + } + assert!(shutdown.lock().await.is_none()); + } + + #[tokio::test] + async fn cancelled_execution_scope_join_keeps_actual_task_owned_for_retry() { + let release = Arc::new(tokio::sync::Notify::new()); + let task_release = Arc::clone(&release); + let scope = RuntimeExecutionTaskScope::new(tokio::spawn(async move { + task_release.notified().await; + })); + let owned = scope.tasks.lock().unwrap()[0].clone(); + let joining_scope = scope.clone(); + let joining = tokio::spawn(async move { joining_scope.join().await }); + tokio::time::timeout(Duration::from_secs(5), async { + while owned.completion.try_lock().is_ok() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + joining.abort(); + assert!(joining.await.unwrap_err().is_cancelled()); + assert!(Arc::ptr_eq(&scope.tasks.lock().unwrap()[0], &owned)); + assert!(owned.completion.lock().await.task.is_some()); + release.notify_one(); scope.join().await.unwrap(); + assert!(scope.tasks.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn execution_scope_timeout_retains_exact_task_until_completion_is_proven() { + let release = Arc::new(tokio::sync::Notify::new()); + let task_release = Arc::clone(&release); + let scope = RuntimeExecutionTaskScope::new(tokio::spawn(async move { + task_release.notified().await; + })); + let owned = scope.tasks.lock().unwrap()[0].clone(); + let error = scope.join().await.unwrap_err(); + assert!( + error.contains("task did not stop before timeout"), + "{error}" + ); + assert!(Arc::ptr_eq(&scope.tasks.lock().unwrap()[0], &owned)); + assert!(owned.completion.lock().await.task.is_some()); + release.notify_one(); + scope.join().await.unwrap(); + assert!(scope.tasks.lock().unwrap().is_empty()); } fn adapter_command( @@ -4623,6 +7314,557 @@ mod tests { } } + #[cfg(feature = "fs-store")] + struct RestartArtifactFailureFactory { + inner: MockFactory, + runtime_store: PathBuf, + restore_id: Arc>>, + } + + #[cfg(feature = "fs-store")] + #[async_trait] + impl RuntimeWorkerFactory for RestartArtifactFailureFactory { + async fn spawn_controller( + &self, + request: WorkerExecutionSpawnRequest, + ) -> Result { + self.inner.spawn_controller(request).await + } + + async fn restore_controller( + &self, + request: WorkerExecutionRestoreRequest, + ) -> Result { + *self.restore_id.lock().unwrap() = Some(request.operation_id); + let runs = self + .runtime_store + .join("workers") + .join(request.worker_ref.worker_id.to_string()) + .join("runs"); + fs::create_dir_all(&runs).unwrap(); + // Deterministic cleanup-validation failure, including for root users. + fs::write( + runs.join(format!("restore-{}", request.operation_id)), + b"incomplete restore artifact", + ) + .unwrap(); + Err("injected restore failure with persisted artifacts".to_string()) + } + + async fn cleanup_failed_restore( + &self, + _: &WorkerExecutionRestoreRequest, + ) -> Result<(), String> { + Err("injected restore artifact cleanup failure".to_string()) + } + } + + #[cfg(feature = "fs-store")] + #[test] + fn fs_restart_stop_cleans_prior_restore_artifacts_and_retains_journal_on_failure() { + let root = tempfile::tempdir().unwrap(); + let runtime_store = root.path().join("runtime"); + let options = crate::fs_store::FsRuntimeStoreOptions { + root: runtime_store.clone(), + runtime_id: "restart-artifact-runtime".to_string(), + display_name: None, + }; + let restore_id = Arc::new(Mutex::new(None)); + let backend = Arc::new( + WorkerRuntimeExecutionBackend::new(RestartArtifactFailureFactory { + inner: MockFactory { + client: MockClient::sequential(vec![]), + runtime_base: root.path().join("controllers"), + cwd: root.path().to_path_buf(), + store_dir: root.path().join("sessions"), + worker_metadata_dir: root.path().join("metadata"), + observed_cwds: Arc::new(Mutex::new(Vec::new())), + observed_workspace_clients: Arc::new(Mutex::new(Vec::new())), + }, + runtime_store: runtime_store.clone(), + restore_id: Arc::clone(&restore_id), + }) + .unwrap(), + ); + let runtime = + EmbeddedRuntime::with_fs_store_and_execution_backend(options.clone(), backend.clone()) + .unwrap(); + runtime.store_config_bundle(test_bundle()).unwrap(); + let worker = runtime + .create_worker(create_request("restart cleanup")) + .unwrap(); + runtime.stop_worker(&worker.worker_ref, None).unwrap(); + assert!(runtime.restore_worker(&worker.worker_ref).is_err()); + assert!(runtime.stop_worker(&worker.worker_ref, None).is_err()); + let prior_restore_id = restore_id.lock().unwrap().unwrap(); + let aggregate = runtime_store + .join("workers") + .join(worker.worker_id.to_string()); + let journal_path = aggregate.join("worker.json"); + let journal_before = fs::read(&journal_path).unwrap(); + let journal: serde_json::Value = serde_json::from_slice(&journal_before).unwrap(); + assert_eq!(journal["execution_state"]["execution"]["operation"], "stop"); + assert_eq!( + journal["execution_state"]["execution"]["intent"]["pending_restore"]["operation_id"], + prior_restore_id.to_string() + ); + let artifact = aggregate + .join("runs") + .join(format!("restore-{prior_restore_id}")); + let create_run = aggregate + .join("runs") + .join(uuid::Uuid::now_v7().to_string()); + fs::create_dir_all(&create_run).unwrap(); + fs::write(create_run.join("keep.log"), "ordinary creation run").unwrap(); + let another_worker = runtime_store + .join("workers") + .join(WorkerId::now_v7().to_string()) + .join("runs") + .join(format!( + "restore-{}", + crate::execution::WorkerLifecycleOperationId::new() + )); + fs::create_dir_all(&another_worker).unwrap(); + fs::write(another_worker.join("keep.log"), "another Worker").unwrap(); + fs::create_dir_all(aggregate.join("session")).unwrap(); + fs::write( + aggregate.join("session").join("keep-authority"), + "saved Session authority", + ) + .unwrap(); + drop(runtime); + drop(backend); + + // All backend maps are empty after process restart. The nested journal + // restore id is not part of StopRequest; artifacts are discovered only + // below this exact Worker's aggregate, after excluding live resources. + let profile_base = root.path().join("must-not-launch-current-profile"); + let backend = Arc::new( + WorkerRuntimeExecutionBackend::new( + ProfileRuntimeWorkerFactory::new(&profile_base) + .with_runtime_store_dir(&runtime_store), + ) + .unwrap(), + ); + let restarted = + EmbeddedRuntime::with_fs_store_and_execution_backend(options.clone(), backend.clone()) + .unwrap(); + assert_eq!(fs::read(&journal_path).unwrap(), journal_before); + assert!(backend.workers.lock().unwrap().is_empty()); + assert!(backend.pending_restores.lock().unwrap().is_empty()); + assert!( + backend + .factory + .failed_restore_sessions + .lock() + .unwrap() + .is_empty() + ); + let failed = restarted.stop_worker(&worker.worker_ref, None).unwrap_err(); + assert!( + failed.to_string().contains("unsafe restore artifact"), + "{failed}" + ); + assert_eq!(fs::read(&journal_path).unwrap(), journal_before); + assert!(artifact.is_file()); + assert!(!profile_base.exists()); + + fs::remove_file(&artifact).unwrap(); + fs::create_dir(&artifact).unwrap(); + fs::write(artifact.join("worker.err.log"), "stale restore diagnostics").unwrap(); + restarted.stop_worker(&worker.worker_ref, None).unwrap(); + assert!(!artifact.exists()); + assert_eq!( + restarted.worker_detail(&worker.worker_ref).unwrap().status, + crate::catalog::WorkerStatus::Stopped + ); + let journal: serde_json::Value = + serde_json::from_slice(&fs::read(&journal_path).unwrap()).unwrap(); + assert_ne!( + journal["execution_state"]["state"], + "reconciliation_required" + ); + assert_eq!( + fs::read_to_string(aggregate.join("session").join("keep-authority")).unwrap(), + "saved Session authority" + ); + assert!(create_run.join("keep.log").is_file()); + assert!(another_worker.join("keep.log").is_file()); + assert!(!profile_base.exists()); + assert!(backend.workers.lock().unwrap().is_empty()); + restarted.stop_worker(&worker.worker_ref, None).unwrap(); + } + + #[tokio::test] + async fn profile_restore_artifact_failure_retains_failed_session_owner_for_retry() { + let root = tempfile::tempdir().unwrap(); + let factory = + ProfileRuntimeWorkerFactory::new(root.path()).with_runtime_store_dir(root.path()); + let worker_ref = WorkerRef::new(WorkerId::now_v7()); + let request = direct_restore_request(worker_ref.clone()); + let run = factory + .worker_restore_run_dir(&worker_ref, request.operation_id) + .unwrap(); + fs::create_dir_all(run.parent().unwrap()).unwrap(); + fs::write(&run, "invalid restore artifact").unwrap(); + let sessions = Arc::new(WorkdirSessionRouter::new()); + factory.failed_restore_sessions.lock().unwrap().insert( + worker_ref.clone(), + (request.operation_id, Arc::clone(&sessions)), + ); + assert!( + factory + .cleanup_failed_restore(&request) + .await + .unwrap_err() + .contains("unsafe restore artifact") + ); + assert!( + factory + .reconcile_stopped_worker(&worker_ref) + .await + .unwrap_err() + .contains("unsafe restore artifact") + ); + let retained = factory + .failed_restore_sessions + .lock() + .unwrap() + .get(&worker_ref) + .cloned() + .unwrap(); + assert_eq!(retained.0, request.operation_id); + assert!(Arc::ptr_eq(&retained.1, &sessions)); + fs::remove_file(&run).unwrap(); + fs::create_dir(&run).unwrap(); + fs::write(run.join("worker.out.log"), "stale restore artifact").unwrap(); + factory.reconcile_stopped_worker(&worker_ref).await.unwrap(); + assert!(!run.exists()); + assert!(factory.failed_restore_sessions.lock().unwrap().is_empty()); + } + + #[cfg(unix)] + #[test] + fn stopped_restore_artifact_cleanup_rejects_symlinks_and_noncanonical_names() { + use std::os::unix::fs::symlink; + let root = tempfile::tempdir().unwrap(); + let outside = tempfile::tempdir().unwrap(); + fs::write(outside.path().join("keep"), "outside authority").unwrap(); + let factory = + ProfileRuntimeWorkerFactory::new(root.path()).with_runtime_store_dir(root.path()); + let worker_ref = WorkerRef::new(WorkerId::now_v7()); + let run = factory + .worker_restore_run_dir( + &worker_ref, + crate::execution::WorkerLifecycleOperationId::new(), + ) + .unwrap(); + let runs = run.parent().unwrap(); + fs::create_dir_all(runs).unwrap(); + symlink(outside.path(), &run).unwrap(); + assert!( + factory + .stopped_restore_run_dirs(&worker_ref) + .unwrap_err() + .contains("unsafe restore artifact") + ); + fs::remove_file(&run).unwrap(); + fs::create_dir(&run).unwrap(); + symlink(outside.path(), run.join("escaped")).unwrap(); + assert!(factory.stopped_restore_run_dirs(&worker_ref).is_err()); + assert!(run.exists()); + fs::remove_file(run.join("escaped")).unwrap(); + let invalid = runs.join("restore-not-an-operation"); + fs::create_dir(&invalid).unwrap(); + assert!( + factory + .stopped_restore_run_dirs(&worker_ref) + .unwrap_err() + .contains("invalid restore artifact operation name") + ); + fs::remove_dir(&invalid).unwrap(); + let noncanonical = runs.join("restore-550E8400-E29B-41D4-A716-446655440000"); + fs::create_dir(&noncanonical).unwrap(); + assert!( + factory + .stopped_restore_run_dirs(&worker_ref) + .unwrap_err() + .contains("noncanonical restore artifact operation name") + ); + fs::remove_dir(&noncanonical).unwrap(); + fs::remove_dir(&run).unwrap(); + fs::remove_dir(runs).unwrap(); + symlink(outside.path(), runs).unwrap(); + assert!( + factory + .stopped_restore_run_dirs(&worker_ref) + .unwrap_err() + .contains("unsafe restore artifact directory") + ); + assert_eq!( + fs::read_to_string(outside.path().join("keep")).unwrap(), + "outside authority" + ); + } + + struct PendingFailureFactory { + restore_calls: Arc, + cleanup_calls: Arc, + cleanup_allowed: Arc, + } + + #[async_trait] + impl RuntimeWorkerFactory for PendingFailureFactory { + async fn spawn_controller( + &self, + _request: WorkerExecutionSpawnRequest, + ) -> Result { + panic!("stop must not spawn a Controller") + } + async fn restore_controller( + &self, + _request: WorkerExecutionRestoreRequest, + ) -> Result { + self.restore_calls.fetch_add(1, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(50)).await; + Err("injected uncertain restore side effect".to_string()) + } + async fn cleanup_failed_restore( + &self, + _request: &WorkerExecutionRestoreRequest, + ) -> Result<(), String> { + self.cleanup_calls.fetch_add(1, Ordering::SeqCst); + if self.cleanup_allowed.load(Ordering::SeqCst) { + Ok(()) + } else { + Err("injected pending restore cleanup failure".to_string()) + } + } + } + + fn direct_restore_request(worker_ref: WorkerRef) -> WorkerExecutionRestoreRequest { + WorkerExecutionRestoreRequest { + operation_id: crate::execution::WorkerLifecycleOperationId::new(), + request: create_request("pending operation"), + workspace_scope: None, + context: test_execution_context(worker_ref.clone()), + worker_ref, + previous_workdir_attachments: Vec::new(), + logical_workdir_attachments: Vec::new(), + workdir_attachments: BTreeMap::new(), + config_bundle: None, + } + } + + #[test] + fn absent_execution_stop_reconciles_pending_restore_cleanup_without_launching() { + let restore_calls = Arc::new(AtomicUsize::new(0)); + let cleanup_calls = Arc::new(AtomicUsize::new(0)); + let cleanup_allowed = Arc::new(AtomicBool::new(false)); + let backend = WorkerRuntimeExecutionBackend::new(PendingFailureFactory { + restore_calls: Arc::clone(&restore_calls), + cleanup_calls: Arc::clone(&cleanup_calls), + cleanup_allowed: Arc::clone(&cleanup_allowed), + }) + .unwrap() + .with_spawn_restore_timeout(Duration::from_millis(1)); + let worker_ref = WorkerRef::new(WorkerId::now_v7()); + let request = direct_restore_request(worker_ref.clone()); + let operation_id = request.operation_id; + assert!(matches!( + backend.restore_worker(request.clone()), + WorkerExecutionSpawnResult::ReconciliationRequired { .. } + )); + assert!(backend.workers.lock().unwrap().is_empty()); + assert_eq!( + backend + .pending_restore(&worker_ref) + .unwrap() + .unwrap() + .request + .operation_id, + operation_id + ); + // Same operation must await/reuse the original future, never start again. + assert!(matches!( + backend.reconcile_restore(request.clone()), + WorkerExecutionSpawnResult::ReconciliationRequired { .. } + )); + let pending = backend.preflight_restore(&request).unwrap_err(); + assert_eq!( + pending.outcome, + crate::execution::WorkerExecutionOutcome::Busy + ); + assert!( + pending.worker_state.is_none(), + "pending resources must not attest AlreadyConnected" + ); + let stop = WorkerExecutionStopRequest { + operation_id: crate::execution::WorkerLifecycleOperationId::new(), + worker_ref: worker_ref.clone(), + }; + let failed = backend.stop_worker_operation(stop.clone()); + assert_eq!( + failed.outcome, + crate::execution::WorkerExecutionOutcome::Errored + ); + assert!( + failed + .message + .unwrap() + .contains("pending restore cleanup failure") + ); + assert_eq!( + backend + .pending_restore(&worker_ref) + .unwrap() + .unwrap() + .request + .operation_id, + operation_id + ); + cleanup_allowed.store(true, Ordering::SeqCst); + assert!(backend.stop_worker_operation(stop.clone()).is_accepted()); + assert!(backend.pending_restore(&worker_ref).unwrap().is_none()); + assert!( + backend.stop_worker_operation(stop).is_accepted(), + "cleanup then catalog-commit retry must converge" + ); + assert_eq!(restore_calls.load(Ordering::SeqCst), 1); + assert_eq!(cleanup_calls.load(Ordering::SeqCst), 2); + } + + struct LateControllerFactory { + inner: MockFactory, + restores: Arc, + release: Arc, + actual_handles: Arc>>, + } + + #[async_trait] + impl RuntimeWorkerFactory for LateControllerFactory { + async fn spawn_controller( + &self, + request: WorkerExecutionSpawnRequest, + ) -> Result { + self.inner.spawn_controller(request).await + } + async fn restore_controller( + &self, + request: WorkerExecutionRestoreRequest, + ) -> Result { + let first = self.restores.fetch_add(1, Ordering::SeqCst) == 0; + if first { + self.release.notified().await; + } + let controller = self.inner.restore_controller(request).await?; + self.actual_handles + .lock() + .unwrap() + .push(controller.handle.clone()); + Ok(controller) + } + } + + #[test] + fn stop_joins_late_restore_controller_and_old_methods_do_not_follow_new_execution() { + let root = tempfile::tempdir().unwrap(); + let restores = Arc::new(AtomicUsize::new(0)); + let release = Arc::new(tokio::sync::Notify::new()); + let actual_handles = Arc::new(Mutex::new(Vec::new())); + let backend = WorkerRuntimeExecutionBackend::new(LateControllerFactory { + inner: MockFactory { + client: MockClient::sequential(vec![]), + runtime_base: root.path().join("runtime"), + cwd: root.path().to_path_buf(), + store_dir: root.path().join("sessions"), + worker_metadata_dir: root.path().join("workers"), + observed_cwds: Arc::new(Mutex::new(Vec::new())), + observed_workspace_clients: Arc::new(Mutex::new(Vec::new())), + }, + restores: Arc::clone(&restores), + release: Arc::clone(&release), + actual_handles: Arc::clone(&actual_handles), + }) + .unwrap() + .with_spawn_restore_timeout(Duration::from_millis(1)); + let worker_ref = WorkerRef::new(WorkerId::now_v7()); + assert!(matches!( + backend.restore_worker(direct_restore_request(worker_ref.clone())), + WorkerExecutionSpawnResult::ReconciliationRequired { .. } + )); + assert!(backend.workers.lock().unwrap().is_empty()); + release.notify_one(); + assert!(backend.stop_worker(&worker_ref).is_accepted()); + assert_eq!( + restores.load(Ordering::SeqCst), + 1, + "stop must await, not relaunch" + ); + assert!(backend.workers.lock().unwrap().is_empty()); + assert!(backend.pending_restore(&worker_ref).unwrap().is_none()); + #[cfg(feature = "ws-server")] + assert!(backend.worker_snapshot(&worker_ref).is_none()); + let old_handle = actual_handles.lock().unwrap()[0].clone(); + // Restore the same identity, but long-lived consumers keep the actual + // old Controller transport rather than acquiring a WorkerRef token. + let mut backend = backend; + backend.spawn_restore_timeout = Duration::from_secs(5); + assert!(matches!( + backend.restore_worker(direct_restore_request(worker_ref.clone())), + WorkerExecutionSpawnResult::Connected { .. } + )); + let old_send = backend.run_on_adapter_runtime(async move { + old_handle + .send(Method::ListRewindTargets) + .await + .map_err(|error| error.to_string()) + }); + assert!( + old_send.is_err(), + "old protocol transport must stay closed after restore" + ); + assert_eq!(restores.load(Ordering::SeqCst), 2); + assert!(backend.stop_worker(&worker_ref).is_accepted()); + #[cfg(feature = "ws-server")] + assert!(backend.worker_snapshot(&worker_ref).is_none()); + } + + #[test] + fn absent_execution_has_typed_input_rejection_and_safe_stop() { + let backend = WorkerRuntimeExecutionBackend::new(FailingFactory).unwrap(); + let worker_ref = WorkerRef::new(WorkerId::now_v7()); + let rejected = backend.dispatch_input(&worker_ref, WorkerInput::user("input")); + assert_eq!(rejected.operation, WorkerExecutionOperation::Input); + assert_eq!( + rejected.outcome, + crate::execution::WorkerExecutionOutcome::Rejected + ); + assert!(rejected.message.unwrap().contains("live Worker execution")); + assert!(backend.stop_worker(&worker_ref).is_accepted()); + assert!(backend.workers.lock().unwrap().is_empty()); + } + + #[test] + fn poisoned_pending_resources_are_not_absence_evidence_for_stop() { + let backend = WorkerRuntimeExecutionBackend::new(FailingFactory).unwrap(); + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _guard = backend.pending_restores.lock().unwrap(); + panic!("injected resource registry poison"); + })); + let stopped = backend.stop_worker(&WorkerRef::new(WorkerId::now_v7())); + assert_eq!( + stopped.outcome, + crate::execution::WorkerExecutionOutcome::Errored + ); + assert!( + stopped + .message + .unwrap() + .contains("pending restore resources lock is poisoned") + ); + } + #[test] fn adapter_runtime_reports_task_panic() { let backend = WorkerRuntimeExecutionBackend::new(FailingFactory).unwrap(); @@ -4904,6 +8146,7 @@ mod tests { let worker_aggregate_dir = runtime_store_dir .join("workers") .join(worker_ref.worker_id.to_string()); + WorkerSessionStore::new(worker_aggregate_dir.join("session")).unwrap(); let worker_name = ProfileRuntimeWorkerFactory::runtime_worker_name_for_ref(&worker_ref); let session_id = session_store::new_session_id(); WorkerAggregateStore::new(&worker_aggregate_dir, &worker_name) @@ -4974,6 +8217,7 @@ mod tests { let worker_aggregate_dir = runtime_store_dir .join("workers") .join(worker_ref.worker_id.to_string()); + WorkerSessionStore::new(worker_aggregate_dir.join("session")).unwrap(); let worker_name = ProfileRuntimeWorkerFactory::runtime_worker_name_for_ref(&worker_ref); let session_id = session_store::new_session_id(); let manifest = manifest::WorkerManifest::from_toml(&format!( @@ -5179,8 +8423,8 @@ mod tests { ); assert!(!first_run_socket.exists()); - let handle = WorkerExecutionHandle::new(worker.worker_ref.clone(), backend.backend_id()); - assert!(backend.stop_worker(&handle).is_accepted()); + let worker_ref = worker.worker_ref.clone(); + assert!(backend.stop_worker(&worker_ref).is_accepted()); drop(runtime); drop(backend); @@ -5262,8 +8506,8 @@ mod tests { }; let first = backend.restore_worker(restore_request.clone()); - let first_handle = match first { - WorkerExecutionSpawnResult::Connected { handle, .. } => handle, + match first { + WorkerExecutionSpawnResult::Connected { .. } => (), other => panic!("initial restore was not connected: {other:?}"), }; let factory_calls_after_restore = observed_workspace_clients.lock().unwrap().len(); @@ -5284,7 +8528,7 @@ mod tests { assert_eq!(workers.len(), 1); drop(workers); - assert!(backend.stop_worker(&first_handle).is_accepted()); + assert!(backend.stop_worker(&worker.worker_ref).is_accepted()); } #[test] @@ -5508,6 +8752,7 @@ mod tests { uuid::Uuid::parse_str(submission_id).expect("opaque submission id is a UUID"); } + #[cfg(feature = "ws-server")] #[test] fn adapter_dispatches_user_input_through_worker_run_lifecycle() { let client = MockClient::new(simple_text_events()); @@ -5933,36 +9178,6 @@ mod tests { let detail = runtime .create_worker(create_request("restore-after-stop")) .unwrap(); - let failed_task = backend - .spawn_on_adapter_runtime(async { panic!("injected owned task failure") }) - .unwrap(); - backend - .workers - .lock() - .unwrap() - .get(&detail.worker_ref) - .unwrap() - .tasks - .push("injected failure", failed_task, false); - - let first_stop = runtime.stop_worker(&detail.worker_ref, None).unwrap_err(); - assert!( - first_stop - .to_string() - .contains("injected failure task failed") - ); - assert_eq!( - runtime.worker_detail(&detail.worker_ref).unwrap().status, - crate::catalog::WorkerStatus::Idle - ); - assert!( - backend - .workers - .lock() - .unwrap() - .contains_key(&detail.worker_ref), - "failed cleanup must retain retry authority" - ); runtime.stop_worker(&detail.worker_ref, None).unwrap(); assert_eq!( runtime.worker_detail(&detail.worker_ref).unwrap().status, diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index 65fbdd8a0..fde6c2427 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -6869,7 +6869,7 @@ mod tests { fn dispatch_input( &self, - _handle: &worker_runtime::execution::WorkerExecutionHandle, + _worker_ref: &EmbeddedWorkerRef, _input: EmbeddedWorkerInput, ) -> worker_runtime::execution::WorkerExecutionResult { worker_runtime::execution::WorkerExecutionResult::rejected( @@ -6898,7 +6898,7 @@ mod tests { fn dispatch_input( &self, - _handle: &worker_runtime::execution::WorkerExecutionHandle, + _worker_ref: &EmbeddedWorkerRef, _input: EmbeddedWorkerInput, ) -> worker_runtime::execution::WorkerExecutionResult { unreachable!("Repository observation test does not dispatch Worker input") @@ -6978,10 +6978,6 @@ mod tests { .unwrap() .insert(request.worker_ref.clone(), request.context); worker_runtime::execution::WorkerExecutionSpawnResult::Connected { - handle: worker_runtime::execution::WorkerExecutionHandle::new( - request.worker_ref, - self.backend_id(), - ), worker_state: protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, @@ -7000,15 +6996,10 @@ mod tests { fn dispatch_input( &self, - handle: &worker_runtime::execution::WorkerExecutionHandle, + worker_ref: &EmbeddedWorkerRef, input: EmbeddedWorkerInput, ) -> worker_runtime::execution::WorkerExecutionResult { - let context = self - .contexts - .lock() - .unwrap() - .get(handle.worker_ref()) - .cloned(); + let context = self.contexts.lock().unwrap().get(worker_ref).cloned(); let Some(context) = context else { return worker_runtime::execution::WorkerExecutionResult::rejected( worker_runtime::execution::WorkerExecutionOperation::Input, diff --git a/crates/workspace-server/src/runtime_subscription_tests.rs b/crates/workspace-server/src/runtime_subscription_tests.rs index b58cd991b..1ed9375b0 100644 --- a/crates/workspace-server/src/runtime_subscription_tests.rs +++ b/crates/workspace-server/src/runtime_subscription_tests.rs @@ -3,10 +3,10 @@ use protocol::subscription::{SubscriptionWorkerIds, SubscriptionWorkerState}; use worker_runtime::Runtime; use worker_runtime::catalog::{CreateWorkerRequest, ProfileSelector, ProfileSourceArchiveSource}; use worker_runtime::execution::{ - WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, WorkerExecutionResult, + WorkerExecutionBackend, WorkerExecutionOperation, WorkerExecutionResult, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, }; -use worker_runtime::identity::WorkerId; +use worker_runtime::identity::{WorkerId, WorkerRef}; #[derive(Debug)] struct TestExecutionBackend; @@ -16,9 +16,8 @@ impl WorkerExecutionBackend for TestExecutionBackend { "runtime-subscription-test" } - fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { + fn spawn_worker(&self, _request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { WorkerExecutionSpawnResult::connected( - WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, @@ -28,7 +27,7 @@ impl WorkerExecutionBackend for TestExecutionBackend { fn dispatch_input( &self, - _handle: &WorkerExecutionHandle, + _worker_ref: &WorkerRef, input: worker_runtime::interaction::WorkerInput, ) -> WorkerExecutionResult { if let Some(submission_request_id) = input.submission_request_id { @@ -43,7 +42,7 @@ impl WorkerExecutionBackend for TestExecutionBackend { } } - fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { + fn stop_worker(&self, _worker_ref: &WorkerRef) -> WorkerExecutionResult { WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop) } } diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index cd4463fb6..ed7afcead 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -175,8 +175,8 @@ use crate::memory_staging::{ list_memory_staging_from_authority, memory_staging_backlog_from_authority, }; use crate::observation::{ - BackendObservationProxy, ObservationProxyError, RuntimeObservationClient, - RuntimeObservationSource, RuntimeObservationSourceConfig, + BackendObservationProxy, ObservationProxyError, RuntimeObservationSource, + RuntimeObservationSourceConfig, }; use crate::records::{ MergeRequestListItem, MergeRequestListResponse, MergeRequestRefDiagnostic, ObjectiveDetail, @@ -29832,8 +29832,9 @@ async fn worker_protocol_ws( } }; let input_source = authenticated_browser_input_source(&actor); + let scope = worker_runtime::RuntimeWorkspaceScope::new(api.workspace_id(), "embedded-backend"); ws.on_upgrade(move |socket| { - worker_protocol_ws_session(source, socket, input_source, job_read_only) + worker_protocol_ws_session(source, socket, input_source, job_read_only, scope) }) } @@ -29866,7 +29867,9 @@ pub(crate) async fn connect_workspace_worker_protocol( connect_remote_worker_protocol(config, input_source).await } RuntimeObservationSource::Embedded(source) => { - connect_embedded_worker_protocol(source).await + let scope = + worker_runtime::RuntimeWorkspaceScope::new(api.workspace_id(), "embedded-backend"); + connect_embedded_worker_protocol(source, scope).await } } } @@ -29948,38 +29951,55 @@ async fn connect_remote_worker_protocol( async fn connect_embedded_worker_protocol( source: crate::observation::EmbeddedRuntimeObservationSource, + scope: worker_runtime::RuntimeWorkspaceScope, ) -> Result { - let mut upstream = - RuntimeObservationClient::connect(&RuntimeObservationSource::Embedded(source.clone())) - .await - .map_err(|error| Error::RuntimeOperationFailed { - runtime_id: source.worker.runtime_id.clone(), - code: error.code().to_string(), - message: error.message().to_string(), - })?; + // Capture one real execution endpoint, including its greeting/snapshot. + // The WorkerRef-wide Workspace observation bus is a separate catalogue + // channel and must not reconnect this operational session after Restore. + let mut transport = source + .runtime + .attach_worker_protocol_scoped(&scope, &source.worker_ref) + .map_err(|error| Error::RuntimeOperationFailed { + runtime_id: source.worker.runtime_id.clone(), + code: "embedded_worker_protocol_connect_failed".to_string(), + message: error.to_string(), + })?; let (methods, mut method_receiver) = tokio::sync::mpsc::channel(256); let (event_sender, events) = tokio::sync::mpsc::channel(512); tokio::spawn(async move { + if event_sender.send(transport.snapshot.clone()).await.is_err() { + return; + } loop { tokio::select! { + _ = event_sender.closed() => break, method = method_receiver.recv() => { let Some(method) = method else { break }; - match source.runtime.send_protocol_method(&source.worker_ref, method) { + let shutdown = matches!(&method, protocol::Method::Shutdown { .. }); + match source.runtime.send_connected_protocol_method_scoped( + &scope, &source.worker_ref, &transport, method, + ) { Ok(direct_events) => { for event in direct_events { - if event_sender.send(event).await.is_err() { return; } + let terminal = matches!(event, protocol::Event::Shutdown); + if event_sender.send(event).await.is_err() || terminal { return; } } } Err(error) => { - if event_sender.send(protocol_error_event(error.to_string())).await.is_err() { return; } + let terminal = embedded_protocol_dispatch_is_terminal(&error, &transport); + if event_sender.send(protocol_error_event(error.to_string())).await.is_err() || terminal { return; } } } + if shutdown { break; } } - event = upstream.next_event() => match event { - Ok(event) => { - if event_sender.send(event.payload).await.is_err() { break; } + event = transport.events.recv() => match event { + Some(event) => { + let terminal = matches!(event, protocol::Event::Shutdown); + if event_sender.send(event).await.is_err() || terminal { break; } } - Err(_) => break, + // The execution relay closes on Stop, disconnect or lag. + // Uncertain streams are terminal, never grounds to reattach. + None => break, } } } @@ -29987,18 +30007,37 @@ async fn connect_embedded_worker_protocol( Ok(WorkspaceWorkerProtocolConnection { methods, events }) } +fn embedded_protocol_dispatch_is_terminal( + error: &worker_runtime::error::RuntimeError, + transport: &worker_runtime::execution::WorkerProtocolTransport, +) -> bool { + // A settled, side-effect-free method rejection can leave the connection + // usable. Uncertain lifecycle/store/transport errors cannot. + !matches!( + error, + worker_runtime::error::RuntimeError::WorkerExecutionRejected { + outcome: worker_runtime::execution::WorkerExecutionOutcome::Busy + | worker_runtime::execution::WorkerExecutionOutcome::Rejected + | worker_runtime::execution::WorkerExecutionOutcome::Unsupported, + .. + } + ) || transport.validate().is_err() +} + async fn worker_protocol_ws_session( source: RuntimeObservationSource, socket: WebSocket, input_source: protocol::AuthenticatedInputSource, read_only: bool, + scope: worker_runtime::RuntimeWorkspaceScope, ) { match source { RuntimeObservationSource::RemoteWs(config) => { remote_worker_protocol_ws_session(config, socket, input_source, read_only).await; } RuntimeObservationSource::Embedded(source) => { - embedded_worker_protocol_ws_session(source, socket, input_source, read_only).await; + embedded_worker_protocol_ws_session(source, socket, input_source, read_only, scope) + .await; } } } @@ -30158,19 +30197,23 @@ async fn embedded_worker_protocol_ws_session( mut socket: WebSocket, input_source: protocol::AuthenticatedInputSource, read_only: bool, + scope: worker_runtime::RuntimeWorkspaceScope, ) { - let mut upstream = match RuntimeObservationClient::connect(&RuntimeObservationSource::Embedded( - source.clone(), - )) - .await + let mut transport = match source + .runtime + .attach_worker_protocol_scoped(&scope, &source.worker_ref) { - Ok(client) => client, + Ok(transport) => transport, Err(error) => { - let event = protocol_error_event(error.message()); + let event = protocol_error_event(error.to_string()); let _ = send_protocol_event(&mut socket, &event).await; + let _ = socket.send(WsMessage::Close(None)).await; return; } }; + if !send_protocol_event(&mut socket, &transport.snapshot).await { + return; + } loop { tokio::select! { @@ -30184,24 +30227,44 @@ async fn embedded_worker_protocol_ws_session( } Some(Ok(WsMessage::Text(text))) => match decode_method(&text) { Ok(method) => match authorize_browser_worker_method(method, &input_source) { - Ok(method) => match source.runtime.send_protocol_method(&source.worker_ref, method) { - Ok(events) => { - for event in events { + Ok(method) => { + let shutdown = matches!(&method, protocol::Method::Shutdown { .. }); + match source.runtime.send_connected_protocol_method_scoped( + &scope, &source.worker_ref, &transport, method, + ) { + Ok(events) => { + for event in events { + let terminal = matches!(event, protocol::Event::Shutdown); + if !send_protocol_event(&mut socket, &event).await { + return; + } + if terminal { + let _ = socket.send(WsMessage::Close(None)).await; + return; + } + } + } + Err(error) => { + let terminal = embedded_protocol_dispatch_is_terminal(&error, &transport); + let event = protocol_error_event(error.to_string()); if !send_protocol_event(&mut socket, &event).await { return; } + if terminal { + let _ = socket.send(WsMessage::Close(None)).await; + return; + } } } - Err(error) => { - let event = protocol_error_event(error.to_string()); - if !send_protocol_event(&mut socket, &event).await { - return; - } + if shutdown { + let _ = socket.send(WsMessage::Close(None)).await; + return; } }, Err(message) => { let event = protocol_error_event(message); let _ = send_protocol_event(&mut socket, &event).await; + let _ = socket.send(WsMessage::Close(None)).await; return; } }, @@ -30228,16 +30291,22 @@ async fn embedded_worker_protocol_ws_session( } } } - upstream_event = upstream.next_event() => { + upstream_event = transport.events.recv() => { match upstream_event { - Ok(event) => { - if !send_protocol_event(&mut socket, &event.payload).await { + Some(event) => { + let terminal = matches!(event, protocol::Event::Shutdown); + if !send_protocol_event(&mut socket, &event).await { + return; + } + if terminal { + let _ = socket.send(WsMessage::Close(None)).await; return; } } - Err(error) => { - let event = protocol_error_event(error.message()); - let _ = send_protocol_event(&mut socket, &event).await; + // Stop/disconnect/lag closes the execution-owned relay. + // Never follow the WorkerRef to a restored Controller. + None => { + let _ = socket.send(WsMessage::Close(None)).await; return; } } @@ -42971,11 +43040,125 @@ mod tests { assert!(!workdir_runtime_miss_is_not_found(&unrelated)); } + // The receiver belongs to one mock execution, not to a Runtime WorkerRecord. + // Bound protocol senders become terminal when Stop drops this execution. + struct DeterministicWorkerExecution { + context: worker_runtime::execution::WorkerExecutionContext, + state: std::sync::Mutex, + method_tx: tokio::sync::mpsc::Sender, + method_rx: std::sync::Mutex>>, + events: tokio::sync::broadcast::Sender, + protocol_methods: + Arc>>, + } + + impl DeterministicWorkerExecution { + fn new( + context: worker_runtime::execution::WorkerExecutionContext, + protocol_methods: Arc< + std::sync::Mutex>, + >, + ) -> Self { + let (method_tx, method_rx) = tokio::sync::mpsc::channel(32); + let (events, _) = tokio::sync::broadcast::channel(32); + Self { + context, + state: std::sync::Mutex::new(protocol::WorkerStateSnapshot::initial()), + method_tx, + method_rx: std::sync::Mutex::new(Some(method_rx)), + events, + protocol_methods, + } + } + + fn unavailable() -> worker_runtime::execution::WorkerExecutionResult { + worker_runtime::execution::WorkerExecutionResult::rejected( + worker_runtime::execution::WorkerExecutionOperation::ProtocolMethod, + "deterministic test execution protocol is closed", + ) + } + + fn dispatch( + &self, + method: protocol::Method, + ) -> std::result::Result< + Vec, + worker_runtime::execution::WorkerExecutionResult, + > { + // Drive the actual mock receiver synchronously; no method is routed + // through a WorkerRef lookup after a protocol connection is attached. + let mut receiver = self.method_rx.lock().unwrap(); + let receiver = receiver.as_mut().ok_or_else(Self::unavailable)?; + self.method_tx + .try_send(method) + .map_err(|_| Self::unavailable())?; + let method = receiver.try_recv().map_err(|_| Self::unavailable())?; + self.protocol_methods + .lock() + .unwrap() + .push((self.context.worker_ref().clone(), method.clone())); + Ok(match method { + protocol::Method::ListCompletions { + kind, + prefix, + request_id, + context, + } => vec![protocol::Event::Completions { + kind, + prefix, + request_id, + context, + entries: Vec::new(), + }], + _ => Vec::new(), + }) + } + + fn publish_event( + &self, + event: protocol::Event, + ) -> std::result::Result<(), worker_runtime::error::RuntimeError> { + if let protocol::Event::WorkerState { snapshot } = &event { + *self.state.lock().unwrap() = snapshot.clone(); + } + let _ = self.events.send(event.clone()); + self.context.publish_protocol_event(event).map(|_| ()) + } + + fn snapshot(&self) -> protocol::Event { + protocol::Event::Snapshot { + session: protocol::SessionSnapshot { + pending_submissions: protocol::PendingSubmissionsSnapshot::default(), + entries: Vec::new(), + }, + greeting: protocol::Greeting { + worker_name: self.context.worker_ref().worker_id.to_string(), + cwd: String::new(), + provider: "deterministic-workspace-server-test".to_string(), + model: "deterministic-workspace-server-test".to_string(), + reasoning: None, + scope_summary: "test execution snapshot".to_string(), + tools: Vec::new(), + context_window: 0, + context_tokens: 0, + context_usage: None, + }, + state: self.state.lock().unwrap().clone(), + in_flight: protocol::InFlightSnapshot { + blocks: Vec::new(), + commands: Vec::new(), + compaction: None, + }, + internal_workers: Vec::new(), + } + } + } + struct DeterministicExecutionBackend { contexts: std::sync::Mutex< std::collections::HashMap< worker_runtime::identity::WorkerRef, - worker_runtime::execution::WorkerExecutionContext, + Arc, >, >, materializer: worker_runtime::working_directory::RuntimeGitMaterializer, @@ -42988,7 +43171,8 @@ mod tests { accept_restores: std::sync::atomic::AtomicBool, inputs: std::sync::Mutex>, protocol_methods: - std::sync::Mutex>, + Arc>>, + stops: std::sync::Mutex>, } impl Default for DeterministicExecutionBackend { @@ -43012,7 +43196,8 @@ mod tests { input_request_ids: std::sync::Mutex::new(Vec::new()), accept_restores: std::sync::atomic::AtomicBool::new(false), inputs: std::sync::Mutex::new(Vec::new()), - protocol_methods: std::sync::Mutex::new(Vec::new()), + protocol_methods: Arc::new(std::sync::Mutex::new(Vec::new())), + stops: std::sync::Mutex::new(Vec::new()), } } } @@ -43039,19 +43224,34 @@ mod tests { worker: &RuntimeWorkerRef, snapshot: protocol::WorkerStateSnapshot, ) { - let context = self + let execution = self .contexts .lock() .unwrap() .iter() .find(|(worker_ref, _)| worker_ref.worker_id.to_string() == worker.worker_id) - .map(|(_, context)| context.clone()) - .expect("execution context"); - context - .publish_protocol_event(protocol::Event::WorkerState { snapshot }) + .map(|(_, execution)| execution.clone()) + .expect("mock execution"); + execution + .publish_event(protocol::Event::WorkerState { snapshot }) .unwrap(); } + fn emit_worker_event( + &self, + worker_ref: &worker_runtime::identity::WorkerRef, + event: protocol::Event, + ) { + let execution = self + .contexts + .lock() + .unwrap() + .get(worker_ref) + .cloned() + .expect("mock execution"); + execution.publish_event(event).unwrap(); + } + fn take_input_request_ids(&self) -> Vec> { std::mem::take( &mut *self @@ -43139,15 +43339,14 @@ mod tests { working_directory: binding.status(), }) .collect(); - self.contexts - .lock() - .unwrap() - .insert(request.worker_ref.clone(), request.context); + self.contexts.lock().unwrap().insert( + request.worker_ref.clone(), + Arc::new(DeterministicWorkerExecution::new( + request.context, + self.protocol_methods.clone(), + )), + ); worker_runtime::execution::WorkerExecutionSpawnResult::Connected { - handle: worker_runtime::execution::WorkerExecutionHandle::new( - request.worker_ref, - self.backend_id(), - ), worker_state: protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, @@ -43155,6 +43354,20 @@ mod tests { } } + fn preflight_restore( + &self, + request: &worker_runtime::execution::WorkerExecutionRestoreRequest, + ) -> std::result::Result<(), worker_runtime::execution::WorkerExecutionResult> { + if let Some(execution) = self.contexts.lock().unwrap().get(&request.worker_ref) { + return Err(worker_runtime::execution::WorkerExecutionResult::busy( + worker_runtime::execution::WorkerExecutionOperation::Restore, + "deterministic test backend already has a live execution", + ) + .with_worker_state(execution.state.lock().unwrap().clone())); + } + Ok(()) + } + fn restore_worker( &self, request: worker_runtime::execution::WorkerExecutionRestoreRequest, @@ -43178,15 +43391,14 @@ mod tests { working_directory: binding.status(), }) .collect(); - self.contexts - .lock() - .unwrap() - .insert(request.worker_ref.clone(), request.context); + self.contexts.lock().unwrap().insert( + request.worker_ref.clone(), + Arc::new(DeterministicWorkerExecution::new( + request.context, + self.protocol_methods.clone(), + )), + ); worker_runtime::execution::WorkerExecutionSpawnResult::Connected { - handle: worker_runtime::execution::WorkerExecutionHandle::new( - request.worker_ref, - self.backend_id(), - ), worker_state: protocol::WorkerStateSnapshot::initial(), workdir_attachments, } @@ -43194,53 +43406,88 @@ mod tests { fn dispatch_method( &self, - handle: &worker_runtime::execution::WorkerExecutionHandle, + worker_ref: &worker_runtime::identity::WorkerRef, method: protocol::Method, ) -> worker_runtime::execution::WorkerExecutionResult { - self.protocol_methods + let execution = self.contexts.lock().unwrap().get(worker_ref).cloned(); + let Some(execution) = execution else { + return DeterministicWorkerExecution::unavailable(); + }; + match execution.dispatch(method) { + Ok(_) => worker_runtime::execution::WorkerExecutionResult::accepted( + worker_runtime::execution::WorkerExecutionOperation::ProtocolMethod, + ), + Err(result) => result, + } + } + + fn attach_worker_protocol( + self: Arc, + worker_ref: &worker_runtime::identity::WorkerRef, + ) -> std::result::Result< + worker_runtime::execution::WorkerProtocolTransport, + worker_runtime::execution::WorkerExecutionResult, + > { + let execution = self + .contexts .lock() .unwrap() - .push((handle.worker_ref().clone(), method)); - worker_runtime::execution::WorkerExecutionResult::accepted( - worker_runtime::execution::WorkerExecutionOperation::ProtocolMethod, - ) + .get(worker_ref) + .cloned() + .ok_or_else(DeterministicWorkerExecution::unavailable)?; + let mut execution_events = execution.events.subscribe(); + let snapshot = execution.snapshot(); + let (event_tx, events) = tokio::sync::mpsc::channel(32); + tokio::spawn(async move { + while let Ok(event) = execution_events.recv().await { + if event_tx.send(event).await.is_err() { + break; + } + } + }); + let bound_execution = Arc::downgrade(&execution); + let method_tx = execution.method_tx.clone(); + Ok(worker_runtime::execution::WorkerProtocolTransport::new( + worker_ref.clone(), + snapshot, + events, + Arc::new(move |method| { + bound_execution + .upgrade() + .ok_or_else(DeterministicWorkerExecution::unavailable)? + .dispatch(method) + }), + Arc::new(move || { + if method_tx.is_closed() { + Err(DeterministicWorkerExecution::unavailable()) + } else { + Ok(()) + } + }), + )) } fn worker_snapshot( &self, - handle: &worker_runtime::execution::WorkerExecutionHandle, + worker_ref: &worker_runtime::identity::WorkerRef, ) -> Option { - Some(protocol::Event::Snapshot { - session: protocol::SessionSnapshot { - pending_submissions: protocol::PendingSubmissionsSnapshot::default(), - entries: Vec::new(), - }, - greeting: protocol::Greeting { - worker_name: handle.worker_ref().worker_id.to_string(), - cwd: String::new(), - provider: "deterministic-workspace-server-test".to_string(), - model: "deterministic-workspace-server-test".to_string(), - reasoning: None, - scope_summary: "test execution snapshot".to_string(), - tools: Vec::new(), - context_window: 0, - context_tokens: 0, - context_usage: None, - }, - state: protocol::WorkerStateSnapshot::initial(), - in_flight: protocol::InFlightSnapshot { - blocks: Vec::new(), - commands: Vec::new(), - compaction: None, - }, - internal_workers: Vec::new(), - }) + self.contexts + .lock() + .unwrap() + .get(worker_ref) + .map(|execution| execution.snapshot()) } fn stop_worker( &self, - _handle: &worker_runtime::execution::WorkerExecutionHandle, + worker_ref: &worker_runtime::identity::WorkerRef, ) -> worker_runtime::execution::WorkerExecutionResult { + self.stops.lock().unwrap().push(worker_ref.clone()); + if let Some(execution) = self.contexts.lock().unwrap().remove(worker_ref) { + // Drop the actual receiver even if an in-flight operation still + // briefly owns the execution. Old protocol senders stay closed. + execution.method_rx.lock().unwrap().take(); + } worker_runtime::execution::WorkerExecutionResult::accepted( worker_runtime::execution::WorkerExecutionOperation::Stop, ) @@ -43248,7 +43495,7 @@ mod tests { fn cancel_worker( &self, - _handle: &worker_runtime::execution::WorkerExecutionHandle, + _worker_ref: &worker_runtime::identity::WorkerRef, ) -> worker_runtime::execution::WorkerExecutionResult { worker_runtime::execution::WorkerExecutionResult::accepted( worker_runtime::execution::WorkerExecutionOperation::Cancel, @@ -43257,7 +43504,7 @@ mod tests { fn dispatch_input( &self, - handle: &worker_runtime::execution::WorkerExecutionHandle, + worker_ref: &worker_runtime::identity::WorkerRef, input: worker_runtime::interaction::WorkerInput, ) -> worker_runtime::execution::WorkerExecutionResult { self.input_request_ids @@ -43267,9 +43514,18 @@ mod tests { self.inputs .lock() .expect("inputs lock") - .push((handle.worker_ref().clone(), input.content.clone())); + .push((worker_ref.clone(), input.content.clone())); + let execution = self.contexts.lock().unwrap().get(worker_ref).cloned(); + let Some(execution) = execution else { + return worker_runtime::execution::WorkerExecutionResult::rejected( + worker_runtime::execution::WorkerExecutionOperation::Input, + "deterministic test backend has no live Worker", + ); + }; + let bound_execution = Arc::downgrade(&execution); + drop(execution); if let Some(hook) = self.input_hook.lock().unwrap().take() { - hook(handle.worker_ref()); + hook(worker_ref); } if let Some(message) = self.input_failure.lock().unwrap().clone() { return worker_runtime::execution::WorkerExecutionResult::errored( @@ -43277,21 +43533,16 @@ mod tests { message, ); } - let context = self - .contexts - .lock() - .unwrap() - .get(handle.worker_ref()) - .cloned() - .expect("execution context"); let submission_request_id = input.submission_request_id.clone(); let input_kind = input.kind.clone(); let content = input.content.clone(); std::thread::spawn(move || { std::thread::sleep(std::time::Duration::from_millis(25)); - let _ = context.publish_protocol_event(protocol::Event::TextDone { - text: format!("server companion echoed: {content}"), - }); + if let Some(execution) = bound_execution.upgrade() { + let _ = execution.publish_event(protocol::Event::TextDone { + text: format!("server companion echoed: {content}"), + }); + } }); if input_kind == worker_runtime::interaction::WorkerInputKind::Notify { worker_runtime::execution::WorkerExecutionResult::accepted_notification( @@ -49739,20 +49990,20 @@ mod tests { fn activate_restored_worker( &self, id: attention_execution::WorkerLifecycleOperationId, - handle: &attention_execution::WorkerExecutionHandle, + worker_ref: &worker_runtime::identity::WorkerRef, ) -> std::result::Result<(), String> { - self.inner.activate_restored_worker(id, handle) + self.inner.activate_restored_worker(id, worker_ref) } fn dispatch_input( &self, - handle: &attention_execution::WorkerExecutionHandle, + worker_ref: &worker_runtime::identity::WorkerRef, input: worker_runtime::interaction::WorkerInput, ) -> attention_execution::WorkerExecutionResult { self.attempts .lock() .unwrap() .push(input.submission_request_id.clone().unwrap()); - let result = self.inner.dispatch_input(handle, input); + let result = self.inner.dispatch_input(worker_ref, input); assert_eq!( result.outcome, attention_execution::WorkerExecutionOutcome::Accepted @@ -49771,10 +50022,10 @@ mod tests { } fn dispatch_method( &self, - handle: &attention_execution::WorkerExecutionHandle, + worker_ref: &worker_runtime::identity::WorkerRef, method: protocol::Method, ) -> attention_execution::WorkerExecutionResult { - self.inner.dispatch_method(handle, method) + self.inner.dispatch_method(worker_ref, method) } fn stop_worker_operation( &self, @@ -49784,9 +50035,9 @@ mod tests { } fn worker_snapshot( &self, - handle: &attention_execution::WorkerExecutionHandle, + worker_ref: &worker_runtime::identity::WorkerRef, ) -> Option { - self.inner.worker_snapshot(handle) + self.inner.worker_snapshot(worker_ref) } } @@ -55277,9 +55528,19 @@ mod tests { } fn runtime_with_worker() -> (worker_runtime::Runtime, worker_runtime::identity::WorkerRef) { + let (runtime, worker_ref, _) = runtime_with_worker_backend(); + (runtime, worker_ref) + } + + fn runtime_with_worker_backend() -> ( + worker_runtime::Runtime, + worker_runtime::identity::WorkerRef, + Arc, + ) { + let execution = Arc::new(DeterministicExecutionBackend::default()); let runtime = worker_runtime::Runtime::with_execution_backend( worker_runtime::RuntimeOptions::default(), - Arc::new(DeterministicExecutionBackend::default()), + execution.clone(), ) .unwrap(); runtime.store_config_bundle(runtime_test_bundle()).unwrap(); @@ -55289,7 +55550,7 @@ mod tests { runtime_create_request(), ) .unwrap(); - (runtime, worker.worker_ref) + (runtime, worker.worker_ref, execution) } #[test] @@ -57289,10 +57550,11 @@ mod tests { tokio::time::sleep(std::time::Duration::from_millis(50)).await; let restored_store = test_control_store(&config); + let restored_execution = Arc::new(DeterministicExecutionBackend::default()); let restored = WorkspaceApi::new_with_execution_backend( config, Arc::new(restored_store), - Arc::new(DeterministicExecutionBackend::default()), + restored_execution.clone(), ) .await .expect("restored fs-backed api starts"); @@ -57300,10 +57562,34 @@ mod tests { .runtime .worker(&worker_ref) .expect("restored worker"); - // The durable projection keeps an observer attached, so the restored Runtime - // catalog remains controllable while its failed execution is represented in - // `worker_state` rather than by removing the catalog Worker. + // A read-only automatic restore rejection retains the active catalog Worker, + // but cannot invent a live execution state or protocol snapshot. assert_eq!(restored_worker.state, "idle"); + assert!(restored_worker.worker_state.is_none()); + let crate::observation::RuntimeObservationSource::Embedded(observation) = restored + .runtime + .observation_source(&worker_ref) + .expect("catalog Worker retains its observation source") + else { + panic!("expected embedded observation source"); + }; + let diagnostics = observation.runtime.diagnostics().unwrap(); + assert!( + diagnostics.iter().any(|diagnostic| { + diagnostic.worker_ref.as_ref() == Some(&observation.worker_ref) + && diagnostic + .message + .contains("deterministic test backend restore disabled") + }), + "rejected automatic restore must retain its concrete backend diagnostic: {diagnostics:?}" + ); + assert!(matches!( + observation + .runtime + .worker_observation_snapshot(&observation.worker_ref), + Err(worker_runtime::error::RuntimeError::WorkerExecutionUnavailable { .. }) + )); + assert!(restored_execution.contexts.lock().unwrap().is_empty()); assert!( serde_json::to_value(&restored_worker) .unwrap() @@ -57326,13 +57612,108 @@ mod tests { &worker_ref, WorkerInputRequest { kind: WorkerInputKind::User, - content: "should not be routed to corrupted handle".to_string(), + content: "reject input without live Worker".to_string(), submission_request_id: None, segments: None, }, ) .expect("stale worker input is projected as an operation result"); assert_eq!(rejected_input.state, InternalWorkerOperationState::Rejected); + assert_eq!(rejected_input.disposition, WorkerInputDisposition::Rejected); + assert!(rejected_input.runtime_run_id.is_none()); + assert!(rejected_input.diagnostics.iter().any(|diagnostic| { + diagnostic.code == "embedded_worker_execution_rejected" + && diagnostic + .message + .contains("deterministic test backend has no live Worker") + })); + assert_eq!( + restored_execution.take_inputs(), + vec![( + observation.worker_ref.clone(), + "reject input without live Worker".to_string() + )] + ); + assert!( + restored + .runtime + .worker(&worker_ref) + .unwrap() + .worker_state + .is_none() + ); + + // Exercise the ordinary guarded removal service, not direct catalog deletion. + let Json(orchestrator) = scoped_start_workspace_orchestrator( + State(restored.clone()), + AxumPath(ScopedWorkspacePath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + }), + ) + .await + .unwrap(); + let source = orchestrator.worker.expect("removal controller"); + let source = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id); + sync_worker_observation(&restored, &restored_worker).unwrap(); + seed_worker_control_grant(&restored, &source, &worker_ref, "rejected-restore-remove"); + let verified_source = || crate::worker_source::VerifiedWorkerMutationSource { + runtime_id: source.runtime_id.clone(), + worker_id: source.worker_id.clone(), + actor_kind: worker_runtime::auth::WorkerMutationActorKind::Worker, + permission: worker_runtime::auth::WORKER_REMOVE_PERMISSION.to_string(), + jti: "rejected-restore-remove".to_string(), + }; + let removal = WorkerRemovalService::new(&restored); + let before_stop = removal + .execute_async( + verified_source(), + &worker_ref.runtime_id, + &worker_ref.worker_id, + "rejected restore is not Stop", + ) + .await + .unwrap(); + assert_eq!(before_stop.status, StatusCode::CONFLICT.as_u16()); + assert!(before_stop.body.contains("worker_not_stopped")); + + let stopped = restored + .runtime + .stop_worker( + &worker_ref, + WorkerLifecycleRequest { + reason: Some("retire Worker after rejected automatic restore".to_string()), + ticket_assignment: None, + }, + ) + .unwrap(); + assert_eq!(stopped.state, InternalWorkerOperationState::Accepted); + assert_eq!( + restored_execution.stops.lock().unwrap().as_slice(), + &[observation.worker_ref.clone()] + ); + let stopped_worker = restored.runtime.worker(&worker_ref).unwrap(); + assert_eq!(stopped_worker.state, "stopped"); + sync_worker_observation(&restored, &stopped_worker).unwrap(); + seed_worker_session_for_cleanup(&dir.path().join("workspace"), &worker_ref); + let removed = removal + .execute_async( + verified_source(), + &worker_ref.runtime_id, + &worker_ref.worker_id, + "remove stopped Worker after rejected automatic restore", + ) + .await + .unwrap(); + assert_eq!(removed.status, StatusCode::OK.as_u16(), "{}", removed.body); + assert!(removed.body.contains("\"removed\":true")); + assert!(restored.runtime.worker(&worker_ref).is_err()); + assert!( + restored + .store + .get_worker_registry(TEST_WORKSPACE_ID, &worker_ref) + .unwrap() + .is_none() + ); } #[tokio::test] @@ -57662,7 +58043,7 @@ mod tests { #[tokio::test] async fn proxies_worker_protocol_ws_as_raw_events() { - let (runtime, worker_ref, endpoint) = spawn_runtime_worker().await; + let (_runtime, worker_ref, endpoint, execution) = spawn_runtime_worker().await; let source = RuntimeObservationSourceConfig { worker: RuntimeWorkerRef::new("runtime-a", "worker-a"), endpoint, @@ -57693,23 +58074,617 @@ mod tests { protocol::Event::Completions { .. } )); - runtime - .observe_worker_event( - &worker_ref, - protocol::Event::TextDelta { - text: "live".into(), - }, - ) - .unwrap(); + execution.emit_worker_event( + &worker_ref, + protocol::Event::TextDelta { + text: "live".into(), + }, + ); assert!(matches!( next_client_frame(&mut stream).await, protocol::Event::TextDelta { .. } )); } + fn embedded_protocol_test_source( + runtime: &worker_runtime::Runtime, + worker_ref: &worker_runtime::identity::WorkerRef, + ) -> crate::observation::EmbeddedRuntimeObservationSource { + crate::observation::EmbeddedRuntimeObservationSource { + worker: RuntimeWorkerRef::new("runtime-a", worker_ref.worker_id.to_string()), + runtime: runtime.clone(), + worker_ref: worker_ref.clone(), + } + } + + fn embedded_protocol_test_scope() -> worker_runtime::RuntimeWorkspaceScope { + worker_runtime::RuntimeWorkspaceScope::new("local", "local-token") + } + + fn protocol_test_completion(request_id: &str) -> protocol::Method { + protocol::Method::ListCompletions { + kind: protocol::CompletionKind::File, + prefix: String::new(), + request_id: Some(request_id.to_string()), + context: None, + } + } + + fn delayed_protocol_test_methods() -> Vec { + vec![ + protocol::Method::Submit { + submission_request_id: "delayed-submit-a".to_string(), + input: vec![protocol::Segment::Text { + content: "must not reach restored B".to_string(), + }], + }, + protocol_test_completion("delayed-completion-a"), + protocol::Method::Shutdown { + command: protocol::WorkerCommandEnvelope::new(700), + }, + ] + } + + fn assert_embedded_protocol_snapshot( + event: protocol::Event, + worker_ref: &worker_runtime::identity::WorkerRef, + ) { + let protocol::Event::Snapshot { + greeting, state, .. + } = event + else { + panic!("protocol connection must start with its attached execution snapshot"); + }; + assert_eq!(greeting.worker_name, worker_ref.worker_id.to_string()); + assert_eq!(greeting.provider, "deterministic-workspace-server-test"); + assert_eq!(state, protocol::WorkerStateSnapshot::initial()); + } + + async fn next_embedded_protocol_event( + connection: &mut WorkspaceWorkerProtocolConnection, + ) -> protocol::Event { + tokio::time::timeout(std::time::Duration::from_secs(2), connection.events.recv()) + .await + .expect("embedded protocol event timeout") + .expect("embedded protocol unexpectedly closed") + } + + async fn assert_embedded_protocol_closed(connection: &mut WorkspaceWorkerProtocolConnection) { + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while let Some(event) = connection.events.recv().await { + // A late method may be rejected before the relay closure is + // consumed. No completion, snapshot or B event is permitted. + assert!(matches!( + event, + protocol::Event::Error { .. } | protocol::Event::Shutdown + )); + } + }) + .await + .expect("old embedded protocol connection did not close"); + assert!(connection.methods.is_closed()); + } + + async fn spawn_embedded_protocol_test_ws( + source: crate::observation::EmbeddedRuntimeObservationSource, + scope: worker_runtime::RuntimeWorkspaceScope, + read_only: bool, + ) -> (TestWebSocket, tokio::task::JoinHandle<()>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let input_source = authenticated_browser_input_source(&test_browser_request_actor()); + let app = Router::new().route( + "/protocol/ws", + get(move |ws: WebSocketUpgrade| { + let source = source.clone(); + let scope = scope.clone(); + let input_source = input_source.clone(); + async move { + ws.on_upgrade(move |socket| { + worker_protocol_ws_session( + RuntimeObservationSource::Embedded(source), + socket, + input_source, + read_only, + scope, + ) + }) + } + }), + ); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let (socket, _) = connect_async(format!("ws://{address}/protocol/ws")) + .await + .unwrap(); + (socket, server) + } + + async fn send_protocol_test_ws_method(socket: &mut TestWebSocket, method: &protocol::Method) { + socket + .send(Message::Text( + protocol::stream::encode_method(method).unwrap().into(), + )) + .await + .unwrap(); + } + + async fn next_embedded_protocol_ws_event(socket: &mut TestWebSocket) -> protocol::Event { + tokio::time::timeout(std::time::Duration::from_secs(2), next_client_frame(socket)) + .await + .expect("embedded protocol WebSocket event timeout") + } + + async fn assert_embedded_protocol_ws_closed(socket: &mut TestWebSocket) { + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + match socket.next().await { + Some(Ok(Message::Text(text))) => { + let event = protocol::stream::decode_event(text.as_ref()).unwrap(); + assert!(matches!( + event, + protocol::Event::Error { .. } | protocol::Event::Shutdown + )); + } + Some(Ok(Message::Close(_))) | None => break, + other => panic!("expected terminal WebSocket frame, got {other:?}"), + } + } + }) + .await + .expect("old embedded protocol WebSocket did not close"); + } + + #[tokio::test] + async fn embedded_internal_protocol_bridge_closes_across_stop_restore_without_retargeting() { + // Each delayed method gets its own A connection: a rejected Submit + // must not hide whether completion or Shutdown can retarget B. + for delayed in delayed_protocol_test_methods() { + let (runtime, worker_ref, execution) = runtime_with_worker_backend(); + let source = embedded_protocol_test_source(&runtime, &worker_ref); + let mut old = + connect_embedded_worker_protocol(source.clone(), embedded_protocol_test_scope()) + .await + .unwrap(); + assert_embedded_protocol_snapshot( + next_embedded_protocol_event(&mut old).await, + &worker_ref, + ); + old.methods + .send(protocol_test_completion("execution-a")) + .await + .unwrap(); + assert!(matches!( + next_embedded_protocol_event(&mut old).await, + protocol::Event::Completions { request_id: Some(id), .. } if id == "execution-a" + )); + execution.emit_worker_event( + &worker_ref, + protocol::Event::TextDelta { + text: "execution-a".into(), + }, + ); + assert!(matches!( + next_embedded_protocol_event(&mut old).await, + protocol::Event::TextDelta { text } if text == "execution-a" + )); + + // No await here: the old loop has not yet observed A's closure when + // B is installed and a delayed method enters its existing queue. + runtime.stop_worker(&worker_ref, None).unwrap(); + execution.accept_restores(); + runtime.restore_worker(&worker_ref).unwrap(); + old.methods.try_send(delayed).unwrap(); + execution.emit_worker_event( + &worker_ref, + protocol::Event::TextDelta { + text: "execution-b".into(), + }, + ); + assert_embedded_protocol_closed(&mut old).await; + assert_eq!( + execution.protocol_methods().len(), + 1, + "delayed A method reached B" + ); + assert_eq!( + runtime.worker_detail(&worker_ref).unwrap().status, + worker_runtime::catalog::WorkerStatus::Idle + ); + assert_eq!( + execution.stops.lock().unwrap().as_slice(), + &[worker_ref.clone()] + ); + + let mut current = + connect_embedded_worker_protocol(source, embedded_protocol_test_scope()) + .await + .unwrap(); + assert_embedded_protocol_snapshot( + next_embedded_protocol_event(&mut current).await, + &worker_ref, + ); + current + .methods + .send(protocol_test_completion("execution-b")) + .await + .unwrap(); + assert!(matches!( + next_embedded_protocol_event(&mut current).await, + protocol::Event::Completions { request_id: Some(id), .. } if id == "execution-b" + )); + execution.emit_worker_event( + &worker_ref, + protocol::Event::TextDelta { + text: "fresh-b".into(), + }, + ); + assert!(matches!( + next_embedded_protocol_event(&mut current).await, + protocol::Event::TextDelta { text } if text == "fresh-b" + )); + } + } + + #[tokio::test] + async fn embedded_browser_protocol_ws_closes_across_stop_restore_without_retargeting() { + for delayed in delayed_protocol_test_methods() { + let (runtime, worker_ref, execution) = runtime_with_worker_backend(); + let source = embedded_protocol_test_source(&runtime, &worker_ref); + let (mut old, server) = spawn_embedded_protocol_test_ws( + source.clone(), + embedded_protocol_test_scope(), + false, + ) + .await; + assert_embedded_protocol_snapshot( + next_embedded_protocol_ws_event(&mut old).await, + &worker_ref, + ); + let submit = protocol::Method::Submit { + submission_request_id: "authenticated-a".to_string(), + input: vec![protocol::Segment::Text { + content: "from browser".into(), + }], + }; + send_protocol_test_ws_method(&mut old, &submit).await; + send_protocol_test_ws_method(&mut old, &protocol_test_completion("execution-a")).await; + assert!(matches!( + next_embedded_protocol_ws_event(&mut old).await, + protocol::Event::Completions { .. } + )); + assert!(matches!( + &execution.protocol_methods()[0].1, + protocol::Method::SubmitTracked { source, .. } + if source == &authenticated_browser_input_source(&test_browser_request_actor()) + )); + execution.emit_worker_event( + &worker_ref, + protocol::Event::TextDelta { + text: "execution-a".into(), + }, + ); + assert!( + matches!(next_embedded_protocol_ws_event(&mut old).await, protocol::Event::TextDelta { text } if text == "execution-a") + ); + + // Stop and restore happen before this socket loop processes closure. + runtime.stop_worker(&worker_ref, None).unwrap(); + execution.accept_restores(); + runtime.restore_worker(&worker_ref).unwrap(); + // A peer close may win this race; otherwise the old connected + // dispatcher must reject the frame without touching B. + let _ = old + .send(Message::Text( + protocol::stream::encode_method(&delayed).unwrap().into(), + )) + .await; + execution.emit_worker_event( + &worker_ref, + protocol::Event::TextDelta { + text: "execution-b".into(), + }, + ); + assert_embedded_protocol_ws_closed(&mut old).await; + assert_eq!( + execution.protocol_methods().len(), + 2, + "delayed A method reached B" + ); + assert_eq!( + runtime.worker_detail(&worker_ref).unwrap().status, + worker_runtime::catalog::WorkerStatus::Idle + ); + assert_eq!( + execution.stops.lock().unwrap().as_slice(), + &[worker_ref.clone()] + ); + server.abort(); + + let (mut current, current_server) = + spawn_embedded_protocol_test_ws(source, embedded_protocol_test_scope(), false) + .await; + assert_embedded_protocol_snapshot( + next_embedded_protocol_ws_event(&mut current).await, + &worker_ref, + ); + send_protocol_test_ws_method(&mut current, &protocol_test_completion("execution-b")) + .await; + assert!(matches!( + next_embedded_protocol_ws_event(&mut current).await, + protocol::Event::Completions { request_id: Some(id), .. } if id == "execution-b" + )); + execution.emit_worker_event( + &worker_ref, + protocol::Event::TextDelta { + text: "fresh-b".into(), + }, + ); + assert!(matches!( + next_embedded_protocol_ws_event(&mut current).await, + protocol::Event::TextDelta { text } if text == "fresh-b" + )); + current.close(None).await.unwrap(); + current_server.abort(); + } + } + + #[tokio::test] + async fn embedded_protocol_paths_close_on_shutdown_event_and_lagged_relay() { + for lagged in [false, true] { + let (runtime, worker_ref, execution) = runtime_with_worker_backend(); + let source = embedded_protocol_test_source(&runtime, &worker_ref); + let mut bridge = + connect_embedded_worker_protocol(source.clone(), embedded_protocol_test_scope()) + .await + .unwrap(); + assert_embedded_protocol_snapshot( + next_embedded_protocol_event(&mut bridge).await, + &worker_ref, + ); + let (mut socket, server) = + spawn_embedded_protocol_test_ws(source, embedded_protocol_test_scope(), false) + .await; + assert_embedded_protocol_snapshot( + next_embedded_protocol_ws_event(&mut socket).await, + &worker_ref, + ); + if lagged { + // Overflow the execution broadcast before either relay runs. + // A lost terminal event is uncertain, never a reason to follow + // the Runtime WorkerRef-wide observation bus instead. + for _ in 0..40 { + execution.emit_worker_event( + &worker_ref, + protocol::Event::TextDelta { + text: "overflow".into(), + }, + ); + } + } else { + execution.emit_worker_event(&worker_ref, protocol::Event::Shutdown); + } + assert_embedded_protocol_closed(&mut bridge).await; + assert_embedded_protocol_ws_closed(&mut socket).await; + server.abort(); + } + } + + #[tokio::test] + async fn embedded_protocol_paths_close_after_connected_shutdown_method() { + let shutdown = protocol::Method::Shutdown { + command: protocol::WorkerCommandEnvelope::new(701), + }; + let (runtime, worker_ref, execution) = runtime_with_worker_backend(); + let mut bridge = connect_embedded_worker_protocol( + embedded_protocol_test_source(&runtime, &worker_ref), + embedded_protocol_test_scope(), + ) + .await + .unwrap(); + assert_embedded_protocol_snapshot( + next_embedded_protocol_event(&mut bridge).await, + &worker_ref, + ); + bridge.methods.try_send(shutdown.clone()).unwrap(); + bridge + .methods + .try_send(protocol_test_completion("after-shutdown")) + .unwrap(); + assert_embedded_protocol_closed(&mut bridge).await; + assert!(matches!( + execution.protocol_methods().as_slice(), + [(_, protocol::Method::Shutdown { .. })] + )); + assert_eq!( + runtime.worker_detail(&worker_ref).unwrap().status, + worker_runtime::catalog::WorkerStatus::Stopped + ); + + let (runtime, worker_ref, execution) = runtime_with_worker_backend(); + let (mut socket, server) = spawn_embedded_protocol_test_ws( + embedded_protocol_test_source(&runtime, &worker_ref), + embedded_protocol_test_scope(), + false, + ) + .await; + assert_embedded_protocol_snapshot( + next_embedded_protocol_ws_event(&mut socket).await, + &worker_ref, + ); + send_protocol_test_ws_method(&mut socket, &shutdown).await; + let _ = socket + .send(Message::Text( + protocol::stream::encode_method(&protocol_test_completion("after-shutdown")) + .unwrap() + .into(), + )) + .await; + assert_embedded_protocol_ws_closed(&mut socket).await; + assert!(matches!( + execution.protocol_methods().as_slice(), + [(_, protocol::Method::Shutdown { .. })] + )); + assert_eq!( + runtime.worker_detail(&worker_ref).unwrap().status, + worker_runtime::catalog::WorkerStatus::Stopped + ); + server.abort(); + } + + #[tokio::test] + async fn embedded_protocol_paths_enforce_workspace_scope_and_browser_method_authority() { + let (runtime, worker_ref, execution) = runtime_with_worker_backend(); + let source = embedded_protocol_test_source(&runtime, &worker_ref); + let wrong_scope = + worker_runtime::RuntimeWorkspaceScope::new("other-workspace", "local-token"); + assert!( + connect_embedded_worker_protocol(source.clone(), wrong_scope.clone()) + .await + .is_err() + ); + let (mut socket, server) = + spawn_embedded_protocol_test_ws(source.clone(), wrong_scope, false).await; + assert!(matches!( + next_embedded_protocol_ws_event(&mut socket).await, + protocol::Event::Error { .. } + )); + assert_embedded_protocol_ws_closed(&mut socket).await; + server.abort(); + + let (mut socket, server) = + spawn_embedded_protocol_test_ws(source.clone(), embedded_protocol_test_scope(), true) + .await; + assert_embedded_protocol_snapshot( + next_embedded_protocol_ws_event(&mut socket).await, + &worker_ref, + ); + send_protocol_test_ws_method(&mut socket, &protocol_test_completion("read-only")).await; + assert!(matches!( + next_embedded_protocol_ws_event(&mut socket).await, + protocol::Event::Error { .. } + )); + socket.close(None).await.unwrap(); + server.abort(); + + let (mut socket, server) = + spawn_embedded_protocol_test_ws(source, embedded_protocol_test_scope(), false).await; + assert_embedded_protocol_snapshot( + next_embedded_protocol_ws_event(&mut socket).await, + &worker_ref, + ); + let forged = protocol::Method::SubmitTracked { + submission_request_id: "forged-source".into(), + input: Vec::new(), + source: authenticated_browser_input_source(&test_browser_request_actor()), + }; + send_protocol_test_ws_method(&mut socket, &forged).await; + assert!(matches!( + next_embedded_protocol_ws_event(&mut socket).await, + protocol::Event::Error { .. } + )); + assert_embedded_protocol_ws_closed(&mut socket).await; + assert!(execution.protocol_methods().is_empty()); + server.abort(); + } + + #[tokio::test] + async fn deterministic_protocol_transport_remains_bound_across_stop_and_restore() { + use worker_runtime::execution::WorkerExecutionBackend; + + let (runtime, worker_ref, execution) = runtime_with_worker_backend(); + let completion = |request_id: &str| protocol::Method::ListCompletions { + kind: protocol::CompletionKind::File, + prefix: String::new(), + request_id: Some(request_id.to_string()), + context: None, + }; + let mut old = execution + .clone() + .attach_worker_protocol(&worker_ref) + .unwrap(); + assert!(matches!(old.snapshot, protocol::Event::Snapshot { .. })); + old.validate().unwrap(); + assert!(matches!( + old.dispatch(completion("execution-a")).unwrap().as_slice(), + [protocol::Event::Completions { .. }] + )); + + // Runtime observations are not this execution's protocol event stream. + runtime + .observe_worker_event( + &worker_ref, + protocol::Event::TextDelta { + text: "runtime-only".to_string(), + }, + ) + .unwrap(); + execution.emit_worker_event( + &worker_ref, + protocol::Event::TextDelta { + text: "execution-a".to_string(), + }, + ); + let event = tokio::time::timeout(std::time::Duration::from_secs(2), old.events.recv()) + .await + .unwrap() + .unwrap(); + assert!(matches!(event, protocol::Event::TextDelta { text } if text == "execution-a")); + + runtime.stop_worker(&worker_ref, None).unwrap(); + assert!(old.validate().is_err()); + assert!(old.dispatch(completion("stopped-a")).is_err()); + assert!( + tokio::time::timeout(std::time::Duration::from_secs(2), old.events.recv()) + .await + .unwrap() + .is_none() + ); + + execution.accept_restores(); + runtime.restore_worker(&worker_ref).unwrap(); + let mut current = execution + .clone() + .attach_worker_protocol(&worker_ref) + .unwrap(); + current.validate().unwrap(); + current.dispatch(completion("execution-b")).unwrap(); + assert!(old.validate().is_err()); + assert!(old.dispatch(completion("must-not-reach-b")).is_err()); + execution.emit_worker_event( + &worker_ref, + protocol::Event::TextDelta { + text: "execution-b".to_string(), + }, + ); + let event = tokio::time::timeout(std::time::Duration::from_secs(2), current.events.recv()) + .await + .unwrap() + .unwrap(); + assert!(matches!(event, protocol::Event::TextDelta { text } if text == "execution-b")); + let delivered: Vec<_> = execution + .protocol_methods() + .into_iter() + .map(|(worker, method)| { + let protocol::Method::ListCompletions { request_id, .. } = method else { + panic!("expected completion method"); + }; + (worker, request_id.unwrap()) + }) + .collect(); + assert_eq!( + delivered, + vec![ + (worker_ref.clone(), "execution-a".to_string()), + (worker_ref, "execution-b".to_string()), + ] + ); + } + #[tokio::test] async fn proxy_maps_runtime_worker_not_found_http_404_to_protocol_error_event() { - let (_runtime, _worker_ref, endpoint) = spawn_runtime_worker().await; + let (_runtime, _worker_ref, endpoint, _execution) = spawn_runtime_worker().await; let endpoint = endpoint.replace("/protocol/ws", "/missing-worker/protocol/ws"); let source = RuntimeObservationSourceConfig { worker: RuntimeWorkerRef::new("runtime-a", "worker-a"), @@ -57804,8 +58779,9 @@ mod tests { worker_runtime::Runtime, worker_runtime::identity::WorkerRef, String, + Arc, ) { - let (runtime, worker_ref) = runtime_with_worker(); + let (runtime, worker_ref, execution) = runtime_with_worker_backend(); let runtime_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let runtime_addr = runtime_listener.local_addr().unwrap(); tokio::spawn({ @@ -57824,7 +58800,7 @@ mod tests { "ws://{runtime_addr}/v1/workers/{}/protocol/ws", worker_ref.worker_id ); - (runtime, worker_ref, endpoint) + (runtime, worker_ref, endpoint, execution) } fn test_browser_request_actor() -> RequestActor { diff --git a/docs/report/T-700-subworker-followup-observation.md b/docs/report/T-700-subworker-followup-observation.md new file mode 100644 index 000000000..3dc6d95d3 --- /dev/null +++ b/docs/report/T-700-subworker-followup-observation.md @@ -0,0 +1,15 @@ +# SubWorker follow-up後のcommitted sessionを確認できなかった + +T-700の実装中、既存のIdleなdirect child `backend`へ`WorkerSendInput`で追加修正を依頼した。toolは成功を返し、その後Idle通知と`WorkerList`のIdle表示を確認した。 + +一方、`ViewSessionOverview`と`SearchSessionEntries`が返した最後のentryは、follow-up前の保存Manifest fixture修正報告（entry range 744）だった。今回送った追加依頼の文字列の検索は0件で、offset 744からもその旧報告1件しか見えなかった。したがって、通知だけから追加作業の完了は認定できなかった。これは最新turnの実行失敗、captureの遅延、観測面のfreshnessのいずれかを切り分ける証拠ではない。 + +親は現在のdiffを直接確認し、子をStopしてdelegated scopeを解放した後、残りのShutdown admission修正を新しいdirect childへ限定して委譲した。新しい子では最初の依頼と修正報告をsession observationから確認できた。 + +改善案: + +- Idle通知に、対象turnとcommitted captureの末尾entryの対応が分かる情報を含める。 +- follow-upの開始またはcaptureに失敗した場合、成功した旧turnだけを返す代わりにbounded diagnosticを表示する。 +- observationで最新のcommitted captureが未取得の場合、古いsnapshotであることを明示する。 + +credential、raw prompt、Backend保存領域の直接読取で観測制限を迂回することは行っていない。