worker: route sub-worker tools through internal sessions

2026-08-06 23:15:30 +09:00
parent d9f399b97b
commit 485918ebe3
2 changed files with 143 additions and 4 deletions
@@ -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::<Vec<_>>();
+        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<ToolOutput, ToolError> {
        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::<Vec<_>>();
+            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<ToolOutput, ToolError> {
        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)
@@ -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<dyn Fn(&SpawnedWorkerRecord) -> 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<ScopeRule>,
+    pub session: InternalWorkerSessionHandle,
+}
+
 pub struct SpawnedWorkerRegistry {
    records: Mutex<Vec<SpawnedWorkerRecord>>,
+    internal_records: std::sync::Mutex<Vec<InternalSpawnedWorkerRecord>>,
    cursors: Mutex<HashMap<String, usize>>,
    mutations: Mutex<()>,
    runtime_dir: Arc<RuntimeDir>,
@@ -58,6 +67,7 @@ impl SpawnedWorkerRegistry {
    pub fn new(runtime_dir: Arc<RuntimeDir>) -> Arc<Self> {
        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<InternalSpawnedWorkerRecord> {
+        self.internal_records
+            .lock()
+            .ok()?
+            .iter()
+            .find(|record| record.worker_name == worker_name)
+            .cloned()
+    }
+
+    pub(crate) fn list_internal(&self) -> Vec<InternalSpawnedWorkerRecord> {
+        self.internal_records
+            .lock()
+            .map(|records| records.clone())
+            .unwrap_or_default()
+    }
+
+    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) {
+            let write_rules = record
+                .scope_delegated
+                .iter()
+                .filter(|rule| rule.permission == Permission::Write)
+                .cloned()
+                .collect::<Vec<_>>();
+            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<()> {