From a9bb8063876e91c8695386e113067568bf860133 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 7 Aug 2026 13:36:44 +0900 Subject: [PATCH] subworker: harden internal session lifecycle --- crates/worker/src/internal_worker.rs | 14 ++- crates/worker/src/spawn/registry.rs | 159 +++++++++++++++++++++------ crates/worker/src/spawn/tool.rs | 113 ++++++++++++++++--- 3 files changed, 239 insertions(+), 47 deletions(-) diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index f6056f10..7f7811db 100644 --- a/crates/worker/src/internal_worker.rs +++ b/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, EphemeralSessionStore>, store: EphemeralSessionStore, - input: String, on_turn_end: Option>, ) -> Result { 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, EphemeralSessionStore>, + store: EphemeralSessionStore, + input: String, + on_turn_end: Option>, +) -> Result { + let handle = prepare_internal_worker_session(worker, store, on_turn_end).await?; handle.send(input).await?; Ok(handle) } diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 372acb60..4eaccb67 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/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, pub session: InternalWorkerSessionHandle, + scope_reclaimed: Arc, +} + +impl InternalSpawnedWorkerRecord { + pub(crate) fn new( + worker_name: String, + scope_delegated: Vec, + 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, + 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>, + internal_names: std::sync::Mutex>, cursors: Mutex>, parent_scope: Option, } @@ -48,6 +110,7 @@ impl SpawnedWorkerRegistry { pub fn new(_runtime_dir: Arc) -> Arc { 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 { 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<()> { - let mut records = self - .internal_records + pub(crate) fn reserve_internal_name( + self: &Arc, + worker_name: String, + ) -> io::Result { + let mut names = self + .internal_names .lock() - .map_err(|_| io::Error::other("internal spawned-worker registry lock poisoned"))?; - if records - .iter() - .any(|existing| existing.worker_name == record.worker_name) - { + .map_err(|_| io::Error::other("internal SubWorker name registry lock poisoned"))?; + if !names.insert(worker_name.clone()) { return Err(io::Error::new( io::ErrorKind::AlreadyExists, - format!( - "spawned worker `{}` is already registered", - record.worker_name - ), + format!("spawned worker `{worker_name}` is already registered"), )); } - records.push(record); - Ok(()) + drop(names); + Ok(InternalSpawnReservation { + registry: Arc::clone(self), + worker_name, + committed: false, + }) } pub(crate) fn get_internal(&self, worker_name: &str) -> Option { @@ -177,26 +243,56 @@ impl SpawnedWorkerRegistry { .unwrap_or_default() } + pub(crate) fn reclaim_internal_scope(&self, worker_name: &str) -> io::Result { + 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 { + 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> { - let removed = { - let mut records = self - .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))?; + if let Some(record) = self.get_internal(worker_name) { + self.reclaim_record_scope(&record)?; } + 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::>(); let _ = parent_scope.update(|current| current.with_removed_deny_rules(write_rules)); diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index c0850f75..6df32771 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/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 { - worker_name: input.name.clone(), - scope_delegated: scope_allow, - session: session.clone(), - }; - if let Err(error) = self.registry.add_internal(record) { + let record = crate::spawn::registry::InternalSpawnedWorkerRecord::new( + input.name.clone(), + scope_allow, + session.clone(), + ); + 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_instruction_override = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let observed_parent_write_revoked = Arc::new(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, observed_instruction_override: Arc, + fail_requests: Arc, } #[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")),