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/tests/ticket_worker_launch_tests.rs b/crates/workspace-server/src/server/tests/ticket_worker_launch_tests.rs index 50e92b0de..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< @@ -134,7 +135,7 @@ impl WorkerExecutionBackend for LaunchExecutionBackend { .materializer .bind_working_directory(id, None)?; let scope = manifest::Scope::writable(binding.root()).unwrap(); - Ok(Arc::new( + let session: WorkdirSessionHandle = Arc::new( workdir::LocalWorkdirSession::materialized_bound_with_environment( workdir::Workdir::new(id), binding.root().to_owned(), @@ -144,7 +145,11 @@ impl WorkerExecutionBackend for LaunchExecutionBackend { 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( @@ -205,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(); } @@ -217,6 +226,17 @@ 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); } @@ -290,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 { @@ -376,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 index e4daae4a2..7221d6430 100644 --- a/crates/workspace-server/src/server/tests/workdir_session_close_tests.rs +++ b/crates/workspace-server/src/server/tests/workdir_session_close_tests.rs @@ -4,7 +4,74 @@ //! 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 { @@ -19,6 +86,9 @@ struct CloseTransport { fault: Mutex>, attempts: Mutex>, runtime_closes: Mutex>, + direct_runtime: Mutex>, + pending_dispatches: Mutex>, + dispatched: Notify, entered: Notify, release: Semaphore, } @@ -42,6 +112,30 @@ impl CloseTransport { // 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 @@ -125,6 +219,9 @@ impl Fixture { 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), }); @@ -313,6 +410,88 @@ async fn failure_restoration(path: ClosePath) { 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;