subworker: harden internal session lifecycle

crates/worker/src/internal_worker.rs

@@ -419,10 +419,9 @@ pub(crate) async fn spawn_internal_worker_session(
     spawn_prepared_internal_worker_session(worker, store, input, None).await
 }
 
-pub(crate) async fn spawn_prepared_internal_worker_session(
+pub(crate) async fn prepare_internal_worker_session(
     mut worker: Worker<Box<dyn LlmClient>, EphemeralSessionStore>,
     store: EphemeralSessionStore,
-    input: String,
     on_turn_end: Option<Arc<dyn Fn(InternalWorkerSessionStatus) + Send + Sync>>,
 ) -> Result<InternalWorkerSessionHandle, InternalWorkerSessionError> {
     let session_id = worker.session_id();
@@ -496,6 +495,17 @@ pub(crate) async fn spawn_prepared_internal_worker_session(
         }
     });
 
+    Ok(handle)
+}
+
+#[cfg(test)]
+pub(crate) async fn spawn_prepared_internal_worker_session(
+    worker: Worker<Box<dyn LlmClient>, EphemeralSessionStore>,
+    store: EphemeralSessionStore,
+    input: String,
+    on_turn_end: Option<Arc<dyn Fn(InternalWorkerSessionStatus) + Send + Sync>>,
+) -> Result<InternalWorkerSessionHandle, InternalWorkerSessionError> {
+    let handle = prepare_internal_worker_session(worker, store, on_turn_end).await?;
     handle.send(input).await?;
     Ok(handle)
 }

crates/worker/src/spawn/registry.rs

@@ -9,9 +9,12 @@
 //! yield only new assistant text. Parent registry drop closes all session handles and synchronously
 //! returns delegated Write deny rules to the parent scope.
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::io;
-use std::sync::Arc;
+use std::sync::{
+    Arc,
+    atomic::{AtomicBool, Ordering},
+};
 
 use manifest::{Permission, ScopeRule, SharedScope};
 use session_store::{
@@ -29,10 +32,69 @@ pub(crate) struct InternalSpawnedWorkerRecord {
     pub worker_name: String,
     pub scope_delegated: Vec<ScopeRule>,
     pub session: InternalWorkerSessionHandle,
+    scope_reclaimed: Arc<AtomicBool>,
+}
+
+impl InternalSpawnedWorkerRecord {
+    pub(crate) fn new(
+        worker_name: String,
+        scope_delegated: Vec<ScopeRule>,
+        session: InternalWorkerSessionHandle,
+    ) -> Self {
+        Self {
+            worker_name,
+            scope_delegated,
+            session,
+            scope_reclaimed: Arc::new(AtomicBool::new(false)),
+        }
+    }
+
+    fn claim_scope_reclaim(&self) -> bool {
+        !self.scope_reclaimed.swap(true, Ordering::AcqRel)
+    }
+
+    fn restore_scope_reclaim(&self) {
+        self.scope_reclaimed.store(false, Ordering::Release);
+    }
+}
+
+pub(crate) struct InternalSpawnReservation {
+    registry: Arc<SpawnedWorkerRegistry>,
+    worker_name: String,
+    committed: bool,
+}
+
+impl InternalSpawnReservation {
+    pub(crate) fn commit(mut self, record: InternalSpawnedWorkerRecord) -> io::Result<()> {
+        if record.worker_name != self.worker_name {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "internal SubWorker reservation name does not match record name",
+            ));
+        }
+        self.registry
+            .internal_records
+            .lock()
+            .map_err(|_| io::Error::other("internal spawned-worker registry lock poisoned"))?
+            .push(record);
+        self.committed = true;
+        Ok(())
+    }
+}
+
+impl Drop for InternalSpawnReservation {
+    fn drop(&mut self) {
+        if !self.committed {
+            if let Ok(mut names) = self.registry.internal_names.lock() {
+                names.remove(&self.worker_name);
+            }
+        }
+    }
 }
 
 pub struct SpawnedWorkerRegistry {
     internal_records: std::sync::Mutex<Vec<InternalSpawnedWorkerRecord>>,
+    internal_names: std::sync::Mutex<HashSet<String>>,
     cursors: Mutex<HashMap<String, usize>>,
     parent_scope: Option<SharedScope>,
 }
