From bb5f64fadd021d60e029b8a14f1373dc11d7547f Mon Sep 17 00:00:00 2001 From: daniel Date: Sat, 5 Sep 2026 19:58:44 +0100 Subject: [PATCH 1/3] fix: make shared-state transitions unwind-safe --- .agents/skills/shared-state/SKILL.md | 14 + AGENTS.md | 4 + src/acp_child.rs | 214 +++-- src/protocols/acp.rs | 1135 ++++++++++++++++++++++---- src/protocols/acp/activity.rs | 287 ++++++- src/protocols/acp/v2.rs | 395 +++++++-- src/provider/openai_auth.rs | 54 +- src/resilient_fs/mod.rs | 266 +++--- src/resilient_fs/tests.rs | 242 +++++- src/runtime.rs | 474 ++++++----- src/runtime/tests.rs | 315 ++++++- src/session.rs | 287 ++++++- src/tools/mcp.rs | 266 +++++- 13 files changed, 3346 insertions(+), 607 deletions(-) create mode 100644 .agents/skills/shared-state/SKILL.md diff --git a/.agents/skills/shared-state/SKILL.md b/.agents/skills/shared-state/SKILL.md new file mode 100644 index 00000000..7db1e31c --- /dev/null +++ b/.agents/skills/shared-state/SKILL.md @@ -0,0 +1,14 @@ +--- +name: shared-state +description: MUST use whenever adding or editing shared state protected by a Mutex, RwLock, or async lock, including guarded transitions, lock acquisition, poison recovery, and coordinated registries or counters. +--- + +# Shared state + +1. Before editing, list the guarded invariants and ALL writers: normal methods, callbacks, destructors, unwind cleanup, cancellation, and coordinated objects or registries. Inspect lock ordering and every await while a guard or ownership claim is live. Keep a bounded inventory of owners, guarantees, findings, and unresolved blockers. +2. Design each transition to leave valid state on success, rejection, error, unwind, and cancellation. Prepare fallible or panicking work before the commit where feasible. Commit complete valid state without exposing intermediate invalid combinations. Use private ownership, rollback, repair, or typed-error isolation when justified; rollback must itself be safe during unwind and must account for external effects. +3. Release guards before dropping replaced values or invoking arbitrary callbacks when those operations can panic, reenter, or acquire other locks. Include implicit drops, trait implementations, backend calls, and wakeups in the analysis. Mutex exclusion and a single assignment do not prove consistency across objects or external effects. Memory safety is not logical consistency. +4. Recover a poisoned lock with `into_inner` ONLY beside an explicit invariant argument covering EVERY guarded transition and any external state it coordinates. Standard-library poisoning is advisory, not a correctness mechanism; async locks may not poison at all. Preserve invariants across unwind and cancellation for both. If they cannot survive failure, redesign, repair, or isolate the owner with a typed error; never assert that recovery is safe without evidence. +5. Do not blanket-replace `expect` with poison recovery or defaults, use production `catch_unwind`, introduce a generic recovery wrapper, or switch to a non-poisoning lock crate to conceal inconsistent state. Do not enable global panic lints as part of a shared-state fix. +6. Add failure-path tests for changed invariants: rejected and failed commits, unwind at real callback/backend boundaries, poison behavior, cancellation at suspension points, stale generations, and cleanup/drop ordering as applicable. Assert subsequent state and cross-owner effects, not just absence of panic. Use test-only modules and fakes at real boundaries; no test-specific production branches, counters, or hooks. Test-only `catch_unwind` may observe unwinding. +7. Run focused tests and formatting/lint checks. Get an independent review of the invariant arguments, code, and failure tests. Report audited coverage and residual blockers honestly; do not claim a repository-wide guarantee from a bounded audit. diff --git a/AGENTS.md b/AGENTS.md index 84f8c138..19f7e62e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -3,3 +3,7 @@ Only when you are using Kit as your agent harness: report issues with Kit's harness at https://github.com/speakeasy-api/kit/issues. Do not report issues with other harnesses there. Do not open an issue on the user's behalf unless the user explicitly requests it; ask the user first when they have not already made that request. Follow [Reporting Kit Issues](docs/user/reporting-kit-issues.md). Do not change release versions in ordinary pull requests. Use a Conventional Commit title for every pull request. Mark a breaking change with `!` after the commit type or scope, or with a `BREAKING CHANGE:` line in the commit body. The release workflow derives the next version from commits since the latest release, advances the version files in a release commit, and applies a minor bump when any commit is breaking or a patch bump otherwise. + +## Shared state + +You MUST load and follow `.agents/skills/shared-state/SKILL.md` whenever adding or editing lock-protected shared state, guarded transitions, lock acquisition, or recovery. diff --git a/src/acp_child.rs b/src/acp_child.rs index cd30eeda..7a389d39 100644 --- a/src/acp_child.rs +++ b/src/acp_child.rs @@ -411,12 +411,15 @@ impl std::fmt::Display for ChildError { } struct Prompt { + // The worker, not the waiting caller, owns serialization until settlement. + serial: tokio::sync::OwnedMutexGuard<()>, session_id: SessionId, text: String, cancellation: TurnCancellation, reply: oneshot::Sender>, } struct Fork { + serial: tokio::sync::OwnedMutexGuard<()>, session_id: SessionId, model: Option, parent: Option<(String, String)>, @@ -689,8 +692,8 @@ impl ChildSession { parent: Option<(String, String)>, cancellation: &TurnCancellation, ) -> Result { - let _serial = tokio::select! { - serial = self.serial.lock() => serial, + let serial = tokio::select! { + serial = self.serial.clone().lock_owned() => serial, () = cancellation.cancelled() => return Err(ChildError::Cancelled), }; if !self.supports_native_fork() { @@ -702,6 +705,7 @@ impl ChildSession { let (reply, response) = oneshot::channel(); tokio::select! { sent = self.tx.send(Request::Fork(Fork { + serial, session_id: self.session_id.clone(), model: model.map(str::to_owned), parent, @@ -728,12 +732,13 @@ impl ChildSession { text: String, cancellation: TurnCancellation, ) -> Result { - let _serial = tokio::select! { - serial = self.serial.lock() => serial, + let serial = tokio::select! { + serial = self.serial.clone().lock_owned() => serial, () = cancellation.cancelled() => return Err(ChildError::Cancelled), }; let (reply, response) = oneshot::channel(); let request = Request::Prompt(Prompt { + serial, session_id: self.session_id.clone(), text, cancellation: cancellation.clone(), @@ -908,6 +913,7 @@ async fn run( let sessions = Arc::clone(&sessions); let root = root.clone(); tasks.spawn(async move { + let serial = fork.serial; let mut request = ForkSessionRequest::new(fork.session_id, root); if let Some((id, name)) = fork.parent { request.meta = Some(serde_json::Map::from_iter([ @@ -967,6 +973,9 @@ async fn run( let cleanup_connection = connection.clone(); let cleanup_sessions = Arc::clone(&sessions); tokio::spawn(async move { + // The remote fork still owns source-session + // serialization until its response and cleanup. + let _serial = serial; let Ok(response) = request.await else { return; }; @@ -992,6 +1001,9 @@ async fn run( let cleanup_connection = connection.clone(); let cleanup_sessions = Arc::clone(&sessions); tokio::spawn(async move { + // The remote fork still owns source-session + // serialization until its response and cleanup. + let _serial = serial; let Ok(response) = request.await else { return; }; @@ -1050,6 +1062,7 @@ async fn run( let routes = Arc::clone(&routes); let fatal = fatal_tx.clone(); tasks.spawn(async move { + let _serial = prompt.serial; let session_id = prompt.session_id.clone(); let output = Arc::new(Mutex::new(ChildOutput::default())); if let Ok(mut routes) = routes.lock() { routes.insert(session_id.clone(), Arc::clone(&output)); } @@ -1235,6 +1248,61 @@ mod tests { serde_json::from_value(value).unwrap() } + #[tokio::test] + async fn cancelled_caller_does_not_release_child_request_serialization() { + for fork in [false, true] { + let (tx, mut rx) = mpsc::channel(1); + let mut capabilities = agentkit_acp::AgentCapabilities::default(); + capabilities.session_capabilities.fork = Some(Default::default()); + let child = ChildSession { + tx, + session_id: "test".into(), + capabilities, + serial: Arc::new(tokio::sync::Mutex::new(())), + closed: watch::channel(false).1, + descendant_parent: None, + }; + let caller = child.clone(); + let task = tokio::spawn(async move { + if fork { + caller + .fork(None, None, &TurnCancellation::default()) + .await + .map(|_| ()) + } else { + caller + .prompt("first".into(), TurnCancellation::default()) + .await + .map(|_| ()) + } + }); + // The actor has accepted the request but has not settled it. + let request = rx.recv().await.unwrap(); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + assert!(child.serial.try_lock().is_err()); + drop(request); + assert!(child.serial.try_lock().is_ok()); + + // Settlement releases the same gate for a subsequent prompt. + let answer = async { + let Some(Request::Prompt(prompt)) = rx.recv().await else { + panic!("expected a prompt"); + }; + prompt.reply.send(Ok(ChildOutput::default())).unwrap(); + }; + let (result, ()) = tokio::time::timeout(Duration::from_secs(1), async { + tokio::join!( + child.prompt("next".into(), TurnCancellation::default()), + answer + ) + }) + .await + .unwrap(); + assert!(result.is_ok()); + } + } + #[test] fn nested_runtime_events_are_not_forwarded_as_parent_events() { let event = crate::events::RuntimeEvent::ChildStarted { @@ -1961,69 +2029,87 @@ mod tests { } #[tokio::test] - async fn cancelling_native_fork_keeps_the_shared_process_usable() { - let root = tempfile::tempdir().unwrap(); - let release = root.path().join("release-fork"); - let mut profiles = BTreeMap::new(); - profiles.insert( - "mock".into(), - AcpHarnessProfile { - command: "python3".into(), - args: vec![ - format!("{}/fixtures/mock-acp.py", env!("CARGO_MANIFEST_DIR")), - format!("--fork-release={}", release.display()), - ], - permissions: AcpPermissionPolicy::Deny, - }, - ); - let config = ChildConfig { - root: root.path().to_path_buf(), - model: "unused".into(), - provider: Default::default(), - reasoning_effort: None, - openrouter_api_key: None, - configured_mcp_config: None, - configured_mcp_config_inherited: false, - legacy_mcp_config: false, - mcp_config: None, - credential_storage: Default::default(), - telemetry: Default::default(), - harnesses: AcpHarnesses::new(profiles).unwrap(), - default_harness: "acp.mock".into(), - parent_id: None, - parent_name: None, - }; - let base = ChildSession::start( - config, - "acp.mock".into(), - None, - None, - 1, - TurnCancellation::default(), - ) - .await - .unwrap(); - let controller = agentkit_core::CancellationController::new(); - let cancellation = controller.handle().checkpoint(); - let child = base.clone(); - let fork = tokio::spawn(async move { child.fork(None, None, &cancellation).await }); - tokio::time::sleep(Duration::from_millis(100)).await; - controller.interrupt(); - - assert!(matches!(fork.await.unwrap(), Err(ChildError::Cancelled))); - assert_eq!( - base.prompt( - "source survives cancellation".into(), + async fn cancelled_or_timed_out_fork_retains_serialization_until_remote_settlement() { + for cancel in [true, false] { + let root = tempfile::tempdir().unwrap(); + let release = root.path().join("release-fork"); + let mut profiles = BTreeMap::new(); + profiles.insert( + "mock".into(), + AcpHarnessProfile { + command: "python3".into(), + args: vec![ + format!("{}/fixtures/mock-acp.py", env!("CARGO_MANIFEST_DIR")), + format!("--fork-release={}", release.display()), + ], + permissions: AcpPermissionPolicy::Deny, + }, + ); + let config = ChildConfig { + root: root.path().to_path_buf(), + model: "unused".into(), + provider: Default::default(), + reasoning_effort: None, + openrouter_api_key: None, + configured_mcp_config: None, + configured_mcp_config_inherited: false, + legacy_mcp_config: false, + mcp_config: None, + credential_storage: Default::default(), + telemetry: Default::default(), + harnesses: AcpHarnesses::new(profiles).unwrap(), + default_harness: "acp.mock".into(), + parent_id: None, + parent_name: None, + }; + let base = ChildSession::start( + config, + "acp.mock".into(), + None, + None, + 1, TurnCancellation::default(), ) .await - .unwrap() - .text, - "source survives cancellation" - ); - std::fs::write(release, b"ready").unwrap(); - tokio::time::sleep(Duration::from_millis(100)).await; - base.close().await.unwrap(); + .unwrap(); + let controller = agentkit_core::CancellationController::new(); + let cancellation = controller.handle().checkpoint(); + let child = base.clone(); + let fork = tokio::spawn(async move { child.fork(None, None, &cancellation).await }); + tokio::time::sleep(Duration::from_millis(100)).await; + if cancel { + controller.interrupt(); + } + let outcome = tokio::time::timeout(HANDSHAKE + Duration::from_secs(5), fork) + .await + .unwrap() + .unwrap(); + if cancel { + assert!(matches!(outcome, Err(ChildError::Cancelled))); + } else { + assert!(matches!(outcome, Err(ChildError::Failed(_)))); + } + assert!(base.serial.try_lock().is_err()); + let mut next = Box::pin(base.prompt( + "source survives cancellation".into(), + TurnCancellation::default(), + )); + assert!( + tokio::time::timeout(Duration::from_millis(20), &mut next) + .await + .is_err() + ); + std::fs::write(release, b"ready").unwrap(); + assert_eq!( + tokio::time::timeout(Duration::from_secs(5), next) + .await + .unwrap() + .unwrap() + .text, + "source survives cancellation" + ); + base.close().await.unwrap(); + } } #[test] diff --git a/src/protocols/acp.rs b/src/protocols/acp.rs index 4893ba68..51cccf3a 100644 --- a/src/protocols/acp.rs +++ b/src/protocols/acp.rs @@ -513,6 +513,19 @@ struct RegistryState { v2_sessions: HashMap, } +impl RegistryState { + fn close_gate(&mut self, permanently: bool) { + self.accepting = false; + self.permanently_closed |= permanently; + if let Some(generation) = self.generation.checked_add(1) { + self.generation = generation; + } else { + // Never let an old admission become current again through wraparound. + self.permanently_closed = true; + } + } +} + struct SessionRegistryInner { next_token: AtomicU64, lifecycle: tokio::sync::Mutex<()>, @@ -520,6 +533,27 @@ struct SessionRegistryInner { state: Mutex, } +impl SessionRegistryInner { + fn lock_state(&self) -> std::sync::MutexGuard<'_, RegistryState> { + self.state.lock().unwrap_or_else(|poisoned| { + // This registry owns only admission bookkeeping and live handle indexes, + // not the durable v2 publication commit. All writers below finish their + // in-memory transitions without callbacks, wakeups, arbitrary drops, or + // awaits: admission counts are checked before incrementing; completion + // consumes a validated admission; insertion rejects existing tokens; + // removal extracts ownership; gate closure cannot wrap its generation. + // Snapshots clone only concrete Arc/channel/handle types. Reopening + // happens only after teardown and the external reset have succeeded; + // cancellation or panic before then leaves admission closed. Thus unwind + // cannot leave a partial transition or imply an external commit succeeded. + // In particular, actor/admission destructors must work after poison too. + let state = poisoned.into_inner(); + self.state.clear_poison(); + state + }) + } +} + pub(super) struct SessionAdmission { generation: u64, active: bool, @@ -530,7 +564,6 @@ impl SessionAdmission { fn complete(&mut self, state: &mut RegistryState) { state.pending_attachments = state.pending_attachments.saturating_sub(1); self.active = false; - self.registry.attachments_changed.notify_waiters(); } } @@ -540,10 +573,7 @@ impl Drop for SessionAdmission { return; } let registry = Arc::clone(&self.registry); - let mut state = registry - .state - .lock() - .expect("ACP session registry poisoned"); + let mut state = registry.lock_state(); state.pending_attachments = state.pending_attachments.saturating_sub(1); self.active = false; drop(state); @@ -577,19 +607,20 @@ impl SessionRegistry { } pub(super) fn next_token(&self) -> u64 { - self.inner.next_token.fetch_add(1, Ordering::Relaxed) + self.inner + .next_token + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |token| { + token.checked_add(1) + }) + .expect("ACP session registry token space exhausted") } fn begin_attachment(&self) -> Result { - let mut state = self - .inner - .state - .lock() - .expect("ACP session registry poisoned"); + let mut state = self.inner.lock_state(); if !state.accepting { return Err(()); } - state.pending_attachments += 1; + state.pending_attachments = state.pending_attachments.checked_add(1).ok_or(())?; Ok(SessionAdmission { generation: state.generation, active: true, @@ -602,16 +633,20 @@ impl SessionRegistry { admission: &mut SessionAdmission, session: RegisteredSession, ) -> Result<(), ()> { - let mut state = self - .inner - .state - .lock() - .expect("ACP session registry poisoned"); - if !state.accepting || state.generation != admission.generation { + let mut state = self.inner.lock_state(); + if !admission.active + || !Arc::ptr_eq(&self.inner, &admission.registry) + || !state.accepting + || state.generation != admission.generation + || state.sessions.contains_key(&session.token) + || state.v2_sessions.contains_key(&session.token) + { return Err(()); } state.sessions.insert(session.token, session); admission.complete(&mut state); + drop(state); + self.inner.attachments_changed.notify_waiters(); Ok(()) } @@ -624,12 +659,14 @@ impl SessionRegistry { actor: AbortHandle, completed: watch::Receiver, ) -> Result<(), ()> { - let mut state = self - .inner - .state - .lock() - .expect("ACP session registry poisoned"); - if !state.accepting || state.generation != admission.generation { + let mut state = self.inner.lock_state(); + if !admission.active + || !Arc::ptr_eq(&self.inner, &admission.registry) + || !state.accepting + || state.generation != admission.generation + || state.sessions.contains_key(&token) + || state.v2_sessions.contains_key(&token) + { return Err(()); } state.v2_sessions.insert( @@ -643,31 +680,27 @@ impl SessionRegistry { }, ); admission.complete(&mut state); + drop(state); + self.inner.attachments_changed.notify_waiters(); Ok(()) } pub(super) fn remove(&self, token: u64) { - let mut state = self - .inner - .state - .lock() - .expect("ACP session registry poisoned"); - state.sessions.remove(&token); - state.v2_sessions.remove(&token); + let mut state = self.inner.lock_state(); + let session = state.sessions.remove(&token); + let v2_session = state.v2_sessions.remove(&token); + drop(state); + // Captured callback state can reenter the registry or panic on final drop. + drop(session); + drop(v2_session); } fn close_gate_and_snapshot( &self, permanently: bool, ) -> (Vec, Vec) { - let mut state = self - .inner - .state - .lock() - .expect("ACP session registry poisoned"); - state.accepting = false; - state.permanently_closed |= permanently; - state.generation = state.generation.wrapping_add(1); + let mut state = self.inner.lock_state(); + state.close_gate(permanently); ( state.sessions.values().cloned().collect(), state.v2_sessions.values().cloned().collect(), @@ -677,14 +710,7 @@ impl SessionRegistry { async fn wait_for_pending_attachments(&self) { loop { let changed = self.inner.attachments_changed.notified(); - if self - .inner - .state - .lock() - .expect("ACP session registry poisoned") - .pending_attachments - == 0 - { + if self.inner.lock_state().pending_attachments == 0 { return; } changed.await; @@ -729,14 +755,8 @@ impl SessionRegistry { Ok(lifecycle) => lifecycle, Err(_) => { if !reopen { - let mut state = self - .inner - .state - .lock() - .expect("ACP session registry poisoned"); - state.accepting = false; - state.permanently_closed = true; - state.generation = state.generation.wrapping_add(1); + let mut state = self.inner.lock_state(); + state.close_gate(true); } return (false, None); } @@ -752,7 +772,6 @@ impl SessionRegistry { let mut closing = JoinSet::new(); for mut session in sessions.iter().cloned() { - let registry = self.clone(); closing.spawn(async move { cancel_background_jobs(&session.tasks, &session.background_jobs).await; if let Some(commands) = session.commands.upgrade() { @@ -761,31 +780,33 @@ impl SessionRegistry { let _ = acknowledged.await; } } - if !*session.completed.borrow() { - let _ = session.completed.changed().await; - } - registry.remove(session.token); + session.completed.wait_for(|done| *done).await.is_ok() }); } for mut session in v2_sessions.iter().cloned() { - let registry = self.clone(); closing.spawn(async move { (session.close)().await; - if !*session.completed.borrow() { - let _ = session.completed.changed().await; - } - registry.remove(session.token); + session.completed.wait_for(|done| *done).await.is_ok() }); } - let timed_out = timeout(limit, async { - while closing.join_next().await.is_some() {} - self.wait_for_pending_attachments().await; - }) - .await - .is_err(); - let teardown_complete = if timed_out { + let graceful = matches!( + timeout(limit, async { + let mut completed = true; + while let Some(result) = closing.join_next().await { + completed &= matches!(result, Ok(true)); + } + self.wait_for_pending_attachments().await; + completed + }) + .await, + Ok(true) + ); + let teardown_complete = if !graceful { + // A panicking close callback is not successful teardown. Abort any + // unfinished actor and require its positive completion signal before + // resetting external credentials or reopening admission. closing.abort_all(); for session in &sessions { if !*session.completed.borrow() { @@ -798,21 +819,24 @@ impl SessionRegistry { } } while closing.join_next().await.is_some() {} - timeout(limit, async { - for mut session in sessions.iter().cloned() { - if !*session.completed.borrow() { - let _ = session.completed.changed().await; + matches!( + timeout(limit, async { + for mut session in sessions.iter().cloned() { + if session.completed.wait_for(|done| *done).await.is_err() { + return false; + } } - } - for mut session in v2_sessions.iter().cloned() { - if !*session.completed.borrow() { - let _ = session.completed.changed().await; + for mut session in v2_sessions.iter().cloned() { + if session.completed.wait_for(|done| *done).await.is_err() { + return false; + } } - } - self.wait_for_pending_attachments().await; - }) - .await - .is_ok() + self.wait_for_pending_attachments().await; + true + }) + .await, + Ok(true) + ) } else { true }; @@ -823,11 +847,7 @@ impl SessionRegistry { self.remove(session.token); } if !teardown_complete { - self.inner - .state - .lock() - .expect("ACP session registry poisoned") - .permanently_closed = true; + self.inner.lock_state().permanently_closed = true; } let output = if teardown_complete { Some(after_close().await) @@ -836,11 +856,7 @@ impl SessionRegistry { }; let mut reopened = false; if reopen && teardown_complete { - let mut state = self - .inner - .state - .lock() - .expect("ACP session registry poisoned"); + let mut state = self.inner.lock_state(); state.accepting = !state.permanently_closed; reopened = state.accepting; } @@ -945,6 +961,28 @@ struct PreparedFork { creation: crate::session::SessionObserver, } +impl PreparedFork { + /// Validation and cleanup ownership precede submission. Only private, + /// infallible publication and actor activation may follow a success response. + fn submit( + self, + respond: impl FnOnce(Result) -> Result<(), E>, + ) -> Result, E> { + let creation = match self.creation.prepare_creation() { + Ok(creation) => creation, + Err(error) => { + respond(Err(sdk_error(AcpRuntimeError::Loop(error))))?; + return Ok(None); + } + }; + let session_id = self.response.session_id.clone(); + respond(Ok(self.response))?; + creation.commit(); + let _ = self.activation.send(()); + Ok(Some(session_id)) + } +} + /// Owns an ACP integration binding for an in-flight request or live actor. /// Dropping either owner must release the durable identity. struct SessionBindingGuard { @@ -1047,6 +1085,35 @@ impl Drop for SessionActorGuard { } } +#[derive(Debug)] +enum LegacyPublicationError { + AdmissionClosed, + Commit(AcpRuntimeError), +} + +// Declared before the map guard: rollback runs after map exclusion is released, +// including on unwind. Consuming a fresh admission proves registry ownership. +struct SessionPublicationRollback<'a> { + registry: &'a SessionRegistry, + admission: &'a mut SessionAdmission, + token: u64, + actor: AbortHandle, + armed: bool, +} + +impl Drop for SessionPublicationRollback<'_> { + fn drop(&mut self) { + if self.armed { + self.actor.abort(); + // Registration consumes admission before notifying waiters. A + // rejection must not unregister a different actor with this token. + if !self.admission.active { + self.registry.remove(self.token); + } + } + } +} + struct Server { runtime: Arc, integration: Arc, @@ -1099,14 +1166,65 @@ impl Server { Ok(LogoutResponse::new()) } + fn publish_session( + &self, + admission: &mut SessionAdmission, + registered: RegisteredSession, + session: SessionHandle, + commit: impl FnOnce() -> Result, + ) -> Result { + if !admission.active || !Arc::ptr_eq(&admission.registry, &self.registry.inner) { + return Err(LegacyPublicationError::AdmissionClosed); + } + let mut rollback = SessionPublicationRollback { + registry: &self.registry, + admission, + token: registered.token, + actor: registered.actor.clone(), + armed: false, + }; + // The commit coordinates durable identity with this local map. An + // unwind leaves that outcome unknown: isolate the connection on poison, + // but still abort and unregister the actor after releasing this guard. + let mut sessions = self + .sessions + .lock() + .map_err(|_| LegacyPublicationError::Commit(AcpRuntimeError::ClientClosed))?; + if sessions.contains_key(®istered.session_id) { + return Err(LegacyPublicationError::Commit(AcpRuntimeError::Loop( + "ACP session is already published".into(), + ))); + } + sessions.try_reserve(1).map_err(|error| { + LegacyPublicationError::Commit(AcpRuntimeError::Loop(error.to_string())) + })?; + let session_id = registered.session_id.clone(); + rollback.armed = true; + self.registry + .register(rollback.admission, registered) + .map_err(|()| LegacyPublicationError::AdmissionClosed)?; + let committed = commit().map_err(LegacyPublicationError::Commit)?; + sessions.insert(session_id, session); + rollback.armed = false; + Ok(committed) + } + fn remove_session(&self, session_id: &agentkit_acp::SessionId, token: u64) { - let mut sessions = self.sessions.lock().expect("ACP session map poisoned"); - if sessions + // Cleanup must still unregister and signal completion after a commit + // unwind poisoned the map. Do not inspect uncertain publication state. + let Ok(mut sessions) = self.sessions.lock() else { + return; + }; + let removed = if sessions .get(session_id) .is_some_and(|session| session.token == token) { - sessions.remove(session_id); - } + sessions.remove(session_id) + } else { + None + }; + drop(sessions); + drop(removed); } async fn initialize(&self, request: InitializeRequest) -> InitializeResponse { @@ -1400,38 +1518,11 @@ impl Server { completed: completion, }; - // Hold the request-scoped map lock across registration, commit, and publication. - // Shutdown either closes the gate before this point or snapshots this actor. - let mut sessions = self.sessions.lock().expect("ACP session map poisoned"); - if self.registry.register(&mut admission, registered).is_err() { - drop(sessions); - drop(activation); - actor_task.abort(); - let _ = actor_task.await; - return Err(AcpRuntimeError::ClientClosed); - } - let pending_fork_creation = if claim.is_fork() { - Some(claim.defer_fork_commit()) - } else { - if let Err(error) = claim.commit() { - self.registry.remove(token); - drop(sessions); - drop(activation); - actor_task.abort(); - let _ = actor_task.await; - return Err(record_acp_runtime_failure( - &session_id, - "session_commit", - error, - )); - } - None - }; - crate::events::emit(&crate::events::RuntimeEvent::SessionStarted { - session_id: session_id.to_string(), - }); - sessions.insert( - session_id.clone(), + // Shutdown either closes the gate before registration or snapshots this + // actor. The helper rolls back after releasing the local map on failure. + let publication = self.publish_session( + &mut admission, + registered, SessionHandle { token, commands: tx, @@ -1439,8 +1530,32 @@ impl Server { structured_completion, tasks, }, + || { + if claim.is_fork() { + Ok(Some(claim.defer_fork_commit())) + } else { + claim.commit()?; + Ok(None) + } + }, ); - drop(sessions); + let pending_fork_creation = match publication { + Ok(creation) => creation, + Err(error) => { + drop(activation); + actor_task.abort(); + let _ = actor_task.await; + return Err(match error { + LegacyPublicationError::AdmissionClosed => AcpRuntimeError::ClientClosed, + LegacyPublicationError::Commit(error) => { + record_acp_runtime_failure(&session_id, "session_commit", error) + } + }); + } + }; + crate::events::emit(&crate::events::RuntimeEvent::SessionStarted { + session_id: session_id.to_string(), + }); drop(actor_task); Ok(AttachedSession { session_id, @@ -1481,7 +1596,7 @@ impl Server { let (sender, background_jobs, tasks, structured_completion) = self .sessions .lock() - .expect("ACP session map poisoned") + .map_err(|_| AcpRuntimeError::ClientClosed)? .get(¬ification.session_id) .map(|session| { ( @@ -1512,7 +1627,7 @@ impl Server { let session = self .sessions .lock() - .expect("ACP session map poisoned") + .map_err(|_| AcpRuntimeError::ClientClosed)? .remove(&request.session_id) .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; cancel_background_jobs(&session.tasks, &session.background_jobs).await; @@ -1534,7 +1649,7 @@ impl Server { ) -> Result, AcpRuntimeError> { self.sessions .lock() - .expect("ACP session map poisoned") + .map_err(|_| AcpRuntimeError::ClientClosed)? .get(session_id) .map(|session| session.commands.clone()) .ok_or_else(|| AcpRuntimeError::SessionNotFound(session_id.to_string())) @@ -1547,7 +1662,7 @@ impl Server { let (background_jobs, tasks) = self .sessions .lock() - .expect("ACP session map poisoned") + .map_err(|_| AcpRuntimeError::ClientClosed)? .get(&request.session_id) .map(|session| (session.background_jobs.clone(), session.tasks.clone())) .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; @@ -1563,7 +1678,7 @@ impl Server { let background_jobs = self .sessions .lock() - .expect("ACP session map poisoned") + .map_err(|_| AcpRuntimeError::ClientClosed)? .get(&request.session_id) .map(|session| session.background_jobs.clone()) .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; @@ -2414,11 +2529,13 @@ fn component( cx.spawn(async move { match state.fork_session(request, connection.clone()).await { Ok(prepared) => { - let session_id = prepared.response.session_id.clone(); - responder.respond(prepared.response)?; - prepared.creation.commit_creation(); - let _ = prepared.activation.send(()); - connection.send_notification(available_commands_update(session_id)) + if let Some(session_id) = prepared + .submit(|response| responder.respond_with_result(response))? + { + connection + .send_notification(available_commands_update(session_id))?; + } + Ok(()) } Err(error) => responder.respond_with_result(Err(sdk_error(error))), } @@ -2916,6 +3033,338 @@ pub(super) mod tests { )) } + fn legacy_pending_publication( + server: &Arc, + id: &str, + ) -> ( + RegisteredSession, + SessionHandle, + tokio::task::JoinHandle<()>, + mpsc::Receiver, + ) { + let session_id = agentkit_acp::SessionId::new(id.to_owned()); + let token = server.registry.next_token(); + let (completed, completion) = watch::channel(false); + let guard = SessionActorGuard { + server: Arc::downgrade(server), + registry: server.registry.clone(), + session_id: session_id.clone(), + token, + completed, + }; + let actor = tokio::spawn(async move { + let _guard = guard; + std::future::pending::<()>().await; + }); + let (commands, received) = mpsc::channel(1); + let session = SessionHandle { + token, + commands, + background_jobs: BackgroundJobs::default(), + structured_completion: false, + tasks: AsyncTaskManager::new().handle(), + }; + let registered = RegisteredSession { + token, + session_id, + integration: Arc::clone(&server.integration), + background_jobs: session.background_jobs.clone(), + tasks: session.tasks.clone(), + commands: session.commands.downgrade(), + actor: actor.abort_handle(), + completed: completion, + }; + (registered, session, actor, received) + } + + #[tokio::test] + async fn legacy_publication_commit_error_rolls_back_and_unwind_isolates_map() { + for unwind in [false, true] { + let root = tempfile::tempdir().unwrap(); + let registry = SessionRegistry::new(); + let server = logout_test_server( + Runtime::new(root.path(), "gpt-5.4").unwrap(), + registry.clone(), + ); + let mut admission = registry.begin_attachment().unwrap(); + let (registered, session, actor, _received) = + legacy_pending_publication(&server, "failed"); + let completed = registered.completed.clone(); + let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + server.publish_session(&mut admission, registered, session, || { + if unwind { + panic!("durable commit unwound"); + } + Err::<(), _>(AcpRuntimeError::ClientClosed) + }) + })); + if unwind { + assert!(outcome.is_err()); + } else { + assert!(outcome.unwrap().is_err()); + } + assert!(registry.inner.lock_state().sessions.is_empty()); + assert_eq!(server.sessions.is_poisoned(), unwind); + let id = agentkit_acp::SessionId::new("failed"); + let error = server.sender(&id).await.unwrap_err(); + if unwind { + assert!(matches!(error, AcpRuntimeError::ClientClosed)); + assert!(matches!( + server.cancel(CancelNotification::new(id.clone())).await, + Err(AcpRuntimeError::ClientClosed) + )); + assert!(matches!( + server.close(CloseSessionRequest::new(id.clone())).await, + Err(AcpRuntimeError::ClientClosed) + )); + assert!(matches!( + server + .detach_compose(DetachComposeRequest { + session_id: id.clone(), + call_id: "call".into(), + }) + .await, + Err(AcpRuntimeError::ClientClosed) + )); + assert!(matches!( + server + .cancel_background(CancelBackgroundRequest { + session_id: id, + call_id: "call".into(), + }) + .await, + Err(AcpRuntimeError::ClientClosed) + )); + } else { + assert!(matches!(error, AcpRuntimeError::SessionNotFound(_))); + } + // The real actor guard runs against the poisoned map without a + // second panic, then unregisters and publishes completion. + assert!( + timeout(Duration::from_secs(1), actor) + .await + .unwrap() + .unwrap_err() + .is_cancelled() + ); + assert!(*completed.borrow()); + assert_eq!(registry.inner.lock_state().pending_attachments, 0); + } + } + + #[tokio::test] + async fn legacy_publication_closed_gate_rejects_without_commit() { + let root = tempfile::tempdir().unwrap(); + let registry = SessionRegistry::new(); + let server = logout_test_server( + Runtime::new(root.path(), "gpt-5.4").unwrap(), + registry.clone(), + ); + let mut admission = registry.begin_attachment().unwrap(); + registry.close_gate_and_snapshot(false); + let (registered, session, actor, _received) = legacy_pending_publication(&server, "closed"); + assert!(matches!( + server.publish_session(&mut admission, registered, session, || -> Result<(), _> { + panic!("rejected commit ran") + }), + Err(LegacyPublicationError::AdmissionClosed) + )); + assert!(server.sessions.lock().unwrap().is_empty()); + assert!(registry.inner.lock_state().sessions.is_empty()); + assert!(admission.active); + assert!( + timeout(Duration::from_secs(1), actor) + .await + .unwrap() + .unwrap_err() + .is_cancelled() + ); + drop(admission); + assert_eq!(registry.inner.lock_state().pending_attachments, 0); + } + + #[tokio::test] + async fn legacy_publication_rejection_preserves_incumbent_registration() { + let root = tempfile::tempdir().unwrap(); + let registry = SessionRegistry::new(); + let server = logout_test_server( + Runtime::new(root.path(), "gpt-5.4").unwrap(), + registry.clone(), + ); + let mut consumed = registry.begin_attachment().unwrap(); + let (registered, session, incumbent, _received) = + legacy_pending_publication(&server, "original"); + let token = registered.token; + server + .publish_session(&mut consumed, registered, session, || Ok(())) + .unwrap(); + for rejection in ["local-id", "token", "consumed", "foreign"] { + let foreign = SessionRegistry::new(); + let mut fresh = if rejection == "foreign" { + foreign.begin_attachment().unwrap() + } else { + registry.begin_attachment().unwrap() + }; + let id = if rejection == "local-id" { + "original" + } else { + "duplicate" + }; + let (mut registered, mut session, actor, _received) = + legacy_pending_publication(&server, id); + if rejection != "local-id" { + registered.token = token; + session.token = token; + } + let admission = if rejection == "consumed" { + &mut consumed + } else { + &mut fresh + }; + assert!( + server + .publish_session(admission, registered, session, || -> Result<(), _> { + panic!("rejected commit ran") + }) + .is_err() + ); + assert_eq!( + registry + .inner + .lock_state() + .sessions + .get(&token) + .unwrap() + .actor + .id(), + incumbent.id() + ); + assert!(!incumbent.is_finished()); + // Pre-registration rejection leaves abort/join to the attachment + // caller. This actor's guard retains its separately allocated token. + actor.abort(); + assert!( + timeout(Duration::from_secs(1), actor) + .await + .unwrap() + .unwrap_err() + .is_cancelled() + ); + assert!(registry.inner.lock_state().sessions.contains_key(&token)); + assert_eq!( + server + .sessions + .lock() + .unwrap() + .get(&agentkit_acp::SessionId::new("original")) + .unwrap() + .token, + token + ); + } + incumbent.abort(); + assert!( + timeout(Duration::from_secs(1), incumbent) + .await + .unwrap() + .unwrap_err() + .is_cancelled() + ); + assert!(server.sessions.lock().unwrap().is_empty()); + assert!(registry.inner.lock_state().sessions.is_empty()); + assert_eq!(registry.inner.lock_state().pending_attachments, 0); + } + + #[tokio::test] + async fn legacy_publication_actor_cleanup_drops_mailbox_outside_map_lock() { + struct CheckUnlocked(Weak, AtomicBool); + impl std::task::Wake for CheckUnlocked { + fn wake(self: Arc) { + let server = self.0.upgrade().unwrap(); + assert!(server.sessions.try_lock().is_ok()); + self.1.store(true, Ordering::SeqCst); + } + } + let root = tempfile::tempdir().unwrap(); + let registry = SessionRegistry::new(); + let server = logout_test_server( + Runtime::new(root.path(), "gpt-5.4").unwrap(), + registry.clone(), + ); + let mut admission = registry.begin_attachment().unwrap(); + let (registered, session, actor, mut received) = + legacy_pending_publication(&server, "mailbox-drop"); + server + .publish_session(&mut admission, registered, session, || Ok(())) + .unwrap(); + let wake = Arc::new(CheckUnlocked( + Arc::downgrade(&server), + AtomicBool::new(false), + )); + let waker = std::task::Waker::from(Arc::clone(&wake)); + let mut context = std::task::Context::from_waker(&waker); + assert!(received.poll_recv(&mut context).is_pending()); + actor.abort(); + assert!( + timeout(Duration::from_secs(1), actor) + .await + .unwrap() + .unwrap_err() + .is_cancelled() + ); + assert!(wake.1.load(Ordering::SeqCst)); + assert!(server.sessions.lock().unwrap().is_empty()); + assert!(registry.inner.lock_state().sessions.is_empty()); + assert!(received.recv().await.is_none()); + } + + #[tokio::test] + async fn legacy_publication_wakeup_unwind_rolls_back_consumed_admission() { + struct PanicOnWake(SessionRegistry); + impl std::task::Wake for PanicOnWake { + fn wake(self: Arc) { + assert!(self.0.inner.state.try_lock().is_ok()); + panic!("registration waiter unwound"); + } + } + let root = tempfile::tempdir().unwrap(); + let registry = SessionRegistry::new(); + let server = logout_test_server( + Runtime::new(root.path(), "gpt-5.4").unwrap(), + registry.clone(), + ); + let mut admission = registry.begin_attachment().unwrap(); + let waker = std::task::Waker::from(Arc::new(PanicOnWake(registry.clone()))); + let mut context = std::task::Context::from_waker(&waker); + let mut waiting = Box::pin(registry.wait_for_pending_attachments()); + assert!(waiting.as_mut().poll(&mut context).is_pending()); + let (registered, session, actor, _received) = legacy_pending_publication(&server, "wake"); + let completed = registered.completed.clone(); + let committed = AtomicBool::new(false); + assert!( + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + server.publish_session(&mut admission, registered, session, || { + committed.store(true, Ordering::SeqCst); + Ok(()) + }) + })) + .is_err() + ); + assert!(!committed.load(Ordering::SeqCst)); + assert!(!admission.active); + assert!(registry.inner.lock_state().sessions.is_empty()); + assert_eq!(registry.inner.lock_state().pending_attachments, 0); + assert!(server.sessions.is_poisoned()); + assert!( + timeout(Duration::from_secs(1), actor) + .await + .unwrap() + .unwrap_err() + .is_cancelled() + ); + assert!(*completed.borrow()); + } + fn register_close_tracking_session( server: &Arc, registry: &SessionRegistry, @@ -3064,6 +3513,336 @@ pub(super) mod tests { late_actor.await.unwrap_err(); } + fn register_test_v2( + registry: &SessionRegistry, + admission: &mut SessionAdmission, + token: u64, + interrupt: Arc, + ) -> Result<(), ()> { + let actor = tokio::spawn(std::future::pending::<()>()); + actor.abort(); + let (_completed, completion) = watch::channel(true); + registry.register_v2( + admission, + token, + interrupt, + Arc::new(|| Box::pin(async {})), + actor.abort_handle(), + completion, + ) + } + + #[tokio::test] + async fn registry_rejects_foreign_consumed_duplicate_and_stale_admissions() { + let registry = SessionRegistry::new(); + let other = SessionRegistry::new(); + let mut admission = registry.begin_attachment().unwrap(); + let token = registry.next_token(); + assert!(register_test_v2(&other, &mut admission, token, Arc::new(|| {})).is_err()); + assert_eq!(registry.inner.lock_state().pending_attachments, 1); + assert_eq!(other.inner.lock_state().pending_attachments, 0); + assert!(other.inner.lock_state().v2_sessions.is_empty()); + register_test_v2(®istry, &mut admission, token, Arc::new(|| {})).unwrap(); + let second_token = registry.next_token(); + assert!( + register_test_v2(®istry, &mut admission, second_token, Arc::new(|| {})).is_err() + ); + let mut duplicate = registry.begin_attachment().unwrap(); + assert!(register_test_v2(®istry, &mut duplicate, token, Arc::new(|| {})).is_err()); + assert!(duplicate.active); + assert_eq!(registry.inner.lock_state().pending_attachments, 1); + assert_eq!(registry.inner.lock_state().v2_sessions.len(), 1); + registry.remove(token); + registry.close_gate_and_snapshot(false); + // Reopening never makes an old pending admission current again. + registry.inner.lock_state().accepting = true; + assert!( + register_test_v2(®istry, &mut duplicate, second_token, Arc::new(|| {})).is_err() + ); + drop(duplicate); + assert_eq!(registry.inner.lock_state().pending_attachments, 0); + assert!(registry.begin_attachment().is_ok()); + } + + #[tokio::test] + async fn registry_legacy_registration_obeys_admission_and_cross_version_token_rules() { + let registry = SessionRegistry::new(); + let other = SessionRegistry::new(); + let token = registry.next_token(); + register_test_v2( + ®istry, + &mut registry.begin_attachment().unwrap(), + token, + Arc::new(|| {}), + ) + .unwrap(); + let actor = tokio::spawn(std::future::pending::<()>()); + actor.abort(); + let (commands, _received) = mpsc::channel(1); + let (_completed, completion) = watch::channel(true); + let mut session = RegisteredSession { + token, + session_id: agentkit_acp::SessionId::new("registry-test"), + integration: shutdown_test_integration(), + background_jobs: BackgroundJobs::default(), + tasks: AsyncTaskManager::new().handle(), + commands: commands.downgrade(), + actor: actor.abort_handle(), + completed: completion, + }; + let mut admission = registry.begin_attachment().unwrap(); + assert!(registry.register(&mut admission, session.clone()).is_err()); + assert!(other.register(&mut admission, session.clone()).is_err()); + assert!(admission.active); + assert_eq!(registry.inner.lock_state().pending_attachments, 1); + session.token = registry.next_token(); + registry.register(&mut admission, session.clone()).unwrap(); + assert!(registry.register(&mut admission, session.clone()).is_err()); + assert_eq!(registry.inner.lock_state().pending_attachments, 0); + assert_eq!(registry.inner.lock_state().sessions.len(), 1); + assert!(other.inner.lock_state().sessions.is_empty()); + registry.remove(token); + registry.remove(session.token); + } + + #[tokio::test] + async fn registry_attachment_wakeup_runs_after_unlock() { + struct ReenterOnWake(SessionRegistry, AtomicBool); + impl std::task::Wake for ReenterOnWake { + fn wake(self: Arc) { + assert!(self.0.inner.state.try_lock().is_ok()); + assert!(self.0.begin_attachment().is_ok()); + self.1.store(true, Ordering::SeqCst); + } + } + let registry = SessionRegistry::new(); + let mut admission = registry.begin_attachment().unwrap(); + let wake = Arc::new(ReenterOnWake(registry.clone(), AtomicBool::new(false))); + let waker = std::task::Waker::from(Arc::clone(&wake)); + let mut context = std::task::Context::from_waker(&waker); + let mut waiting = Box::pin(registry.wait_for_pending_attachments()); + assert!(waiting.as_mut().poll(&mut context).is_pending()); + let token = registry.next_token(); + register_test_v2(®istry, &mut admission, token, Arc::new(|| {})).unwrap(); + assert!(wake.1.load(Ordering::SeqCst)); + assert!(waiting.as_mut().poll(&mut context).is_ready()); + assert_eq!(registry.inner.lock_state().pending_attachments, 0); + registry.remove(token); + } + + #[tokio::test] + async fn registry_interrupt_unwind_leaves_gate_closed_without_poison() { + let registry = SessionRegistry::new(); + let token = registry.next_token(); + let callback_registry = registry.clone(); + register_test_v2( + ®istry, + &mut registry.begin_attachment().unwrap(), + token, + Arc::new(move || { + assert!(callback_registry.inner.state.try_lock().is_ok()); + panic!("interrupt callback"); + }), + ) + .unwrap(); + let resetting = registry.clone(); + let reset = tokio::spawn(async move { resetting.reset_authentication().await }); + assert!(reset.await.unwrap_err().is_panic()); + assert!(registry.begin_attachment().is_err()); + assert!(!registry.inner.state.is_poisoned()); + assert!(registry.inner.lock_state().v2_sessions.contains_key(&token)); + registry.remove(token); + assert!(registry.reset_authentication().await); + } + + #[tokio::test] + async fn registry_callback_destructors_run_after_unlock_on_remove_and_rejection() { + struct ReenterOnDrop(SessionRegistry); + impl Drop for ReenterOnDrop { + fn drop(&mut self) { + assert!(self.0.inner.state.try_lock().is_ok()); + assert!(self.0.begin_attachment().is_ok()); + panic!("callback capture destructor"); + } + } + let registry = SessionRegistry::new(); + let token = registry.next_token(); + let capture = ReenterOnDrop(registry.clone()); + register_test_v2( + ®istry, + &mut registry.begin_attachment().unwrap(), + token, + Arc::new(move || { + let _ = &capture; + }), + ) + .unwrap(); + assert!( + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| registry.remove(token))) + .is_err() + ); + assert!(!registry.inner.state.is_poisoned()); + assert!(registry.inner.lock_state().v2_sessions.is_empty()); + let other = SessionRegistry::new(); + let capture = ReenterOnDrop(registry.clone()); + assert!( + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _ = register_test_v2( + ®istry, + &mut other.begin_attachment().unwrap(), + token, + Arc::new(move || { + let _ = &capture; + }), + ); + })) + .is_err() + ); + assert!(!registry.inner.state.is_poisoned()); + assert_eq!(other.inner.lock_state().pending_attachments, 0); + } + + #[tokio::test] + async fn registry_poison_does_not_break_unwind_cleanup_or_reopen_closed_gate() { + struct RemoveOnDrop(SessionRegistry, u64); + impl Drop for RemoveOnDrop { + fn drop(&mut self) { + self.0.remove(self.1); + } + } + let registry = SessionRegistry::new(); + let token = registry.next_token(); + register_test_v2( + ®istry, + &mut registry.begin_attachment().unwrap(), + token, + Arc::new(|| {}), + ) + .unwrap(); + let admission = registry.begin_attachment().unwrap(); + registry.close_gate_and_snapshot(true); + let cleanup = RemoveOnDrop(registry.clone(), token); + // Poison injection represents an unwind with the audited state intact. + assert!( + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _admission = admission; + let _cleanup = cleanup; + let _state = registry.inner.state.lock().unwrap(); + panic!("unwind while registry guard is held"); + })) + .is_err() + ); + assert!(!registry.inner.state.is_poisoned()); + assert!(registry.inner.lock_state().v2_sessions.is_empty()); + assert_eq!(registry.inner.lock_state().pending_attachments, 0); + assert!(registry.begin_attachment().is_err()); + assert!(!registry.reset_authentication().await); + } + + #[tokio::test] + async fn registry_counter_exhaustion_does_not_wrap_or_poison() { + let registry = SessionRegistry::new(); + registry + .inner + .next_token + .store(u64::MAX - 1, Ordering::Relaxed); + assert_eq!(registry.next_token(), u64::MAX - 1); + for _ in 0..2 { + assert!( + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| registry.next_token())) + .is_err() + ); + assert_eq!(registry.inner.next_token.load(Ordering::Relaxed), u64::MAX); + } + registry.inner.lock_state().pending_attachments = usize::MAX; + assert!(registry.begin_attachment().is_err()); + assert!(!registry.inner.state.is_poisoned()); + assert_eq!(registry.inner.lock_state().pending_attachments, usize::MAX); + registry.inner.lock_state().pending_attachments = 0; + registry.inner.lock_state().generation = u64::MAX; + assert!(!registry.reset_authentication().await); + assert_eq!(registry.inner.lock_state().generation, u64::MAX); + assert!(registry.begin_attachment().is_err()); + } + + #[tokio::test] + async fn registry_cancelled_reset_keeps_admission_closed() { + let registry = SessionRegistry::new(); + let started = Arc::new(Notify::new()); + let resetting = registry.clone(); + let action_started = Arc::clone(&started); + let reset = tokio::spawn(async move { + resetting + .reset_authentication_with(async move { + action_started.notify_one(); + std::future::pending::<()>().await; + }) + .await + }); + started.notified().await; + let generation = registry.inner.lock_state().generation; + reset.abort(); + assert!(reset.await.unwrap_err().is_cancelled()); + assert!(registry.begin_attachment().is_err()); + assert_eq!(registry.inner.lock_state().generation, generation); + // A later completed reset can reopen; cancellation itself cannot. + assert!(registry.reset_authentication().await); + assert!(registry.begin_attachment().is_ok()); + assert!(registry.inner.lock_state().generation > generation); + } + + #[tokio::test] + async fn registry_close_panic_requires_positive_actor_completion_before_reset() { + struct Completion { + alive: Arc, + completed: watch::Sender, + acknowledge: bool, + } + impl Drop for Completion { + fn drop(&mut self) { + self.alive.store(false, Ordering::SeqCst); + if self.acknowledge { + self.completed.send_replace(true); + } + } + } + for acknowledge in [true, false] { + let registry = SessionRegistry::new(); + let alive = Arc::new(AtomicBool::new(true)); + let (completed, completion) = watch::channel(false); + let guard = Completion { + alive: Arc::clone(&alive), + completed, + acknowledge, + }; + let actor = tokio::spawn(async move { + let _guard = guard; + std::future::pending::<()>().await; + }); + registry + .register_v2( + &mut registry.begin_attachment().unwrap(), + registry.next_token(), + Arc::new(|| {}), + Arc::new(|| Box::pin(async { panic!("close callback unwound") })), + actor.abort_handle(), + completion, + ) + .unwrap(); + let reset_ran = AtomicBool::new(false); + let (reopened, _) = registry + .close_sessions_with_timeout(Duration::from_millis(100), true, || async { + assert!(!alive.load(Ordering::SeqCst)); + reset_ran.store(true, Ordering::SeqCst); + }) + .await; + assert_eq!(reopened, acknowledge); + assert_eq!(reset_ran.load(Ordering::SeqCst), acknowledge); + assert_eq!(registry.begin_attachment().is_ok(), acknowledge); + assert!(actor.await.unwrap_err().is_cancelled()); + } + } + #[tokio::test] async fn authentication_reset_closes_shared_v2_sessions_and_reopens_registration() { let registry = SessionRegistry::new(); @@ -4966,6 +5745,78 @@ pub(super) mod tests { ); } + #[test] + fn fork_submission_keeps_cleanup_until_response_succeeds() { + for delivery in ["success", "failure", "unwind", "rejected"] { + let root = tempfile::tempdir().unwrap(); + let id = crate::session::new_id(); + let opened = crate::session::open_uncommitted( + root.path(), + &id, + false, + vec![Item::text(ItemKind::System, "system")], + ) + .unwrap(); + // A creation already owned by a prepared publication is rejected by + // normal APIs, just as poison/write fencing is rejected in session tests. + let held = if delivery == "rejected" { + Some(opened.observer.prepare_creation().unwrap()) + } else { + None + }; + let (activation, mut activated) = oneshot::channel(); + let prepared = PreparedFork { + response: ForkSessionResponse::new(id.clone()), + activation, + creation: opened.observer.clone(), + }; + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + prepared.submit(|response| { + if delivery == "rejected" { + assert!(response.is_err(), "invalid creation cannot report success"); + return Ok(()); + } + assert!(response.is_ok()); + // Response callbacks cannot invalidate the prepared writer. + assert!( + opened + .observer + .replace(&[Item::text(ItemKind::User, "reentry")]) + .is_err() + ); + assert!(crate::session::load(root.path(), &id).is_ok()); + match delivery { + "failure" => Err("response closed"), + "unwind" => panic!("response interrupted"), + _ => Ok(()), + } + }) + })); + match delivery { + "success" => { + assert!(result.unwrap().unwrap().is_some()); + assert!(activated.try_recv().is_ok()); + } + "failure" => assert_eq!(result.unwrap().unwrap_err(), "response closed"), + "unwind" => assert!(result.is_err()), + "rejected" => assert!(result.unwrap().unwrap().is_none()), + _ => unreachable!(), + } + if delivery != "success" { + assert!(matches!( + activated.try_recv(), + Err(oneshot::error::TryRecvError::Closed) + )); + } + drop(held); + drop(opened); + assert_eq!( + crate::session::load(root.path(), &id).is_ok(), + delivery == "success" + ); + } + } + #[tokio::test] async fn kit_server_advertises_supported_session_discovery_restoration_and_forking() { let root = tempfile::tempdir().unwrap(); diff --git a/src/protocols/acp/activity.rs b/src/protocols/acp/activity.rs index e7672a3e..3ca8c08c 100644 --- a/src/protocols/acp/activity.rs +++ b/src/protocols/acp/activity.rs @@ -17,6 +17,7 @@ enum State { Idle, Running, Settling, + Unavailable, } /// An ordered snapshot, not another reducer. Projections select only wire format. @@ -35,15 +36,18 @@ struct Activity { next_id: u64, origin: ExecutionOrigin, current: Option, + projecting: bool, + executing: bool, } /// Session-owned lifecycle instrument shared by the actor and its observer. /// Admission alone is silent; TurnStarted allocates the activity identity. All /// logical turns drained by an execution share that identity until settlement. -/// Projection runs synchronously under the state lock: Running precedes content, -/// and the actor must flush final content/diagnostics before settling to Idle. -/// Projections only enqueue wire notifications; they must not reenter this -/// instrument while its ordered transition is being projected. +/// Projection runs synchronously outside the state lock. An in-flight claim +/// prevents another projection from overtaking it, including callback reentry. +/// The actor must flush final content/diagnostics before settling to Idle. +/// Abandoned executions/projections isolate this owner: external delivery cannot +/// be rolled back or safely retried after unwind or cancellation. #[derive(Clone)] pub(super) struct SessionActivity { state: Arc>, @@ -60,33 +64,47 @@ impl SessionActivity { } } + #[cfg(test)] pub(super) fn begin(&self, origin: ExecutionOrigin) { - let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner()); + let Ok(mut state) = self.state.lock() else { + return; // A poisoned owner is never reused. + }; // Continuations cannot redefine the origin of an existing interval. - if state.state == State::Idle { + if state.state == State::Idle && !state.projecting && !state.executing { state.origin = origin; } } pub(super) fn observe(&self, event: &AgentEvent) { - let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner()); - match event { + let Ok(mut state) = self.state.lock() else { + return; + }; + if state.state == State::Unavailable || state.projecting { + return; + } + let transition = match event { AgentEvent::TurnStarted { .. } => { let was_idle = state.state == State::Idle; - state.state = State::Running; if was_idle { - state.next_id = state.next_id.wrapping_add(1); + let Some(id) = state.next_id.checked_add(1) else { + state.state = State::Unavailable; + return; + }; let transition = Transition { - id: state.next_id, + id, origin: state.origin, active: true, reason: FinishReason::Completed, error: None, }; state.current = Some(transition.clone()); - if let Err(error) = (self.project)(transition) { - tracing::debug!(%error, "failed to project session activity"); - } + state.next_id = id; + state.state = State::Running; + state.projecting = true; + Some(transition) + } else { + state.state = State::Running; + None } } AgentEvent::TurnFinished(result) if state.state != State::Idle => { @@ -94,20 +112,48 @@ impl SessionActivity { if let Some(current) = &mut state.current { current.reason = result.finish_reason.clone(); } + None } - _ => {} + _ => None, + }; + drop(state); + if let Some(transition) = transition + && let Err(error) = self.project_transition(transition) + { + tracing::debug!(%error, "failed to project session activity"); } } + fn project_transition(&self, transition: Transition) -> Result<(), AcpRuntimeError> { + // The claim is already installed. Drop isolates the owner if the client + // callback unwinds; it never calls client code during unwind. + let mut claim = ActivityClaim { + activity: self, + projection: true, + completed: false, + }; + let result = (self.project)(transition); + claim.completed = true; + result + } + pub(super) fn settle( &self, reason: Option, error: Option, ) -> Result<(), AcpRuntimeError> { - let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner()); + let mut state = self + .state + .lock() + .map_err(|_| AcpRuntimeError::Loop("session activity is poisoned".into()))?; + if state.state == State::Unavailable || state.projecting { + return Err(AcpRuntimeError::Loop( + "session activity is unavailable".into(), + )); + } state.origin = ExecutionOrigin::Prompt; state.state = State::Idle; - if let Some(mut terminal) = state.current.take() { + let terminal = state.current.take().map(|mut terminal| { terminal.active = false; if let Some(reason) = reason { terminal.reason = reason; @@ -116,7 +162,14 @@ impl SessionActivity { terminal.reason = FinishReason::Error; } terminal.error = error; - (self.project)(terminal)?; + terminal + }); + state.projecting = terminal.is_some(); + drop(state); + if let Some(terminal) = terminal { + // A returned transport error is an at-most-once projection attempt, + // not a reason to replay an already consumed terminal transition. + self.project_transition(terminal)?; } Ok(()) } @@ -129,10 +182,30 @@ impl SessionActivity { operation: impl std::future::Future>, reason: impl FnOnce(&T) -> Option, ) -> Result { - self.begin(origin); + { + let mut state = self + .state + .lock() + .map_err(|_| AcpRuntimeError::Loop("session activity is poisoned".into()))?; + if state.state == State::Unavailable || state.projecting || state.executing { + return Err(AcpRuntimeError::Loop( + "session activity is unavailable".into(), + )); + } + if state.state == State::Idle { + state.origin = origin; + } + state.executing = true; + } + let mut claim = ActivityClaim { + activity: self, + projection: false, + completed: false, + }; let result = operation.await; let terminal = result.as_ref().ok().and_then(reason); let settled = self.settle(terminal, result.as_ref().err().map(ToString::to_string)); + claim.completed = true; match result { Err(error) => Err(error), Ok(value) => settled.map(|()| value), @@ -140,6 +213,31 @@ impl SessionActivity { } } +/// No callbacks, awaits, or external effects occur while releasing a claim. +/// Poison is left isolated rather than recovering potentially incomplete state. +struct ActivityClaim<'a> { + activity: &'a SessionActivity, + projection: bool, + completed: bool, +} + +impl Drop for ActivityClaim<'_> { + fn drop(&mut self) { + if let Ok(mut state) = self.activity.state.lock() { + if self.projection { + state.projecting = false; + } else { + state.executing = false; + } + if !self.completed || (!self.projection && state.state != State::Idle) { + // A failed settlement must not hand an unterminated interval + // to another execution, even if a concurrent projection won. + state.state = State::Unavailable; + } + } + } +} + impl LoopObserver for SessionActivity { fn handle_event(&self, event: ObservedEvent) { self.observe(&event.event); @@ -365,4 +463,155 @@ mod tests { activity.settle(None, None).unwrap(); assert_eq!(*calls.lock().unwrap(), 2); } + + #[test] + fn projection_unwind_isolates_without_poison_or_replay() { + use std::sync::atomic::{AtomicUsize, Ordering}; + for fail_active in [true, false] { + let calls = Arc::new(AtomicUsize::new(0)); + let output = calls.clone(); + let activity = SessionActivity::new(move |transition| { + output.fetch_add(1, Ordering::Relaxed); + assert_ne!(transition.active, fail_active, "projection failed"); + Ok(()) + }); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + activity.observe(&started("first")); + activity.settle(None, None).unwrap(); + })); + assert!(result.is_err()); + assert!(!activity.state.is_poisoned()); + activity.begin(ExecutionOrigin::Autonomous); + activity.observe(&started("later")); + assert!(matches!( + activity.settle(None, None), + Err(AcpRuntimeError::Loop(_)) + )); + assert_eq!( + calls.load(Ordering::Relaxed), + if fail_active { 1 } else { 2 } + ); + } + } + + #[test] + fn projection_reentry_is_rejected_without_deadlock_or_reordering() { + let owner = Arc::new(std::sync::OnceLock::::new()); + let callback_owner = Arc::downgrade(&owner); + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let output = calls.clone(); + let activity = SessionActivity::new(move |_| { + let owner = callback_owner.upgrade().unwrap(); + let activity = owner.get().unwrap(); + assert!(activity.state.try_lock().is_ok()); + assert!(matches!( + activity.settle(None, None), + Err(AcpRuntimeError::Loop(_)) + )); + activity.observe(&started("reentrant")); + output.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + Ok(()) + }); + assert!(owner.set(activity.clone()).is_ok()); + activity.observe(&started("first")); + activity.settle(None, None).unwrap(); + assert_eq!(calls.load(std::sync::atomic::Ordering::Relaxed), 2); + assert_eq!(activity.state.lock().unwrap().next_id, 1); + } + + #[tokio::test] + async fn dropped_execution_isolates_before_another_operation_can_run() { + let transitions = Arc::new(Mutex::new(Vec::new())); + let output = transitions.clone(); + let activity = SessionActivity::new(move |transition| { + output.lock().unwrap().push(transition); + Ok(()) + }); + let mut execution = Box::pin(activity.execute( + ExecutionOrigin::Autonomous, + async { + activity.observe(&started("first")); + std::future::pending::>().await + }, + |_| None, + )); + assert!(futures_util::poll!(execution.as_mut()).is_pending()); + // A concurrent execution must not steal the live interval either. + assert!( + activity + .execute( + ExecutionOrigin::Prompt, + async { panic!("must not run") }, + |_: &()| None + ) + .await + .is_err() + ); + drop(execution); + activity.observe(&started("stale")); + assert!(matches!( + activity + .execute( + ExecutionOrigin::Prompt, + async { panic!("must not run") }, + |_: &()| None + ) + .await, + Err(AcpRuntimeError::Loop(_)) + )); + assert_eq!(transitions.lock().unwrap().len(), 1); + assert!(!activity.state.is_poisoned()); + } + + #[test] + fn exhausted_identity_and_poison_never_resume_projection() { + let activity = SessionActivity::new(|_| panic!("must not project")); + activity.state.lock().unwrap().next_id = u64::MAX; + activity.observe(&started("overflow")); + assert!(activity.settle(None, None).is_err()); + let poisoned = activity.clone(); + assert!( + std::thread::spawn(move || { + let _state = poisoned.state.lock().unwrap(); + panic!("state mutation interrupted"); + }) + .join() + .is_err() + ); + activity.begin(ExecutionOrigin::Prompt); + activity.observe(&started("after-poison")); + assert!(matches!( + activity.settle(None, None), + Err(AcpRuntimeError::Loop(_)) + )); + } + + #[tokio::test] + async fn outcome_callback_unwind_isolates_the_execution_owner() { + use futures_util::FutureExt; + let activity = SessionActivity::new(|_| Ok(())); + let result = std::panic::AssertUnwindSafe(activity.execute( + ExecutionOrigin::Prompt, + async { + activity.observe(&started("first")); + Ok(()) + }, + |_| panic!("outcome callback failed"), + )) + .catch_unwind() + .await; + assert!(result.is_err()); + assert!(!activity.state.is_poisoned()); + assert!(activity.settle(None, None).is_err()); + assert!( + activity + .execute( + ExecutionOrigin::Prompt, + async { panic!("must not run") }, + |_: &()| None + ) + .await + .is_err() + ); + } } diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index eb923cb5..003ce3ab 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -171,6 +171,13 @@ impl AcpSessionUpdateSink for ConnectionSink { } } +// All writers below manipulate only owned IDs and standard collections: no +// user callbacks, backend calls, awaits, or fallible external commits under this +// guard. Each ID is optional and each collection remains valid independently; +// pending_thought is only a hint for the next notification. Clear/reset retire a +// generation before sending best-effort clears outside the lock. Sink failure +// does not resurrect it. Thus poison recovery preserves these local invariants, +// not a promise of atomic delivery to the remote client. #[derive(Default)] struct CurrentReplacementMessages { agent: Option, @@ -528,6 +535,30 @@ enum SessionPublicationError { Commit(AcpRuntimeError), } +// Created before the local map guard so rollback runs only after that guard is +// released, including during unwind. Registry callbacks may own arbitrary drops. +struct PublicationRollback<'a> { + registry: &'a SessionRegistry, + admission: &'a mut super::SessionAdmission, + token: u64, + actor: tokio::task::AbortHandle, + armed: bool, +} + +impl Drop for PublicationRollback<'_> { + fn drop(&mut self) { + if self.armed { + self.actor.abort(); + // Only successful insertion consumes this fresh admission. That + // happens before notifying waiters, which can themselves unwind. + // A rejected duplicate token still belongs to its original actor. + if !self.admission.active { + self.registry.remove(self.token); + } + } + } +} + struct PendingSessionPublication { token: u64, interrupt: Arc, @@ -596,33 +627,65 @@ impl Server { publication: PendingSessionPublication, commit: impl FnOnce() -> Result<(), AcpRuntimeError>, ) -> Result<(), SessionPublicationError> { - let mut sessions = self.sessions.lock().expect("ACP v2 session map poisoned"); + if !admission.active || !Arc::ptr_eq(&admission.registry, &self.registry.inner) { + return Err(SessionPublicationError::AdmissionClosed); + } + let mut rollback = PublicationRollback { + registry: &self.registry, + admission, + token: publication.token, + actor: publication.actor.clone(), + armed: false, + }; + // A commit can affect durable identity as well as this map. If it + // unwinds, rollback removes/aborts the registered actor, but cannot prove + // the external commit complete. Isolate this connection on poison rather + // than interpreting the map's memory safety as successful publication. + let mut sessions = self + .sessions + .lock() + .map_err(|_| SessionPublicationError::Commit(AcpRuntimeError::ClientClosed))?; + if sessions.contains_key(&publication.session_id) { + return Err(SessionPublicationError::Commit(AcpRuntimeError::Loop( + "ACP v2 session is already published".into(), + ))); + } + sessions.try_reserve(1).map_err(|error| { + SessionPublicationError::Commit(AcpRuntimeError::Loop(error.to_string())) + })?; + rollback.armed = true; self.registry .register_v2( - admission, + rollback.admission, publication.token, - publication.interrupt, - publication.close, - publication.actor, - publication.completed, + Arc::clone(&publication.interrupt), + Arc::clone(&publication.close), + publication.actor.clone(), + publication.completed.clone(), ) .map_err(|()| SessionPublicationError::AdmissionClosed)?; - if let Err(error) = commit() { - self.registry.remove(publication.token); - return Err(SessionPublicationError::Commit(error)); - } + commit().map_err(SessionPublicationError::Commit)?; sessions.insert(publication.session_id, publication.session); + rollback.armed = false; Ok(()) } fn remove_session(&self, session_id: &wire::SessionId, token: u64) { - let mut sessions = self.sessions.lock().expect("ACP v2 session map poisoned"); - if sessions + // Actor cleanup must still remove its registry entry and signal + // completion when publication isolated the connection after an unwind. + let Ok(mut sessions) = self.sessions.lock() else { + return; + }; + let removed = if sessions .get(session_id) .is_some_and(|session| session.token == token) { - sessions.remove(session_id); - } + sessions.remove(session_id) + } else { + None + }; + drop(sessions); + drop(removed); } fn initialize( @@ -936,22 +999,22 @@ impl Server { &self, request: wire::PromptRequest, ) -> Result, AcpRuntimeError> { - let (sender, busy, handle, cancellation_generation) = - self.prompt_sender(&request.session_id)?; - let (reply, response) = oneshot::channel(); - if sender - .send(Command::Prompt(PromptCommand { - request, - cancellation_generation, - reply, - })) + // Wait for capacity before claiming busy/injection state. Cancellation + // while a full mailbox is pending must not strand a prompt claim. + let (sender, busy, handle) = self.prompt_route(&request.session_id)?; + let permit = sender + .reserve() .await - .is_err() - { - handle.stop_injection_turn(); - busy.store(false, Ordering::Release); - return Err(AcpRuntimeError::ClientClosed); - } + .map_err(|_| AcpRuntimeError::ClientClosed)?; + claim_prompt(&busy)?; + handle.prepare_injection_turn(); + let cancellation_generation = handle.cancellation_handle().generation(); + let (reply, response) = oneshot::channel(); + permit.send(Command::Prompt(PromptCommand { + request, + cancellation_generation, + reply, + })); match response.await { Ok(response) => response, Err(_) => { @@ -962,32 +1025,21 @@ impl Server { } } - fn prompt_sender( + fn prompt_route( &self, session_id: &wire::SessionId, - ) -> Result< - ( - mpsc::Sender, - Arc, - AcpSessionHandle, - u64, - ), - AcpRuntimeError, - > { - let sessions = self.sessions.lock().expect("ACP v2 session map poisoned"); + ) -> Result<(mpsc::Sender, Arc, AcpSessionHandle), AcpRuntimeError> { + let sessions = self + .sessions + .lock() + .map_err(|_| AcpRuntimeError::ClientClosed)?; let session = sessions .get(session_id) .ok_or_else(|| AcpRuntimeError::SessionNotFound(session_id.to_string()))?; - claim_prompt(&session.busy)?; + let commands = session.commands.clone(); + let busy = Arc::clone(&session.busy); let handle = session.integration.clone(); - handle.prepare_injection_turn(); - let generation = handle.cancellation_handle().generation(); - Ok(( - session.commands.clone(), - Arc::clone(&session.busy), - handle, - generation, - )) + Ok((commands, busy, handle)) } async fn set_config( @@ -1010,7 +1062,7 @@ impl Server { let session = self .sessions .lock() - .expect("ACP v2 session map poisoned") + .map_err(|_| AcpRuntimeError::ClientClosed)? .get(¬ification.session_id) .map(|session| { ( @@ -1036,7 +1088,7 @@ impl Server { let session = self .sessions .lock() - .expect("ACP v2 session map poisoned") + .map_err(|_| AcpRuntimeError::ClientClosed)? .remove(&request.session_id) .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; super::cancel_background_jobs(&session.tasks, &session.background_jobs).await; @@ -1060,7 +1112,7 @@ impl Server { ) -> Result, AcpRuntimeError> { self.sessions .lock() - .expect("ACP v2 session map poisoned") + .map_err(|_| AcpRuntimeError::ClientClosed)? .get(session_id) .map(|session| session.commands.clone()) .ok_or_else(|| AcpRuntimeError::SessionNotFound(session_id.to_string())) @@ -1074,7 +1126,7 @@ impl Server { let (jobs, tasks) = self .sessions .lock() - .expect("ACP v2 session map poisoned") + .map_err(|_| AcpRuntimeError::ClientClosed)? .get(&id) .map(|session| (session.background_jobs.clone(), session.tasks.clone())) .ok_or_else(|| AcpRuntimeError::SessionNotFound(id.to_string()))?; @@ -1091,7 +1143,7 @@ impl Server { let jobs = self .sessions .lock() - .expect("ACP v2 session map poisoned") + .map_err(|_| AcpRuntimeError::ClientClosed)? .get(&id) .map(|session| session.background_jobs.clone()) .ok_or_else(|| AcpRuntimeError::SessionNotFound(id.to_string()))?; @@ -3911,6 +3963,241 @@ mod tests { } } + fn pending_publication( + server: &Arc, + id: &str, + ) -> ( + PendingSessionPublication, + tokio::task::JoinHandle<()>, + mpsc::Receiver, + ) { + let session_id = wire::SessionId::new(id.to_owned()); + let integration = server + .integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + SessionId::new(id.to_owned()), + RecordingSink::default(), + )) + .unwrap(); + let token = server.registry.next_token(); + let (completed, completion) = watch::channel(false); + let guard = ActorGuard { + server: Arc::downgrade(server), + registry: server.registry.clone(), + session_id: session_id.clone(), + token, + completed, + }; + let actor = tokio::spawn(async move { + let _guard = guard; + std::future::pending::<()>().await; + }); + let (commands, received) = mpsc::channel(1); + ( + PendingSessionPublication { + token, + interrupt: Arc::new(|| {}), + close: Arc::new(|| Box::pin(async {})), + actor: actor.abort_handle(), + completed: completion, + session_id, + session: SessionHandle { + token, + commands, + integration, + busy: Arc::new(AtomicBool::new(false)), + background_jobs: BackgroundJobs::default(), + structured_completion: false, + tasks: AsyncTaskManager::new().handle(), + }, + }, + actor, + received, + ) + } + + #[tokio::test] + async fn failed_publication_rolls_back_registration_and_unwind_isolates_connection() { + for unwind in [false, true] { + let root = tempfile::tempdir().unwrap(); + let registry = SessionRegistry::new(); + let server = Arc::new(Server::new( + Runtime::new(root.path(), "gpt-5.4").unwrap(), + registry.clone(), + )); + let mut admission = registry.begin_attachment().unwrap(); + let (publication, actor, _received) = + pending_publication(&server, "failed-publication"); + let mut completed = publication.completed.clone(); + let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + server.publish_session(&mut admission, publication, || { + if unwind { + panic!("commit callback unwound"); + } + Err(AcpRuntimeError::ClientClosed) + }) + })); + if unwind { + assert!(outcome.is_err()); + } else { + assert!(outcome.unwrap().is_err()); + } + assert!(registry.inner.state.lock().unwrap().v2_sessions.is_empty()); + assert_eq!(server.sessions.is_poisoned(), unwind); + let error = server + .sender(&wire::SessionId::new("failed-publication")) + .unwrap_err(); + if unwind { + assert!(matches!(error, AcpRuntimeError::ClientClosed)); + } else { + assert!(matches!(error, AcpRuntimeError::SessionNotFound(_))); + } + assert!( + timeout(Duration::from_secs(1), actor) + .await + .unwrap() + .unwrap_err() + .is_cancelled() + ); + if !*completed.borrow() { + completed.changed().await.unwrap(); + } + assert!(*completed.borrow()); + assert_eq!(registry.inner.state.lock().unwrap().pending_attachments, 0); + } + } + + #[tokio::test] + async fn rejected_publication_does_not_commit_or_register() { + let root = tempfile::tempdir().unwrap(); + let registry = SessionRegistry::new(); + let server = Arc::new(Server::new( + Runtime::new(root.path(), "gpt-5.4").unwrap(), + registry.clone(), + )); + let mut admission = registry.begin_attachment().unwrap(); + registry.close_gate_and_snapshot(false); + let (publication, actor, _received) = pending_publication(&server, "rejected-publication"); + let result = server.publish_session(&mut admission, publication, || { + panic!("rejected commit ran") + }); + assert!(matches!( + result, + Err(SessionPublicationError::AdmissionClosed) + )); + assert!(server.sessions.lock().unwrap().is_empty()); + assert!(registry.inner.state.lock().unwrap().v2_sessions.is_empty()); + drop(admission); + assert_eq!(registry.inner.state.lock().unwrap().pending_attachments, 0); + assert!( + timeout(Duration::from_secs(1), actor) + .await + .unwrap() + .unwrap_err() + .is_cancelled() + ); + } + + #[tokio::test] + async fn duplicate_publication_preserves_the_original_registration() { + let root = tempfile::tempdir().unwrap(); + let registry = SessionRegistry::new(); + let server = Arc::new(Server::new( + Runtime::new(root.path(), "gpt-5.4").unwrap(), + registry.clone(), + )); + let mut original_admission = registry.begin_attachment().unwrap(); + let (original, original_actor, _received) = pending_publication(&server, "original"); + let token = original.token; + server + .publish_session(&mut original_admission, original, || Ok(())) + .unwrap(); + let mut duplicate_admission = registry.begin_attachment().unwrap(); + let (mut duplicate, duplicate_actor, _duplicate_received) = + pending_publication(&server, "duplicate"); + duplicate.token = token; + duplicate.session.token = token; + assert!(matches!( + server.publish_session(&mut duplicate_admission, duplicate, || panic!( + "duplicate commit ran" + )), + Err(SessionPublicationError::AdmissionClosed) + )); + assert_eq!( + registry + .inner + .state + .lock() + .unwrap() + .v2_sessions + .get(&token) + .unwrap() + .actor + .id(), + original_actor.id() + ); + assert!(!original_actor.is_finished()); + assert!(server.sender(&wire::SessionId::new("original")).is_ok()); + assert!(server.sender(&wire::SessionId::new("duplicate")).is_err()); + assert!( + timeout(Duration::from_secs(1), duplicate_actor) + .await + .unwrap() + .unwrap_err() + .is_cancelled() + ); + // The rejected actor's guard uses its independently allocated original + // token; it must not remove the successfully published actor either. + assert!( + registry + .inner + .state + .lock() + .unwrap() + .v2_sessions + .contains_key(&token) + ); + original_actor.abort(); + let _ = original_actor.await; + } + + #[tokio::test] + async fn cancelled_mailbox_wait_does_not_claim_prompt() { + let root = tempfile::tempdir().unwrap(); + let registry = SessionRegistry::new(); + let server = Arc::new(Server::new( + Runtime::new(root.path(), "gpt-5.4").unwrap(), + registry.clone(), + )); + let mut admission = registry.begin_attachment().unwrap(); + let (publication, actor, mut received) = pending_publication(&server, "queued-prompt"); + let busy = Arc::clone(&publication.session.busy); + let session_id = publication.session_id.clone(); + server + .publish_session(&mut admission, publication, || Ok(())) + .unwrap(); + let (reply, _ack) = oneshot::channel(); + server + .sender(&session_id) + .unwrap() + .try_send(Command::Close { reply }) + .unwrap(); + assert!( + timeout( + Duration::from_millis(20), + server.prepare_prompt(wire::PromptRequest::new(session_id.clone(), Vec::new(),)) + ) + .await + .is_err() + ); + assert!(!busy.load(Ordering::Acquire)); + received.recv().await.unwrap(); + assert!(claim_prompt(&busy).is_ok()); + actor.abort(); + let _ = actor.await; + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn logout_reset_waits_for_registered_session_publication() { let root = tempfile::tempdir().unwrap(); diff --git a/src/provider/openai_auth.rs b/src/provider/openai_auth.rs index bda6a713..90606d14 100644 --- a/src/provider/openai_auth.rs +++ b/src/provider/openai_auth.rs @@ -592,7 +592,7 @@ pub(crate) fn access_token( ) })?; if !valid_generation(&record.generation) { - let _thread = refresh_guard(deadline)?; + let _thread = refresh_guard(&REFRESH_LOCK, deadline)?; let _process = process_lock_scoped(deadline, store.lock_scope())?; record = store.load()?.ok_or_else(|| { AuthError::invalid( @@ -635,16 +635,27 @@ fn refresh_locked( rejected_access_token: Option<&str>, token_url: &str, ) -> Result { - let _thread = refresh_guard(deadline)?; + let _thread = refresh_guard(&REFRESH_LOCK, deadline)?; let _process = process_lock_scoped(deadline, store.lock_scope())?; refresh_current(store, deadline, rejected_access_token, token_url) } -fn refresh_guard(deadline: Instant) -> Result, AuthError> { +fn refresh_guard( + lock: &Mutex<()>, + deadline: Instant, +) -> Result, AuthError> { loop { - match REFRESH_LOCK.try_lock() { + match lock.try_lock() { Ok(guard) => return Ok(guard), - Err(std::sync::TryLockError::Poisoned(error)) => return Ok(error.into_inner()), + // This guard also covers backend calls and remote token rotation. + // After unwind, the unit value says nothing about whether those + // external effects committed; do not silently resume refreshing. + Err(std::sync::TryLockError::Poisoned(_)) => { + return Err(AuthError::unavailable( + "credential_refresh_poisoned", + "credential refresh state is unavailable after a panic; restart Kit", + )); + } Err(std::sync::TryLockError::WouldBlock) if Instant::now() < deadline => { std::thread::sleep(Duration::from_millis(10)); } @@ -1673,6 +1684,39 @@ mod tests { } } + #[test] + fn refresh_guard_rejects_backend_unwind_without_retrying() { + struct PanickingStore(std::sync::atomic::AtomicUsize); + + impl CredentialStore for PanickingStore { + fn load(&self) -> Result, AuthError> { + self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + panic!("credential backend panicked"); + } + + fn save(&self, _: &TokenRecord) -> Result<(), AuthError> { + unreachable!() + } + + fn delete(&self) -> Result { + unreachable!() + } + } + + // Isolate the coordinator so this test cannot poison other auth tests. + let lock = Mutex::new(()); + let store = PanickingStore(std::sync::atomic::AtomicUsize::new(0)); + let deadline = Instant::now() + Duration::from_secs(1); + let attempt = || { + let _guard = refresh_guard(&lock, deadline)?; + store.load() + }; + assert!(std::panic::catch_unwind(attempt).is_err()); + assert!(lock.is_poisoned()); + assert_eq!(attempt().unwrap_err().code, "credential_refresh_poisoned"); + assert_eq!(store.0.load(std::sync::atomic::Ordering::Relaxed), 1); + } + struct FailingStore; impl CredentialStore for FailingStore { diff --git a/src/resilient_fs/mod.rs b/src/resilient_fs/mod.rs index 428119fe..78b66ab6 100644 --- a/src/resilient_fs/mod.rs +++ b/src/resilient_fs/mod.rs @@ -58,8 +58,21 @@ fn allocation_oom() -> io::Error { fn oom() -> io::Error { io::ErrorKind::OutOfMemory.into() } -fn lock(m: &Mutex) -> std::sync::MutexGuard<'_, T> { - m.lock().unwrap_or_else(|e| e.into_inner()) +/// An unwind interrupted guarded filesystem state. The affected service or +/// lease is isolated rather than inferring success from a partial transition. +#[derive(Debug)] +pub struct PoisonedState; +impl std::fmt::Display for PoisonedState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("resilient filesystem state poisoned by an interrupted operation") + } +} +impl std::error::Error for PoisonedState {} + +// Never recover a guard: backend callbacks can unwind after changing disk but +// before the corresponding image, cursor, replay stage, or lease is updated. +fn lock(m: &Mutex) -> io::Result> { + m.lock().map_err(|_| io::Error::other(PoisonedState)) } fn bytes(data: &[u8]) -> io::Result>> { let mut v = Vec::new(); @@ -239,7 +252,7 @@ impl Image { } Source::Native(file) => { if base_n > 0 { - let mut file = lock(file); + let mut file = lock(file)?; file.seek(SeekFrom::Start(offset))?; file.read_exact(&mut buf[..base_n])?; } @@ -308,12 +321,20 @@ pub struct Fs { service: Arc, lease: Option>, } +// Lock order: service state -> cursor (when present) -> object -> native file. +// Lease fencing is acquired under service state, or independently by observers; +// it never acquires service state. No guard crosses an await. File and namespace +// callbacks run under state so an unwind also isolates sibling handles. +// Backends must not reenter this service while a synchronous call is active. struct Service { backend: Arc, state: Mutex, max_bytes: usize, max_operations: usize, } +// Fs and File transitions coordinate namespace entries/redirects, live object +// images, budget accounting, replay stages, and retained lease authority here. +// Disk effects cannot be rolled back generically after an unwind. struct State { entries: Vec, redirects: Vec<(PathBuf, PathBuf)>, @@ -329,13 +350,16 @@ struct State { } #[derive(Debug)] pub struct Status { + /// `usize::MAX` when poison prevents a trustworthy snapshot. pub pending_operations: usize, + /// `usize::MAX` when poison prevents a trustworthy snapshot. pub retained_bytes: usize, pub exhausted: bool, } #[derive(Debug)] pub struct RecoveryReport { pub completed_operations: usize, + /// `usize::MAX` when poison prevents inspecting the pending queue. pub remaining_operations: usize, pub blocked: Option, } @@ -359,7 +383,7 @@ struct LeaseInner { } impl LeaseInner { fn check(&self) -> io::Result<()> { - let mut fenced = lock(&self.fenced); + let mut fenced = lock(&self.fenced)?; if let Some(kind) = *fenced { if kind == io::ErrorKind::NotFound { // Keep Missing while the name is absent, but report a replaced @@ -504,7 +528,7 @@ impl Fs { ) -> io::Result { let path = self.norm(path.as_ref())?; let scope = self.norm(scope.as_ref())?; - let mut state = lock(&self.service.state); + let mut state = lock(&self.service.state)?; state .leases .retain(|(_, authority, _)| authority.strong_count() > 0); @@ -541,7 +565,7 @@ impl Fs { } // Retained authority precedes recovery: parent sync touches the lock. let report = self.recover_locked(&mut state); - self.rebase(&mut state); + self.rebase(&mut state)?; Self::prune(&mut state); if state.pending.iter().any(|p| p.action.touches(&path)) { return Err(report.blocked.unwrap_or_else(|| { @@ -573,6 +597,7 @@ impl Fs { Ok(Lease { inner: caller }) } fn norm(&self, path: &Path) -> io::Result { + let _state = lock(&self.service.state)?; // Canonicalize a real ancestor, never collapse `..` through a symlink. let absolute = if path.is_absolute() { path.to_path_buf() @@ -629,7 +654,7 @@ impl Fs { } if let Some(e) = s.entries.iter().find(|e| e.path == cur) { if let Some(o) = &e.object { - if lock(o).meta.kind.symlink { + if lock(o)?.meta.kind.symlink { return Err(error( io::ErrorKind::PermissionDenied, "symlink in managed path", @@ -657,27 +682,38 @@ impl Fs { } Ok(()) } - fn retained(s: &State) -> usize { + fn retained(s: &State) -> io::Result { let objects = s .objects .iter() .filter_map(|o| o.upgrade()) - .map(|o| lock(&o).image.payload_bytes()) - .sum::(); - objects.saturating_add(s.pending.iter().map(|p| p.action.bytes()).sum::()) - } + .try_fold(0usize, |total, o| { + Ok::<_, io::Error>(total.saturating_add(lock(&o)?.image.payload_bytes())) + })?; + Ok(objects.saturating_add(s.pending.iter().map(|p| p.action.bytes()).sum::())) + } + /// Poison makes accounting unknowable. The infallible snapshot reports + /// saturated counts and exhaustion rather than inspecting interrupted state. + /// Use `recover().blocked` for the typed poison error. pub fn status(&self) -> Status { - let s = lock(&self.service.state); - Status { - pending_operations: s.pending.len(), - retained_bytes: Self::retained(&s), - exhausted: s.exhausted - || ALLOCATION_EXHAUSTED.load(std::sync::atomic::Ordering::Acquire), - } + let snapshot = || -> io::Result { + let s = lock(&self.service.state)?; + Ok(Status { + pending_operations: s.pending.len(), + retained_bytes: Self::retained(&s)?, + exhausted: s.exhausted + || ALLOCATION_EXHAUSTED.load(std::sync::atomic::Ordering::Acquire), + }) + }; + snapshot().unwrap_or(Status { + pending_operations: usize::MAX, + retained_bytes: usize::MAX, + exhausted: true, + }) } fn reserve(&self, s: &mut State, additional: usize, entries: usize) -> io::Result<()> { if s.pending.len() >= self.service.max_operations - || Self::retained(s).saturating_add(additional) > self.service.max_bytes + || Self::retained(s)?.saturating_add(additional) > self.service.max_bytes || s.entries.len().saturating_add(entries) > self.service.max_operations.saturating_mul(4) { @@ -745,8 +781,8 @@ impl Fs { return e .object .as_ref() - .map(|o| lock(o).meta.clone()) - .ok_or_else(|| error(io::ErrorKind::NotFound, "removed path")); + .ok_or_else(|| error(io::ErrorKind::NotFound, "removed path")) + .and_then(|o| Ok(lock(o)?.meta.clone())); } if s.entries.iter().any(|e| { e.object.is_none() @@ -783,7 +819,7 @@ impl Fs { // Share logical identity only while the named inode still matches. // Explicit and external replacements must remain distinct. for object in s.objects.iter().filter_map(|w| w.upgrade()) { - let held = lock(&object); + let held = lock(&object)?; if held.path.as_deref() == Some(path) && same_disk_identity(held.meta.disk_identity, opened.meta.disk_identity) { @@ -805,14 +841,19 @@ impl Fs { Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(None), Err(e) => return Err(e), }; - Ok(s.objects.iter().filter_map(|w| w.upgrade()).find(|o| { - let o = lock(o); - o.path.as_deref() == Some(path) && same_disk_identity(o.meta.disk_identity, identity) - })) + for object in s.objects.iter().filter_map(|w| w.upgrade()) { + let o = lock(&object)?; + if o.path.as_deref() == Some(path) && same_disk_identity(o.meta.disk_identity, identity) + { + drop(o); + return Ok(Some(object)); + } + } + Ok(None) } - fn entry(s: &mut State, path: PathBuf, object: Option) { + fn entry(s: &mut State, path: PathBuf, object: Option) -> io::Result<()> { if let Some(o) = &object { - lock(o).path = Some(path.clone()); + lock(o)?.path = Some(path.clone()); } if let Some(o) = &object && !s.objects.iter().any(|w| w.ptr_eq(&Arc::downgrade(o))) @@ -824,6 +865,7 @@ impl Fs { } else { s.entries.push(Entry { path, object }); } + Ok(()) } fn new_permissions( &self, @@ -842,9 +884,20 @@ impl Fs { let mut parent = path .parent() .ok_or_else(|| error(io::ErrorKind::InvalidInput, "no parent"))?; - while s.entries.iter().any(|e| { - e.path == parent && e.object.as_ref().is_some_and(|o| lock(o).meta.is_dir()) - }) && matches!(self.service.backend.metadata(parent, false), + while s + .entries + .iter() + .filter(|e| e.path == parent) + .try_fold(false, |found, e| { + Ok::<_, io::Error>( + found + || match &e.object { + Some(o) => lock(o)?.meta.is_dir(), + None => false, + }, + ) + })? + && matches!(self.service.backend.metadata(parent, false), Err(e) if e.kind() == io::ErrorKind::NotFound) { parent = parent @@ -1050,22 +1103,25 @@ impl Fs { } } fn abandon(&self, action: &Action) { + // Best-effort cleanup needs a trustworthy descriptor identity. A + // poisoned descriptor cannot authorize deleting a temporary pathname. if let Action::Put { temp, temp_file: Some(file), stage: 1 | 2, .. } = action - && let Ok(held) = lock(file).identity() + && let Ok(file) = lock(file) + && let Ok(held) = file.identity() && let Ok(named) = self.service.backend.identity(temp, false) && same_disk_identity(held, named) { let _ = self.service.backend.remove_file(temp); } } - fn rebase(&self, s: &mut State) { + fn rebase(&self, s: &mut State) -> io::Result<()> { for object in s.objects.iter().filter_map(|w| w.upgrade()) { - let mut object = lock(&object); + let mut object = lock(&object)?; let Some(path) = object.path.as_ref() else { continue; }; @@ -1089,7 +1145,7 @@ impl Fs { }); if let Some(file) = published { let snapshot = { - let held = lock(&file); + let held = lock(&file)?; held.metadata() .and_then(|meta| Ok((meta, held.identity()?))) }; @@ -1123,6 +1179,7 @@ impl Fs { object.dirty = false; } } + Ok(()) } fn enqueue(&self, s: &mut State, action: Action) -> (bool, io::Result<()>) { s.pending.push_back(Pending { @@ -1157,9 +1214,20 @@ impl Fs { } } pub fn recover(&self) -> RecoveryReport { - let mut s = lock(&self.service.state); - let report = self.recover_locked(&mut s); - self.rebase(&mut s); + let mut s = match lock(&self.service.state) { + Ok(s) => s, + Err(e) => { + return RecoveryReport { + completed_operations: 0, + remaining_operations: usize::MAX, + blocked: Some(e), + }; + } + }; + let mut report = self.recover_locked(&mut s); + if let Err(e) = self.rebase(&mut s) { + report.blocked = Some(e); + } Self::prune(&mut s); report } @@ -1251,7 +1319,7 @@ impl Fs { let file = temp_file.as_ref().ok_or_else(|| { error(io::ErrorKind::InvalidData, "missing temporary descriptor") })?; - let mut file = lock(file); + let mut file = lock(file)?; #[cfg(unix)] let named = b.metadata(temp, false)?; if !same_disk_identity(file.identity()?, b.identity(temp, false)?) @@ -1332,9 +1400,9 @@ impl Fs { } pub fn require_disk>(&self, path: P) -> io::Result<()> { let path = self.norm(path.as_ref())?; - let mut s = lock(&self.service.state); + let mut s = lock(&self.service.state)?; let report = self.recover_locked(&mut s); - self.rebase(&mut s); + self.rebase(&mut s)?; Self::prune(&mut s); if s.pending.iter().any(|p| p.action.touches(&path)) { return Err(report.blocked.unwrap_or_else(|| { @@ -1348,13 +1416,13 @@ impl Fs { } pub fn metadata>(&self, path: P) -> io::Result { let path = self.norm(path.as_ref())?; - let s = lock(&self.service.state); + let s = lock(&self.service.state)?; self.secure_path(&s, &path, false)?; self.lookup(&s, &path) } pub fn symlink_metadata>(&self, path: P) -> io::Result { let path = self.norm(path.as_ref())?; - let s = lock(&self.service.state); + let s = lock(&self.service.state)?; self.secure_path(&s, &path, true)?; self.lookup(&s, &path) } @@ -1367,7 +1435,8 @@ impl Fs { } pub fn read>(&self, path: P) -> io::Result> { let file = self.open(path)?; - let image = lock(&file.object).image.clone(); + let _state = lock(&self.service.state)?; + let image = lock(&file.object)?.image.clone(); let len = usize::try_from(image.len).map_err(|_| allocation_oom())?; let mut data = Zeroizing::new(Vec::new()); data.try_reserve_exact(len).map_err(|_| allocation_oom())?; @@ -1396,7 +1465,7 @@ impl Fs { new_object: bool, ) -> io::Result<()> { let path = self.norm(path)?; - let mut s = lock(&self.service.state); + let mut s = lock(&self.service.state)?; self.recover_before(&mut s)?; self.preflight(&s, &path)?; s.entries.try_reserve(1).map_err(|_| allocation_oom())?; @@ -1445,20 +1514,20 @@ impl Fs { // Acceptance replaces the logical name even when publication is queued. // Old handles must not publish writes or chmod through that name. for object in s.objects.iter().filter_map(|w| w.upgrade()) { - let mut object = lock(&object); + let mut object = lock(&object)?; if object.path.as_ref() == Some(&path) { object.path = None; } } } let object = if let Some(o) = object { - *lock(&o) = Object::memory(data.clone(), meta); + *lock(&o)? = Object::memory(data.clone(), meta); o } else { Arc::new(Mutex::new(Object::memory(data.clone(), meta))) }; - Self::entry(&mut s, path.clone(), Some(object)); - self.rebase(&mut s); + Self::entry(&mut s, path.clone(), Some(object))?; + self.rebase(&mut s)?; result } fn prune(s: &mut State) { @@ -1470,7 +1539,7 @@ impl Fs { } fn recover_before(&self, s: &mut State) -> io::Result<()> { let report = self.recover_locked(s); - self.rebase(s); + self.rebase(s)?; Self::prune(s); match report.blocked { Some(e) if !capacity(&e) => Err(e), @@ -1489,7 +1558,7 @@ impl Fs { } pub fn read_link>(&self, path: P) -> io::Result { let p = self.norm(path.as_ref())?; - let s = lock(&self.service.state); + let s = lock(&self.service.state)?; self.secure_path(&s, &p, true)?; if !self.lookup(&s, &p)?.file_type().is_symlink() { return Err(error(io::ErrorKind::InvalidInput, "not a symlink")); @@ -1498,10 +1567,12 @@ impl Fs { } pub fn canonicalize>(&self, path: P) -> io::Result { let p = self.norm(path.as_ref())?; + let s = lock(&self.service.state)?; match self.service.backend.canonicalize(&p) { Ok(p) => Ok(p), Err(e) if e.kind() == io::ErrorKind::NotFound => { - self.metadata(&p)?; + self.secure_path(&s, &p, false)?; + self.lookup(&s, &p)?; Ok(p) } Err(e) => Err(e), @@ -1588,7 +1659,7 @@ impl Fs { } pub fn read_dir>(&self, path: P) -> io::Result { let p = self.norm(path.as_ref())?; - let s = lock(&self.service.state); + let s = lock(&self.service.state)?; self.secure_path(&s, &p, false)?; let paths = self.list(&s, &p)?; let mut entries = Vec::new(); @@ -1612,7 +1683,7 @@ impl Fs { } fn mkdir(&self, path: &Path, private: bool) -> io::Result<()> { let p = self.norm(path)?; - let mut s = lock(&self.service.state); + let mut s = lock(&self.service.state)?; self.recover_before(&mut s)?; self.authority(&p)?; self.secure_path(&s, &p, false)?; @@ -1654,7 +1725,7 @@ impl Fs { Arc::new(Zeroizing::new(Vec::new())), meta, )))), - ); + )?; result } pub fn create_dir_all>(&self, path: P) -> io::Result<()> { @@ -1699,7 +1770,7 @@ impl Fs { } fn unlink(&self, path: &Path, dir: bool) -> io::Result<()> { let p = self.norm(path)?; - let mut s = lock(&self.service.state); + let mut s = lock(&self.service.state)?; self.recover_before(&mut s)?; self.authority(&p)?; self.secure_path(&s, &p, false)?; @@ -1724,12 +1795,12 @@ impl Fs { ); if accepted { for object in s.objects.iter().filter_map(|w| w.upgrade()) { - let mut object = lock(&object); + let mut object = lock(&object)?; if object.path.as_ref() == Some(&p) { object.path = None; } } - Self::entry(&mut s, p, None); + Self::entry(&mut s, p, None)?; } result } @@ -1748,7 +1819,7 @@ impl Fs { pub fn rename, Q: AsRef>(&self, from: P, to: Q) -> io::Result<()> { let from = self.norm(from.as_ref())?; let to = self.norm(to.as_ref())?; - let mut s = lock(&self.service.state); + let mut s = lock(&self.service.state)?; self.recover_before(&mut s)?; self.authority(&from)?; self.authority(&to)?; @@ -1820,7 +1891,7 @@ impl Fs { ); if accepted { for object in s.objects.iter().filter_map(|w| w.upgrade()) { - let mut object = lock(&object); + let mut object = lock(&object)?; if let Some(path) = &object.path { if path.starts_with(&from) { object.path = Some(if *path == from { @@ -1838,9 +1909,9 @@ impl Fs { e.object = None; } } - Self::entry(&mut s, from.clone(), None); + Self::entry(&mut s, from.clone(), None)?; for (path, object) in moved { - Self::entry(&mut s, path, object); + Self::entry(&mut s, path, object)?; } s.redirects .retain(|(path, _)| !path.starts_with(&from) && !path.starts_with(&to)); @@ -1863,7 +1934,7 @@ impl Fs { permissions: Permissions, ) -> io::Result<()> { let p = self.norm(path.as_ref())?; - let mut s = lock(&self.service.state); + let mut s = lock(&self.service.state)?; self.recover_before(&mut s)?; self.authority(&p)?; self.secure_path(&s, &p, false)?; @@ -1885,7 +1956,7 @@ impl Fs { .find(|e| e.path == p) .and_then(|e| e.object.as_ref()) { - let mut o = lock(o); + let mut o = lock(o)?; o.meta.permissions = permissions.clone(); o.meta.disk = None; } @@ -1904,13 +1975,13 @@ impl Fs { Arc::new(Zeroizing::new(Vec::new())), meta, )))), - ); + )?; } Ok(()) } pub fn sync_directory>(&self, path: P) -> io::Result<()> { let p = self.norm(path.as_ref())?; - let mut s = lock(&self.service.state); + let mut s = lock(&self.service.state)?; self.recover_before(&mut s)?; self.authority(&p)?; self.secure_path(&s, &p, false)?; @@ -1938,9 +2009,9 @@ impl Fs { } let root = self.norm(root.as_ref())?; let p = root.join(relative); - let mut s = lock(&self.service.state); + let mut s = lock(&self.service.state)?; let _ = self.recover_locked(&mut s); - self.rebase(&mut s); + self.rebase(&mut s)?; Self::prune(&mut s); self.secure_path(&s, &p, false)?; if s.entries.iter().any(|e| e.path == p) { @@ -2023,12 +2094,12 @@ impl OpenOptions { { return Err(error(io::ErrorKind::InvalidInput, "invalid open options")); } - let mut s = lock(&fs.service.state); + let mut s = lock(&fs.service.state)?; if writable { fs.recover_before(&mut s)?; } else { let _ = fs.recover_locked(&mut s); - fs.rebase(&mut s); + fs.rebase(&mut s)?; Fs::prune(&mut s); } fs.secure_path(&s, &p, false)?; @@ -2090,15 +2161,15 @@ impl OpenOptions { unreachable!(); } let object = if let Some(obj) = obj { - *lock(&obj) = Object::memory(data.clone(), meta); + *lock(&obj)? = Object::memory(data.clone(), meta); obj } else { Arc::new(Mutex::new(Object::memory(data.clone(), meta))) }; - Fs::entry(&mut s, p.clone(), Some(object)); + Fs::entry(&mut s, p.clone(), Some(object))?; result?; } - fs.rebase(&mut s); + fs.rebase(&mut s)?; let object = fs.object(&mut s, &p)?; Ok(File { fs: fs.clone(), @@ -2142,11 +2213,16 @@ impl File { }) } pub fn metadata(&self) -> io::Result { - let object = lock(&self.object); + let _state = lock(&self.fs.service.state)?; + self.metadata_locked() + } + // Caller holds service state; seek also holds the cursor in lock order. + fn metadata_locked(&self) -> io::Result { + let object = lock(&self.object)?; if !object.dirty && let Source::Native(file) = &object.image.source { - let file = lock(file); + let file = lock(file)?; let mut meta = Metadata::disk(file.metadata()?, file.identity()?); meta.identity = object.meta.identity; Ok(meta) @@ -2164,15 +2240,15 @@ impl File { if data.is_empty() && size.is_none() { return Ok(0); } - let mut s = lock(&self.fs.service.state); + let mut s = lock(&self.fs.service.state)?; self.fs.recover_before(&mut s)?; - let mut cursor = lock(&self.cursor); - let mut object = lock(&self.object); + let mut cursor = lock(&self.cursor)?; + let mut object = lock(&self.object)?; if !object.dirty && let Source::Native(native) = &object.image.source { let (meta, disk_identity) = { - let native = lock(native); + let native = lock(native)?; (native.metadata()?, native.identity()?) }; if disk_identity.is_none() { @@ -2235,7 +2311,7 @@ impl File { unreachable!(); } { - let mut object = lock(&self.object); + let mut object = lock(&self.object)?; object.image = image; object.meta.len = len; object.meta.modified = SystemTime::now(); @@ -2246,12 +2322,12 @@ impl File { object.dirty = true; } if let Some(path) = path { - Fs::entry(&mut s, path, Some(self.object.clone())); + Fs::entry(&mut s, path, Some(self.object.clone()))?; } if size.is_none() { *cursor = end; } - self.fs.rebase(&mut s); + self.fs.rebase(&mut s)?; result?; Ok(data.len()) } @@ -2269,15 +2345,15 @@ impl File { self.sync_data() } pub fn set_permissions(&self, p: Permissions) -> io::Result<()> { - let mut s = lock(&self.fs.service.state); + let mut s = lock(&self.fs.service.state)?; self.fs.recover_before(&mut s)?; let path = { - let mut object = lock(&self.object); + let mut object = lock(&self.object)?; if let Some(path) = &object.path && !s.pending.iter().any(|p| p.action.touches(path)) && let Source::Native(native) = &object.image.source { - let held = lock(native).identity()?; + let held = lock(native)?.identity()?; match self.fs.service.backend.identity(path, false) { Ok(named) if same_disk_identity(held, named) => {} Ok(None) => { @@ -2311,15 +2387,15 @@ impl File { (true, Ok(())) }; if accepted { - let mut object = lock(&self.object); + let mut object = lock(&self.object)?; object.meta.permissions = p; object.meta.disk = None; object.dirty = true; drop(object); if let Some(path) = path { - Fs::entry(&mut s, path, Some(self.object.clone())); + Fs::entry(&mut s, path, Some(self.object.clone()))?; } - self.fs.rebase(&mut s); + self.fs.rebase(&mut s)?; } result } @@ -2332,13 +2408,14 @@ impl Read for File { "handle is not readable", )); } - let mut cursor = lock(&self.cursor); + let _state = lock(&self.fs.service.state)?; + let mut cursor = lock(&self.cursor)?; let (image, dirty) = { - let object = lock(&self.object); + let object = lock(&self.object)?; (object.image.clone(), object.dirty) }; let n = if !dirty && let Source::Native(file) = &image.source { - let mut file = lock(file); + let mut file = lock(file)?; file.seek(SeekFrom::Start(*cursor))?; file.read(buf)? } else { @@ -2358,11 +2435,12 @@ impl Write for File { } impl Seek for File { fn seek(&mut self, from: SeekFrom) -> io::Result { - let mut cursor = lock(&self.cursor); + let _state = lock(&self.fs.service.state)?; + let mut cursor = lock(&self.cursor)?; let next = match from { SeekFrom::Start(n) => n as i128, SeekFrom::Current(n) => *cursor as i128 + n as i128, - SeekFrom::End(n) => self.metadata()?.len() as i128 + n as i128, + SeekFrom::End(n) => self.metadata_locked()?.len() as i128 + n as i128, }; if !(0..=u64::MAX as i128).contains(&next) { return Err(error(io::ErrorKind::InvalidInput, "invalid seek")); diff --git a/src/resilient_fs/tests.rs b/src/resilient_fs/tests.rs index a1dfd663..0b838511 100644 --- a/src/resilient_fs/tests.rs +++ b/src/resilient_fs/tests.rs @@ -15,10 +15,29 @@ enum Point { Remove, ProbeCollision, ProbeSwap, + Read, + Seek, + Metadata, + AfterWrite, + AfterRename, + LeaseCheck, + FileDrop, } #[derive(Default)] -struct Faults(Mutex>, AtomicUsize); +struct Faults(Mutex>, AtomicUsize, AtomicUsize); impl Faults { + fn arm_panic(&self, point: Point) { + self.2.store(point as usize + 1, Ordering::SeqCst); + } + fn panic_at(&self, point: Point) { + if self + .2 + .compare_exchange(point as usize + 1, 0, Ordering::SeqCst, Ordering::SeqCst) + .is_ok() + { + panic!("backend panic at {point:?}"); + } + } fn arm(&self, point: Point, code: i32) { *self.0.lock().unwrap() = Some((point, code)); } @@ -43,12 +62,18 @@ struct InjectedFile { } impl Read for InjectedFile { fn read(&mut self, b: &mut [u8]) -> io::Result { - self.disk.read(b) + let n = self.disk.read(b)?; + self.faults.panic_at(Point::Read); + self.faults.check(Point::Read)?; + Ok(n) } } impl Seek for InjectedFile { fn seek(&mut self, p: SeekFrom) -> io::Result { - self.disk.seek(p) + let position = self.disk.seek(p)?; + self.faults.panic_at(Point::Seek); + self.faults.check(Point::Seek)?; + Ok(position) } } impl Write for InjectedFile { @@ -60,7 +85,9 @@ impl Write for InjectedFile { self.wrote_prefix = true; return self.disk.write(&b[..b.len().min(3)]); } - self.disk.write(b) + let n = self.disk.write(b)?; + self.faults.panic_at(Point::AfterWrite); + Ok(n) } fn flush(&mut self) -> io::Result<()> { self.disk.flush() @@ -71,7 +98,9 @@ impl BackendFile for InjectedFile { self.disk.identity() } fn metadata(&self) -> io::Result { - self.disk.metadata() + let meta = self.disk.metadata()?; + self.faults.panic_at(Point::Metadata); + Ok(meta) } fn set_len(&self, n: u64) -> io::Result<()> { self.disk.set_len(n) @@ -88,6 +117,22 @@ impl BackendFile for InjectedFile { self.disk.set_permissions(p) } } +impl Drop for InjectedFile { + fn drop(&mut self) { + self.faults.panic_at(Point::FileDrop); + } +} +struct InjectedLease { + disk: Box, + faults: Arc, +} +impl BackendLease for InjectedLease { + fn check(&self) -> io::Result<()> { + self.disk.check()?; + self.faults.panic_at(Point::LeaseCheck); + Ok(()) + } +} impl Backend for Injected { fn identity(&self, path: &Path, follow: bool) -> io::Result> { self.disk.identity(path, follow) @@ -139,7 +184,9 @@ impl Backend for Injected { } fn rename(&self, a: &Path, b: &Path) -> io::Result<()> { self.faults.check(Point::Rename)?; - self.disk.rename(a, b) + self.disk.rename(a, b)?; + self.faults.panic_at(Point::AfterRename); + Ok(()) } fn set_permissions(&self, p: &Path, mode: Permissions) -> io::Result<()> { self.faults.check(Point::Chmod)?; @@ -150,7 +197,10 @@ impl Backend for Injected { self.disk.sync_directory(p) } fn acquire_lease(&self, r: &LeaseRequest) -> io::Result> { - self.disk.acquire_lease(r) + Ok(Box::new(InjectedLease { + disk: self.disk.acquire_lease(r)?, + faults: self.faults.clone(), + })) } fn open_beneath(&self, root: &Path, p: &Path) -> io::Result> { self.disk.open_beneath(root, p) @@ -1442,3 +1492,181 @@ fn capacity_during_permission_probe_cleanup_does_not_reject_write() { assert_eq!(native::read(t.path("value")).unwrap(), b"data"); t.settle(); } + +fn assert_poison(e: io::Error) { + assert_eq!(e.kind(), io::ErrorKind::Other); + assert!(e.get_ref().unwrap().is::()); +} + +fn assert_service_isolated(t: &Fixture) { + let report = t.fs.recover(); + assert_eq!(report.completed_operations, 0); + assert_eq!(report.remaining_operations, usize::MAX); + assert_poison(report.blocked.unwrap()); + let status = t.fs.status(); + assert!(status.exhausted); + assert_eq!(status.pending_operations, usize::MAX); + assert_eq!(status.retained_bytes, usize::MAX); + assert_poison(t.fs.require_disk(&t.root).unwrap_err()); + assert_poison(t.fs.write(t.path("must-not-appear"), b"no").unwrap_err()); + assert!(!t.path("must-not-appear").exists()); + assert!(t.fs.service.state.is_poisoned()); + // Poison is scoped to this service, not a process-wide recovery default. + let other = Fixture::new(); + other.fs.write(other.path("healthy"), b"yes").unwrap(); + other.settle(); +} + +#[test] +fn backend_panic_after_disk_effect_is_not_replayed_or_hidden() { + for point in [Point::AfterWrite, Point::AfterRename, Point::FileDrop] { + let t = Fixture::new(); + let path = t.path("value"); + native::write(&path, b"old").unwrap(); + let held = t.fs.open(&path).unwrap(); + t.faults.arm_panic(point); + let fs = t.fs.clone(); + let target = path.clone(); + assert!( + std::thread::spawn(move || fs.write(target, b"replacement")) + .join() + .is_err() + ); + let expected = if point == Point::AfterRename { + b"replacement".as_slice() + } else { + b"old".as_slice() + }; + assert_eq!(native::read(&path).unwrap(), expected); + let before: Vec<_> = native::read_dir(&t.root) + .unwrap() + .map(|e| e.unwrap().path()) + .collect(); + assert_service_isolated(&t); + assert_poison(held.metadata().unwrap_err()); + assert_eq!(native::read(&path).unwrap(), expected); + let after: Vec<_> = native::read_dir(&t.root) + .unwrap() + .map(|e| e.unwrap().path()) + .collect(); + assert_eq!( + before, after, + "recovery must not clean or replay unknown disk state" + ); + } +} + +#[test] +fn handle_backend_panics_fence_clones_and_service() { + for point in [Point::Read, Point::Seek, Point::Metadata] { + let t = Fixture::new(); + let path = t.path("value"); + native::write(&path, b"contents").unwrap(); + let mut held = t.fs.open(&path).unwrap(); + let mut worker = held.try_clone().unwrap(); + t.faults.arm_panic(point); + assert!( + std::thread::spawn(move || { + if point == Point::Metadata { + worker.seek(SeekFrom::End(0)).unwrap(); + } else { + worker.read_exact(&mut [0; 2]).unwrap(); + } + }) + .join() + .is_err() + ); + assert!(held.cursor.is_poisoned()); + { + let object = held.object.lock(); + if point == Point::Metadata { + assert!(object.is_err()); + drop(object); + assert_poison(lock(&held.object).err().unwrap()); + } else { + let object = object.unwrap(); + let Source::Native(native) = &object.image.source else { + panic!("expected native source") + }; + assert!(native.is_poisoned()); + assert_poison(lock(native).err().unwrap()); + } + } + assert_poison(lock(&held.cursor).err().unwrap()); + assert_poison(held.read(&mut [0; 2]).unwrap_err()); + assert_poison(held.seek(SeekFrom::Start(0)).unwrap_err()); + assert_poison(held.sync_all().unwrap_err()); + assert_service_isolated(&t); + assert_eq!(native::read(&path).unwrap(), b"contents"); + } +} + +#[test] +fn lease_callback_panic_permanently_fences_authority() { + let t = Fixture::new(); + let lease = + t.fs.acquire_lease(t.path("lock"), &t.root, LeaseMode::CreateNew) + .unwrap(); + let guarded = t.fs.guarded(&lease).unwrap(); + let authority = lease.inner.clone(); + t.faults.arm_panic(Point::LeaseCheck); + assert!( + std::thread::spawn(move || authority.check()) + .join() + .is_err() + ); + assert_poison(lease.check().unwrap_err()); + assert_poison(guarded.write(t.path("value"), b"no").unwrap_err()); + assert!(!t.path("value").exists()); + assert!(lease.inner.fenced.is_poisoned()); + assert!(!t.fs.service.state.is_poisoned()); + // A rejected authority check must not corrupt service accounting. + assert_eq!(t.fs.status().pending_operations, 0); +} + +#[test] +#[cfg(unix)] +fn failed_native_read_and_seek_do_not_commit_logical_cursor() { + for point in [Point::Read, Point::Seek] { + let t = Fixture::new(); + let path = t.path("value"); + native::write(&path, b"abcdef").unwrap(); + let mut file = t.fs.open(&path).unwrap(); + file.seek(SeekFrom::Start(2)).unwrap(); + t.faults.arm(point, libc::EIO); + assert_eq!( + file.read(&mut [0; 2]).unwrap_err().raw_os_error(), + Some(libc::EIO) + ); + assert!(!t.fs.service.state.is_poisoned()); + t.faults.clear(); + assert_eq!(file.stream_position().unwrap(), 2); + let mut clone = file.try_clone().unwrap(); + let mut data = [0; 2]; + clone.read_exact(&mut data).unwrap(); + assert_eq!(&data, b"cd"); + assert_eq!(file.stream_position().unwrap(), 4); + t.settle(); + } +} + +#[test] +#[cfg(unix)] +fn panic_during_pending_publication_blocks_further_recovery() { + let t = Fixture::new(); + let path = t.path("value"); + native::write(&path, b"old").unwrap(); + t.faults.arm(Point::Open, libc::ENOSPC); + t.fs.write(&path, b"queued").unwrap(); + assert_eq!(t.fs.status().pending_operations, 1); + assert_eq!(native::read(&path).unwrap(), b"old"); + t.faults.clear(); + t.faults.arm_panic(Point::AfterRename); + let fs = t.fs.clone(); + assert!(std::thread::spawn(move || fs.recover()).join().is_err()); + // Native publication happened, but the stage and completion count were not + // committed. Replaying this action or claiming durability would be unsound. + assert_eq!(native::read(&path).unwrap(), b"queued"); + assert_service_isolated(&t); + assert_eq!(native::read(&path).unwrap(), b"queued"); +} diff --git a/src/runtime.rs b/src/runtime.rs index 6134f640..3055de57 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -210,17 +210,29 @@ impl SessionClaim { pub(crate) fn commit(mut self) -> Result<(), AcpRuntimeError> { if matches!(self.kind, SessionClaimKind::Fork) { - if let Some(observer) = self.uncommitted_observer.take() { - observer.commit_creation(); + if let Some(observer) = &self.uncommitted_observer { + observer.commit_creation().map_err(AcpRuntimeError::Loop)?; } self.committed = true; return Ok(()); } - self.mark_opened(); - let mut selection = - self.runtime.session.lock().map_err(|_| { - AcpRuntimeError::Loop("runtime session selection is poisoned".into()) - })?; + // Already-committed transcripts survive rejection. A guarded creation + // does not: its cleanup still owns the file until commit succeeds. + if self.uncommitted_observer.is_none() { + self.mark_opened(); + } + let runtime = Arc::clone(&self.runtime); + let mut selection = runtime + .session + .lock() + .map_err(|_| AcpRuntimeError::Loop("runtime session selection is poisoned".into()))?; + // Validate transcript publication before consuming the reserved identity. + // Retain the observer until success so rejected commits still run the + // claim's normal uncommitted-transcript cleanup and retry bookkeeping. + if let Some(observer) = &self.uncommitted_observer { + observer.commit_creation().map_err(AcpRuntimeError::Loop)?; + self.mark_opened(); + } match self.kind { SessionClaimKind::New { configured, @@ -231,10 +243,8 @@ impl SessionClaim { } SessionClaimKind::Fork => unreachable!("fork claims commit before locking selection"), } - if let Some(observer) = self.uncommitted_observer.take() { - observer.commit_creation(); - } self.committed = true; + drop(selection); Ok(()) } } @@ -1101,7 +1111,7 @@ impl Runtime { { Ok(driver) => { if let Some(observer) = pending_creation { - observer.commit_creation(); + observer.commit_creation()?; } driver } @@ -1669,7 +1679,9 @@ impl ToolSource for ComposeOnly { } struct BackgroundJob { + registration: Arc<()>, controller: CancellationController, + cancellation_requested: bool, foreground_cancellation: Option, cancellation_relay: Option, detached: bool, @@ -1718,27 +1730,62 @@ impl Default for BackgroundJobs { } } +impl BackgroundJob { + fn request_cancellation(&mut self) -> Option { + // This private controller is signalled only once. In agentkit-core 0.10.5 + // interrupt is an atomic generation increment (waiters poll); selecting + // at most one signal also prevents overflow before the out-of-lock effect. + if std::mem::replace(&mut self.cancellation_requested, true) { + None + } else { + Some(self.controller.clone()) + } + } +} + +impl BackgroundJobState { + fn changed(&mut self) -> u64 { + self.generation = self.generation.wrapping_add(1); + self.generation + } +} + impl BackgroundJobs { - fn changed(&self, jobs: &mut BackgroundJobState) { - jobs.generation = jobs.generation.wrapping_add(1); - self.activity.send_replace(jobs.generation); + fn lock_jobs(&self) -> std::sync::MutexGuard<'_, BackgroundJobState> { + self.state.lock().unwrap_or_else(|poisoned| { + // Every writer commits only owned maps/sets, booleans and wrapping + // counters here. Registration rejects live IDs before changing state; + // finish moves the job out and records its terminal debt atomically; + // acknowledgement handles both sides of that transition. Relay handles, + // removed jobs, aborts, cancellation signals and watch wakeups all leave + // this lock before running/dropping. No await or user code runs guarded. + // Cancellation decisions are sticky per job, so detach cannot overtake + // a signal selected under this lock; cancel_all also covers later jobs. + // Signals target cloned controllers; relays also check registration + // identity, so an old relay racing its abort cannot cancel a reused ID. + // Registration establishes RAII cleanup before any wakeup can unwind. + // Thus poison cannot mean a partly published job or lost terminal debt; + // recover the actual activity, never substitute a fictitious idle state. + let jobs = poisoned.into_inner(); + self.state.clear_poison(); + jobs + }) + } + + fn notify_changed(&self, generation: u64) { + // Notifications may arrive out of commit order after unlocking. The watch + // value is only a wakeup hint; activity_after always reads authoritative state. + self.activity.send_replace(generation); } pub(crate) fn activity(&self) -> BackgroundActivity { - self.state.lock().map_or( - BackgroundActivity { - generation: *self.activity.borrow(), - active: false, - background_started: 0, - unacknowledged_terminals: false, - }, - |jobs| BackgroundActivity { - generation: jobs.generation, - active: !jobs.running.is_empty(), - background_started: jobs.background_started, - unacknowledged_terminals: !jobs.unacknowledged_terminals.is_empty(), - }, - ) + let jobs = self.lock_jobs(); + BackgroundActivity { + generation: jobs.generation, + active: !jobs.running.is_empty(), + background_started: jobs.background_started, + unacknowledged_terminals: !jobs.unacknowledged_terminals.is_empty(), + } } pub(crate) async fn activity_after(&self, generation: u64) -> BackgroundActivity { @@ -1765,86 +1812,106 @@ impl BackgroundJobs { } pub(crate) fn acknowledge_terminal(&self, call_id: &agentkit_core::ToolCallId) { - if let Ok(mut jobs) = self.state.lock() { - let changed = if let Some(job) = jobs.running.get_mut(call_id) { - !std::mem::replace(&mut job.terminal_published, true) - } else { - jobs.unacknowledged_terminals.remove(call_id) - }; - if changed { - self.changed(&mut jobs); - } + let mut jobs = self.lock_jobs(); + let changed = if let Some(job) = jobs.running.get_mut(call_id) { + !std::mem::replace(&mut job.terminal_published, true) + } else { + jobs.unacknowledged_terminals.remove(call_id) + }; + let generation = changed.then(|| jobs.changed()); + drop(jobs); + if let Some(generation) = generation { + self.notify_changed(generation); } } pub(crate) fn begin_turn(&self) { - if let Ok(mut jobs) = self.state.lock() { - jobs.cancel_all = false; - } + self.lock_jobs().cancel_all = false; } pub(crate) fn cancel_all(&self) { - if let Ok(mut jobs) = self.state.lock() { - jobs.cancel_all = true; - for job in jobs.running.values() { - job.controller.interrupt(); - } + let mut jobs = self.lock_jobs(); + jobs.cancel_all = true; + let controllers: Vec<_> = jobs + .running + .values_mut() + .filter_map(BackgroundJob::request_cancellation) + .collect(); + drop(jobs); + for controller in controllers { + controller.interrupt(); } } fn cancel_running(&self, call_id: &str) -> bool { + self.cancel_registration(call_id, None) + } + + fn cancel_registration(&self, call_id: &str, registration: Option<&Arc<()>>) -> bool { let call_id = agentkit_core::ToolCallId::new(call_id); - let Ok(jobs) = self.state.lock() else { - return false; - }; - let Some(job) = jobs.running.get(&call_id) else { + let mut jobs = self.lock_jobs(); + let Some(job) = jobs.running.get_mut(&call_id).filter(|job| { + registration.is_none_or(|registration| Arc::ptr_eq(&job.registration, registration)) + }) else { return false; }; - job.controller.interrupt(); + let controller = job.request_cancellation(); + drop(jobs); + if let Some(controller) = controller { + controller.interrupt(); + } true } pub fn cancel(&self, call_id: &str) -> bool { let call_id = agentkit_core::ToolCallId::new(call_id); - let Ok(mut jobs) = self.state.lock() else { - return false; - }; - if let Some(job) = jobs.running.get(&call_id) { - job.controller.interrupt(); + let mut jobs = self.lock_jobs(); + let controller = if let Some(job) = jobs.running.get_mut(&call_id) { + job.request_cancellation() } else { - // ACP can expose the call just before its execution future registers. - // Remember the request so registration and cancellation are atomic - // from the user's perspective. + // Registration consumes the pending request in the same state commit. jobs.pending_cancellations.insert(call_id); + None + }; + drop(jobs); + if let Some(controller) = controller { + controller.interrupt(); } true } pub(crate) fn detach(&self, call_id: &str) -> Option { let call_id = agentkit_core::ToolCallId::new(call_id); - let Ok(mut jobs) = self.state.lock() else { - return None; - }; + let mut jobs = self.lock_jobs(); if let Some(job) = jobs.running.get_mut(&call_id) { if job.manual_detach { return Some(DetachRegistration::AlreadyDetached); } - let newly_detached = !job.detached; - job.detached = true; - job.manual_detach = true; - if job - .foreground_cancellation - .as_ref() - .is_some_and(agentkit_core::TurnCancellation::is_cancelled) + if job.cancellation_requested + || job + .foreground_cancellation + .as_ref() + .is_some_and(agentkit_core::TurnCancellation::is_cancelled) { - job.detached = false; - job.manual_detach = false; - job.controller.interrupt(); + let controller = job.request_cancellation(); + drop(jobs); + if let Some(controller) = controller { + controller.interrupt(); + } return None; } - if newly_detached { + let newly_detached = !job.detached; + job.detached = true; + job.manual_detach = true; + let generation = if newly_detached { jobs.background_started = jobs.background_started.wrapping_add(1); - self.changed(&mut jobs); + Some(jobs.changed()) + } else { + None + }; + drop(jobs); + if let Some(generation) = generation { + self.notify_changed(generation); } return Some(DetachRegistration::Registered); } @@ -1857,9 +1924,7 @@ impl BackgroundJobs { pub(crate) fn restore_foreground(&self, call_id: &str) { let call_id = agentkit_core::ToolCallId::new(call_id); - let Ok(mut jobs) = self.state.lock() else { - return; - }; + let mut jobs = self.lock_jobs(); let Some(job) = jobs.running.get_mut(&call_id) else { jobs.pending_detaches.remove(&call_id); return; @@ -1868,67 +1933,92 @@ impl BackgroundJobs { if job.foreground_cancellation.is_some() { job.detached = false; } - if job + let controller = if job .foreground_cancellation .as_ref() .is_some_and(agentkit_core::TurnCancellation::is_cancelled) { - job.controller.interrupt(); + job.request_cancellation() + } else { + None + }; + drop(jobs); + if let Some(controller) = controller { + controller.interrupt(); } } - fn propagate_foreground_cancellation(&self, call_id: &agentkit_core::ToolCallId) { - let Ok(jobs) = self.state.lock() else { - return; - }; - if let Some(job) = jobs.running.get(call_id) - && !job.detached - { - job.controller.interrupt(); + fn propagate_foreground_cancellation( + &self, + call_id: &agentkit_core::ToolCallId, + registration: &Arc<()>, + ) { + let mut jobs = self.lock_jobs(); + let controller = jobs + .running + .get_mut(call_id) + .filter(|job| !job.detached && Arc::ptr_eq(&job.registration, registration)) + .and_then(BackgroundJob::request_cancellation); + drop(jobs); + if let Some(controller) = controller { + controller.interrupt(); } } fn finish(&self, call_id: &agentkit_core::ToolCallId) { - if let Ok(mut jobs) = self.state.lock() { - if let Some(job) = jobs.running.remove(call_id) { - if job.detached && !job.terminal_published { - jobs.unacknowledged_terminals.insert(call_id.clone()); - } - if let Some(relay) = job.cancellation_relay { - relay.abort(); - } - } - jobs.pending_cancellations.remove(call_id); - jobs.pending_detaches.remove(call_id); - self.changed(&mut jobs); + let mut jobs = self.lock_jobs(); + let job = jobs.running.remove(call_id); + if job + .as_ref() + .is_some_and(|job| job.detached && !job.terminal_published) + { + jobs.unacknowledged_terminals.insert(call_id.clone()); + } + jobs.pending_cancellations.remove(call_id); + jobs.pending_detaches.remove(call_id); + let generation = jobs.changed(); + drop(jobs); + if let Some(job) = &job + && let Some(relay) = &job.cancellation_relay + { + relay.abort(); } + drop(job); + self.notify_changed(generation); } #[cfg(test)] pub(crate) fn register_foreground_for_test(&self, call_id: &str) { - if let Ok(mut jobs) = self.state.lock() { - let call_id = agentkit_core::ToolCallId::new(call_id); - let manual_detach = jobs.pending_detaches.remove(&call_id); - let controller = CancellationController::new(); - if jobs.cancel_all { - controller.interrupt(); - } - jobs.running.insert( - call_id, - BackgroundJob { - controller, - foreground_cancellation: None, - cancellation_relay: None, - detached: manual_detach, - manual_detach, - terminal_published: false, - }, - ); - if manual_detach { - jobs.background_started = jobs.background_started.wrapping_add(1); - } - self.changed(&mut jobs); + let call_id = agentkit_core::ToolCallId::new(call_id); + let controller = CancellationController::new(); + let mut jobs = self.lock_jobs(); + if jobs.running.contains_key(&call_id) { + return; + } + let cancelled = jobs.pending_cancellations.remove(&call_id) || jobs.cancel_all; + let manual_detach = jobs.pending_detaches.remove(&call_id); + jobs.running.insert( + call_id, + BackgroundJob { + registration: Arc::new(()), + controller: controller.clone(), + cancellation_requested: cancelled, + foreground_cancellation: None, + cancellation_relay: None, + detached: manual_detach, + manual_detach, + terminal_published: false, + }, + ); + if manual_detach { + jobs.background_started = jobs.background_started.wrapping_add(1); + } + let generation = jobs.changed(); + drop(jobs); + if cancelled { + controller.interrupt(); } + self.notify_changed(generation); } #[cfg(test)] @@ -1938,21 +2028,18 @@ impl BackgroundJobs { #[cfg(test)] pub(crate) fn is_cancelled_for_test(&self, call_id: &str) -> bool { - let call_id = agentkit_core::ToolCallId::new(call_id); - self.state.lock().is_ok_and(|jobs| { - jobs.running - .get(&call_id) - .is_some_and(|job| job.controller.handle().is_cancelled_since(0)) - }) + self.lock_jobs() + .running + .get(&agentkit_core::ToolCallId::new(call_id)) + .is_some_and(|job| job.controller.handle().is_cancelled_since(0)) } #[cfg(test)] pub(crate) fn is_detached_for_test(&self, call_id: &str) -> bool { let call_id = agentkit_core::ToolCallId::new(call_id); - self.state.lock().is_ok_and(|jobs| { - jobs.running.get(&call_id).is_some_and(|job| job.detached) - || jobs.pending_detaches.contains(&call_id) - }) + let jobs = self.lock_jobs(); + jobs.running.get(&call_id).is_some_and(|job| job.detached) + || jobs.pending_detaches.contains(&call_id) } } @@ -2049,7 +2136,7 @@ impl Tool for BackgroundableCompose { let artifact_directory = crate::artifacts::directory(&self.root, &request.session_id.0, &call_id.0); let request = Self::sanitized(request)?; - let _job = self.begin_background(background, &call_id, ctx); + let _job = self.begin_background(background, &call_id, ctx)?; match self.inner.invoke(request, ctx).await { Ok(mut result) => { match crate::compose_output::guard(&artifact_directory, result.result.output).await @@ -2078,7 +2165,10 @@ impl Tool for BackgroundableCompose { Ok(request) => request, Err(error) => return ToolExecutionOutcome::Failed(error), }; - let _job = self.begin_background(background, &call_id, ctx); + let _job = match self.begin_background(background, &call_id, ctx) { + Ok(job) => job, + Err(error) => return ToolExecutionOutcome::FailedBeforeInvocation(error), + }; match self.inner.invoke_outcome(request, ctx).await { ToolExecutionOutcome::Completed(mut result) => { match crate::compose_output::guard(&artifact_directory, result.result.output).await @@ -2101,67 +2191,83 @@ impl BackgroundableCompose { background: bool, call_id: &agentkit_core::ToolCallId, ctx: &mut ToolContext<'_>, - ) -> BackgroundJobGuard { + ) -> Result { let foreground_cancellation = (!background).then(|| ctx.cancellation.clone()).flatten(); + let registration = Arc::new(()); let controller = CancellationController::new(); let cancellation = controller.handle().checkpoint(); - ctx.cancellation = Some(cancellation.clone()); - if let Some(scope) = &mut ctx.execution_scope { - scope.cancellation = Some(cancellation); + let mut jobs = self.background_jobs.lock_jobs(); + if jobs.running.contains_key(call_id) || jobs.unacknowledged_terminals.contains(call_id) { + drop(jobs); + return Err(ToolError::InvalidInput( + "compose call ID is already registered".into(), + )); } - if let Ok(mut jobs) = self.background_jobs.state.lock() { - if jobs.pending_cancellations.remove(call_id) || jobs.cancel_all { - controller.interrupt(); - } - let manual_detach = jobs.pending_detaches.remove(call_id); - let detached = background || manual_detach; - if !detached + let manual_detach = jobs.pending_detaches.remove(call_id); + let detached = background || manual_detach; + let cancelled = jobs.pending_cancellations.remove(call_id) + || jobs.cancel_all + || (!detached && foreground_cancellation .as_ref() - .is_some_and(agentkit_core::TurnCancellation::is_cancelled) - { - controller.interrupt(); - } - jobs.running.insert( - call_id.clone(), - BackgroundJob { - controller, - foreground_cancellation: foreground_cancellation.clone(), - cancellation_relay: None, - detached, - manual_detach, - terminal_published: false, - }, - ); - if detached { - jobs.background_started = jobs.background_started.wrapping_add(1); - } - self.background_jobs.changed(&mut jobs); - } - { - let shutdown = crate::resilient_fs::shutdown_token().child_token(); - let jobs = self.background_jobs.clone(); - let relay_call_id = call_id.clone(); - let relay = tokio::spawn(async move { - tokio::select! { - _ = shutdown.cancelled() => { jobs.cancel_running(&relay_call_id.0); }, - _ = async { - if let Some(cancellation) = foreground_cancellation { cancellation.cancelled().await; } - else { std::future::pending::<()>().await; } - } => jobs.propagate_foreground_cancellation(&relay_call_id), - } - }) - .abort_handle(); - if let Ok(mut jobs) = self.background_jobs.state.lock() - && let Some(job) = jobs.running.get_mut(call_id) - { - job.cancellation_relay = Some(relay); - } - } - BackgroundJobGuard { + .is_some_and(agentkit_core::TurnCancellation::is_cancelled)); + jobs.running.insert( + call_id.clone(), + BackgroundJob { + registration: Arc::clone(®istration), + controller: controller.clone(), + cancellation_requested: cancelled, + foreground_cancellation: foreground_cancellation.clone(), + cancellation_relay: None, + detached, + manual_detach, + terminal_published: false, + }, + ); + if detached { + jobs.background_started = jobs.background_started.wrapping_add(1); + } + let generation = jobs.changed(); + drop(jobs); + // The guard must exist before notification, spawning, or context drops can + // unwind. It owns this unique registration until execution completes/cancels. + let guard = BackgroundJobGuard { jobs: self.background_jobs.clone(), call_id: call_id.clone(), + }; + if cancelled { + controller.interrupt(); + } + ctx.cancellation = Some(cancellation.clone()); + if let Some(scope) = &mut ctx.execution_scope { + scope.cancellation = Some(cancellation); } + self.background_jobs.notify_changed(generation); + let shutdown = crate::resilient_fs::shutdown_token().child_token(); + let relay_jobs = self.background_jobs.clone(); + let relay_call_id = call_id.clone(); + let relay = tokio::spawn(async move { + tokio::select! { + _ = shutdown.cancelled() => { relay_jobs.cancel_registration(&relay_call_id.0, Some(®istration)); }, + _ = async { + if let Some(cancellation) = foreground_cancellation { cancellation.cancelled().await; } + else { std::future::pending::<()>().await; } + } => relay_jobs.propagate_foreground_cancellation(&relay_call_id, ®istration), + } + }).abort_handle(); + let mut jobs = self.background_jobs.lock_jobs(); + // This guard is the sole removal owner and has not yet left this function. + let old_relay = jobs + .running + .get_mut(call_id) + .map(|job| job.cancellation_relay.replace(relay.clone())); + drop(jobs); + match old_relay { + Some(Some(old_relay)) => old_relay.abort(), + None => relay.abort(), + Some(None) => {} + } + Ok(guard) } } diff --git a/src/runtime/tests.rs b/src/runtime/tests.rs index 37dc1498..dbb8bc77 100644 --- a/src/runtime/tests.rs +++ b/src/runtime/tests.rs @@ -445,6 +445,57 @@ fn dropped_uncommitted_session_claim_retries_as_new() { assert!(!retry.request.resume); } +#[test] +fn rejected_creation_commit_recreates_configured_and_generated_claims_as_new() { + for configured in [true, false] { + let root = tempfile::tempdir().unwrap(); + let runtime = if configured { + Runtime::with_session( + root.path(), + "gpt-5.4", + SessionRequest { + id: "selected".into(), + resume: false, + force: false, + }, + ) + .unwrap() + } else { + Runtime::new(root.path(), "gpt-5.4").unwrap() + }; + let mut failed = runtime.claim_session().unwrap(); + let id = failed.id().to_owned(); + let opened = crate::session::open_uncommitted( + root.path(), + &id, + false, + vec![agentkit_core::Item::text(ItemKind::System, "system")], + ) + .unwrap(); + failed.guard_uncommitted_transcript(&opened.observer); + // Abandon normal response publication: this removes the uncommitted + // file and freezes the writer, so its later creation commit must fail. + drop(opened.observer.prepare_creation().unwrap()); + drop(opened); + assert!(failed.commit().is_err()); + assert!(crate::session::load(root.path(), &id).is_err()); + let mut retry = runtime.claim_session().unwrap(); + assert_eq!(retry.id(), id); + assert!(!retry.request.resume); + let recreated = crate::session::open_uncommitted( + root.path(), + retry.id(), + false, + vec![agentkit_core::Item::text(ItemKind::System, "recreated")], + ) + .unwrap(); + retry.guard_uncommitted_transcript(&recreated.observer); + retry.commit().unwrap(); + drop(recreated); + assert!(crate::session::load(root.path(), &id).is_ok()); + } +} + #[test] fn uncommitted_claim_rolls_back_before_waiting_to_publish_retry() { let root = tempfile::tempdir().unwrap(); @@ -1274,7 +1325,8 @@ async fn close_tool_can_cancel_a_detached_compose() { let job = compose .backgroundable - .begin_background(true, &call_id, &mut context); + .begin_background(true, &call_id, &mut context) + .unwrap(); let cancellation = context.cancellation.clone().expect("job cancellation"); assert!(!cancellation.is_cancelled()); assert_eq!( @@ -1334,7 +1386,8 @@ async fn foreground_compose_can_detach_from_turn_cancellation_and_still_be_kille let mut context = owned.borrowed(); let job = compose .backgroundable - .begin_background(false, &call_id, &mut context); + .begin_background(false, &call_id, &mut context) + .unwrap(); let cancellation = context.cancellation.clone().expect("compose cancellation"); assert_eq!( @@ -1352,11 +1405,10 @@ async fn foreground_compose_can_detach_from_turn_cancellation_and_still_be_kille let already_cancelled_id = ToolCallId::new("already-cancelled"); let mut already_cancelled_context = owned.borrowed(); - let already_cancelled_job = compose.backgroundable.begin_background( - false, - &already_cancelled_id, - &mut already_cancelled_context, - ); + let already_cancelled_job = compose + .backgroundable + .begin_background(false, &already_cancelled_id, &mut already_cancelled_context) + .unwrap(); assert!( already_cancelled_context .cancellation @@ -1375,11 +1427,10 @@ async fn foreground_compose_can_detach_from_turn_cancellation_and_still_be_kille ); let pending_detach_id = ToolCallId::new("pending-detach"); let mut pending_detach_context = owned.borrowed(); - let pending_detach_job = compose.backgroundable.begin_background( - false, - &pending_detach_id, - &mut pending_detach_context, - ); + let pending_detach_job = compose + .backgroundable + .begin_background(false, &pending_detach_id, &mut pending_detach_context) + .unwrap(); assert!( !pending_detach_context .cancellation @@ -1621,6 +1672,246 @@ fn cancel_all_covers_running_and_late_background_registration() { jobs.finish_for_test("next-turn"); } +fn background_job_test_context( + cancellation: Option, +) -> OwnedToolContext { + OwnedToolContext { + session_id: SessionId::new("background-test"), + turn_id: TurnId::new("turn"), + metadata: MetadataMap::new(), + permissions: Arc::new(AllowAllPermissions), + resources: Arc::new(()), + cancellation, + execution_scope: None, + approved_request: None, + } +} + +#[test] +fn background_jobs_poison_preserves_activity_cancellation_and_terminal_debt() { + let jobs = BackgroundJobs::default(); + jobs.register_foreground_for_test("active"); + jobs.register_foreground_for_test("finished"); + jobs.detach("finished").unwrap(); + jobs.finish_for_test("finished"); + let before = jobs.activity(); + assert!( + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _state = jobs.state.lock().unwrap(); + panic!("poison intact background bookkeeping"); + })) + .is_err() + ); + assert_eq!(jobs.activity(), before); + assert!(jobs.activity().active); + assert!(jobs.activity().unacknowledged_terminals); + assert!(!jobs.state.is_poisoned()); + jobs.cancel_all(); + assert!(jobs.is_cancelled_for_test("active")); + jobs.register_foreground_for_test("late"); + assert!(jobs.is_cancelled_for_test("late")); + jobs.acknowledge_terminal(&ToolCallId::new("finished")); + jobs.finish_for_test("active"); + jobs.finish_for_test("late"); + assert!(!jobs.activity().active); + assert!(!jobs.activity().unacknowledged_terminals); +} + +#[tokio::test] +async fn background_registration_waker_unwind_cleans_up_without_poison() { + use std::{ + future::Future, + sync::atomic::{AtomicBool, Ordering}, + task::{Context, Wake, Waker}, + }; + struct ReenterAndPanic(BackgroundJobs, AtomicBool); + impl Wake for ReenterAndPanic { + fn wake(self: Arc) { + assert!( + self.0.state.try_lock().is_ok(), + "activity wake held jobs lock" + ); + let _ = self.0.activity(); + if !self.1.swap(true, Ordering::SeqCst) { + panic!("activity waker"); + } + } + } + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); + let compose = runtime.compose(0); + let jobs = &compose.backgroundable.background_jobs; + let wake = Arc::new(ReenterAndPanic(jobs.clone(), AtomicBool::new(false))); + let waker = Waker::from(Arc::clone(&wake)); + let mut waiting = Box::pin(jobs.activity_after(jobs.activity().generation)); + assert!( + waiting + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + let owned = background_job_test_context(None); + let mut context = owned.borrowed(); + let call_id = ToolCallId::new("wake-unwind"); + assert!( + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _job = compose + .backgroundable + .begin_background(true, &call_id, &mut context) + .unwrap(); + })) + .is_err() + ); + assert!(wake.1.load(Ordering::SeqCst)); + assert!(!jobs.state.is_poisoned()); + assert!( + !jobs.activity().active, + "registration must establish cleanup before waking" + ); + assert!(jobs.activity().unacknowledged_terminals); + jobs.acknowledge_terminal(&call_id); + assert!(!jobs.activity().unacknowledged_terminals); + assert!( + waiting + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_ready() + ); + let guard = compose + .backgroundable + .begin_background(false, &call_id, &mut context) + .unwrap(); + assert!(jobs.activity().active); + drop(guard); + assert!(!jobs.activity().active); +} + +#[tokio::test] +async fn background_finish_and_acknowledgement_wake_with_committed_state() { + use std::{ + future::Future, + sync::atomic::{AtomicBool, Ordering}, + task::{Context, Wake, Waker}, + }; + struct ObserveDebt(BackgroundJobs, AtomicBool); + impl Wake for ObserveDebt { + fn wake(self: Arc) { + assert!(self.0.state.try_lock().is_ok()); + self.1 + .store(self.0.activity().unacknowledged_terminals, Ordering::SeqCst); + } + } + let jobs = BackgroundJobs::default(); + jobs.register_foreground_for_test("terminal"); + jobs.detach("terminal").unwrap(); + let observe = Arc::new(ObserveDebt(jobs.clone(), AtomicBool::new(false))); + let waker = Waker::from(Arc::clone(&observe)); + let mut context = Context::from_waker(&waker); + let mut waiting = Box::pin(jobs.activity_after(jobs.activity().generation)); + assert!(waiting.as_mut().poll(&mut context).is_pending()); + jobs.finish_for_test("terminal"); + assert!(observe.1.load(Ordering::SeqCst)); + assert!(waiting.as_mut().poll(&mut context).is_ready()); + let mut waiting = Box::pin(jobs.activity_after(jobs.activity().generation)); + assert!(waiting.as_mut().poll(&mut context).is_pending()); + jobs.acknowledge_terminal(&ToolCallId::new("terminal")); + assert!(!observe.1.load(Ordering::SeqCst)); + assert!(waiting.as_mut().poll(&mut context).is_ready()); + assert!(!jobs.activity().active); + assert!(!jobs.state.is_poisoned()); +} + +#[tokio::test] +async fn background_cancellation_rejects_detach_and_stale_relays_ignore_reused_ids() { + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); + let compose = runtime.compose(0); + let jobs = &compose.backgroundable.background_jobs; + let parent = CancellationController::new(); + let owned = background_job_test_context(Some(parent.handle().checkpoint())); + let mut context = owned.borrowed(); + let call_id = ToolCallId::new("relay"); + let guard = compose + .backgroundable + .begin_background(false, &call_id, &mut context) + .unwrap(); + let old_registration = jobs.lock_jobs().running[&call_id].registration.clone(); + let cancellation = context.cancellation.clone().unwrap(); + parent.interrupt(); + tokio::time::timeout(Duration::from_secs(1), cancellation.cancelled()) + .await + .unwrap(); + assert_eq!(jobs.detach(&call_id.0), None); + assert!(jobs.cancel(&call_id.0)); + assert!(jobs.cancel(&call_id.0)); + assert_eq!( + cancellation.handle().generation(), + 1, + "signal a job only once" + ); + drop(guard); + let owned = background_job_test_context(None); + let mut context = owned.borrowed(); + let guard = compose + .backgroundable + .begin_background(false, &call_id, &mut context) + .unwrap(); + let cancellation = context.cancellation.clone().unwrap(); + // The old relay can race abort after finish unlocks; exercise its real handlers. + assert!(!jobs.cancel_registration(&call_id.0, Some(&old_registration))); + jobs.propagate_foreground_cancellation(&call_id, &old_registration); + assert!(!cancellation.is_cancelled()); + assert!(jobs.cancel(&call_id.0)); + assert!(cancellation.is_cancelled()); + drop(guard); +} + +#[tokio::test] +async fn background_registration_rejection_preserves_owner_and_unwind_cleanup() { + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); + let compose = runtime.compose(0); + let jobs = &compose.backgroundable.background_jobs; + let owned = background_job_test_context(None); + let call_id = ToolCallId::new("duplicate"); + let guard = compose + .backgroundable + .begin_background(false, &call_id, &mut owned.borrowed()) + .unwrap(); + let before = jobs.activity(); + assert!( + compose + .backgroundable + .begin_background(false, &call_id, &mut owned.borrowed()) + .is_err() + ); + assert_eq!(jobs.activity(), before); + assert!( + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _guard = guard; + let _state = jobs.state.lock().unwrap(); + panic!("execution unwind with poisoned registry"); + })) + .is_err() + ); + assert!(!jobs.state.is_poisoned()); + assert!(!jobs.activity().active); + let guard = compose + .backgroundable + .begin_background(true, &call_id, &mut owned.borrowed()) + .unwrap(); + drop(guard); + assert!(jobs.activity().unacknowledged_terminals); + assert!( + compose + .backgroundable + .begin_background(true, &call_id, &mut owned.borrowed()) + .is_err() + ); + jobs.acknowledge_terminal(&call_id); + assert!(!jobs.activity().unacknowledged_terminals); +} + #[tokio::test] async fn persistent_startup_failure_does_not_commit_new_session() { let root = tempfile::tempdir().unwrap(); diff --git a/src/session.rs b/src/session.rs index aff39381..a749a7ca 100644 --- a/src/session.rs +++ b/src/session.rs @@ -10,7 +10,7 @@ use std::{ path::{Path, PathBuf}, sync::{ Arc, Mutex, - atomic::{AtomicU64, Ordering}, + atomic::{AtomicBool, AtomicU64, Ordering}, }, time::{SystemTime, UNIX_EPOCH}, }; @@ -79,17 +79,26 @@ impl CatalogEntry { } } +// The writer serializes generation, open-file identity, and its ownership lease. +// No guard crosses an await. Backend I/O and replaced file/lease destruction can +// unwind after external effects, so poison is isolation, never recovery. The +// infallible observer refuses the mutation; result-bearing callers get an error. #[derive(Clone)] pub struct SessionObserver(Arc>); struct Writer { session_id: String, generation: u64, + // An uncertain append/reconstruction must never reuse a generation. + write_failed: bool, path: PathBuf, workspace_root: PathBuf, file: File, lock: SessionLock, created: Option, + // A prepared response privately owns cleanup. Writers remain frozen until + // that owner publishes success, or forever if submission is abandoned. + publication: Option>, } struct SessionLock { @@ -98,12 +107,31 @@ struct SessionLock { } /// Removes an incompletely bootstrapped new transcript unless opening commits. +/// The guarded filesystem retains the lease through cleanup, independently of +/// Writer's field drop order, and cannot delete a replacement owner's file. struct CreatedTranscript { filesystem: Fs, path: PathBuf, keep: bool, } +/// Exclusive ownership of a validated, not-yet-published creation. No writer +/// mutation is admitted while this exists. Dropping it before commit retains +/// the original cleanup behavior, including failed response delivery/unwind. +pub(crate) struct PreparedCreation { + created: CreatedTranscript, + published: Arc, +} + +impl PreparedCreation { + /// Called only after response submission. Both steps are infallible and + /// require no shared lock or callback; readers resume after cleanup is kept. + pub(crate) fn commit(mut self) { + self.created.keep = true; + self.published.store(true, Ordering::Release); + } +} + #[derive(Clone, Copy)] struct InitialTranscriptOptions { stamp_items: bool, @@ -402,11 +430,13 @@ fn open_with_initial_timestamps_in( let mut writer = Writer { session_id: session_id.into(), generation, + write_failed: false, path, workspace_root, file, lock, created, + publication: None, }; if resume && stored_workspace.is_none() { writer.replace(&transcript)?; @@ -435,7 +465,7 @@ fn open_with_initial_timestamps_in( writer.append(&item)?; } if initial_options.commit_creation { - writer.commit_creation(); + writer.commit_creation()?; } Ok(OpenSession { transcript, @@ -444,11 +474,31 @@ fn open_with_initial_timestamps_in( } impl SessionObserver { - pub(crate) fn commit_creation(&self) { + pub(crate) fn prepare_creation(&self) -> Result { + let published = Arc::new(AtomicBool::new(false)); + let mut writer = self + .0 + .lock() + .map_err(|_| "session transcript writer poisoned".to_string())?; + // Validate/reconstruct before extracting cleanup. From this point until + // commit, all shared writer APIs reject mutation without touching I/O. + writer.ensure_lock()?; + let created = writer + .created + .take() + .ok_or_else(|| "session creation is not available for publication".to_string())?; + writer.publication = Some(Arc::clone(&published)); + Ok(PreparedCreation { created, published }) + } + + pub(crate) fn commit_creation(&self) -> Result<(), String> { + // Creation publication cannot bless an uncertain transcript. All + // result-bearing APIs isolate poison; the observer trait below has no + // error channel and must instead refuse an unpersisted mutation. self.0 .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .commit_creation(); + .map_err(|_| "session transcript writer poisoned".to_string())? + .commit_creation() } /// Durably records a complete transcript replacement produced by a mutator. @@ -467,8 +517,14 @@ impl SessionObserver { impl TranscriptObserver for SessionObserver { fn on_transcript_event(&self, event: TranscriptEvent<'_>) { - let mut writer = self.0.lock().expect("session transcript writer poisoned"); - if let Err(error) = writer.append(event.item) { + // Release the writer before shutdown/logging/panic paths. A rejected + // input does not poison a healthy writer; uncertain I/O fences it below. + let result = self + .0 + .lock() + .expect("session transcript writer poisoned") + .append(event.item); + if let Err(error) = result { if fs::global().status().exhausted || fs::shutdown_token().is_cancelled() { fs::request_shutdown(); return; @@ -482,10 +538,21 @@ impl TranscriptObserver for SessionObserver { } impl Writer { - fn commit_creation(&mut self) { + fn commit_creation(&mut self) -> Result<(), String> { + if self + .publication + .as_ref() + .is_some_and(|published| !published.load(Ordering::Acquire)) + { + return Err("session creation is awaiting publication".into()); + } + if self.write_failed { + return Err("session transcript writer is isolated after an uncertain write".into()); + } if let Some(created) = self.created.take() { created.keep(); } + Ok(()) } fn append(&mut self, item: &Item) -> Result<(), String> { @@ -537,6 +604,10 @@ impl Writer { let mut encoded = serde_json::to_vec(&record) .map_err(|error| format!("could not encode transcript record: {error}"))?; encoded.push(b'\n'); + // write_all or sync_data can fail after accepting some or all bytes. + // Without a verified rollback, neither retry nor generation reuse is + // safe. Only a complete successful append re-enables this writer. + self.write_failed = true; self.file .write_all(&encoded) .and_then(|_| self.file.sync_data()) @@ -547,10 +618,21 @@ impl Writer { format!("could not persist transcript record: {error}") })?; self.generation = generation; + self.write_failed = false; Ok(()) } fn ensure_lock(&mut self) -> Result<(), String> { + if self + .publication + .as_ref() + .is_some_and(|published| !published.load(Ordering::Acquire)) + { + return Err("session creation is awaiting publication".into()); + } + if self.write_failed { + return Err("session transcript writer is isolated after an uncertain write".into()); + } match self.lock.check() { Ok(()) => { if fs::try_exists(&self.path).map_err(|error| { @@ -594,6 +676,9 @@ impl Writer { .create_new(true) .open_in(&self.lock.filesystem()?, &self.path) .map_err(|error| format!("could not reconstruct {}: {error}", self.path.display()))?; + // Cleanup can itself fail. Do not later mistake an incomplete rebuilt + // path for the complete history still held by the original open file. + self.write_failed = true; if let Err(error) = io::copy(&mut source, &mut file).and_then(|_| file.sync_all()) { let _ = self.lock.filesystem()?.remove_file(&self.path); return Err(format!( @@ -602,6 +687,7 @@ impl Writer { )); } self.file = file; + self.write_failed = false; Ok(()) } } @@ -3925,4 +4011,189 @@ mod tests { assert!(!transcript_path(root.path(), "abc").exists()); drop(other); } + + #[test] + fn rejected_observer_input_does_not_poison_or_advance_generation() { + let root = tempfile::tempdir().unwrap(); + let opened = open( + root.path(), + "abc", + false, + false, + vec![Item::text(ItemKind::System, "system")], + ) + .unwrap(); + let id = agentkit_core::SessionId::new("abc"); + assert!( + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + opened.observer.on_transcript_event(TranscriptEvent { + session_id: &id, + item: &Item::text(ItemKind::User, "unstamped"), + }); + })) + .is_err() + ); + assert!(!opened.observer.0.is_poisoned()); + assert_eq!(opened.observer.0.lock().unwrap().generation, 1); + let item = Item::text(ItemKind::User, "accepted").with_created_at(Timestamp(123)); + opened.observer.on_transcript_event(TranscriptEvent { + session_id: &id, + item: &item, + }); + assert_eq!(opened.observer.0.lock().unwrap().generation, 2); + assert_eq!(load(root.path(), "abc").unwrap().len(), 2); + } + + #[test] + fn failed_write_isolates_replacement_append_and_creation_commit() { + let root = tempfile::tempdir().unwrap(); + let opened = open( + root.path(), + "abc", + false, + false, + vec![Item::text(ItemKind::System, "system")], + ) + .unwrap(); + let path = { + let mut writer = opened.observer.0.lock().unwrap(); + // A real file boundary rejects writes, without production hooks. + writer.file = File::open(&writer.path).unwrap(); + writer.path.clone() + }; + let original = fs::read(&path).unwrap(); + let item = Item::text(ItemKind::User, "not accepted").with_created_at(Timestamp(123)); + assert!( + opened + .observer + .replace(std::slice::from_ref(&item)) + .is_err() + ); + assert!(!opened.observer.0.is_poisoned()); + // Even repairing the handle cannot prove what a failed write accepted. + opened.observer.0.lock().unwrap().file = + OpenOptions::new().append(true).open(&path).unwrap(); + let error = opened + .observer + .replace(std::slice::from_ref(&item)) + .unwrap_err(); + assert!(error.contains("isolated"), "{error}"); + assert!( + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + opened.observer.on_transcript_event(TranscriptEvent { + session_id: &agentkit_core::SessionId::new("abc"), + item: &item, + }); + })) + .is_err() + ); + assert!(opened.observer.commit_creation().is_err()); + assert!(opened.observer.prepare_creation().is_err()); + assert_eq!(opened.observer.0.lock().unwrap().generation, 1); + assert_eq!(fs::read(&path).unwrap(), original); + } + + #[test] + fn poisoned_writer_cannot_publish_creation_or_replace() { + let root = tempfile::tempdir().unwrap(); + let opened = open_with_initial_timestamps_in( + &project_root(root.path()), + &session_directory(root.path()), + "abc", + false, + false, + vec![Item::text(ItemKind::System, "system")], + InitialTranscriptOptions { + stamp_items: true, + commit_creation: false, + }, + ) + .unwrap(); + let path = opened.observer.0.lock().unwrap().path.clone(); + let observer = opened.observer.clone(); + assert!( + std::thread::spawn(move || { + let _writer = observer.0.lock().unwrap(); + panic!("writer interrupted"); + }) + .join() + .is_err() + ); + assert!( + opened + .observer + .replace(&[Item::text(ItemKind::User, "replacement")]) + .unwrap_err() + .contains("poisoned") + ); + assert!(opened.observer.commit_creation().is_err()); + assert!(opened.observer.prepare_creation().is_err()); + drop(opened); + assert!(!path.exists(), "failed creation must still roll back"); + } + + #[test] + fn prepared_creation_freezes_writers_until_private_commit_or_cleanup() { + for commit in [true, false] { + let root = tempfile::tempdir().unwrap(); + let opened = open_with_initial_timestamps_in( + &project_root(root.path()), + &session_directory(root.path()), + "abc", + false, + false, + vec![Item::text(ItemKind::System, "system")], + InitialTranscriptOptions { + stamp_items: true, + commit_creation: false, + }, + ) + .unwrap(); + let path = opened.observer.0.lock().unwrap().path.clone(); + let original = fs::read(&path).unwrap(); + let prepared = opened.observer.prepare_creation().unwrap(); + let item = Item::text(ItemKind::User, "later").with_created_at(Timestamp(123)); + assert!(opened.observer.prepare_creation().is_err()); + assert!(opened.observer.commit_creation().is_err()); + assert!( + opened + .observer + .replace(std::slice::from_ref(&item)) + .is_err() + ); + assert!( + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + opened.observer.on_transcript_event(TranscriptEvent { + session_id: &agentkit_core::SessionId::new("abc"), + item: &item, + }); + })) + .is_err() + ); + assert!(!opened.observer.0.is_poisoned()); + assert_eq!(fs::read(&path).unwrap(), original); + if commit { + prepared.commit(); + opened + .observer + .replace(std::slice::from_ref(&item)) + .unwrap(); + drop(opened); + assert!(path.exists()); + } else { + drop(prepared); + assert!(!path.exists()); + assert!( + opened + .observer + .replace(std::slice::from_ref(&item)) + .is_err() + ); + assert!( + !path.exists(), + "abandoned publication must never reconstruct" + ); + } + } + } } diff --git a/src/tools/mcp.rs b/src/tools/mcp.rs index 1d8f5569..89ac2ec5 100644 --- a/src/tools/mcp.rs +++ b/src/tools/mcp.rs @@ -1372,11 +1372,6 @@ impl McpRuntime { .filter(|name| current_entries.get(*name) != entries.get(*name)) .cloned() .collect::>(); - let deleted = current_entries - .keys() - .filter(|name| !entries.contains_key(*name)) - .cloned() - .collect::>(); let gates = if changed.is_empty() { Vec::new() @@ -1470,13 +1465,16 @@ impl McpRuntime { } drop(generation_writer); drop(initialization_guard); - if !deleted.is_empty() { - let mut operations = self.inner.operations.lock().await; - for name in deleted { - operations.remove(&name); - } - } drop(operation_guards); + drop(gates); + // A deleted name can still have callers holding or waiting on its + // gate. Keep that identity until its last owner leaves, including + // across deletion and re-addition of the same server name. + self.inner + .operations + .lock() + .await + .retain(|_, gate| gate.strong_count() > 0); return Ok(()); } } @@ -1788,11 +1786,43 @@ impl McpRuntime { ) .await .map_err(ToolError::Unavailable)?; + self.start_authorization( + name.to_string(), + record.fingerprint, + request, + pending, + session_id, + ) + .await + } + + async fn start_authorization( + &self, + name: String, + fingerprint: Vec, + request: AuthRequest, + pending: auth::PendingAuthorization, + session_id: String, + ) -> Result { + // The caller retains the server operation gate and auth_setup guard. + // Acquire both publication guards before changing either registry. No + // writer holds servers while waiting for pending (reload releases its + // servers guard first). Cancellation at either await only drops the + // uninstalled OAuth listener; it needs no generation-sensitive rollback. + let mut registrations = self.inner.pending.lock().await; + let mut servers = self.inner.servers.write().await; + let record = servers + .get_mut(&name) + .filter(|record| record.fingerprint == fingerprint) + .ok_or_else(|| { + ToolError::Unavailable( + "MCP configuration changed during authentication setup; retry auth".into(), + ) + })?; let url = pending.url.clone(); - self.set_status(name, ServerStatus::Pending).await; let runtime = self.clone(); - let server = name.to_string(); - let fingerprint = record.fingerprint.clone(); + let server = name.clone(); + let task_fingerprint = fingerprint.clone(); let event_generation = self.event_generation(&session_id); let (start, started) = oneshot::channel(); let task = tokio::spawn(async move { @@ -1800,7 +1830,7 @@ impl McpRuntime { runtime .complete_authorization( server, - fingerprint, + task_fingerprint, request, pending, session_id, @@ -1809,16 +1839,21 @@ impl McpRuntime { .await; } }); - self.inner.pending.lock().await.insert( - name.to_string(), + registrations.insert( + name.clone(), PendingRecord { url: url.clone(), expires: Instant::now() + auth::FLOW_TIMEOUT, - fingerprint: record.fingerprint, + fingerprint, abort: task.abort_handle(), }, ); + record.status = ServerStatus::Pending; + // Publish the registered worker and its status without another await. + // The start barrier prevents completion before registration is visible. let _ = start.send(()); + drop(servers); + drop(registrations); Ok(json!({ "server": name, "status": "pending", @@ -3094,6 +3129,174 @@ mod tests { ) } + async fn pending_authorization() -> super::auth::PendingAuthorization { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let server_base = base.clone(); + let server = tokio::spawn(async move { + for _ in 0..2 { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + let read = stream.read(&mut request).await.unwrap(); + let request = std::str::from_utf8(&request[..read]).unwrap(); + let path = request.split_whitespace().nth(1).unwrap(); + let body = if path == "/resource-metadata" { + json!({ + "resource": format!("{server_base}/mcp"), + "authorization_servers": [server_base] + }) + } else { + json!({ + "issuer": server_base, + "authorization_endpoint": format!("{server_base}/authorize"), + "token_endpoint": format!("{server_base}/token"), + "response_types_supported": ["code"], + "code_challenge_methods_supported": ["S256"] + }) + } + .to_string(); + stream.write_all(format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ).as_bytes()).await.unwrap(); + } + }); + let config = serde_json::from_value(json!({ + "type": "oauth", "clientId": "kit-client" + })) + .unwrap(); + let challenge = format!(r#"Bearer resource_metadata="{base}/resource-metadata""#); + let pending = super::auth::begin( + &format!("{base}/mcp"), + &config, + &CredentialStorage::Memory, + None, + Some(&challenge), + ) + .await + .unwrap(); + server.await.unwrap(); + pending + } + + #[tokio::test] + async fn cancelled_authorization_registration_leaves_no_pending_status_or_listener() { + for block_pending in [true, false] { + let runtime = super::empty(); + let mut record = connected_oauth_record(); + record.status = ServerStatus::AuthenticationRequired; + let fingerprint = record.fingerprint.clone(); + runtime + .inner + .servers + .write() + .await + .insert("remote".into(), record); + let pending = pending_authorization().await; + let url = url::Url::parse(&pending.url).unwrap(); + let callback = url + .query_pairs() + .find(|(key, _)| key == "redirect_uri") + .unwrap() + .1; + let callback = url::Url::parse(&callback).unwrap(); + let pending_guard = if block_pending { + Some(runtime.inner.pending.lock().await) + } else { + None + }; + let servers_guard = if block_pending { + None + } else { + Some(runtime.inner.servers.write().await) + }; + let mut registration = Box::pin(runtime.start_authorization( + "remote".into(), + fingerprint.clone(), + tool_auth_request(), + pending, + "session".into(), + )); + // Poll at the real registry boundary, then cancel the caller while + // the selected lock is unavailable. No sleeps or production hooks. + assert!(futures_util::poll!(registration.as_mut()).is_pending()); + drop(registration); + drop(servers_guard); + drop(pending_guard); + let records = runtime.inner.servers.read().await; + assert!(matches!( + records["remote"].status, + ServerStatus::AuthenticationRequired + )); + assert_eq!(records["remote"].fingerprint, fingerprint); + drop(records); + assert!(runtime.inner.pending.lock().await.is_empty()); + assert!( + tokio::net::TcpStream::connect(( + callback.host_str().unwrap(), + callback.port().unwrap() + )) + .await + .is_err() + ); + } + } + + #[tokio::test] + async fn authorization_registration_rejects_stale_generation_before_publication() { + let runtime = super::empty(); + let mut record = connected_oauth_record(); + let stale = record.fingerprint.clone(); + record.fingerprint.push(1); + let current = record.fingerprint.clone(); + runtime + .inner + .servers + .write() + .await + .insert("remote".into(), record); + let result = runtime + .start_authorization( + "remote".into(), + stale, + tool_auth_request(), + pending_authorization().await, + "session".into(), + ) + .await; + assert!( + result + .unwrap_err() + .to_string() + .contains("configuration changed") + ); + assert!(runtime.inner.pending.lock().await.is_empty()); + assert!(matches!( + runtime.inner.servers.read().await["remote"].status, + ServerStatus::Connected + )); + + let result = runtime + .start_authorization( + "remote".into(), + current.clone(), + tool_auth_request(), + pending_authorization().await, + "session".into(), + ) + .await + .unwrap(); + assert_eq!(result["status"], "pending"); + let registrations = runtime.inner.pending.lock().await; + assert_eq!(registrations["remote"].fingerprint, current); + assert_eq!(registrations["remote"].url, result["url"].as_str().unwrap()); + assert!(matches!( + runtime.inner.servers.read().await["remote"].status, + ServerStatus::Pending + )); + registrations["remote"].abort.abort(); + } + #[test] fn mcp_tool_timeout_schema_is_optional_and_has_no_schema_default() { let tool = McpTool::new(super::empty()); @@ -4521,6 +4724,33 @@ mod tests { assert_eq!(record.url.as_deref(), Some("https://example.com/mcp")); } + #[tokio::test] + async fn removed_server_operation_gate_remains_shared_with_existing_owners() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("mcp.json"); + let config = r#"{"mcpServers":{"removed":{"command":"unused"}}}"#; + std::fs::write(&path, config).unwrap(); + let runtime = super::connect(Some(&path), &[], true, CredentialStorage::Memory) + .await + .unwrap(); + // Model a caller that fetched the gate but has not acquired it yet. + let existing = runtime.operation_gate("removed").await; + + std::fs::write(&path, r#"{"mcpServers":{}}"#).unwrap(); + runtime.reload_config().await.unwrap(); + assert!(Arc::ptr_eq( + &existing, + &runtime.operation_gate("removed").await + )); + + std::fs::write(&path, config).unwrap(); + runtime.reload_config().await.unwrap(); + let replacement = runtime.operation_gate("removed").await; + assert!(Arc::ptr_eq(&existing, &replacement)); + let _held = existing.read().await; + assert!(replacement.try_write().is_err()); + } + #[tokio::test] async fn removed_server_operation_gate_is_reclaimed() { let directory = tempfile::tempdir().unwrap(); From dccd0b371f700b89008d5ac7f61f2d0088716243 Mon Sep 17 00:00:00 2001 From: daniel Date: Sat, 5 Sep 2026 22:34:22 +0100 Subject: [PATCH 2/3] fix(acp): address shared-state review feedback --- AGENTS.md | 4 - src/protocols/acp.rs | 175 ++++++++++++++++++++++++++-------- src/protocols/acp/activity.rs | 52 ++++++---- src/protocols/acp/v2.rs | 89 ++++++++++++++++- 4 files changed, 257 insertions(+), 63 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 19f7e62e..84f8c138 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -3,7 +3,3 @@ Only when you are using Kit as your agent harness: report issues with Kit's harness at https://github.com/speakeasy-api/kit/issues. Do not report issues with other harnesses there. Do not open an issue on the user's behalf unless the user explicitly requests it; ask the user first when they have not already made that request. Follow [Reporting Kit Issues](docs/user/reporting-kit-issues.md). Do not change release versions in ordinary pull requests. Use a Conventional Commit title for every pull request. Mark a breaking change with `!` after the commit type or scope, or with a `BREAKING CHANGE:` line in the commit body. The release workflow derives the next version from commits since the latest release, advances the version files in a release commit, and applies a minor bump when any commit is breaking or a patch bump otherwise. - -## Shared state - -You MUST load and follow `.agents/skills/shared-state/SKILL.md` whenever adding or editing lock-protected shared state, guarded transitions, lock acquisition, or recovery. diff --git a/src/protocols/acp.rs b/src/protocols/acp.rs index 51cccf3a..29eb71f1 100644 --- a/src/protocols/acp.rs +++ b/src/protocols/acp.rs @@ -606,13 +606,13 @@ impl SessionRegistry { } } - pub(super) fn next_token(&self) -> u64 { + pub(super) fn next_token(&self) -> Result { self.inner .next_token .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |token| { token.checked_add(1) }) - .expect("ACP session registry token space exhausted") + .map_err(|_| AcpRuntimeError::Loop("ACP session registry token space exhausted".into())) } fn begin_attachment(&self) -> Result { @@ -1387,6 +1387,8 @@ impl Server { mut claim: crate::runtime::SessionClaim, forked: Option, ) -> Result { + // Reject exhaustion before admission or any binding, driver, or actor effects. + let token = self.registry.next_token()?; let mut admission = self .registry .begin_attachment() @@ -1490,7 +1492,6 @@ impl Server { activity, mcp_events, }; - let token = self.registry.next_token(); let (activation, activated) = oneshot::channel(); let (completed, completion) = watch::channel(false); let guard = SessionActorGuard { @@ -3043,7 +3044,7 @@ pub(super) mod tests { mpsc::Receiver, ) { let session_id = agentkit_acp::SessionId::new(id.to_owned()); - let token = server.registry.next_token(); + let token = server.registry.next_token().unwrap(); let (completed, completion) = watch::channel(false); let guard = SessionActorGuard { server: Arc::downgrade(server), @@ -3369,7 +3370,7 @@ pub(super) mod tests { server: &Arc, registry: &SessionRegistry, ) -> (mpsc::Sender, Arc) { - let token = registry.next_token(); + let token = registry.next_token().unwrap(); let (commands, mut received) = mpsc::channel(1); let (completed, completion) = watch::channel(false); let closed = Arc::new(AtomicBool::new(false)); @@ -3433,7 +3434,7 @@ pub(super) mod tests { let integration = shutdown_test_integration(); let (background_jobs, tasks) = start_non_cooperative_background("shutdown-call").await; let session_id = agentkit_acp::SessionId::new("close-me"); - let token = registry.next_token(); + let token = registry.next_token().unwrap(); let (commands, mut received) = mpsc::channel(1); let weak_commands = commands.downgrade(); let (completed, completion) = watch::channel(false); @@ -3480,7 +3481,7 @@ pub(super) mod tests { tokio::task::yield_now().await; } - let late_token = registry.next_token(); + let late_token = registry.next_token().unwrap(); let (late_commands, _late_received) = mpsc::channel(1); let (late_completed, late_completion) = watch::channel(false); let late_actor = tokio::spawn(async move { @@ -3537,13 +3538,13 @@ pub(super) mod tests { let registry = SessionRegistry::new(); let other = SessionRegistry::new(); let mut admission = registry.begin_attachment().unwrap(); - let token = registry.next_token(); + let token = registry.next_token().unwrap(); assert!(register_test_v2(&other, &mut admission, token, Arc::new(|| {})).is_err()); assert_eq!(registry.inner.lock_state().pending_attachments, 1); assert_eq!(other.inner.lock_state().pending_attachments, 0); assert!(other.inner.lock_state().v2_sessions.is_empty()); register_test_v2(®istry, &mut admission, token, Arc::new(|| {})).unwrap(); - let second_token = registry.next_token(); + let second_token = registry.next_token().unwrap(); assert!( register_test_v2(®istry, &mut admission, second_token, Arc::new(|| {})).is_err() ); @@ -3568,7 +3569,7 @@ pub(super) mod tests { async fn registry_legacy_registration_obeys_admission_and_cross_version_token_rules() { let registry = SessionRegistry::new(); let other = SessionRegistry::new(); - let token = registry.next_token(); + let token = registry.next_token().unwrap(); register_test_v2( ®istry, &mut registry.begin_attachment().unwrap(), @@ -3595,7 +3596,7 @@ pub(super) mod tests { assert!(other.register(&mut admission, session.clone()).is_err()); assert!(admission.active); assert_eq!(registry.inner.lock_state().pending_attachments, 1); - session.token = registry.next_token(); + session.token = registry.next_token().unwrap(); registry.register(&mut admission, session.clone()).unwrap(); assert!(registry.register(&mut admission, session.clone()).is_err()); assert_eq!(registry.inner.lock_state().pending_attachments, 0); @@ -3622,7 +3623,7 @@ pub(super) mod tests { let mut context = std::task::Context::from_waker(&waker); let mut waiting = Box::pin(registry.wait_for_pending_attachments()); assert!(waiting.as_mut().poll(&mut context).is_pending()); - let token = registry.next_token(); + let token = registry.next_token().unwrap(); register_test_v2(®istry, &mut admission, token, Arc::new(|| {})).unwrap(); assert!(wake.1.load(Ordering::SeqCst)); assert!(waiting.as_mut().poll(&mut context).is_ready()); @@ -3633,7 +3634,7 @@ pub(super) mod tests { #[tokio::test] async fn registry_interrupt_unwind_leaves_gate_closed_without_poison() { let registry = SessionRegistry::new(); - let token = registry.next_token(); + let token = registry.next_token().unwrap(); let callback_registry = registry.clone(); register_test_v2( ®istry, @@ -3666,7 +3667,7 @@ pub(super) mod tests { } } let registry = SessionRegistry::new(); - let token = registry.next_token(); + let token = registry.next_token().unwrap(); let capture = ReenterOnDrop(registry.clone()); register_test_v2( ®istry, @@ -3711,7 +3712,7 @@ pub(super) mod tests { } } let registry = SessionRegistry::new(); - let token = registry.next_token(); + let token = registry.next_token().unwrap(); register_test_v2( ®istry, &mut registry.begin_attachment().unwrap(), @@ -3739,6 +3740,85 @@ pub(super) mod tests { assert!(!registry.reset_authentication().await); } + #[tokio::test] + async fn legacy_token_exhaustion_rejects_attachment_before_publication() { + let root = tempfile::tempdir().unwrap(); + let session_id = crate::session::new_id(); + let runtime = Runtime::with_session_provider_and_credentials( + root.path(), + "gpt-5.4", + crate::ProviderKind::OpenAiSubscription, + crate::runtime::SessionRequest { + id: session_id.clone(), + resume: false, + force: false, + }, + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(); + let registry = SessionRegistry::new(); + registry + .inner + .next_token + .store(u64::MAX, std::sync::atomic::Ordering::Relaxed); + let (client_transport, agent_transport) = agent_client_protocol::Channel::duplex(); + let router = component(runtime.clone(), registry.clone()).unwrap(); + let server = tokio::spawn(async move { router.connect_to(agent_transport).await }); + let workspace = root.path().to_path_buf(); + let client = agent_client_protocol::Client.builder().connect_with( + client_transport, + async move |connection| { + connection + .send_request(InitializeRequest::new(ProtocolVersion::V1)) + .block_task() + .await?; + for _ in 0..2 { + let error = connection + .send_request(NewSessionRequest::new(workspace.clone())) + .block_task() + .await + .expect_err("exhausted registry must reject attachment"); + assert!( + error + .data + .unwrap() + .to_string() + .contains("ACP session registry token space exhausted") + ); + assert_eq!( + registry + .inner + .next_token + .load(std::sync::atomic::Ordering::Relaxed), + u64::MAX + ); + assert!(!registry.inner.state.is_poisoned()); + { + let state = registry.inner.lock_state(); + assert!(state.accepting); + assert_eq!(state.pending_attachments, 0); + assert!(state.sessions.is_empty()); + assert!(state.v2_sessions.is_empty()); + } + // Failed attachment releases its claim without creating a transcript. + let claim = runtime.claim_session().unwrap(); + assert_eq!(claim.id(), session_id); + assert!(claim.is_configured()); + drop(claim); + assert_eq!( + crate::session::load(&workspace, &session_id).unwrap_err(), + format!("session {session_id:?} does not exist") + ); + } + Ok(()) + }, + ); + let result = timeout(Duration::from_secs(2), client).await; + server.abort(); + let _ = server.await; + result.expect("exhaustion client timed out").unwrap(); + } + #[tokio::test] async fn registry_counter_exhaustion_does_not_wrap_or_poison() { let registry = SessionRegistry::new(); @@ -3746,14 +3826,24 @@ pub(super) mod tests { .inner .next_token .store(u64::MAX - 1, Ordering::Relaxed); - assert_eq!(registry.next_token(), u64::MAX - 1); + assert_eq!(registry.next_token().unwrap(), u64::MAX - 1); for _ in 0..2 { - assert!( - std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| registry.next_token())) - .is_err() - ); + assert!(matches!( + registry.next_token(), + Err(AcpRuntimeError::Loop(message)) + if message == "ACP session registry token space exhausted" + )); assert_eq!(registry.inner.next_token.load(Ordering::Relaxed), u64::MAX); + assert!(!registry.inner.state.is_poisoned()); + let state = registry.inner.lock_state(); + assert!(state.accepting); + assert_eq!(state.pending_attachments, 0); + assert!(state.sessions.is_empty()); + assert!(state.v2_sessions.is_empty()); } + // Allocation failure does not corrupt admission or its drop cleanup. + drop(registry.begin_attachment().unwrap()); + assert_eq!(registry.inner.lock_state().pending_attachments, 0); registry.inner.lock_state().pending_attachments = usize::MAX; assert!(registry.begin_attachment().is_err()); assert!(!registry.inner.state.is_poisoned()); @@ -3822,7 +3912,7 @@ pub(super) mod tests { registry .register_v2( &mut registry.begin_attachment().unwrap(), - registry.next_token(), + registry.next_token().unwrap(), Arc::new(|| {}), Arc::new(|| Box::pin(async { panic!("close callback unwound") })), actor.abort_handle(), @@ -3846,7 +3936,7 @@ pub(super) mod tests { #[tokio::test] async fn authentication_reset_closes_shared_v2_sessions_and_reopens_registration() { let registry = SessionRegistry::new(); - let token = registry.next_token(); + let token = registry.next_token().unwrap(); let closed = Arc::new(AtomicBool::new(false)); let close_flag = Arc::clone(&closed); let (completed, completion) = watch::channel(false); @@ -3949,7 +4039,7 @@ pub(super) mod tests { #[tokio::test] async fn failed_authentication_reset_cannot_reopen_on_retry() { let registry = SessionRegistry::new(); - let token = registry.next_token(); + let token = registry.next_token().unwrap(); let (_completed, completion) = watch::channel(false); let actor = tokio::spawn(std::future::pending::<()>()); registry @@ -4011,7 +4101,7 @@ pub(super) mod tests { #[tokio::test] async fn shutdown_aborts_actor_at_the_shared_deadline() { let registry = SessionRegistry::new(); - let token = registry.next_token(); + let token = registry.next_token().unwrap(); let (commands, _received) = mpsc::channel(1); let (_completed, completion) = watch::channel(false); let actor = tokio::spawn(std::future::pending::<()>()); @@ -5117,24 +5207,33 @@ pub(super) mod tests { // Multiple logical turns drained within one autonomous interval do not // publish an intermediate terminal state or allocate another turn ID. - activity.begin(activity::ExecutionOrigin::Autonomous); - for text in ["first continuation", "second continuation"] { - driver.submit_input(vec![Item::notification(text)]).unwrap(); - drive_autonomous(&session_id, &integration, &mut driver) - .await - .unwrap(); - } - let started = states.try_recv().unwrap(); - assert!(started.active); - assert_eq!(started.turn_id, 3); - assert!(states.try_recv().is_err()); + let mut started = None; activity - .settle(None, Some("terminal error".into())) - .unwrap(); + .execute( + activity::ExecutionOrigin::Autonomous, + async { + for text in ["first continuation", "second continuation"] { + driver.submit_input(vec![Item::notification(text)]).unwrap(); + drive_autonomous(&session_id, &integration, &mut driver) + .await + .unwrap(); + } + let event = states.try_recv().unwrap(); + assert!(event.active); + assert_eq!(event.turn_id, 3); + assert!(states.try_recv().is_err()); + started = Some(event); + Err::<(), _>(AcpRuntimeError::Loop("terminal error".into())) + }, + |_| None, + ) + .await + .unwrap_err(); + let started = started.unwrap(); let ended = states.try_recv().unwrap(); assert!(!ended.active); assert_eq!(ended.turn_id, started.turn_id); - assert_eq!(ended.error.as_deref(), Some("terminal error")); + assert!(ended.error.as_deref().unwrap().contains("terminal error")); activity.settle(None, None).unwrap(); assert!(states.try_recv().is_err()); assert_eq!(turns.load(Ordering::SeqCst), 6); diff --git a/src/protocols/acp/activity.rs b/src/protocols/acp/activity.rs index 3ca8c08c..1c218efe 100644 --- a/src/protocols/acp/activity.rs +++ b/src/protocols/acp/activity.rs @@ -64,17 +64,6 @@ impl SessionActivity { } } - #[cfg(test)] - pub(super) fn begin(&self, origin: ExecutionOrigin) { - let Ok(mut state) = self.state.lock() else { - return; // A poisoned owner is never reused. - }; - // Continuations cannot redefine the origin of an existing interval. - if state.state == State::Idle && !state.projecting && !state.executing { - state.origin = origin; - } - } - pub(super) fn observe(&self, event: &AgentEvent) { let Ok(mut state) = self.state.lock() else { return; @@ -345,7 +334,16 @@ mod tests { assert!(transitions.lock().unwrap().last().unwrap().active); activity.observe(&finished(FinishReason::ToolCall)); // Steering and background synthesis cannot redefine the interval. - activity.begin(ExecutionOrigin::Autonomous); + assert!( + activity + .execute( + ExecutionOrigin::Autonomous, + async { panic!("overlapping operation must not run") }, + |_: &()| None, + ) + .await + .is_err() + ); activity.observe(&started("continuation")); activity.observe(&finished(reason.clone())); assert_eq!(transitions.lock().unwrap().len(), index * 2 + 1); @@ -464,8 +462,8 @@ mod tests { assert_eq!(*calls.lock().unwrap(), 2); } - #[test] - fn projection_unwind_isolates_without_poison_or_replay() { + #[tokio::test] + async fn projection_unwind_isolates_without_poison_or_replay() { use std::sync::atomic::{AtomicUsize, Ordering}; for fail_active in [true, false] { let calls = Arc::new(AtomicUsize::new(0)); @@ -481,7 +479,16 @@ mod tests { })); assert!(result.is_err()); assert!(!activity.state.is_poisoned()); - activity.begin(ExecutionOrigin::Autonomous); + assert!( + activity + .execute( + ExecutionOrigin::Autonomous, + async { panic!("isolated operation must not run") }, + |_: &()| None, + ) + .await + .is_err() + ); activity.observe(&started("later")); assert!(matches!( activity.settle(None, None), @@ -563,8 +570,8 @@ mod tests { assert!(!activity.state.is_poisoned()); } - #[test] - fn exhausted_identity_and_poison_never_resume_projection() { + #[tokio::test] + async fn exhausted_identity_and_poison_never_resume_projection() { let activity = SessionActivity::new(|_| panic!("must not project")); activity.state.lock().unwrap().next_id = u64::MAX; activity.observe(&started("overflow")); @@ -578,7 +585,16 @@ mod tests { .join() .is_err() ); - activity.begin(ExecutionOrigin::Prompt); + assert!( + activity + .execute( + ExecutionOrigin::Prompt, + async { panic!("isolated operation must not run") }, + |_: &()| None, + ) + .await + .is_err() + ); activity.observe(&started("after-poison")); assert!(matches!( activity.settle(None, None), diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index 003ce3ab..195127e3 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -847,6 +847,8 @@ impl Server { connection: V2ConnectionTo, mut claim: crate::runtime::SessionClaim, ) -> Result { + // Reject exhaustion before admission or any binding, driver, or actor effects. + let token = self.registry.next_token()?; let mut admission = self .registry .begin_attachment() @@ -912,7 +914,6 @@ impl Server { commands: rx, mcp_events, }; - let token = self.registry.next_token(); let (activation, activated) = oneshot::channel(); let (completed, completion) = watch::channel(false); let guard = ActorGuard { @@ -3980,7 +3981,7 @@ mod tests { RecordingSink::default(), )) .unwrap(); - let token = server.registry.next_token(); + let token = server.registry.next_token().unwrap(); let (completed, completion) = watch::channel(false); let guard = ActorGuard { server: Arc::downgrade(server), @@ -4215,7 +4216,7 @@ mod tests { RecordingSink::default(), )) .unwrap(); - let token = registry.next_token(); + let token = registry.next_token().unwrap(); let (completed, completion) = watch::channel(false); let release = Arc::new(Notify::new()); let actor_release = Arc::clone(&release); @@ -4360,6 +4361,88 @@ mod tests { } } + #[tokio::test] + async fn v2_token_exhaustion_rejects_attachment_before_publication() { + let root = tempfile::tempdir().unwrap(); + let session_id = crate::session::new_id(); + let runtime = Runtime::with_session_provider_and_credentials( + root.path(), + "gpt-5.4", + crate::ProviderKind::OpenAiSubscription, + crate::runtime::SessionRequest { + id: session_id.clone(), + resume: false, + force: false, + }, + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(); + let registry = SessionRegistry::new(); + registry + .inner + .next_token + .store(u64::MAX, std::sync::atomic::Ordering::Relaxed); + let (client_transport, agent_transport) = agent_client_protocol::Channel::duplex(); + let router = v2_router(runtime.clone(), registry.clone()).unwrap(); + let server = tokio::spawn(async move { router.connect_to(agent_transport).await }); + let workspace = root.path().to_path_buf(); + let client = agent_client_protocol::Client.v2().connect_with( + client_transport, + async move |connection| { + connection + .send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("exhaustion-test", "0"), + )) + .block_task() + .await?; + for _ in 0..2 { + let error = connection + .send_request(wire::NewSessionRequest::new(workspace.clone())) + .block_task() + .await + .expect_err("exhausted registry must reject attachment"); + assert!( + error + .data + .unwrap() + .to_string() + .contains("ACP session registry token space exhausted") + ); + assert_eq!( + registry + .inner + .next_token + .load(std::sync::atomic::Ordering::Relaxed), + u64::MAX + ); + assert!(!registry.inner.state.is_poisoned()); + { + let state = registry.inner.lock_state(); + assert!(state.accepting); + assert_eq!(state.pending_attachments, 0); + assert!(state.sessions.is_empty()); + assert!(state.v2_sessions.is_empty()); + } + // Failed attachment releases its claim without creating a transcript. + let claim = runtime.claim_session().unwrap(); + assert_eq!(claim.id(), session_id); + assert!(claim.is_configured()); + drop(claim); + assert_eq!( + crate::session::load(&workspace, &session_id).unwrap_err(), + format!("session {session_id:?} does not exist") + ); + } + Ok(()) + }, + ); + let result = timeout(Duration::from_secs(2), client).await; + server.abort(); + let _ = server.await; + result.expect("exhaustion client timed out").unwrap(); + } + #[tokio::test] async fn v2_router_advertises_and_routes_pending_injection_replacement() { let root = tempfile::tempdir().unwrap(); From 86a1a3619275fb5e316a09c25d0cd912f6b590a0 Mon Sep 17 00:00:00 2001 From: daniel Date: Sat, 5 Sep 2026 23:17:10 +0100 Subject: [PATCH 3/3] fix(acp): reject busy prompts before mailbox waits --- src/protocols/acp/v2.rs | 65 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 63 insertions(+), 2 deletions(-) diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index 195127e3..db6710e4 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -1000,9 +1000,15 @@ impl Server { &self, request: wire::PromptRequest, ) -> Result, AcpRuntimeError> { - // Wait for capacity before claiming busy/injection state. Cancellation - // while a full mailbox is pending must not strand a prompt claim. let (sender, busy, handle) = self.prompt_route(&request.session_id)?; + // Reject overlaps before waiting: mailbox pressure must not queue another turn. + if busy.load(Ordering::Acquire) { + return Err(AcpRuntimeError::Unsupported( + "session is already running a prompt".into(), + )); + } + // Claim only after capacity is available so cancellation cannot strand + // ownership. Recheck atomically afterward in case another request now owns it. let permit = sender .reserve() .await @@ -4163,6 +4169,61 @@ mod tests { let _ = original_actor.await; } + #[tokio::test] + async fn busy_prompt_is_rejected_before_or_after_mailbox_wait() { + for initially_busy in [true, false] { + let root = tempfile::tempdir().unwrap(); + let registry = SessionRegistry::new(); + let server = Arc::new(Server::new( + Runtime::new(root.path(), "gpt-5.4").unwrap(), + registry.clone(), + )); + let mut admission = registry.begin_attachment().unwrap(); + let (publication, actor, mut received) = pending_publication(&server, "busy-prompt"); + let busy = Arc::clone(&publication.session.busy); + let session_id = publication.session_id.clone(); + server + .publish_session(&mut admission, publication, || Ok(())) + .unwrap(); + let sender = server.sender(&session_id).unwrap(); + let (reply, _ack) = oneshot::channel(); + sender.try_send(Command::Close { reply }).unwrap(); + assert_eq!(sender.capacity(), 0); + busy.store(initially_busy, Ordering::Release); + + let mut request = + Box::pin(server.prepare_prompt(wire::PromptRequest::new(session_id, Vec::new()))); + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + if !initially_busy { + assert!(std::future::Future::poll(request.as_mut(), &mut context).is_pending()); + assert!(!busy.load(Ordering::Acquire)); + // Another request wins admission while this one waits for capacity. + claim_prompt(&busy).unwrap(); + assert!(matches!(received.try_recv(), Ok(Command::Close { .. }))); + } + assert!(matches!( + std::future::Future::poll(request.as_mut(), &mut context), + std::task::Poll::Ready(Err(AcpRuntimeError::Unsupported(message))) + if message == "session is already running a prompt" + )); + drop(request); + assert!(busy.load(Ordering::Acquire)); + if initially_busy { + // Rejection must not wait for or consume the occupied mailbox slot. + assert_eq!(sender.capacity(), 0); + assert!(matches!(received.try_recv(), Ok(Command::Close { .. }))); + } + busy.store(false, Ordering::Release); + assert!(matches!( + received.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + assert_eq!(sender.capacity(), 1); + actor.abort(); + let _ = actor.await; + } + } + #[tokio::test] async fn cancelled_mailbox_wait_does_not_claim_prompt() { let root = tempfile::tempdir().unwrap();