Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

subworker: harden internal session lifecycle

Browse Source
This commit is contained in:
Hare
2026-08-07 13:36:44 +09:00
parent 8fe0525295
commit a9bb806387
3 changed files with 239 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/internal_worker.rs
+12 -2
View File
@@ -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
+128 -31
View File
@@ -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<()> {
let mut records = self
.internal_records
pub(crate) fn reserve_internal_name(
self: &Arc<Self>,
worker_name: String,
) -> io::Result<InternalSpawnReservation> {
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<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 = {
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::<Vec<_>>();
let _ = parent_scope.update(|current| current.with_removed_deny_rules(write_rules));
crates/worker/src/spawn/tool.rs
+99 -14
View File
@@ -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<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")),