fix: confirm remote Workdir close before releasing session ownership

This commit is contained in:
2026-10-30 21:44:51 +09:00
parent 0f82d6fa8f
commit 1c03758fc1
5 changed files with 1018 additions and 32 deletions
+19 -2
View File
@@ -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(_))
));
}
}
}
+287
View File
@@ -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<TcpStream>,
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::<usize>().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<F: Future>(&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<dyn WorkdirHttpAuthorization>) -> 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<u8>,
) -> 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<Option<tokio::sync::oneshot::Receiver<()>>>,
}
impl WorkdirHttpAuthorization for GatedAuthorization {
fn bearer_token(
&self,
method: String,
_: String,
_: Vec<u8>,
) -> 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());
}
}
+56 -25
View File
@@ -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<Mutex<WorkdirSessionRegistry>>,
registered: Vec<RegisteredWorkdirSession>,
}
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<Mutex<WorkdirSessionRegistry>>,
registered: Vec<RegisteredWorkdirSession>,
) -> 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();
@@ -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 {
@@ -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<Option<CloseFault>>,
attempts: Mutex<Vec<String>>,
runtime_closes: Mutex<Vec<(String, StatusCode)>>,
entered: Notify,
release: Semaphore,
}
impl CloseTransport {
async fn forward(State(state): State<Arc<Self>>, request: Request<Body>) -> 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<String> {
self.attempts.lock().unwrap().clone()
}
}
struct Fixture {
source: LaunchFixture,
transport: Arc<CloseTransport>,
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;
}