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/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index 4a8d0b375..fd87a6a04 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -11,6 +11,8 @@ //! filesystem/Git inspection and unlink; conflicting target use is rejected promptly. mod runtime_management_api; +#[cfg(test)] +mod workdir_close_tests; use crate::auth::{RuntimeAuthContext, new_token_id, unix_now_seconds}; use crate::catalog::{ @@ -373,9 +375,12 @@ pub struct WorkspaceRuntimeHttpAuth { pub runtime_id: String, } +type RuntimeHttpWorkdirCloseResult = Result<(), RuntimeHttpWorkdirError>; + struct RuntimeHttpWorkdirSession { owner: RuntimeWorkspaceScope, session: WorkdirSessionHandle, + close_attempt: Option>>, } /// `GET /v1/runtime` response. @@ -878,7 +883,11 @@ async fn open_workdir_session( .map_err(|_| RuntimeHttpWorkdirError::internal())? .insert( session_id.as_str().to_string(), - RuntimeHttpWorkdirSession { owner, session }, + RuntimeHttpWorkdirSession { + owner, + session, + close_attempt: None, + }, ); Ok(Json(response)) } @@ -898,7 +907,7 @@ async fn run_workdir_session_operation( .map_err(|_| RuntimeHttpWorkdirError::internal())?; let record = sessions .get(&session_id) - .filter(|record| record.owner == owner) + .filter(|record| record.owner == owner && record.close_attempt.is_none()) .ok_or_else(RuntimeHttpWorkdirError::not_found)?; record.session.clone() }; @@ -934,24 +943,55 @@ async fn close_workdir_session( auth: Option>, ) -> Result { let owner = required_workdir_owner(auth)?; - let session = { + let mut completion = { let mut sessions = state .workdir_sessions .lock() .map_err(|_| RuntimeHttpWorkdirError::internal())?; - if sessions - .get(&session_id) - .is_some_and(|record| record.owner == owner) - { - sessions.remove(&session_id).map(|record| record.session) + let Some(record) = sessions + .get_mut(&session_id) + .filter(|record| record.owner == owner) + else { + return Ok(StatusCode::NO_CONTENT); + }; + if let Some(completion) = &record.close_attempt { + completion.clone() } else { - None + let (result_tx, completion) = tokio::sync::watch::channel(None); + record.close_attempt = Some(completion.clone()); + let session = record.session.clone(); + let registry = state.workdir_sessions.clone(); + // The task, not an HTTP waiter, owns provider teardown. In particular, + // Local close marks itself closed before awaiting command teardown; + // cancelling that future would make a later provider close a false ack. + tokio::spawn(async move { + let result = session.close().await.map_err(RuntimeHttpWorkdirError::from); + let result = match registry.lock() { + Ok(mut sessions) => { + if result.is_ok() { + sessions.remove(&session_id); + } else if let Some(record) = sessions.get_mut(&session_id) { + record.close_attempt = None; + } + result + } + Err(_) => Err(RuntimeHttpWorkdirError::internal()), + }; + result_tx.send_replace(Some(result)); + }); + completion } }; - if let Some(session) = session { - session.close().await?; + loop { + if let Some(result) = completion.borrow_and_update().clone() { + result?; + return Ok(StatusCode::NO_CONTENT); + } + completion + .changed() + .await + .map_err(|_| RuntimeHttpWorkdirError::internal())?; } - Ok(StatusCode::NO_CONTENT) } fn required_workdir_owner( @@ -2149,7 +2189,7 @@ fn required_runtime_permission(method: &Method, path: &str) -> Option<&'static s None } -#[derive(Debug)] +#[derive(Clone, Debug)] struct RuntimeHttpWorkdirError { status: StatusCode, payload: WorkdirTransportError, @@ -2890,6 +2930,7 @@ mod tests { RuntimeHttpWorkdirSession { owner: RuntimeWorkspaceScope::new("workspace-a", "https://backend.test"), session, + close_attempt: None, }, )]))), ssh_keyscan_program: Arc::new(PathBuf::from("ssh-keyscan")), @@ -3194,6 +3235,7 @@ mod tests { RuntimeHttpWorkdirSession { owner: owner.clone(), session: session.clone(), + close_attempt: None, }, )]))), ssh_keyscan_program: Arc::new(PathBuf::from("ssh-keyscan")), @@ -3314,6 +3356,41 @@ mod tests { .expect_err("cross-workspace session access must fail"); assert_eq!(error.status, StatusCode::NOT_FOUND); + for unauthorized in [ + None, + Some(Extension(RuntimeAuthContext { + workspace_id: "workspace-b".into(), + ..auth.clone() + })), + Some(Extension(RuntimeAuthContext { + server_id: "server-b".into(), + ..auth.clone() + })), + ] { + let missing_auth = unauthorized.is_none(); + let result = + close_workdir_session(State(state.clone()), Path("session-1".into()), unauthorized) + .await; + if missing_auth { + assert_eq!(result.unwrap_err().status, StatusCode::FORBIDDEN); + } else { + assert_eq!(result.unwrap(), StatusCode::NO_CONTENT); + } + assert!( + state + .workdir_sessions + .lock() + .unwrap() + .contains_key("session-1") + ); + session + .stat(StatRequest { + path: WorkdirPath::new("hello.txt").unwrap(), + }) + .await + .expect("unauthorized close must not reach provider"); + } + assert_eq!( close_workdir_session( State(state.clone()), diff --git a/crates/worker-runtime/src/http_server/tests/command_transport_tests.rs b/crates/worker-runtime/src/http_server/tests/command_transport_tests.rs index cdaa8ce40..f368ce42c 100644 --- a/crates/worker-runtime/src/http_server/tests/command_transport_tests.rs +++ b/crates/worker-runtime/src/http_server/tests/command_transport_tests.rs @@ -99,6 +99,7 @@ impl Harness { RuntimeHttpWorkdirSession { owner: RuntimeWorkspaceScope::new("workspace-a", "https://backend.test"), session: session.clone(), + close_attempt: None, }, )]))), ssh_keyscan_program: Arc::new(PathBuf::from("ssh-keyscan")), diff --git a/crates/worker-runtime/src/http_server/workdir_close_tests.rs b/crates/worker-runtime/src/http_server/workdir_close_tests.rs new file mode 100644 index 000000000..da8b47646 --- /dev/null +++ b/crates/worker-runtime/src/http_server/workdir_close_tests.rs @@ -0,0 +1,415 @@ +//! Runtime HTTP close integration: production local-token authentication, session registry, +//! Local command teardown, and Runtime active-use guard. The execution-backend mapping is +//! a fixture, not a materializer; a provider wrapper gates a genuinely pending Local close. +//! Run with `cargo test -p worker-runtime --lib http_server::workdir_close_tests`. + +use super::*; +use crate::RuntimeOptions; +use crate::execution::{ + WorkerExecutionBackend, WorkerExecutionResult, WorkerExecutionSpawnRequest, + WorkerExecutionSpawnResult, +}; +use crate::working_directory::WorkingDirectoryDiagnostic; +use axum::body::to_bytes; +use std::future::Future; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::task::Poll; +use tokio::sync::{Notify, Semaphore}; +use tower::ServiceExt; +use workdir::{LocalWorkdirSession, Workdir, WorkdirError, WorkdirSession}; + +#[derive(Debug)] +struct CloseProvider { + local: LocalWorkdirSession, + calls: AtomicUsize, + cancelled: AtomicUsize, + pause_first: bool, + fail_first: bool, + entered: Notify, + release: Semaphore, +} + +struct AttemptGuard<'a> { + cancelled: &'a AtomicUsize, + finished: bool, +} + +impl Drop for AttemptGuard<'_> { + fn drop(&mut self) { + if !self.finished { + self.cancelled.fetch_add(1, Ordering::SeqCst); + } + } +} + +// Only close scheduling/error injection is substituted. All session operations and +// actual teardown remain Local; in particular its early closed.swap is exercised. +#[async_trait::async_trait] +impl WorkdirSession for CloseProvider { + fn workdir(&self) -> &Workdir { + self.local.workdir() + } + + fn capabilities(&self) -> workdir::WorkdirSessionCapabilities { + self.local.capabilities() + } + + async fn stat( + &self, + request: workdir::StatRequest, + ) -> Result { + self.local.stat(request).await + } + async fn read( + &self, + request: workdir::ReadRequest, + ) -> Result { + self.local.read(request).await + } + async fn write( + &self, + request: workdir::WriteRequest, + ) -> Result { + self.local.write(request).await + } + async fn edit( + &self, + request: workdir::EditRequest, + ) -> Result { + self.local.edit(request).await + } + async fn list( + &self, + request: workdir::ListRequest, + ) -> Result { + self.local.list(request).await + } + async fn glob( + &self, + request: workdir::GlobRequest, + ) -> Result { + self.local.glob(request).await + } + async fn grep( + &self, + request: workdir::GrepRequest, + ) -> Result { + self.local.grep(request).await + } + async fn start_command( + &self, + request: workdir::CommandRequest, + ) -> Result { + self.local.start_command(request).await + } + async fn command_status( + &self, + request: workdir::CommandHandle, + ) -> Result { + self.local.command_status(request).await + } + async fn command_output( + &self, + request: workdir::CommandOutputRequest, + ) -> Result { + self.local.command_output(request).await + } + async fn cancel_command(&self, request: workdir::CommandHandle) -> Result<(), WorkdirError> { + self.local.cancel_command(request).await + } + + async fn close(&self) -> Result<(), WorkdirError> { + let attempt = self.calls.fetch_add(1, Ordering::SeqCst); + if attempt == 0 && self.fail_first { + return Err(WorkdirError::Unavailable("injected close failure".into())); + } + let mut guard = AttemptGuard { + cancelled: &self.cancelled, + finished: false, + }; + if attempt == 0 && self.pause_first { + self.local + .start_command(workdir::CommandRequest { + command: "printf close-probe".into(), + timeout_secs: 5, + transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES, + cwd: workdir::WorkdirPath::root(), + tool_call_id: None, + }) + .await?; + } + let mut closing = Box::pin(self.local.close()); + if attempt == 0 && self.pause_first { + // On the current-thread executor the command task cannot run before + // this poll. Local close sets closed, drains/cancels the command, + // then is genuinely Pending on command.task.await (reviewer's probe). + assert!( + std::future::poll_fn(|cx| Poll::Ready(closing.as_mut().poll(cx))) + .await + .is_pending() + ); + self.entered.notify_one(); + // Hold the already-polled provider future, not a fake release flag. + // Cancelling this future would strand Local's unconfirmed teardown. + let _permit = self.release.acquire().await.unwrap(); + } + let result = closing.await; + guard.finished = true; + result + } +} + +struct Backend { + session: Arc, +} + +impl WorkerExecutionBackend for Backend { + fn backend_id(&self) -> &str { + "close-confirmation-fixture" + } + + fn spawn_worker(&self, _: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { + panic!("not used") + } + + fn dispatch_input(&self, _: &WorkerRef, _: WorkerInput) -> WorkerExecutionResult { + panic!("not used") + } + + fn working_directory( + &self, + _: &str, + ) -> Result { + Ok(serde_json::from_value(serde_json::json!({ + "summary": { + "working_directory_id": "w", "repository_id": "r", + "materializer_kind": "runtime_git_clone", "status": "active" + } + })) + .unwrap()) + } + + fn open_workdir_session( + &self, + _: &str, + ) -> Result { + Ok(self.session.clone()) + } + + fn cleanup_working_directory( + &self, + id: &str, + ) -> Result { + self.working_directory(id) + } +} + +struct Fixture { + _temp: tempfile::TempDir, + provider: Arc, + runtime: Runtime, + router: Router, + path: String, +} + +impl Fixture { + async fn new(pause_first: bool, fail_first: bool) -> Self { + let temp = tempfile::tempdir().unwrap(); + std::fs::write(temp.path().join("hello.txt"), "hello").unwrap(); + let provider = Arc::new(CloseProvider { + local: LocalWorkdirSession::materialized_bound( + Workdir::new("w"), + temp.path().to_owned(), + temp.path().to_owned(), + manifest::SharedScope::new(manifest::Scope::writable(temp.path()).unwrap()), + workdir::WorkdirSessionCapabilities::ALL, + ), + calls: AtomicUsize::new(0), + cancelled: AtomicUsize::new(0), + pause_first, + fail_first, + entered: Notify::new(), + release: Semaphore::new(0), + }); + let runtime = Runtime::with_execution_backend( + RuntimeOptions::default(), + Arc::new(Backend { + session: provider.clone(), + }), + ) + .unwrap(); + let router = runtime_http_router(runtime.clone(), "probe-token".into()); + let opened = router + .clone() + .oneshot(request("POST", "/v1/working-directories/w/sessions", "{}")) + .await + .unwrap(); + assert_eq!(opened.status(), StatusCode::OK); + let opened: OpenWorkdirSessionResponse = + serde_json::from_slice(&to_bytes(opened.into_body(), 65536).await.unwrap()).unwrap(); + let path = format!("/v1/workdir-sessions/{}", opened.session_id.as_str()); + Self { + _temp: temp, + provider, + runtime, + router, + path, + } + } + + fn close(&self) -> impl Future + use<> { + let router = self.router.clone(); + let request = request("DELETE", &self.path, ""); + async move { router.oneshot(request).await.unwrap() } + } + + fn assert_busy(&self) { + let error = self + .runtime + .cleanup_working_directory("w") + .expect_err("unconfirmed close must retain the real Runtime use guard"); + assert!( + matches!(error, RuntimeError::WorkingDirectory(ref d) + if d.code == "working_directory_cleanup_resource_busy"), + "{error:?}" + ); + } + + async fn assert_operations_unavailable(&self) { + let response = self + .router + .clone() + .oneshot(request( + "POST", + &format!("{}/operations", self.path), + r#"{"operation":{"operation":"stat","request":{"path":"hello.txt"}}}"#, + )) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + } +} + +impl Drop for Fixture { + fn drop(&mut self) { + // Let a detached close finish even if a test assertion unwinds. + self.provider.release.add_permits(1); + } +} + +fn request(method: &str, path: &str, body: &str) -> Request { + Request::builder() + .method(method) + .uri(path) + .header("authorization", "Bearer probe-token") + .header("content-type", "application/json") + .body(Body::from(body.to_owned())) + .unwrap() +} + +async fn bounded(future: impl Future) -> T { + tokio::time::timeout(std::time::Duration::from_secs(10), future) + .await + .expect("close probe must not hang") +} + +async fn assert_pending(mut future: std::pin::Pin<&mut impl Future>) { + assert!( + std::future::poll_fn(|cx| Poll::Ready(future.as_mut().poll(cx))) + .await + .is_pending(), + "close acknowledged before teardown confirmation" + ); +} + +#[tokio::test(flavor = "current_thread")] +async fn duplicate_delete_waits_for_actual_local_teardown_and_runtime_guard_release() { + let fixture = Fixture::new(true, false).await; + fixture.assert_busy(); + let mut first = Box::pin(fixture.close()); + assert_pending(first.as_mut()).await; + bounded(fixture.provider.entered.notified()).await; + let mut duplicate = Box::pin(fixture.close()); + assert_pending(duplicate.as_mut()).await; + fixture.assert_busy(); + fixture.assert_operations_unavailable().await; + assert_eq!(fixture.provider.calls.load(Ordering::SeqCst), 1); + + fixture.provider.release.add_permits(1); + assert_eq!(bounded(first).await.status(), StatusCode::NO_CONTENT); + assert_eq!(bounded(duplicate).await.status(), StatusCode::NO_CONTENT); + fixture + .runtime + .cleanup_working_directory("w") + .expect("confirmed close releases use guard"); + assert_eq!( + bounded(fixture.close()).await.status(), + StatusCode::NO_CONTENT + ); + assert_eq!(fixture.provider.calls.load(Ordering::SeqCst), 1); +} + +#[tokio::test(flavor = "current_thread")] +async fn cancelled_http_waiter_does_not_cancel_provider_or_let_retries_ack_early() { + let fixture = Fixture::new(true, false).await; + let mut first = Box::pin(fixture.close()); + assert_pending(first.as_mut()).await; + bounded(fixture.provider.entered.notified()).await; + drop(first); + let mut retry = Box::pin(fixture.close()); + let mut duplicate = Box::pin(fixture.close()); + assert_pending(retry.as_mut()).await; + assert_pending(duplicate.as_mut()).await; + fixture.assert_busy(); + assert_eq!(fixture.provider.calls.load(Ordering::SeqCst), 1); + assert_eq!(fixture.provider.cancelled.load(Ordering::SeqCst), 0); + + fixture.provider.release.add_permits(1); + assert_eq!(bounded(retry).await.status(), StatusCode::NO_CONTENT); + assert_eq!(bounded(duplicate).await.status(), StatusCode::NO_CONTENT); + assert_eq!(fixture.provider.cancelled.load(Ordering::SeqCst), 0); + fixture + .runtime + .cleanup_working_directory("w") + .expect("task completes real provider close"); + assert_eq!( + bounded(fixture.close()).await.status(), + StatusCode::NO_CONTENT + ); + assert_eq!(fixture.provider.calls.load(Ordering::SeqCst), 1); +} + +#[tokio::test(flavor = "current_thread")] +async fn provider_close_error_retains_session_and_runtime_guard_for_retry() { + let fixture = Fixture::new(false, true).await; + let mut first = Box::pin(fixture.close()); + let mut duplicate = Box::pin(fixture.close()); + assert_pending(first.as_mut()).await; + assert_pending(duplicate.as_mut()).await; + let response = bounded(first).await; + assert_eq!( + bounded(duplicate).await.status(), + StatusCode::SERVICE_UNAVAILABLE + ); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + let error: WorkdirTransportError = + serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap(); + assert_eq!(error.code, WorkdirTransportErrorCode::Unavailable); + fixture.assert_busy(); + assert_eq!(fixture.provider.calls.load(Ordering::SeqCst), 1); + + assert_eq!( + bounded(fixture.close()).await.status(), + StatusCode::NO_CONTENT + ); + assert_eq!(fixture.provider.calls.load(Ordering::SeqCst), 2); + fixture + .runtime + .cleanup_working_directory("w") + .expect("retry releases the retained Runtime guard"); + assert_eq!( + bounded(fixture.close()).await.status(), + StatusCode::NO_CONTENT + ); + assert_eq!(fixture.provider.calls.load(Ordering::SeqCst), 2); +} diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 6f7ab46a3..b8b362da6 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() { @@ -21407,17 +21431,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(()) } @@ -39310,6 +39336,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; @@ -51253,6 +51280,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..802589f2a 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 @@ -13,6 +13,7 @@ use worker_runtime::execution::{ #[derive(Default)] struct LaunchExecutionBackend { execution: DeterministicExecutionBackend, + provider_close_gate: Mutex>>, bundles: Mutex>, creates: Mutex>, submits: Mutex< @@ -122,6 +123,35 @@ 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(); + let session: WorkdirSessionHandle = 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()?, + ), + ); + Ok(match self.provider_close_gate.lock().unwrap().as_ref() { + Some(gate) => gate.wrap(session), + None => session, + }) + } + fn create_working_directory( &self, request: &worker_runtime::catalog::WorkingDirectoryRequest, @@ -180,11 +210,15 @@ pub(super) struct LaunchFixture { backend: Arc, runtime_id: String, runtime_server: Option>, + runtime_router: Option, runtime_identity: Option, } impl Drop for LaunchFixture { fn drop(&mut self) { + if let Some(gate) = self.backend.provider_close_gate.lock().unwrap().as_ref() { + gate.release(); + } if let Some(server) = &self.runtime_server { server.abort(); } @@ -192,13 +226,38 @@ impl Drop for LaunchFixture { } impl LaunchFixture { + pub(super) fn hold_provider_close( + &self, + gate: Arc, + ) { + *self.backend.provider_close_gate.lock().unwrap() = Some(gate); + } + + pub(super) fn runtime_router(&self) -> Router { + self.runtime_router.clone().unwrap() + } + + 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 +283,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(); @@ -247,6 +310,7 @@ impl LaunchFixture { backend, runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), runtime_server: None, + runtime_router: None, runtime_identity: None, }; if remote { @@ -282,11 +346,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 { @@ -318,18 +397,17 @@ impl LaunchFixture { .unwrap() .unwrap(); let runtime_id = fixture.runtime_id.clone(); + let router = worker_runtime::http_server::runtime_http_router_with_workspace_auth( + runtime, + None, + worker_runtime::http_server::WorkspaceRuntimeHttpAuth { + verifier, + runtime_id, + }, + ); + fixture.runtime_router = Some(router.clone()); fixture.runtime_server = Some(tokio::spawn(async move { - worker_runtime::http_server::serve_runtime_http_with_workspace_auth( - runtime, - listener, - None, - worker_runtime::http_server::WorkspaceRuntimeHttpAuth { - verifier, - runtime_id, - }, - ) - .await - .unwrap(); + axum::serve(listener, router).await.unwrap(); })); let mut remote_config = RemoteRuntimeConfig { runtime_id: fixture.runtime_id.clone(), 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..7221d6430 --- /dev/null +++ b/crates/workspace-server/src/server/tests/workdir_session_close_tests.rs @@ -0,0 +1,772 @@ +//! 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 std::{future::Future, task::Poll}; +use tokio::sync::{Notify, Semaphore}; +use tower::ServiceExt; + +#[derive(Debug)] +pub(super) struct ProviderCloseGate { + entered: Notify, + release: Semaphore, +} + +impl ProviderCloseGate { + pub(super) fn release(&self) { + self.release.add_permits(1); + } + + pub(super) fn wrap(self: &Arc, source: WorkdirSessionHandle) -> WorkdirSessionHandle { + Arc::new(GatedProviderClose { + source, + gate: self.clone(), + }) + } +} + +#[derive(Debug)] +struct GatedProviderClose { + source: WorkdirSessionHandle, + gate: Arc, +} + +// Forward the real provider operations; only suspend an already-pending close. +macro_rules! gated_provider { + ($($name:ident($request:ident: $input:ty) -> $output:ty;)*) => { + #[async_trait::async_trait] + impl WorkdirSession for GatedProviderClose { + fn workdir(&self) -> &workdir::Workdir { self.source.workdir() } + fn capabilities(&self) -> workdir::WorkdirSessionCapabilities { self.source.capabilities() } + fn environment(&self) -> workdir::WorkdirSessionEnvironment { self.source.environment() } + $(async fn $name(&self, $request: $input) -> std::result::Result<$output, workdir::WorkdirError> { + self.source.$name($request).await + })* + async fn close(&self) -> std::result::Result<(), workdir::WorkdirError> { + let mut closing = Box::pin(self.source.close()); + match std::future::poll_fn(|cx| Poll::Ready(closing.as_mut().poll(cx))).await { + Poll::Ready(result) => result, + Poll::Pending => { + self.gate.entered.notify_one(); + self.gate.release.acquire().await.unwrap().forget(); + closing.await + } + } + } + } + }; +} + +gated_provider! { + stat(request: workdir::StatRequest) -> workdir::StatResult; + read(request: workdir::ReadRequest) -> workdir::ReadResult; + write(request: workdir::WriteRequest) -> workdir::WriteResult; + edit(request: workdir::EditRequest) -> workdir::EditResult; + list(request: workdir::ListRequest) -> workdir::ListResult; + glob(request: workdir::GlobRequest) -> workdir::GlobResult; + grep(request: workdir::GrepRequest) -> workdir::GrepResult; + start_command(request: workdir::CommandRequest) -> workdir::CommandHandle; + command_status(handle: workdir::CommandHandle) -> workdir::CommandStatus; + command_output(request: workdir::CommandOutputRequest) -> workdir::CommandOutput; + cancel_command(handle: workdir::CommandHandle) -> (); +} + +#[derive(Clone, Copy)] +enum CloseFault { + Reject, + HoldThenReject, + RejectAfterRuntimeClose, +} + +struct CloseTransport { + upstream: String, + client: reqwest::Client, + fault: Mutex>, + attempts: Mutex>, + runtime_closes: Mutex>, + direct_runtime: Mutex>, + pending_dispatches: Mutex>, + dispatched: Notify, + 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 direct = state.direct_runtime.lock().unwrap().clone(); + if let Some(router) = direct.filter(|_| is_close) { + // Poll the production authenticated handler without another socket + // scheduling boundary, so a duplicate's Pending/204 is observable. + let mut response = Box::pin(router.oneshot(request)); + let first_poll = + std::future::poll_fn(|cx| Poll::Ready(response.as_mut().poll(cx))).await; + state + .pending_dispatches + .lock() + .unwrap() + .push(first_poll.is_pending()); + state.dispatched.notify_one(); + let response = match first_poll { + Poll::Ready(result) => result.unwrap(), + Poll::Pending => response.await.unwrap(), + }; + state + .runtime_closes + .lock() + .unwrap() + .push((path, response.status())); + return 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()), + direct_runtime: Mutex::new(None), + pending_dispatches: Mutex::new(Vec::new()), + dispatched: Notify::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] +async fn cancelled_backend_close_retry_waits_for_actual_provider_teardown_before_confirmation() { + let fixture = Fixture::new().await; + let gate = Arc::new(ProviderCloseGate { + entered: Notify::new(), + release: Semaphore::new(0), + }); + fixture.source.hold_provider_close(gate.clone()); + *fixture.transport.direct_runtime.lock().unwrap() = Some(fixture.source.runtime_router()); + let worker = fixture.registry_owner(); + let session = fixture.open().await; + session + .start_command(workdir::CommandRequest { + command: "sleep 60".into(), + timeout_secs: 60, + transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES, + cwd: workdir::WorkdirPath::root(), + tool_call_id: None, + }) + .await + .unwrap(); + fixture + .api() + .workdir_sessions + .lock() + .unwrap() + .insert_attachment(worker.clone(), session.clone()); + fixture.assert_busy().await; + let api = fixture.api().clone(); + let first_worker = worker.clone(); + let first = tokio::spawn(async move { close(&api, &first_worker, ClosePath::Registry).await }); + tokio::time::timeout(std::time::Duration::from_secs(5), gate.entered.notified()) + .await + .expect("actual Local provider teardown reached Pending"); + fixture.transport.dispatched.notified().await; + first.abort(); + assert!(first.await.unwrap_err().is_cancelled()); + fixture.assert_restored(&worker, &session); + fixture.assert_busy().await; + + let api = fixture.api().clone(); + let retry_worker = worker.clone(); + let retry = tokio::spawn(async move { close(&api, &retry_worker, ClosePath::Registry).await }); + tokio::time::timeout( + std::time::Duration::from_secs(5), + fixture.transport.dispatched.notified(), + ) + .await + .expect("retry reached production Runtime handler"); + assert_eq!( + *fixture.transport.pending_dispatches.lock().unwrap(), + vec![true, true], + "duplicate DELETE must wait, not acknowledge an extracted record" + ); + assert!( + !retry.is_finished(), + "Backend cannot confirm while provider close is held" + ); + assert!(fixture.transport.runtime_closes.lock().unwrap().is_empty()); + fixture.assert_busy().await; + gate.release.add_permits(1); + assert!( + tokio::time::timeout(std::time::Duration::from_secs(5), retry) + .await + .unwrap() + .unwrap() + ); + assert!( + !fixture + .api() + .workdir_sessions + .lock() + .unwrap() + .attachments + .contains_key(&worker) + ); + let attempts = fixture.transport.attempts(); + assert_eq!(attempts.len(), 2); + assert_eq!(attempts[0], attempts[1]); + 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; +}