fix: retain Runtime close ownership until provider teardown completes

Resolve MR runtime-close-ack-before-release: keep sessions registered while a task-owned close is pending; all DELETE waiters share its result, successful teardown alone removes the entry, and provider errors permit retry. Cancelling HTTP waiters cannot cancel Local provider command teardown.

Add deterministic Local teardown/Runtime active-use tests plus authenticated Backend registry cancellation and same-session retry through actual materializer busy/removal. Focused dependent Server 26 pass, Runtime new 3 and existing close tests pass; root check, fmt, diff, names pass. Full Runtime attempt: 505 pass/63 restore-environment failures; corrected exact 63 pass separately, not a full-suite green claim. Validation logs: target/t805-validation. Base remains 0f82; no T795/T796 import.
This commit is contained in:
2026-10-30 22:25:27 +09:00
parent 1c03758fc1
commit 5d0e03dc53
5 changed files with 718 additions and 26 deletions
+90 -13
View File
@@ -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<tokio::sync::watch::Receiver<Option<RuntimeHttpWorkdirCloseResult>>>,
}
/// `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<Extension<RuntimeAuthContext>>,
) -> Result<StatusCode, RuntimeHttpWorkdirError> {
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()),
@@ -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")),
@@ -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<workdir::StatResult, WorkdirError> {
self.local.stat(request).await
}
async fn read(
&self,
request: workdir::ReadRequest,
) -> Result<workdir::ReadResult, WorkdirError> {
self.local.read(request).await
}
async fn write(
&self,
request: workdir::WriteRequest,
) -> Result<workdir::WriteResult, WorkdirError> {
self.local.write(request).await
}
async fn edit(
&self,
request: workdir::EditRequest,
) -> Result<workdir::EditResult, WorkdirError> {
self.local.edit(request).await
}
async fn list(
&self,
request: workdir::ListRequest,
) -> Result<workdir::ListResult, WorkdirError> {
self.local.list(request).await
}
async fn glob(
&self,
request: workdir::GlobRequest,
) -> Result<workdir::GlobResult, WorkdirError> {
self.local.glob(request).await
}
async fn grep(
&self,
request: workdir::GrepRequest,
) -> Result<workdir::GrepResult, WorkdirError> {
self.local.grep(request).await
}
async fn start_command(
&self,
request: workdir::CommandRequest,
) -> Result<workdir::CommandHandle, WorkdirError> {
self.local.start_command(request).await
}
async fn command_status(
&self,
request: workdir::CommandHandle,
) -> Result<workdir::CommandStatus, WorkdirError> {
self.local.command_status(request).await
}
async fn command_output(
&self,
request: workdir::CommandOutputRequest,
) -> Result<workdir::CommandOutput, WorkdirError> {
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<CloseProvider>,
}
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<WorkingDirectoryStatus, WorkingDirectoryDiagnostic> {
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<WorkdirSessionHandle, WorkingDirectoryDiagnostic> {
Ok(self.session.clone())
}
fn cleanup_working_directory(
&self,
id: &str,
) -> Result<WorkingDirectoryStatus, WorkingDirectoryDiagnostic> {
self.working_directory(id)
}
}
struct Fixture {
_temp: tempfile::TempDir,
provider: Arc<CloseProvider>,
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<Output = Response> + 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<Body> {
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<T>(future: impl Future<Output = T>) -> T {
tokio::time::timeout(std::time::Duration::from_secs(10), future)
.await
.expect("close probe must not hang")
}
async fn assert_pending<T>(mut future: std::pin::Pin<&mut impl Future<Output = T>>) {
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);
}
@@ -13,6 +13,7 @@ use worker_runtime::execution::{
#[derive(Default)]
struct LaunchExecutionBackend {
execution: DeterministicExecutionBackend,
provider_close_gate: Mutex<Option<Arc<super::workdir_session_close_tests::ProviderCloseGate>>>,
bundles: Mutex<std::collections::HashMap<String, worker_runtime::config_bundle::ConfigBundle>>,
creates: Mutex<Vec<worker_runtime::catalog::CreateWorkerRequest>>,
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<LaunchExecutionBackend>,
runtime_id: String,
runtime_server: Option<tokio::task::JoinHandle<()>>,
runtime_router: Option<Router>,
runtime_identity: Option<RuntimeIdentityMaterial>,
}
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<super::workdir_session_close_tests::ProviderCloseGate>,
) {
*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(),
@@ -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<Self>, source: WorkdirSessionHandle) -> WorkdirSessionHandle {
Arc::new(GatedProviderClose {
source,
gate: self.clone(),
})
}
}
#[derive(Debug)]
struct GatedProviderClose {
source: WorkdirSessionHandle,
gate: Arc<ProviderCloseGate>,
}
// 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<Option<CloseFault>>,
attempts: Mutex<Vec<String>>,
runtime_closes: Mutex<Vec<(String, StatusCode)>>,
direct_runtime: Mutex<Option<Router>>,
pending_dispatches: Mutex<Vec<bool>>,
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;