diff --git a/crates/worker/src/spawn/comm_tools.rs b/crates/worker/src/spawn/comm_tools.rs index e86b105d..03ec206a 100644 --- a/crates/worker/src/spawn/comm_tools.rs +++ b/crates/worker/src/spawn/comm_tools.rs @@ -62,15 +62,23 @@ impl Tool for SubWorkerListTool { let _input: SubWorkerListInput = serde_json::from_str(input_json).map_err(|error| { ToolError::InvalidArgument(format!("invalid SubWorkerList input: {error}")) })?; - let items = self + let mut items = self .registry - .list() - .await + .list_internal() .into_iter() .map(|record| SubWorkerListItem { name: record.worker_name, }) .collect::>(); + items.extend( + self.registry + .list() + .await + .into_iter() + .map(|record| SubWorkerListItem { + name: record.worker_name, + }), + ); let count = items.len(); let content = serde_json::to_string_pretty(&serde_json::json!({ "sub_workers": items })) .map_err(|error| ToolError::ExecutionFailed(error.to_string()))?; @@ -125,6 +133,15 @@ impl Tool for SubWorkerSendTool { ) -> Result { let input: SubWorkerSendInput = serde_json::from_str(input_json) .map_err(|e| ToolError::InvalidArgument(format!("invalid SubWorkerSend input: {e}")))?; + if let Some(record) = self.registry.get_internal(&input.name) { + record.session.send(input.message).await.map_err(|error| { + ToolError::ExecutionFailed(format!("send to `{}`: {error}", input.name)) + })?; + return Ok(ToolOutput { + summary: format!("sent message to `{}`", input.name), + content: None, + }); + } let record = self .registry .get(&input.name) @@ -191,6 +208,35 @@ impl Tool for SubWorkerReadOutputTool { let input: NameInput = serde_json::from_str(input_json).map_err(|e| { ToolError::InvalidArgument(format!("invalid SubWorkerReadOutput input: {e}")) })?; + if let Some(record) = self.registry.get_internal(&input.name) { + let entries = record.session.entries(); + let cursor = self.registry.cursor(&input.name).await; + let new_entries = if cursor >= entries.len() { + &[] as &[LogEntry] + } else { + &entries[cursor..] + }; + let values = new_entries + .iter() + .filter_map(|entry| serde_json::to_value(entry).ok()) + .collect::>(); + let new_text = extract_assistant_text(&values); + self.registry.set_cursor(&input.name, entries.len()).await; + let status = format!("{:?}", record.session.status()).to_lowercase(); + let summary = if new_text.is_empty() { + format!("worker `{}` {status}; no new assistant text", input.name) + } else { + format!( + "worker `{}` {status}: {} new line(s) of assistant text", + input.name, + new_text.lines().count() + ) + }; + return Ok(ToolOutput { + summary, + content: (!new_text.is_empty()).then_some(new_text), + }); + } let record = self .registry .get(&input.name) @@ -269,6 +315,22 @@ impl Tool for SubWorkerStopTool { ) -> Result { let input: NameInput = serde_json::from_str(input_json) .map_err(|e| ToolError::InvalidArgument(format!("invalid SubWorkerStop input: {e}")))?; + if let Some(record) = self.registry.get_internal(&input.name) { + record.session.stop().await.map_err(|error| { + ToolError::ExecutionFailed(format!("stop `{}`: {error}", input.name)) + })?; + self.registry + .remove_internal(&input.name) + .await + .map_err(|error| ToolError::ExecutionFailed(error.to_string()))?; + return Ok(ToolOutput { + summary: format!( + "stopped worker `{}` and reclaimed delegated scope", + input.name + ), + content: None, + }); + } let record = self .registry .get(&input.name) diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 900cb99f..b079ebfb 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -29,6 +29,7 @@ use tokio::net::UnixStream; use tokio::sync::Mutex; use tracing::warn; +use crate::internal_worker::InternalWorkerSessionHandle; use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord}; use crate::runtime::worker_allocation; @@ -38,8 +39,16 @@ type RegistryReclaimWriter = Arc io::Result<()> const RESTORE_REACHABILITY_TIMEOUT: Duration = Duration::from_millis(500); const REGISTRY_CLEANUP_TIMEOUT: Duration = Duration::from_secs(15); +#[derive(Clone)] +pub(crate) struct InternalSpawnedWorkerRecord { + pub worker_name: String, + pub scope_delegated: Vec, + pub session: InternalWorkerSessionHandle, +} + pub struct SpawnedWorkerRegistry { records: Mutex>, + internal_records: std::sync::Mutex>, cursors: Mutex>, mutations: Mutex<()>, runtime_dir: Arc, @@ -58,6 +67,7 @@ impl SpawnedWorkerRegistry { pub fn new(runtime_dir: Arc) -> Arc { Arc::new(Self { records: Mutex::new(Vec::new()), + internal_records: std::sync::Mutex::new(Vec::new()), cursors: Mutex::new(HashMap::new()), mutations: Mutex::new(()), runtime_dir, @@ -170,6 +180,7 @@ impl SpawnedWorkerRegistry { Ok(SpawnedWorkerRegistryLoad { registry: Arc::new(Self { records: Mutex::new(records), + internal_records: std::sync::Mutex::new(Vec::new()), cursors: Mutex::new(HashMap::new()), mutations: Mutex::new(()), runtime_dir, @@ -182,7 +193,73 @@ impl SpawnedWorkerRegistry { }) } - /// Append a new record and persist the full list. Returns an I/O + pub(crate) fn add_internal(&self, record: InternalSpawnedWorkerRecord) -> io::Result<()> { + let mut records = self + .internal_records + .lock() + .map_err(|_| io::Error::other("internal spawned-worker registry lock poisoned"))?; + if records + .iter() + .any(|existing| existing.worker_name == record.worker_name) + { + return Err(io::Error::new( + io::ErrorKind::AlreadyExists, + format!( + "spawned worker `{}` is already registered", + record.worker_name + ), + )); + } + records.push(record); + Ok(()) + } + + pub(crate) fn get_internal(&self, worker_name: &str) -> Option { + self.internal_records + .lock() + .ok()? + .iter() + .find(|record| record.worker_name == worker_name) + .cloned() + } + + pub(crate) fn list_internal(&self) -> Vec { + self.internal_records + .lock() + .map(|records| records.clone()) + .unwrap_or_default() + } + + 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) { + let write_rules = record + .scope_delegated + .iter() + .filter(|rule| rule.permission == Permission::Write) + .cloned() + .collect::>(); + parent_scope + .update(|current| current.with_removed_deny_rules(write_rules)) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))?; + } + Ok(removed) + } + + /// Append a new legacy process record and persist the full list. /// error if either persisted write fails; the in-memory state is still /// updated in that case — the next successful write will reconcile. pub async fn add(&self, record: SpawnedWorkerRecord) -> io::Result<()> {