fix: fence late command events and spawn rollback

crates/workdir/src/scope.rs

@@ -75,6 +75,7 @@ impl WorkdirToolBroker {
             owned_commands: Arc::new(Mutex::new(HashSet::new())),
             pending_command_events: Arc::new(Mutex::new(HashMap::new())),
             starting_tool_calls: Arc::new(Mutex::new(HashSet::new())),
+            forwarded_starts: Arc::new(Mutex::new(HashSet::new())),
             forwarded_terminals: Arc::new(Mutex::new(HashSet::new())),
             command_events,
             closes_source: true,
@@ -314,6 +315,7 @@ struct ScopedWorkdirSession {
     owned_commands: Arc<Mutex<HashSet<String>>>,
     pending_command_events: Arc<Mutex<HashMap<String, Vec<CommandEvent>>>>,
     starting_tool_calls: Arc<Mutex<HashSet<String>>>,
+    forwarded_starts: Arc<Mutex<HashSet<String>>>,
     forwarded_terminals: Arc<Mutex<HashSet<String>>>,
     command_events: broadcast::Sender<CommandEvent>,
     closes_source: bool,
@@ -675,6 +677,7 @@ impl ScopedWorkdirSession {
         let owned_commands = Arc::new(Mutex::new(HashSet::new()));
         let pending_command_events = Arc::new(Mutex::new(HashMap::new()));
         let starting_tool_calls = Arc::new(Mutex::new(HashSet::new()));
+        let forwarded_starts = Arc::new(Mutex::new(HashSet::new()));
         let forwarded_terminals = Arc::new(Mutex::new(HashSet::new()));
         let (command_events, _) = broadcast::channel(64);
         let event_forwarder = forward_owned_command_events(
@@ -682,6 +685,7 @@ impl ScopedWorkdirSession {
             owned_commands.clone(),
             pending_command_events.clone(),
             starting_tool_calls.clone(),
+            forwarded_starts.clone(),
             forwarded_terminals.clone(),
             command_events.clone(),
         )
@@ -699,6 +703,7 @@ impl ScopedWorkdirSession {
             owned_commands,
             pending_command_events,
             starting_tool_calls,
+            forwarded_starts,
             forwarded_terminals,
             command_events,
             closes_source: false,
@@ -862,6 +867,7 @@ impl WorkdirSession for ScopedWorkdirSession {
         {
             publish_owned_command_event(
                 &self.command_events,
+                &self.forwarded_starts,
                 &self.forwarded_terminals,
                 CommandEvent::Started {
                     command_id: handle.0.clone(),
@@ -871,7 +877,12 @@ impl WorkdirSession for ScopedWorkdirSession {
             );
         }
         for event in pending {
-            publish_owned_command_event(&self.command_events, &self.forwarded_terminals, event);
+            publish_owned_command_event(
+                &self.command_events,
+                &self.forwarded_starts,
+                &self.forwarded_terminals,
+                event,
+            );
         }
         Ok(handle)
     }
@@ -1033,6 +1044,7 @@ fn forward_owned_command_events(
     owned_commands: Arc<Mutex<HashSet<String>>>,
     pending_command_events: Arc<Mutex<HashMap<String, Vec<CommandEvent>>>>,
     starting_tool_calls: Arc<Mutex<HashSet<String>>>,
+    forwarded_starts: Arc<Mutex<HashSet<String>>>,
     forwarded_terminals: Arc<Mutex<HashSet<String>>>,
     sender: broadcast::Sender<CommandEvent>,
 ) -> Option<tokio::task::JoinHandle<()>> {
@@ -1082,7 +1094,7 @@ fn forward_owned_command_events(
             }
             drop(pending);
             drop(owned);
-            publish_owned_command_event(&sender, &forwarded_terminals, event);
+            publish_owned_command_event(&sender, &forwarded_starts, &forwarded_terminals, event);
         }
     }))
 }
@@ -1097,6 +1109,7 @@ fn command_event_id(event: &CommandEvent) -> &str {
 
 fn publish_owned_command_event(
     sender: &broadcast::Sender<CommandEvent>,
+    forwarded_starts: &Mutex<HashSet<String>>,
     forwarded_terminals: &Mutex<HashSet<String>>,
     event: CommandEvent,
 ) {
@@ -1106,11 +1119,16 @@ fn publish_owned_command_event(
         .expect("forwarded terminal command mutex poisoned");
     match &event {
         CommandEvent::Terminal { .. } if !terminals.insert(command_id.to_string()) => return,
-        CommandEvent::Started { .. } | CommandEvent::Output { .. }
-            if terminals.contains(command_id) =>
+        CommandEvent::Started { .. } if terminals.contains(command_id) => return,
+        CommandEvent::Started { .. }
+            if !forwarded_starts
+                .lock()
+                .expect("forwarded command start mutex poisoned")
+                .insert(command_id.to_string()) =>
         {
             return;
         }
+        CommandEvent::Output { .. } if terminals.contains(command_id) => return,
         _ => {}
     }
     drop(terminals);
@@ -1657,8 +1675,16 @@ mod tests {
         }
         assert_eq!(kinds.first(), Some(&"started"));
         assert_eq!(kinds.last(), Some(&"terminal"));
+        assert_eq!(kinds.iter().filter(|kind| **kind == "started").count(), 1);
         assert!(kinds.contains(&"output"));
         assert!(streamed.contains("fast-output"));
+        assert!(
+            tokio::time::timeout(std::time::Duration::from_millis(100), events.recv())
+                .await
+                .is_err(),
+            "no provider event may follow the terminal event"
+        );
+        assert!(child.command_snapshot().is_empty());
         child.close().await.unwrap();
     }
 

crates/worker/src/spawn/registry.rs

@@ -12,7 +12,7 @@ use std::collections::{BTreeMap, HashSet};
 use std::io;
 use std::sync::{
     Arc, Mutex,
-    atomic::{AtomicBool, AtomicU64, Ordering},
+    atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
 };
 use std::time::Instant;
 
@@ -23,7 +23,7 @@ use protocol::{Event, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnaps
 use session_store::{
     LoggedItem, WorkerMetadataStore, WorkerReclaimedChild, WorkerSpawnedChild, WorkerStoreError,
 };
-use tokio::sync::broadcast;
+use tokio::sync::{Notify, broadcast};
 use tracing::warn;
 use workdir::WorkdirScopeLease;
 
@@ -238,39 +238,56 @@ pub(crate) struct InternalSpawnReservation {
 }
 
 impl InternalSpawnReservation {
-    pub(crate) fn commit(
-        mut self,
-        record: InternalSpawnedWorkerRecord,
-    ) -> Result<(), (io::Error, InternalSpawnedWorkerRecord)> {
-        if record.worker_name != self.worker_name {
-            return Err((
-                io::Error::new(
-                    io::ErrorKind::InvalidInput,
-                    "internal SubWorker reservation name does not match record name",
-                ),
-                record,
-            ));
-        }
-        let mut records = match self.registry.internal_records.lock() {
-            Ok(records) => records,
-            Err(_) => {
-                return Err((
-                    io::Error::other("internal spawned-worker registry lock poisoned"),
-                    record,
-                ));
+    pub(crate) async fn commit(mut self, record: InternalSpawnedWorkerRecord) -> io::Result<()> {
+        let rejection = if record.worker_name != self.worker_name {
+            Some(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "internal SubWorker reservation name does not match record name",
+            ))
+        } else {
+            match self.registry.internal_records.lock() {
+                Ok(mut records) => {
+                    if self.registry.internal_shutting_down.load(Ordering::Acquire) {
+                        Some(io::Error::new(
+                            io::ErrorKind::Interrupted,
+                            "internal SubWorker registry is shutting down",
+                        ))
+                    } else {
+                        records.push(record.clone());
+                        None
+                    }
+                }
+                Err(_) => Some(io::Error::other(
+                    "internal spawned-worker registry lock poisoned",
+                )),
             }
         };
-        if self.registry.internal_shutting_down.load(Ordering::Acquire) {
-            return Err((
-                io::Error::new(
-                    io::ErrorKind::Interrupted,
-                    "internal SubWorker registry is shutting down",
-                ),
-                record,
-            ));
+        if let Some(error) = rejection {
+            let mut cleanup_failures = Vec::new();
+            if let Err(cleanup) = record.session.stop().await {
+                cleanup_failures.push(format!("stop rejected Internal SubWorker: {cleanup}"));
+            }
+            if let Err(cleanup) = Box::pin(record.child_registry.shutdown_internal()).await {
+                cleanup_failures.push(format!(
+                    "stop rejected Internal SubWorker descendants: {cleanup}"
+                ));
+            }
+            if let Err(cleanup) = record.workdir_tool_scope.close().await {
+                cleanup_failures.push(format!(
+                    "close rejected Internal SubWorker Workdir tools: {cleanup}"
+                ));
+            }
+            if cleanup_failures.is_empty() {
+                return Err(error);
+            }
+            self.registry
+                .internal_spawn_cleanup_failed
+                .store(true, Ordering::Release);
+            return Err(io::Error::other(format!(
+                "{error}; {}",
+                cleanup_failures.join("; ")
+            )));
         }
-        records.push(record.clone());
-        drop(records);
         self.registry.start_protocol_forwarding(record);
         self.committed = true;
         Ok(())
@@ -284,6 +301,10 @@ impl Drop for InternalSpawnReservation {
                 names.remove(&self.worker_name);
             }
         }
+        self.registry
+            .pending_internal_spawns
+            .fetch_sub(1, Ordering::AcqRel);
+        self.registry.pending_internal_notify.notify_waiters();
     }
 }
 
@@ -292,6 +313,9 @@ pub struct SpawnedWorkerRegistry {
     service_records: std::sync::Mutex<Vec<InternalServiceWorkerRecord>>,
     internal_names: std::sync::Mutex<HashSet<String>>,
     internal_shutting_down: AtomicBool,
+    pending_internal_spawns: AtomicUsize,
+    pending_internal_notify: Notify,
+    internal_spawn_cleanup_failed: AtomicBool,
     parent_scope: Option<SharedScope>,
     parent_protocol: Mutex<Option<(broadcast::Sender<Event>, String)>>,
 }
@@ -309,6 +333,9 @@ impl SpawnedWorkerRegistry {
             service_records: std::sync::Mutex::new(Vec::new()),
             internal_names: std::sync::Mutex::new(HashSet::new()),
             internal_shutting_down: AtomicBool::new(false),
+            pending_internal_spawns: AtomicUsize::new(0),
+            pending_internal_notify: Notify::new(),
+            internal_spawn_cleanup_failed: AtomicBool::new(false),
             parent_scope: None,
             parent_protocol: Mutex::new(None),
         })
@@ -321,6 +348,9 @@ impl SpawnedWorkerRegistry {
             service_records: std::sync::Mutex::new(Vec::new()),
             internal_names: std::sync::Mutex::new(HashSet::new()),
             internal_shutting_down: AtomicBool::new(false),
+            pending_internal_spawns: AtomicUsize::new(0),
+            pending_internal_notify: Notify::new(),
+            internal_spawn_cleanup_failed: AtomicBool::new(false),
             parent_scope: None,
             parent_protocol: Mutex::new(None),
         })
@@ -332,6 +362,9 @@ impl SpawnedWorkerRegistry {
             service_records: std::sync::Mutex::new(Vec::new()),
             internal_names: std::sync::Mutex::new(HashSet::new()),
             internal_shutting_down: AtomicBool::new(false),
+            pending_internal_spawns: AtomicUsize::new(0),
+            pending_internal_notify: Notify::new(),
+            internal_spawn_cleanup_failed: AtomicBool::new(false),
             parent_scope: Some(parent_scope),
             parent_protocol: Mutex::new(None),
         })
@@ -412,6 +445,9 @@ impl SpawnedWorkerRegistry {
                 service_records: std::sync::Mutex::new(Vec::new()),
                 internal_names: std::sync::Mutex::new(HashSet::new()),
                 internal_shutting_down: AtomicBool::new(false),
+                pending_internal_spawns: AtomicUsize::new(0),
+                pending_internal_notify: Notify::new(),
+                internal_spawn_cleanup_failed: AtomicBool::new(false),
                 parent_scope,
                 parent_protocol: Mutex::new(None),
             }),
@@ -423,6 +459,10 @@ impl SpawnedWorkerRegistry {
         self: &Arc<Self>,
         worker_name: String,
     ) -> io::Result<InternalSpawnReservation> {
+        let records = self
+            .internal_records
+            .lock()
+            .map_err(|_| io::Error::other("internal Worker registry lock poisoned"))?;
         if self.internal_shutting_down.load(Ordering::Acquire) {
             return Err(io::Error::new(
                 io::ErrorKind::Interrupted,
@@ -439,7 +479,9 @@ impl SpawnedWorkerRegistry {
                 format!("spawned worker `{worker_name}` is already registered"),
             ));
         }
+        self.pending_internal_spawns.fetch_add(1, Ordering::AcqRel);
         drop(names);
+        drop(records);
         Ok(InternalSpawnReservation {
             registry: Arc::clone(self),
             worker_name,
@@ -758,17 +800,31 @@ impl SpawnedWorkerRegistry {
                 .map(|record| record.worker_name.clone())
                 .collect::<Vec<_>>()
         };
+        loop {
+            let notified = self.pending_internal_notify.notified();
+            if self.pending_internal_spawns.load(Ordering::Acquire) == 0 {
+                break;
+            }
+            notified.await;
+        }
         let mut first_error = None;
         for name in names {
             if let Err(error) = self.remove_internal(&name).await {
                 first_error.get_or_insert(error);
             }
         }
+        if first_error.is_none() && self.internal_spawn_cleanup_failed.load(Ordering::Acquire) {
+            first_error = Some(io::Error::other(
+                "an in-flight Internal SubWorker failed cleanup during shutdown",
+            ));
+        }
         first_error.map_or(Ok(()), Err)
     }
 
     pub(crate) fn reopen_internal(&self) {
         self.internal_shutting_down.store(false, Ordering::Release);
+        self.internal_spawn_cleanup_failed
+            .store(false, Ordering::Release);
     }
 
     /// Stop one direct Internal SubWorker and discard its registry/scope state.
@@ -1342,24 +1398,22 @@ mod tests {
         let (record, _events) =
             record("racing-child", InternalWorkerVisibility::ParentClient).await;
         let scope = record.workdir_tool_scope.clone();
-        let barrier = Arc::new(std::sync::Barrier::new(2));
+        let barrier = Arc::new(tokio::sync::Barrier::new(2));
         let commit_barrier = barrier.clone();
-        let commit = tokio::task::spawn_blocking(move || {
-            commit_barrier.wait();
-            reservation.commit(record)
+        let commit = tokio::spawn(async move {
+            commit_barrier.wait().await;
+            reservation.commit(record).await
         });
         let shutdown_registry = registry.clone();
         let shutdown = tokio::spawn(async move {
-            barrier.wait();
+            barrier.wait().await;
             shutdown_registry.shutdown_internal().await
         });
 
         let commit = commit.await.unwrap();
         shutdown.await.unwrap().unwrap();
-        if let Err((_error, record)) = commit {
-            record.session.stop().await.unwrap();
-            record.child_registry.shutdown_internal().await.unwrap();
-            record.workdir_tool_scope.close().await.unwrap();
+        if let Err(error) = commit {
+            assert_eq!(error.kind(), io::ErrorKind::Interrupted);
         }
 
         assert!(registry.list_internal().is_empty());
@@ -1375,12 +1429,54 @@ mod tests {
         let (record, _events) =
             record("racing-child", InternalWorkerVisibility::ParentClient).await;
 
-        registry.shutdown_internal().await.unwrap();
-        let (error, record) = reservation.commit(record).unwrap_err();
+        let mut shutdown = {
+            let registry = registry.clone();
+            tokio::spawn(async move { registry.shutdown_internal().await })
+        };
+        while !registry.internal_shutting_down.load(Ordering::Acquire) {
+            tokio::task::yield_now().await;
+        }
+        assert!(
+            tokio::time::timeout(std::time::Duration::from_millis(50), &mut shutdown)
+                .await
+                .is_err(),
+            "shutdown must wait for the pending spawn to roll back"
+        );
+        let error = reservation.commit(record).await.unwrap_err();
         assert_eq!(error.kind(), io::ErrorKind::Interrupted);
-        record.session.stop().await.unwrap();
-        record.child_registry.shutdown_internal().await.unwrap();
-        record.workdir_tool_scope.close().await.unwrap();
+        shutdown.await.unwrap().unwrap();
+    }
+
+    #[tokio::test]
+    async fn rejected_spawn_cleanup_failure_keeps_shutdown_failed_closed() {
+        let registry = registry();
+        let reservation = registry
+            .reserve_internal_name("cleanup-failure".into())
+            .unwrap();
+        let (record, _events) =
+            record("cleanup-failure", InternalWorkerVisibility::ParentClient).await;
+        record.session.force_stop_failure();
+        let shutdown = {
+            let registry = registry.clone();
+            tokio::spawn(async move { registry.shutdown_internal().await })
+        };
+        while !registry.internal_shutting_down.load(Ordering::Acquire) {
+            tokio::task::yield_now().await;
+        }
+
+        let error = reservation.commit(record).await.unwrap_err();
+        assert!(
+            error
+                .to_string()
+                .contains("stop rejected Internal SubWorker")
+        );
+        let shutdown_error = shutdown.await.unwrap().unwrap_err();
+        assert!(
+            shutdown_error
+                .to_string()
+                .contains("failed cleanup during shutdown")
+        );
+        assert!(registry.internal_shutting_down.load(Ordering::Acquire));
     }
 
     #[tokio::test]

crates/worker/src/spawn/tool.rs

@@ -630,10 +630,7 @@ impl Tool for SubWorkerSpawnTool {
             child_registry,
             child_change_tracker,
         );
-        if let Err((error, record)) = name_reservation.commit(record) {
-            let _ = session.stop().await;
-            let _ = record.child_registry.shutdown_internal().await;
-            let _ = record.workdir_tool_scope.close().await;
+        if let Err(error) = name_reservation.commit(record).await {
             return Err(ToolError::ExecutionFailed(format!(
                 "register Internal Worker session: {error}"
             )));