From 1c03758fc1f9ed6d7a2a4f428368cff57eaffef4 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 30 Oct 2026 21:44:51 +0900 Subject: [PATCH] fix: confirm remote Workdir close before releasing session ownership --- crates/workdir/src/http.rs | 21 +- crates/workdir/tests/remote_close.rs | 287 +++++++++ crates/workspace-server/src/server.rs | 81 ++- .../tests/ticket_worker_launch_tests.rs | 68 +- .../tests/workdir_session_close_tests.rs | 593 ++++++++++++++++++ 5 files changed, 1018 insertions(+), 32 deletions(-) create mode 100644 crates/workdir/tests/remote_close.rs create mode 100644 crates/workspace-server/src/server/tests/workdir_session_close_tests.rs diff --git a/crates/workdir/src/http.rs b/crates/workdir/src/http.rs index 20a57a696..5155962a0 100644 --- a/crates/workdir/src/http.rs +++ b/crates/workdir/src/http.rs @@ -556,6 +556,7 @@ mod client { capabilities: WorkdirSessionCapabilities, environment: crate::WorkdirSessionEnvironment, closed: AtomicBool, + close_lock: tokio::sync::Mutex<()>, } impl RemoteWorkdirSession { @@ -616,6 +617,7 @@ mod client { capabilities: opened.capabilities, environment: opened.environment, closed: AtomicBool::new(false), + close_lock: tokio::sync::Mutex::new(()), }) } @@ -1014,7 +1016,11 @@ mod client { } async fn close(&self) -> Result<(), WorkdirError> { - if self.closed.swap(true, Ordering::AcqRel) { + // In-flight is not closed. Dropping this guard on error or cancellation + // leaves confirmation unset, so the next caller can retry Runtime's + // idempotent DELETE even when the previous response was lost. + let _close = self.close_lock.lock().await; + if self.closed.load(Ordering::Acquire) { return Ok(()); } let url = endpoint( @@ -1033,6 +1039,7 @@ mod client { .await .map_err(http_unavailable)?; if response.status() == StatusCode::NO_CONTENT || response.status().is_success() { + self.closed.store(true, Ordering::Release); Ok(()) } else { Err(decode_error(response).await) @@ -1211,7 +1218,7 @@ mod client { // Seed an already-open session to exercise the public operation and // close boundaries independently of the rejected open request. - let remote = RemoteWorkdirSession { + let mut remote = RemoteWorkdirSession { client, base_url, authorization: authorization.clone(), @@ -1222,6 +1229,7 @@ mod client { WorkdirSessionCapabilities::EMPTY, ), closed: AtomicBool::new(false), + close_lock: tokio::sync::Mutex::new(()), }; let stat = StatRequest { path: crate::WorkdirPath::new("file").unwrap(), @@ -1263,6 +1271,15 @@ mod client { listener.accept().unwrap_err().kind(), std::io::ErrorKind::WouldBlock ); + // An invalid endpoint cannot be produced by a successful open, but + // must not turn a pre-dispatch error into a confirmed close either. + remote.base_url = Url::parse("data:invalid-base").unwrap(); + for _ in 0..2 { + assert!(matches!( + remote.close().await, + Err(WorkdirError::InvalidArgument(_)) + )); + } } } diff --git a/crates/workdir/tests/remote_close.rs b/crates/workdir/tests/remote_close.rs new file mode 100644 index 000000000..f00c40252 --- /dev/null +++ b/crates/workdir/tests/remote_close.rs @@ -0,0 +1,287 @@ +#![cfg(feature = "http-client")] +//! Remote close protocol tests with a loopback HTTP boundary. Requests and responses +//! are explicitly gated; Runtime lease integration belongs to workspace-server. +use std::{ + future::Future, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + task::Poll, + time::Duration, +}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + sync::mpsc, +}; +use workdir::{ + Workdir, WorkdirError, WorkdirSession, + http::{ + OpenWorkdirSessionRequest, OpenWorkdirSessionResponse, RemoteWorkdirSession, + WorkdirHttpAuthorization, WorkdirHttpAuthorizationFuture, WorkdirSessionId, + }, +}; + +struct Provider { + url: reqwest::Url, + requests: mpsc::UnboundedReceiver, + task: tokio::task::JoinHandle<()>, +} +impl Drop for Provider { + fn drop(&mut self) { + self.task.abort(); + } +} +impl Provider { + async fn new() -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()) + .parse() + .unwrap(); + let (tx, requests) = mpsc::unbounded_channel(); + let task = tokio::spawn(async move { + loop { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut bytes = Vec::new(); + while !bytes.ends_with(b"\r\n\r\n") { + bytes.push(stream.read_u8().await.unwrap()); + assert!(bytes.len() < 16384); + } + let headers = String::from_utf8(bytes).unwrap(); + let length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().unwrap()) + }) + .unwrap_or(0); + assert!(length < 16384); + stream.read_exact(&mut vec![0; length]).await.unwrap(); + assert!( + headers.starts_with("POST /v1/working-directories/w/sessions ") + || headers.starts_with("DELETE /v1/workdir-sessions/s ") + ); + if tx.send(stream).is_err() { + break; + } + } + }); + Self { + url, + requests, + task, + } + } + async fn request_while(&mut self, future: std::pin::Pin<&mut F>) -> TcpStream { + tokio::time::timeout(Duration::from_secs(5), async { + tokio::select! { + request = self.requests.recv() => request.expect("provider stopped"), + _ = future => panic!("close/open completed without the expected HTTP request"), + } + }) + .await + .expect("request timed out") + } + async fn open(&mut self, auth: Arc) -> RemoteWorkdirSession { + let future = RemoteWorkdirSession::open_with_authorization( + reqwest::Client::builder() + .no_proxy() + .timeout(Duration::from_secs(5)) + .build() + .unwrap(), + self.url.clone(), + auth, + Workdir::new("w").id().clone(), + OpenWorkdirSessionRequest::default(), + ); + tokio::pin!(future); + let stream = self.request_while(future.as_mut()).await; + let body = serde_json::to_string(&OpenWorkdirSessionResponse { + session_id: WorkdirSessionId::new("s").unwrap(), + workdir_id: Workdir::new("w").id().clone(), + capabilities: workdir::WorkdirSessionCapabilities::ALL, + environment: workdir::WorkdirSessionEnvironment::root( + workdir::WorkdirSessionCapabilities::ALL, + ), + }) + .unwrap(); + reply(stream, "200 OK", &body).await; + future.await.unwrap() + } +} +async fn reply(mut stream: TcpStream, status: &str, body: &str) { + stream.write_all(format!("HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap(); +} +#[derive(Debug, Default)] +struct Authorization { + reject_delete: AtomicBool, +} +impl WorkdirHttpAuthorization for Authorization { + fn bearer_token( + &self, + method: String, + _: String, + _: Vec, + ) -> WorkdirHttpAuthorizationFuture { + let reject = method == "DELETE" && self.reject_delete.swap(false, Ordering::SeqCst); + Box::pin(async move { + if reject { + Err(WorkdirError::Unavailable( + "fixture authorization rejected".into(), + )) + } else { + Ok("fixture-token".into()) + } + }) + } +} + +#[tokio::test] +async fn close_authorization_failure_retries_same_session_and_confirmed_close_skips_http() { + let mut provider = Provider::new().await; + let auth = Arc::new(Authorization::default()); + let remote = provider.open(auth.clone()).await; + auth.reject_delete.store(true, Ordering::SeqCst); + assert!(remote.close().await.is_err()); + assert!(provider.requests.try_recv().is_err()); + let retry = remote.close(); + tokio::pin!(retry); + let request = provider.request_while(retry.as_mut()).await; + reply(request, "204 No Content", "").await; + retry.await.unwrap(); + remote.close().await.unwrap(); + assert!(provider.requests.try_recv().is_err()); +} + +#[derive(Debug)] +struct GatedAuthorization { + entered: mpsc::UnboundedSender<()>, + release: std::sync::Mutex>>, +} +impl WorkdirHttpAuthorization for GatedAuthorization { + fn bearer_token( + &self, + method: String, + _: String, + _: Vec, + ) -> WorkdirHttpAuthorizationFuture { + let gate = if method == "DELETE" { + self.release.lock().unwrap().take() + } else { + None + }; + let entered = self.entered.clone(); + Box::pin(async move { + if let Some(gate) = gate { + entered.send(()).unwrap(); + gate.await.unwrap(); + } + Ok("fixture-token".into()) + }) + } +} + +#[tokio::test] +async fn close_cancel_during_authorization_allows_an_actual_retry() { + let mut provider = Provider::new().await; + let (entered, mut events) = mpsc::unbounded_channel(); + let (_release, gate) = tokio::sync::oneshot::channel(); + let remote = provider + .open(Arc::new(GatedAuthorization { + entered, + release: std::sync::Mutex::new(Some(gate)), + })) + .await; + let mut first = Box::pin(remote.close()); + tokio::time::timeout(Duration::from_secs(5), async { + tokio::select! { + _ = events.recv() => (), + _ = &mut first => panic!("close completed while authorization was gated"), + } + }) + .await + .unwrap(); + drop(first); + assert!(provider.requests.try_recv().is_err()); + let retry = remote.close(); + tokio::pin!(retry); + let request = provider.request_while(retry.as_mut()).await; + reply(request, "204 No Content", "").await; + retry.await.unwrap(); +} + +#[tokio::test] +async fn close_http_rejection_or_lost_response_requires_another_delete() { + for response in [ + Some(( + "503 Service Unavailable", + r#"{"code":"unavailable","message":"fixture"}"#, + )), + None, + ] { + let mut provider = Provider::new().await; + let remote = provider.open(Arc::new(Authorization::default())).await; + let first = remote.close(); + tokio::pin!(first); + let request = provider.request_while(first.as_mut()).await; + if let Some((status, body)) = response { + reply(request, status, body).await; + } else { + drop(request); + } // Runtime may have deleted; no response confirms it. + assert!(first.await.is_err()); + let retry = remote.close(); + tokio::pin!(retry); + let request = provider.request_while(retry.as_mut()).await; + reply(request, "204 No Content", "").await; + retry.await.unwrap(); + } +} + +#[tokio::test] +async fn close_cancel_releases_serialization_and_waiter_must_confirm_delete() { + let mut provider = Provider::new().await; + let remote = provider.open(Arc::new(Authorization::default())).await; + let mut first = Box::pin(remote.close()); + let request = provider.request_while(first.as_mut()).await; + let mut waiter = Box::pin(remote.close()); + assert!( + std::future::poll_fn(|cx| Poll::Ready(waiter.as_mut().poll(cx))) + .await + .is_pending() + ); + drop(first); + drop(request); + let request = provider.request_while(waiter.as_mut()).await; + reply(request, "204 No Content", "").await; + waiter.await.unwrap(); +} + +#[tokio::test] +async fn concurrent_close_waits_for_confirmation_and_retries_after_failure() { + for success in [false, true] { + let mut provider = Provider::new().await; + let remote = provider.open(Arc::new(Authorization::default())).await; + let mut first = Box::pin(remote.close()); + let request = provider.request_while(first.as_mut()).await; + let mut waiter = Box::pin(remote.close()); + assert!( + std::future::poll_fn(|cx| Poll::Ready(waiter.as_mut().poll(cx))) + .await + .is_pending() + ); + if success { + reply(request, "204 No Content", "").await; + first.await.unwrap(); + } else { + drop(request); + assert!(first.await.is_err()); + let request = provider.request_while(waiter.as_mut()).await; + reply(request, "204 No Content", "").await; + } + waiter.await.unwrap(); + assert!(provider.requests.try_recv().is_err()); + } +} diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 0b8faaa2f..8abd78c7a 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -704,28 +704,52 @@ impl WorkdirSessionRegistry { } } +// Extraction is not release confirmation. Retain every unconfirmed entry across +// the await, including entries not yet attempted when the close future is dropped. +// Callers keep their Worker session fence until this restoration has completed. +struct UnconfirmedWorkdirSessionCloses<'a> { + registry: &'a Arc>, + registered: Vec, +} + +impl Drop for UnconfirmedWorkdirSessionCloses<'_> { + fn drop(&mut self) { + if self.registered.is_empty() { + return; + } + let mut sessions = self + .registry + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for item in self.registered.drain(..) { + sessions.restore(item); + } + } +} + async fn close_registered_workdir_sessions( registry: &Arc>, registered: Vec, ) -> std::result::Result<(), String> { - let mut failed = Vec::new(); + let mut unconfirmed = UnconfirmedWorkdirSessionCloses { + registry, + registered, + }; let mut errors = Vec::new(); - for item in registered { - let session = match &item { + let mut index = 0; + while let Some(item) = unconfirmed.registered.get(index) { + let session = match item { RegisteredWorkdirSession::Attachment { session, .. } => session, RegisteredWorkdirSession::Command { session, .. } => &session.source, }; - if let Err(error) = session.close().await { - errors.push(error.to_string()); - failed.push(item); - } - } - if !failed.is_empty() { - let mut sessions = registry - .lock() - .map_err(|_| "Workdir session registry was poisoned".to_string())?; - for item in failed { - sessions.restore(item); + match session.close().await { + Ok(()) => { + unconfirmed.registered.remove(index); + } + Err(error) => { + errors.push(error.to_string()); + index += 1; + } } } if errors.is_empty() { @@ -21380,17 +21404,19 @@ async fn close_current_worker_attachment_session_locked( .expect("Workdir session registry lock poisoned") .remove_attachment(worker); if let Some(session) = session { - if let Err(error) = session.close().await { - api.workdir_sessions - .lock() - .expect("Workdir session registry lock poisoned") - .insert_attachment(worker.clone(), session); - return Err(Error::RuntimeOperationFailed { - runtime_id: worker.runtime_id.clone(), - code: "workdir_session_close_failed".to_string(), - message: error.to_string(), - }); - } + close_registered_workdir_sessions( + &api.workdir_sessions, + vec![RegisteredWorkdirSession::Attachment { + worker: worker.clone(), + session, + }], + ) + .await + .map_err(|message| Error::RuntimeOperationFailed { + runtime_id: worker.runtime_id.clone(), + code: "workdir_session_close_failed".to_string(), + message, + })?; } Ok(()) } @@ -39283,6 +39309,7 @@ mod tests { mod ticket_worker_launch_tests; mod value_profiles_tests; mod workdir_delegation_tests; + mod workdir_session_close_tests; mod worker_drive_tests; mod worker_operations_tests; mod worker_protocol_tests; @@ -51226,6 +51253,10 @@ mod tests { let worker_root = root .join(".test-embedded-runtime-store/workers") .join(&worker.worker_id); + seed_worker_session_at_root(&worker_root); + } + + fn seed_worker_session_at_root(worker_root: &Path) { use session_store::Store; let store = session_store::WorkerSessionStore::new(worker_root.join("session")).unwrap(); let session = session_store::new_session_id(); diff --git a/crates/workspace-server/src/server/tests/ticket_worker_launch_tests.rs b/crates/workspace-server/src/server/tests/ticket_worker_launch_tests.rs index 83ce14322..50e92b0de 100644 --- a/crates/workspace-server/src/server/tests/ticket_worker_launch_tests.rs +++ b/crates/workspace-server/src/server/tests/ticket_worker_launch_tests.rs @@ -122,6 +122,31 @@ impl WorkerExecutionBackend for LaunchExecutionBackend { self.execution.working_directory(id) } + fn open_workdir_session( + &self, + id: &str, + ) -> std::result::Result< + WorkdirSessionHandle, + worker_runtime::working_directory::WorkingDirectoryDiagnostic, + > { + let binding = self + .execution + .materializer + .bind_working_directory(id, None)?; + let scope = manifest::Scope::writable(binding.root()).unwrap(); + Ok(Arc::new( + workdir::LocalWorkdirSession::materialized_bound_with_environment( + workdir::Workdir::new(id), + binding.root().to_owned(), + binding.root().to_owned(), + manifest::SharedScope::new(scope), + workdir::WorkdirSessionCapabilities::ALL, + binding.command_environment(), + binding.session_resources()?, + ), + )) + } + fn create_working_directory( &self, request: &worker_runtime::catalog::WorkingDirectoryRequest, @@ -192,13 +217,27 @@ impl Drop for LaunchFixture { } impl LaunchFixture { + pub(super) fn publish_profile(&self, profile: &str) { + self.backend.publish_profile(&self.api, profile); + } + + pub(super) fn seed_remote_session_for_cleanup(&self, worker: &RuntimeWorkerRef) { + seed_worker_session_at_root( + &self + ._temp + .path() + .join(".test-embedded-runtime-store/source-runtime/workers") + .join(&worker.worker_id), + ); + } + // Reuse the authenticated Runtime HTTP harness for Drive/source lifetime // assertions without duplicating its trust/bootstrap/listener lifecycle. pub(super) async fn with_source_workdir() -> ( Self, worker_runtime::working_directory::WorkingDirectoryBinding, ) { - let fixture = Self::new(true).await; + let fixture = Self::new_with_runtime_store(true, true).await; init_clean_git_workspace(fixture._temp.path()); let (_, Json(created)) = create_workspace_working_directory( &fixture.api, @@ -224,6 +263,10 @@ impl LaunchFixture { } async fn new(remote: bool) -> Self { + Self::new_with_runtime_store(remote, false).await + } + + async fn new_with_runtime_store(remote: bool, persistent: bool) -> Self { let temp = tempfile::tempdir().unwrap(); let config = test_server_config(temp.path()); let store = SqliteWorkspaceStore::open_for_test(config.database_path.clone()).unwrap(); @@ -282,11 +325,26 @@ impl LaunchFixture { Arc::new(InMemoryWorkspaceClaimReplayProtection::default()), ) .unwrap(); - let runtime = worker_runtime::Runtime::with_execution_backend( - worker_runtime::RuntimeOptions::default(), - fixture.backend.clone(), - ) + let runtime = if persistent { + worker_runtime::Runtime::with_fs_store_and_execution_backend( + worker_runtime::fs_store::FsRuntimeStoreOptions { + root: fixture + ._temp + .path() + .join(".test-embedded-runtime-store/source-runtime"), + runtime_id: fixture.runtime_id.clone(), + display_name: None, + }, + fixture.backend.clone(), + ) + } else { + worker_runtime::Runtime::with_execution_backend( + worker_runtime::RuntimeOptions::default(), + fixture.backend.clone(), + ) + } .unwrap(); + runtime.bind_runtime_identity(&fixture.runtime_id).unwrap(); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let endpoint = format!("http://{}", listener.local_addr().unwrap()); let binding = WorkspaceRuntimeBinding { diff --git a/crates/workspace-server/src/server/tests/workdir_session_close_tests.rs b/crates/workspace-server/src/server/tests/workdir_session_close_tests.rs new file mode 100644 index 000000000..e4daae4a2 --- /dev/null +++ b/crates/workspace-server/src/server/tests/workdir_session_close_tests.rs @@ -0,0 +1,593 @@ +//! Backend registry/current-attachment/WorkerRemove integration with authenticated +//! Runtime HTTP sessions and real Git materializer resource leases. Only Worker +//! execution and the failing/held close transport are replaced. Ordinary library +//! tests use LaunchFixture's isolated SQLite, materializations and loopback Runtime. +use super::ticket_worker_launch_tests::LaunchFixture; +use super::*; +use tokio::sync::{Notify, Semaphore}; + +#[derive(Clone, Copy)] +enum CloseFault { + Reject, + HoldThenReject, + RejectAfterRuntimeClose, +} + +struct CloseTransport { + upstream: String, + client: reqwest::Client, + fault: Mutex>, + attempts: Mutex>, + runtime_closes: Mutex>, + entered: Notify, + release: Semaphore, +} + +impl CloseTransport { + async fn forward(State(state): State>, request: Request) -> Response { + let is_close = request.method() == Method::DELETE + && request.uri().path().starts_with("/v1/workdir-sessions/"); + let path = request.uri().path().to_owned(); + let fault = if is_close { + state.attempts.lock().unwrap().push(path.clone()); + state.fault.lock().unwrap().take() + } else { + None + }; + if matches!(fault, Some(CloseFault::Reject | CloseFault::HoldThenReject)) { + if matches!(fault, Some(CloseFault::HoldThenReject)) { + state.entered.notify_one(); + state.release.acquire().await.unwrap().forget(); + } + // The actual Runtime has not seen this request or released its lease. + return StatusCode::SERVICE_UNAVAILABLE.into_response(); + } + let (parts, body) = request.into_parts(); + let body = to_bytes(body, 1024 * 1024).await.unwrap(); + let response = state + .client + .request(parts.method, format!("{}{}", state.upstream, parts.uri)) + .headers(parts.headers) + .body(body) + .send() + .await + .unwrap(); + let status = response.status(); + let headers = response.headers().clone(); + let bytes = response.bytes().await.unwrap(); + if is_close { + state.runtime_closes.lock().unwrap().push((path, status)); + } + if matches!(fault, Some(CloseFault::RejectAfterRuntimeClose)) { + assert_eq!(status, StatusCode::NO_CONTENT); + // Lose only the successful acknowledgement, after Runtime has closed + // the actual session and released its active-use/materializer leases. + return StatusCode::BAD_GATEWAY.into_response(); + } + let mut response = Response::new(Body::from(bytes)); + *response.status_mut() = status; + *response.headers_mut() = headers; + response + } + + fn fail_next(&self, fault: CloseFault) { + assert!(self.fault.lock().unwrap().replace(fault).is_none()); + } + + fn attempts(&self) -> Vec { + self.attempts.lock().unwrap().clone() + } +} + +struct Fixture { + source: LaunchFixture, + transport: Arc, + proxy: tokio::task::JoinHandle<()>, + runtime_id: String, + workdir_id: String, + materialization_root: PathBuf, +} + +impl Drop for Fixture { + fn drop(&mut self) { + self.proxy.abort(); + } +} + +impl Fixture { + async fn new() -> Self { + let (source, binding) = LaunchFixture::with_source_workdir().await; + let workdir_id = binding.working_directory.id.clone(); + let materialization_root = binding.working_directory_root().to_owned(); + // Do not accidentally retain a second materializer lease in the fixture. + drop(binding); + let record = source + .api + .store + .get_workdir_registry(TEST_WORKSPACE_ID, &workdir_id) + .unwrap() + .unwrap(); + let runtime_id = record.source.runtime_id().unwrap().to_owned(); + let binding = source + .api + .store + .get_workspace_runtime_binding(TEST_WORKSPACE_ID, &runtime_id) + .await + .unwrap() + .unwrap(); + let transport = Arc::new(CloseTransport { + upstream: binding.base_url.clone(), + client: reqwest::Client::builder() + .no_proxy() + .timeout(std::time::Duration::from_secs(10)) + .build() + .unwrap(), + fault: Mutex::new(None), + attempts: Mutex::new(Vec::new()), + runtime_closes: Mutex::new(Vec::new()), + entered: Notify::new(), + release: Semaphore::new(0), + }); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let router = Router::new() + .fallback(axum::routing::any(CloseTransport::forward)) + .with_state(transport.clone()); + let proxy = tokio::spawn(async move { + axum::serve(listener, router).await.unwrap(); + }); + let backend_url = source.api.config.backend_base_url.clone().unwrap(); + let remote = RemoteWorkerRuntime::new( + RemoteRuntimeConfig { + runtime_id: runtime_id.clone(), + workspace_id: Some(TEST_WORKSPACE_ID.into()), + display_name: binding.display_name.clone(), + base_url: endpoint, + bearer_token: None, + workspace_authorization: Some(crate::hosts::WorkspaceRuntimeAuthorization::new( + source.api.store.clone(), + source.api.signing_identities.clone(), + backend_url.clone(), + binding, + )), + strict_public_egress: false, + cached_worker_creation_available: true, + cached_os: "linux".into(), + cached_arch: "x86_64".into(), + cached_status: "active".into(), + timeout: std::time::Duration::from_secs(10), + }, + TEST_WORKSPACE_ID.into(), + backend_url, + ) + .unwrap(); + source.api.runtime.register_or_replace(remote); + Self { + source, + transport, + proxy, + runtime_id, + workdir_id, + materialization_root, + } + } + + fn api(&self) -> &WorkspaceApi { + &self.source.api + } + + fn registry_owner(&self) -> RuntimeWorkerRef { + RuntimeWorkerRef::new(&self.runtime_id, "registry-close-owner") + } + + async fn open(&self) -> WorkdirSessionHandle { + self.api() + .runtime + .open_workdir_session(&self.runtime_id, &self.workdir_id, None) + .await + .unwrap() + } + + fn assert_restored(&self, worker: &RuntimeWorkerRef, expected: &WorkdirSessionHandle) { + let registry = self.api().workdir_sessions.lock().unwrap(); + assert!(Arc::ptr_eq( + registry + .attachments + .get(worker) + .expect("unconfirmed session retained"), + expected + )); + } + + async fn assert_live(&self, session: &WorkdirSessionHandle) { + session + .stat(workdir::StatRequest { + path: workdir::WorkdirPath::new("README.md").unwrap(), + }) + .await + .expect("same Runtime session still serves filesystem operations"); + } + + async fn remove_workdir(&self) -> WorkingDirectoryRemovalResponse { + let api = self.api().clone(); + let workdir_id = self.workdir_id.clone(); + tokio::task::spawn_blocking(move || { + execute_workdir_removal( + &api, + &workdir_id, + "account:close-test", + "remove normally after session close", + ) + }) + .await + .unwrap() + .unwrap() + } + + async fn assert_busy(&self) { + let result = self.remove_workdir().await; + assert_eq!( + result.disposition, + WorkingDirectoryRemovalDisposition::AttentionRequired, + "{result:?}" + ); + assert_eq!( + result.failure_category.as_deref(), + Some("resource_busy"), + "{result:?}" + ); + assert!(result.retryable); + assert!(self.materialization_root.exists()); + } + + async fn assert_removed(&self) { + let result = self.remove_workdir().await; + assert_eq!( + result.disposition, + WorkingDirectoryRemovalDisposition::Removed, + "{result:?}" + ); + assert!(!result.retryable); + assert!(!self.materialization_root.exists()); + assert!( + self.api() + .store + .get_workdir_registry(TEST_WORKSPACE_ID, &self.workdir_id) + .unwrap() + .is_none() + ); + } +} + +#[derive(Clone, Copy)] +enum ClosePath { + Registry, + Attachment, +} + +async fn close(api: &WorkspaceApi, worker: &RuntimeWorkerRef, path: ClosePath) -> bool { + let lock = current_worker_session_lock(api, worker); + let _guard = lock.lock().await; + match path { + ClosePath::Registry => close_worker_workdir_sessions(&api.workdir_sessions, worker) + .await + .is_ok(), + ClosePath::Attachment => close_current_worker_attachment_session_locked(api, worker) + .await + .is_ok(), + } +} + +async fn failure_restoration(path: ClosePath) { + let fixture = Fixture::new().await; + let worker = fixture.registry_owner(); + let session = fixture.open().await; + fixture + .api() + .workdir_sessions + .lock() + .unwrap() + .insert_attachment(worker.clone(), session.clone()); + fixture.assert_busy().await; + fixture.transport.fail_next(CloseFault::Reject); + assert!(!close(fixture.api(), &worker, path).await); + fixture.assert_restored(&worker, &session); + fixture.assert_live(&session).await; + fixture.assert_busy().await; + assert!(close(fixture.api(), &worker, path).await); + assert!( + !fixture + .api() + .workdir_sessions + .lock() + .unwrap() + .attachments + .contains_key(&worker) + ); + let attempts = fixture.transport.attempts(); + assert_eq!(attempts.len(), 2, "retry must actually contact Runtime"); + assert_eq!( + attempts[0], attempts[1], + "retry closes the SAME Runtime session" + ); + fixture.assert_removed().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn registry_close_failure_restores_same_remote_session_and_retry_releases_runtime_lease() { + failure_restoration(ClosePath::Registry).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn attachment_close_failure_restores_same_remote_session_and_retry_releases_runtime_lease() { + failure_restoration(ClosePath::Attachment).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn lost_close_response_restores_handle_and_retry_confirms_actual_runtime_idempotent_release() +{ + let fixture = Fixture::new().await; + let worker = fixture.registry_owner(); + let session = fixture.open().await; + fixture + .api() + .workdir_sessions + .lock() + .unwrap() + .insert_attachment(worker.clone(), session.clone()); + fixture.assert_busy().await; + fixture + .transport + .fail_next(CloseFault::RejectAfterRuntimeClose); + + assert!(!close(fixture.api(), &worker, ClosePath::Registry).await); + fixture.assert_restored(&worker, &session); + let first_attempt = fixture.transport.attempts(); + assert_eq!(first_attempt.len(), 1); + assert_eq!( + *fixture.transport.runtime_closes.lock().unwrap(), + vec![(first_attempt[0].clone(), StatusCode::NO_CONTENT)], + "fault must follow the real authenticated Runtime's successful DELETE" + ); + let absent = session + .stat(workdir::StatRequest { + path: workdir::WorkdirPath::new("README.md").unwrap(), + }) + .await + .expect_err("original session must already be absent on Runtime"); + assert!( + matches!( + absent.classification_source(), + workdir::WorkdirError::NotFound(_) + ), + "must observe Runtime absence, not a local closed flag: {absent:?}" + ); + + assert!(close(fixture.api(), &worker, ClosePath::Registry).await); + let attempts = fixture.transport.attempts(); + assert_eq!( + attempts, + vec![first_attempt[0].clone(), first_attempt[0].clone()], + "retry must send DELETE to the same already-absent Runtime session" + ); + assert_eq!( + *fixture.transport.runtime_closes.lock().unwrap(), + vec![ + (attempts[0].clone(), StatusCode::NO_CONTENT), + (attempts[1].clone(), StatusCode::NO_CONTENT) + ], + "Runtime itself must acknowledge the idempotent retry" + ); + assert!( + !fixture + .api() + .workdir_sessions + .lock() + .unwrap() + .attachments + .contains_key(&worker) + ); + // Runtime has released its leases, but only now has Backend confirmed it. + fixture.assert_removed().await; +} + +async fn cancellation_restoration(path: ClosePath) { + let fixture = Fixture::new().await; + let worker = fixture.registry_owner(); + let session = fixture.open().await; + fixture + .api() + .workdir_sessions + .lock() + .unwrap() + .insert_attachment(worker.clone(), session.clone()); + // A second real Runtime session proves the batch retains not-yet-attempted + // command sources as well as the close currently awaiting a response. + let command = if matches!(path, ClosePath::Registry) { + let source = fixture.open().await; + let handle = source + .start_command(workdir::CommandRequest { + command: "printf close-test".into(), + timeout_secs: 5, + transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES, + cwd: workdir::WorkdirPath::root(), + tool_call_id: None, + }) + .await + .unwrap(); + let external = fixture + .api() + .workdir_sessions + .lock() + .unwrap() + .register_command(worker.clone(), "checkout".into(), source.clone(), handle); + Some((external, source)) + } else { + None + }; + fixture.assert_busy().await; + fixture.transport.fail_next(CloseFault::HoldThenReject); + let api = fixture.api().clone(); + let closing_worker = worker.clone(); + let closing = tokio::spawn(async move { close(&api, &closing_worker, path).await }); + tokio::time::timeout( + std::time::Duration::from_secs(5), + fixture.transport.entered.notified(), + ) + .await + .expect("close reached pre-Runtime hold"); + closing.abort(); + assert!(closing.await.unwrap_err().is_cancelled()); + fixture.assert_restored(&worker, &session); + fixture.assert_live(&session).await; + if let Some((external, source)) = &command { + let restored = fixture + .api() + .workdir_sessions + .lock() + .unwrap() + .command(&worker, "checkout", external) + .expect("unattempted command session retained"); + assert!(Arc::ptr_eq(&restored.source, source)); + } + fixture.transport.release.add_permits(1); + fixture.assert_busy().await; + assert!(close(fixture.api(), &worker, path).await); + let attempts = fixture.transport.attempts(); + assert_eq!(attempts.len(), if command.is_some() { 3 } else { 2 }); + assert_eq!( + attempts[0], attempts[1], + "retry closes cancelled session, not a replacement" + ); + fixture.assert_removed().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn registry_close_cancellation_restores_inflight_and_unattempted_remote_sessions() { + cancellation_restoration(ClosePath::Registry).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn attachment_close_cancellation_restores_remote_session_for_retry() { + cancellation_restoration(ClosePath::Attachment).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn worker_removal_close_failure_retains_attachment_and_retry_releases_runtime_lease() { + let fixture = Fixture::new().await; + let api = fixture.api(); + fixture.source.publish_profile("builtin:companion"); + let Json(created) = create_workspace_worker( + State(api.clone()), + HeaderMap::new(), + Json(CreateWorkspaceWorkerRequest { + runtime_id: fixture.runtime_id.clone(), + display_name: "session-close removal".into(), + singleton_key: None, + profile: Some("builtin:companion".into()), + ticket_assignment: None, + initial_submit: Vec::new(), + workdir_attachments: Vec::new(), + feature_connections: Default::default(), + control_operation_id: None, + execution_grant: None, + }), + ) + .await + .unwrap(); + let worker = RuntimeWorkerRef::new(created.runtime_id, created.worker_id); + api.runtime + .stop_worker( + &worker, + WorkerLifecycleRequest { + reason: Some("prepare removal".into()), + ticket_assignment: None, + }, + ) + .unwrap(); + let summary = api.runtime.worker(&worker).unwrap(); + assert_eq!(summary.state, "stopped"); + sync_worker_observation(api, &summary).unwrap(); + fixture.source.seed_remote_session_for_cleanup(&worker); + let session = fixture.open().await; + api.workdir_sessions + .lock() + .unwrap() + .insert_attachment(worker.clone(), session.clone()); + fixture.assert_busy().await; + api.store + .attach_worker_workdir(&WorkerWorkdirLinkRecord { + connection_id: String::new(), + workspace_id: TEST_WORKSPACE_ID.into(), + worker: worker.clone(), + workdir_id: fixture.workdir_id.clone(), + alias: "checkout".into(), + capabilities: workdir::WorkdirSessionCapabilities::ALL, + linked_at: TEST_CREATED_AT.into(), + unlinked_at: None, + }) + .unwrap(); + fixture.transport.fail_next(CloseFault::Reject); + let removal = WorkerRemovalService::new(api); + let failed = removal + .execute_target_removal(&api.runtime, &worker, "remove stopped Worker", None) + .await + .unwrap(); + assert_eq!( + failed.status, + StatusCode::SERVICE_UNAVAILABLE.as_u16(), + "{}", + failed.body + ); + assert!( + failed.body.contains("attachment_close_failed"), + "{}", + failed.body + ); + fixture.assert_restored(&worker, &session); + fixture.assert_live(&session).await; + assert!( + api.store + .list_worker_workdir_links(TEST_WORKSPACE_ID, &worker) + .unwrap() + .iter() + .any(|link| link.unlinked_at.is_none()) + ); + assert!( + api.store + .get_worker_registry(TEST_WORKSPACE_ID, &worker) + .unwrap() + .is_some() + ); + let operation = api + .store + .recover_worker_removal_execution(TEST_WORKSPACE_ID, &worker) + .unwrap() + .unwrap(); + assert_eq!( + operation.plan.state, + crate::retention::WorkerRemovalPlanState::Failed + ); + let retried = removal + .execute_target_removal(&api.runtime, &worker, "remove stopped Worker", None) + .await + .unwrap(); + assert_eq!(retried.status, StatusCode::OK.as_u16(), "{}", retried.body); + assert!( + api.store + .get_worker_registry(TEST_WORKSPACE_ID, &worker) + .unwrap() + .is_none() + ); + assert!( + api.store + .list_worker_workdir_links(TEST_WORKSPACE_ID, &worker) + .unwrap() + .iter() + .all(|link| link.unlinked_at.is_some()) + ); + let attempts = fixture.transport.attempts(); + assert_eq!(attempts.len(), 2); + assert_eq!(attempts[0], attempts[1]); + fixture.assert_removed().await; +}