@@ -48,6 +110,7 @@ impl SpawnedWorkerRegistry {
     pub fn new(_runtime_dir: Arc<RuntimeDir>) -> Arc<Self> {
         Arc::new(Self {
             internal_records: std::sync::Mutex::new(Vec::new()),
+            internal_names: std::sync::Mutex::new(HashSet::new()),
             cursors: Mutex::new(HashMap::new()),
             parent_scope: None,
         })
@@ -56,6 +119,7 @@ impl SpawnedWorkerRegistry {
     pub(crate) fn new_internal(_parent_name: String, parent_scope: SharedScope) -> Arc<Self> {
         Arc::new(Self {
             internal_records: std::sync::Mutex::new(Vec::new()),
+            internal_names: std::sync::Mutex::new(HashSet::new()),
             cursors: Mutex::new(HashMap::new()),
             parent_scope: Some(parent_scope),
         })
@@ -133,6 +197,7 @@ impl SpawnedWorkerRegistry {
         Ok(SpawnedWorkerRegistryLoad {
             registry: Arc::new(Self {
                 internal_records: std::sync::Mutex::new(Vec::new()),
+                internal_names: std::sync::Mutex::new(HashSet::new()),
                 cursors: Mutex::new(HashMap::new()),
                 parent_scope,
             }),
@@ -140,25 +205,26 @@ impl SpawnedWorkerRegistry {
         })
     }
 
-    pub(crate) fn add_internal(&self, record: InternalSpawnedWorkerRecord) -> io::Result<()> {
+    pub(crate) fn reserve_internal_name(
-        let mut records = self
+        self: &Arc<Self>,
-            .internal_records
+        worker_name: String,
+    ) -> io::Result<InternalSpawnReservation> {
+        let mut names = self
+            .internal_names
             .lock()
-            .map_err(|_| io::Error::other("internal spawned-worker registry lock poisoned"))?;
+            .map_err(|_| io::Error::other("internal SubWorker name registry lock poisoned"))?;
-        if records
+        if !names.insert(worker_name.clone()) {
-            .iter()
-            .any(|existing| existing.worker_name == record.worker_name)
-        {
             return Err(io::Error::new(
                 io::ErrorKind::AlreadyExists,
-                format!(
+                format!("spawned worker `{worker_name}` is already registered"),
-                    "spawned worker `{}` is already registered",
-                    record.worker_name
-                ),
             ));
         }
-        records.push(record);
+        drop(names);
-        Ok(())
+        Ok(InternalSpawnReservation {
+            registry: Arc::clone(self),
+            worker_name,
+            committed: false,
+        })
     }
 
     pub(crate) fn get_internal(&self, worker_name: &str) -> Option<InternalSpawnedWorkerRecord> {
@@ -177,26 +243,56 @@ impl SpawnedWorkerRegistry {
             .unwrap_or_default()
     }
 
+    pub(crate) fn reclaim_internal_scope(&self, worker_name: &str) -> io::Result<bool> {
+        let record = self.get_internal(worker_name).ok_or_else(|| {
+            io::Error::new(io::ErrorKind::NotFound, "internal SubWorker not found")
+        })?;
+        self.reclaim_record_scope(&record)
+    }
+
+    fn reclaim_record_scope(&self, record: &InternalSpawnedWorkerRecord) -> io::Result<bool> {
+        if !record.claim_scope_reclaim() {
+            return Ok(false);
+        }
+        let result = if let Some(parent_scope) = &self.parent_scope {
+            parent_scope
+                .update(|current| current.with_removed_deny_rules(delegated_write_rules(record)))
+                .map(|_| true)
+                .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))
+        } else {
+            Ok(true)
+        };
+        if result.is_err() {
+            record.restore_scope_reclaim();
+        }
+        result
+    }
+
     pub(crate) async fn remove_internal(
         &self,
         worker_name: &str,
     ) -> io::Result<Option<InternalSpawnedWorkerRecord>> {
-        let removed = {
+        if let Some(record) = self.get_internal(worker_name) {
-            let mut records = self
+            self.reclaim_record_scope(&record)?;
-                .internal_records
-                .lock()
-                .map_err(|_| io::Error::other("internal spawned-worker registry lock poisoned"))?;
-            records
-                .iter()
-                .position(|record| record.worker_name == worker_name)
-                .map(|index| records.remove(index))
-        };
-        self.cursors.lock().await.remove(worker_name);
-        if let (Some(record), Some(parent_scope)) = (&removed, &self.parent_scope) {
-            parent_scope
-                .update(|current| current.with_removed_deny_rules(delegated_write_rules(record)))
-                .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))?;
         }
+        let removed =
+            {
+                let mut records = self.internal_records.lock().map_err(|_| {
+                    io::Error::other("internal spawned-worker registry lock poisoned")
+                })?;
+                let mut names = self.internal_names.lock().map_err(|_| {
+                    io::Error::other("internal SubWorker name registry lock poisoned")
+                })?;
+                let removed = records
+                    .iter()
+                    .position(|record| record.worker_name == worker_name)
+                    .map(|index| records.remove(index));
+                if removed.is_some() {
+                    names.remove(worker_name);
+                }
+                removed
+            };
+        self.cursors.lock().await.remove(worker_name);
         Ok(removed)
     }
 
@@ -222,6 +318,7 @@ impl Drop for SpawnedWorkerRegistry {
         };
         let write_rules = records
             .iter()
+            .filter(|record| !record.scope_reclaimed.load(Ordering::Acquire))
             .flat_map(delegated_write_rules)
             .collect::<Vec<_>>();
         let _ = parent_scope.update(|current| current.with_removed_deny_rules(write_rules));

crates/worker/src/spawn/tool.rs

@@ -21,7 +21,9 @@ use serde::Deserialize;
 
 use crate::PromptLoader;
 use crate::controller::register_worker_tools;
-use crate::internal_worker::{EphemeralSessionStore, spawn_prepared_internal_worker_session};
+use crate::internal_worker::{
+    EphemeralSessionStore, InternalWorkerSessionStatus, prepare_internal_worker_session,
+};
 use crate::prompt::catalog::PromptCatalog;
 use crate::spawn::registry::SpawnedWorkerRegistry;
 use crate::worker::{Worker, WorkerFilesystemAuthority};
@@ -303,6 +305,10 @@ impl Tool for SubWorkerSpawnTool {
                 input.name
             )));
         }
+        let name_reservation = self
+            .registry
+            .reserve_internal_name(input.name.clone())
+            .map_err(|error| ToolError::InvalidArgument(error.to_string()))?;
 
         let scope_allow = parse_scope(&input.scope)?;
         self.validate_delegation_scope(&scope_allow)?;
@@ -385,15 +391,26 @@ impl Tool for SubWorkerSpawnTool {
         }
 
         let child_name = input.name.clone();
+        let registry = Arc::downgrade(&self.registry);
         let parent_notifies = self.parent_notifies.clone();
-        let session_result = spawn_prepared_internal_worker_session(
+        let session_result = prepare_internal_worker_session(
             child,
             store,
-            input.task.clone(),
             Some(Arc::new(move |status| {
+                if status == InternalWorkerSessionStatus::Failed {
+                    if let Some(registry) = registry.upgrade() {
+                        if let Err(error) = registry.reclaim_internal_scope(&child_name) {
+                            tracing::warn!(
+                                child_name,
+                                %error,
+                                "failed to reclaim delegated scope after Internal SubWorker failure"
+                            );
+                        }
+                    }
+                }
                 parent_notifies.push_notify(
                     format!("SubWorker `{child_name}` turn ended with status {status:?}. Read its output before making completion decisions."),
-                    false,
+                    true,
                 );
             })),
         )
@@ -407,17 +424,17 @@ impl Tool for SubWorkerSpawnTool {
                         .update(|current| current.with_removed_deny_rules(revoke_write.clone()));
                 }
                 return Err(ToolError::ExecutionFailed(format!(
-                    "start Internal Worker session: {error}"
+                    "prepare Internal Worker session: {error}"
                 )));
             }
         };
 
-        let record = crate::spawn::registry::InternalSpawnedWorkerRecord {
+        let record = crate::spawn::registry::InternalSpawnedWorkerRecord::new(
-            worker_name: input.name.clone(),
+            input.name.clone(),
-            scope_delegated: scope_allow,
+            scope_allow,
-            session: session.clone(),
+            session.clone(),
-        };
+        );
-        if let Err(error) = self.registry.add_internal(record) {
+        if let Err(error) = name_reservation.commit(record) {
             let _ = session.stop().await;
             if !revoke_write.is_empty() {
                 let _ = self
@@ -428,6 +445,13 @@ impl Tool for SubWorkerSpawnTool {
                 "register Internal Worker session: {error}"
             )));
         }
+        if let Err(error) = session.send(input.task).await {
+            let _ = session.stop().await;
+            let _ = self.registry.remove_internal(&input.name).await;
+            return Err(ToolError::ExecutionFailed(format!(
+                "start Internal Worker session: {error}"
+            )));
+        }
 
         Ok(ToolOutput {
             summary: format!("spawned internal worker `{}`", input.name),
@@ -820,7 +844,7 @@ fn sub_worker_spawn_tool_impl(
 mod tests {
     use super::*;
     use std::pin::Pin;
-    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 
     use crate::WorkspaceId;
     use async_trait::async_trait;
@@ -890,8 +914,9 @@ extract_threshold = 4000
             Arc::new(AvailableWorkspaceClient),
         );
         let calls = Arc::new(AtomicUsize::new(0));
-        let observed_parent_write_revoked = Arc::new(std::sync::atomic::AtomicBool::new(false));
+        let observed_parent_write_revoked = Arc::new(AtomicBool::new(false));
-        let observed_instruction_override = Arc::new(std::sync::atomic::AtomicBool::new(false));
+        let observed_instruction_override = Arc::new(AtomicBool::new(false));
+        let fail_requests = Arc::new(AtomicBool::new(false));
         let workspace_prompts = runtime.path().join("workspace-prompts");
         std::fs::create_dir_all(&workspace_prompts).unwrap();
         std::fs::write(
@@ -921,6 +946,7 @@ extract_threshold = 4000
             delegated_path: workspace_root.clone(),
             observed_parent_write_revoked: observed_parent_write_revoked.clone(),
             observed_instruction_override: observed_instruction_override.clone(),
+            fail_requests: fail_requests.clone(),
         }));
         let input = serde_json::json!({
             "name": "reviewer-child",
@@ -936,6 +962,18 @@ extract_threshold = 4000
 
         assert!(spawner_scope.snapshot().is_writable(&workspace_root));
 
+        let mut invalid_input = input.clone();
+        invalid_input["scope"][0]["target"] =
+            serde_json::json!(runtime.path().join("outside-parent-scope"));
+        tool.execute(
+            &serde_json::to_string(&invalid_input).unwrap(),
+            llm_engine::tool::ToolExecutionContext::direct(),
+        )
+        .await
+        .expect_err("invalid delegation must fail before child preparation");
+        assert_eq!(calls.load(Ordering::SeqCst), 0);
+        assert!(spawner_scope.snapshot().is_writable(&workspace_root));
+
         let output = tool
             .execute(
                 &serde_json::to_string(&input).unwrap(),
@@ -956,8 +994,30 @@ extract_threshold = 4000
         assert!(observed_parent_write_revoked.load(Ordering::SeqCst));
         assert!(observed_instruction_override.load(Ordering::SeqCst));
         assert_eq!(parent_notifies.len(), 1);
+        assert!(
+            parent_notifies.has_auto_run_pending(),
+            "SubWorker completion must auto-invoke the parent"
+        );
         assert!(!runtime.path().join("reviewer-child/sock").exists());
 
+        let duplicate_error = tool
+            .execute(
+                &serde_json::to_string(&input).unwrap(),
+                llm_engine::tool::ToolExecutionContext::direct(),
+            )
+            .await
+            .expect_err("duplicate child name must be rejected before a first turn starts");
+        assert!(
+            format!("{duplicate_error:?}").contains("already registered"),
+            "unexpected duplicate error: {duplicate_error:?}"
+        );
+        assert_eq!(
+            calls.load(Ordering::SeqCst),
+            1,
+            "duplicate rejection must not invoke the child provider"
+        );
+        assert_eq!(parent_notifies.len(), 1);
+
         let context = llm_engine::tool::ToolExecutionContext::direct();
         let list = (crate::spawn::comm_tools::sub_worker_list_tool(registry.clone()))().1;
         let listed = list.execute("{}", context.clone()).await.unwrap();
@@ -998,6 +1058,24 @@ extract_threshold = 4000
         );
         assert_eq!(calls.load(Ordering::SeqCst), 2);
 
+        fail_requests.store(true, Ordering::SeqCst);
+        send.execute(
+            r#"{"name":"reviewer-child","message":"trigger terminal failure"}"#,
+            context.clone(),
+        )
+        .await
+        .unwrap();
+        assert_eq!(
+            record.session.wait_until_idle().await,
+            InternalWorkerSessionStatus::Failed
+        );
+        assert_eq!(calls.load(Ordering::SeqCst), 3);
+        assert!(
+            spawner_scope.snapshot().is_writable(&workspace_root),
+            "Failed terminal child must automatically reclaim its delegated write scope"
+        );
+        assert!(registry.get_internal("reviewer-child").is_some());
+
         let stop = (crate::spawn::comm_tools::sub_worker_stop_tool(registry.clone()))().1;
         stop.execute(r#"{"name":"reviewer-child"}"#, context)
             .await
@@ -1005,6 +1083,7 @@ extract_threshold = 4000
         assert!(registry.get_internal("reviewer-child").is_none());
         assert!(spawner_scope.snapshot().is_writable(&workspace_root));
 
+        fail_requests.store(false, Ordering::SeqCst);
         let mut teardown_input = input;
         teardown_input["name"] = serde_json::json!("reviewer-child-parent-drop");
         tool.execute(
@@ -1127,6 +1206,7 @@ extract_threshold = 4000
         delegated_path: PathBuf,
         observed_parent_write_revoked: Arc<std::sync::atomic::AtomicBool>,
         observed_instruction_override: Arc<std::sync::atomic::AtomicBool>,
+        fail_requests: Arc<AtomicBool>,
     }
 
     #[async_trait]
@@ -1155,6 +1235,11 @@ extract_threshold = 4000
                     .is_some_and(|prompt| prompt.contains("WORKSPACE REVIEWER OVERRIDE")),
                 Ordering::SeqCst,
             );
+            if self.fail_requests.load(Ordering::SeqCst) {
+                return Err(ClientError::Config(
+                    "scripted Internal Worker failure".into(),
+                ));
+            }
             Ok(Box::pin(futures::stream::iter(vec![
                 Ok(LlmEvent::text_block_start(0)),
                 Ok(LlmEvent::text_delta(0, "reviewed")),