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

worker: run sub-workers as internal sessions

Browse Source
This commit is contained in:
Hare
2026-08-07 01:44:04 +09:00
parent 485918ebe3
commit c79db24016
13 changed files with 531 additions and 2260 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+5 -13
View File
@@ -336,7 +336,6 @@ impl WorkerController {
let fs_for_view = register_worker_tools(
&mut worker,
bash_output_dir,
runtime_dir.socket_path(),
runtime_base.to_path_buf(),
spawned_registry.clone(),
)
@@ -585,10 +584,9 @@ fn wire_event_bridges_on_engine<C, St>(
/// and the Worker-orchestration tools (SubWorkerSpawn + comm) on the Worker's
/// Engine. Returns the WorkdirSession handle used to attach a `WorkerFsView` to
/// the shared state.
async fn register_worker_tools<C, St>(
pub(crate) async fn register_worker_tools<C, St>(
worker: &mut Worker<C, St>,
bash_output_dir: PathBuf,
spawner_socket: PathBuf,
runtime_base: PathBuf,
spawned_registry: Arc<SpawnedWorkerRegistry>,
) -> std::io::Result<Option<workdir::WorkdirSessionHandle>>
@@ -620,9 +618,9 @@ where
let mcp_config = worker.manifest().mcp.clone();
let spawner_name = worker.manifest().worker.name.clone();
let spawner_manifest = worker.manifest().clone();
let spawner_workspace_context = worker.workspace_context_handle();
let parent_notifies = worker.notify_buffer_handle();
let prompts = worker.prompts().clone();
let self_parent_socket = worker.callback_socket().cloned();
// Resolve the existing Worker–Workdir binding into the domain provider.
// Tools only consume the provider handle; they do not own its root, cwd,
// scope, or lifecycle. No-workdir Workers expose no local tools.
@@ -787,12 +785,6 @@ where
// profile feature and require delegation authority up front so enabling
// the surface cannot imply broad child scope by accident.
if feature_config.sub_worker.enabled {
if spawner_manifest.delegation_scope.allow.is_empty() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"[feature.sub_worker].enabled = true requires non-empty [[delegation_scope.allow]]",
));
}
let spawner_cwd = local_filesystem
.as_ref()
.map(|local| local.cwd.clone())
@@ -810,12 +802,12 @@ where
})?;
engine.register_tool(sub_worker_spawn_tool(
spawner_name.clone(),
spawner_socket,
spawner_workspace_context,
parent_notifies,
runtime_base.clone(),
spawner_workspace_root,
spawner_cwd.clone(),
spawned_registry.clone(),
self_parent_socket,
spawner_manifest,
scope_handle,
prompts,
crates/worker/src/discovery.rs
+1 -26
View File
@@ -29,7 +29,6 @@ use session_store::{
use tokio::net::UnixStream;
use tokio::process::Command;
use crate::runtime::dir::SpawnedWorkerRecord;
use crate::runtime::worker_allocation;
use crate::spawn::comm_tools::connect_and_send;
use crate::spawn::registry::SpawnedWorkerRegistry;
@@ -44,7 +43,6 @@ pub struct WorkerDiscovery<St> {
runtime_base: PathBuf,
cwd: Option<PathBuf>,
store_dir: Option<PathBuf>,
spawned_registry: Arc<SpawnedWorkerRegistry>,
}
impl<St> WorkerDiscovery<St>
@@ -56,7 +54,7 @@ where
self_worker_name: String,
runtime_base: PathBuf,
cwd: Option<PathBuf>,
spawned_registry: Arc<SpawnedWorkerRegistry>,
_spawned_registry: Arc<SpawnedWorkerRegistry>,
) -> Self {
let store_dir = store.root_dir();
Self {
@@ -65,7 +63,6 @@ where
runtime_base,
cwd,
store_dir,
spawned_registry,
}
}
@@ -250,20 +247,6 @@ where
}
}
// The live in-memory registry covers just-spawned children even if a
// state write failed after the process became reachable. It is an
// additive visibility hint, not the source of Worker metadata.
for record in self.spawned_registry.list().await {
visible
.entry(record.worker_name.clone())
.or_insert(VisibilityReason::SpawnedChild);
child_sockets.insert(record.worker_name.clone(), record.socket_path.clone());
comm_registry.insert(
record.worker_name.clone(),
CommRegistryInfo::from_record(&record),
);
}
Ok(VisibilitySet {
visible,
child_sockets,
@@ -569,14 +552,6 @@ impl CommRegistryInfo {
scope_delegated: Vec::new(),
}
}
fn from_record(record: &SpawnedWorkerRecord) -> Self {
Self {
registered: true,
socket_path: Some(record.socket_path.clone()),
scope_delegated: record.scope_delegated.clone(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
crates/worker/src/internal_worker.rs
+80 -17
View File
@@ -228,6 +228,7 @@ impl InternalWorkerSessionStatus {
#[derive(Debug, thiserror::Error)]
pub(crate) enum InternalWorkerSessionError {
#[cfg(test)]
#[error("failed to build internal Worker session: {message}")]
Build { message: String },
#[error("internal Worker session is busy")]
@@ -254,7 +255,6 @@ pub(crate) struct InternalWorkerSessionHandle {
store: EphemeralSessionStore,
session_id: SessionId,
segment_id: SegmentId,
last_error: Arc<Mutex<Option<String>>>,
state_changed: Arc<tokio::sync::Notify>,
}
@@ -269,10 +269,6 @@ impl InternalWorkerSessionHandle {
.unwrap_or_default()
}
pub(crate) fn last_error(&self) -> Option<String> {
self.last_error.lock().ok().and_then(|error| error.clone())
}
pub(crate) async fn send(
&self,
input: impl Into<String>,
@@ -310,6 +306,7 @@ impl InternalWorkerSessionHandle {
Ok(())
}
#[cfg(test)]
pub(crate) async fn wait_until_idle(&self) -> InternalWorkerSessionStatus {
loop {
let notified = self.state_changed.notified();
@@ -346,6 +343,7 @@ impl InternalWorkerSessionHandle {
}
/// Start a reusable Internal Worker session and accept its first turn.
#[cfg(test)]
pub(crate) async fn spawn_internal_worker_session(
spec: InternalWorkerSpec,
) -> Result<InternalWorkerSessionHandle, InternalWorkerSessionError> {
@@ -418,13 +416,21 @@ 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(
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();
let segment_id = worker.segment_id();
let (command_tx, mut command_rx) = tokio::sync::mpsc::channel(8);
let status = Arc::new(std::sync::atomic::AtomicU8::new(
InternalWorkerSessionStatus::Idle.encode(),
));
let last_error = Arc::new(Mutex::new(None));
let state_changed = Arc::new(tokio::sync::Notify::new());
let handle = InternalWorkerSessionHandle {
command_tx,
@@ -432,7 +438,6 @@ pub(crate) async fn spawn_internal_worker_session(
store,
session_id,
segment_id,
last_error: last_error.clone(),
state_changed: state_changed.clone(),
};
@@ -445,14 +450,13 @@ pub(crate) async fn spawn_internal_worker_session(
loop {
tokio::select! {
result = &mut run => {
match result {
Ok(_) => {
status.store(InternalWorkerSessionStatus::Idle.encode(), std::sync::atomic::Ordering::Release);
}
Err(error) => {
*last_error.lock().expect("internal Worker session error lock") = Some(error.to_string());
status.store(InternalWorkerSessionStatus::Failed.encode(), std::sync::atomic::Ordering::Release);
}
let turn_status = match result {
Ok(_) => InternalWorkerSessionStatus::Idle,
Err(_) => InternalWorkerSessionStatus::Failed,
};
status.store(turn_status.encode(), std::sync::atomic::Ordering::Release);
if let Some(callback) = &on_turn_end {
callback(turn_status);
}
state_changed.notify_waiters();
break;
@@ -501,9 +505,10 @@ pub(crate) async fn spawn_internal_worker_session(
/// Keeping the normal Store contract makes history/lifecycle/error records identical to a normal
/// Worker while avoiding a second public persistence/catalog policy for helper executions.
#[derive(Clone, Default)]
struct EphemeralSessionStore {
pub(crate) struct EphemeralSessionStore {
entries: Arc<Mutex<HashMap<(SessionId, SegmentId), Vec<LogEntry>>>>,
traces: Arc<Mutex<HashMap<(SessionId, SegmentId), Vec<TraceEntry>>>>,
worker_metadata: Arc<Mutex<HashMap<String, session_store::WorkerMetadata>>>,
}
impl EphemeralSessionStore {
@@ -631,6 +636,65 @@ impl Store for EphemeralSessionStore {
}
}
impl session_store::WorkerMetadataStore for EphemeralSessionStore {
fn write(
&self,
metadata: &session_store::WorkerMetadata,
) -> Result<(), session_store::WorkerStoreError> {
self.worker_metadata
.lock()
.map_err(|_| {
session_store::WorkerStoreError::Io(std::io::Error::other(
"ephemeral metadata lock poisoned",
))
})?
.insert(metadata.worker_name.clone(), metadata.clone());
Ok(())
}
fn read_by_name(
&self,
worker_name: &str,
) -> Result<Option<session_store::WorkerMetadata>, session_store::WorkerStoreError> {
Ok(self
.worker_metadata
.lock()
.map_err(|_| {
session_store::WorkerStoreError::Io(std::io::Error::other(
"ephemeral metadata lock poisoned",
))
})?
.get(worker_name)
.cloned())
}
fn list_names(&self) -> Result<Vec<String>, session_store::WorkerStoreError> {
Ok(self
.worker_metadata
.lock()
.map_err(|_| {
session_store::WorkerStoreError::Io(std::io::Error::other(
"ephemeral metadata lock poisoned",
))
})?
.keys()
.cloned()
.collect())
}
fn delete_by_name(&self, worker_name: &str) -> Result<(), session_store::WorkerStoreError> {
self.worker_metadata
.lock()
.map_err(|_| {
session_store::WorkerStoreError::Io(std::io::Error::other(
"ephemeral metadata lock poisoned",
))
})?
.remove(worker_name);
Ok(())
}
}
#[cfg(test)]
mod tests {
use std::pin::Pin;
@@ -802,7 +866,6 @@ permission = "write"
);
assert_eq!(calls.load(Ordering::SeqCst), 2);
assert!(handle.entries().len() > entries_after_first);
assert!(handle.last_error().is_none());
handle.stop().await.expect("stop Internal Worker session");
assert_eq!(handle.status(), InternalWorkerSessionStatus::Stopped);
crates/worker/src/spawn/comm_tools.rs
+11 -373
View File
@@ -1,13 +1,9 @@
//! Worker-to-Worker communication tools.
//! Parent-facing tools for in-process Internal SubWorker sessions.
//!
//! Three tools in one module: `SubWorkerSend`, `SubWorkerReadOutput`, `SubWorkerStop`,
//! all built on the same `SpawnedWorkerRegistry` handed in by
//! the controller. Each operation is request-response: connect to the
//! target's Unix socket, perform one method exchange, disconnect.
//!
//! These tools only touch Workers listed in the spawner's
//! `SpawnedWorkerRegistry`; there is no machine-wide directory lookup, so
//! the spawner can only reach its own descendants.
//! All five tools share the same parent-owned `SpawnedWorkerRegistry` of typed session handles.
//! There is no Runtime catalog lookup or child socket transport, so a Worker can operate only on
//! its direct Internal children. The socket helper at the bottom remains solely for the legacy
//! top-level Worker callback protocol and is not part of SubWorker communication.
use std::path::Path;
use std::sync::Arc;
@@ -17,12 +13,11 @@ use async_trait::async_trait;
use llm_engine::llm_client::types::{ContentPart, Item, Role};
use llm_engine::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use protocol::stream::{JsonLineReader, JsonLineWriter};
use protocol::{ErrorCode, Event, InvokeKind, Method};
use protocol::{Event, Method};
use serde::{Deserialize, Serialize};
use session_store::LogEntry;
use tokio::net::UnixStream;
use crate::runtime::dir::SpawnedWorkerRecord;
use crate::spawn::registry::SpawnedWorkerRegistry;
/// Timeout applied to each socket-level operation — connect, write,
@@ -62,7 +57,7 @@ impl Tool for SubWorkerListTool {
let _input: SubWorkerListInput = serde_json::from_str(input_json).map_err(|error| {
ToolError::InvalidArgument(format!("invalid SubWorkerList input: {error}"))
})?;
let mut items = self
let items = self
.registry
.list_internal()
.into_iter()
@@ -70,15 +65,6 @@ impl Tool for SubWorkerListTool {
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()))?;
@@ -142,32 +128,7 @@ impl Tool for SubWorkerSendTool {
content: None,
});
}
let record = self
.registry
.get(&input.name)
.await
.ok_or_else(|| unknown_worker_err(&input.name))?;
send_run_and_confirm(&record.socket_path, input.message)
.await
.map_err(|e| match e {
SendRunError::AlreadyRunning => ToolError::ExecutionFailed(format!(
"worker `{}` is already running a turn; wait for it to finish and retry",
input.name
)),
SendRunError::Rejected { code, message } => ToolError::ExecutionFailed(format!(
"worker `{}` rejected the run with {code:?}: {message}",
input.name
)),
SendRunError::Io(msg) => {
ToolError::ExecutionFailed(format!("send to `{}`: {msg}", input.name))
}
})?;
Ok(ToolOutput {
summary: format!("sent message to `{}`", input.name),
content: None,
})
Err(unknown_worker_err(&input.name))
}
}
@@ -237,46 +198,7 @@ impl Tool for SubWorkerReadOutputTool {
content: (!new_text.is_empty()).then_some(new_text),
});
}
let record = self
.registry
.get(&input.name)
.await
.ok_or_else(|| unknown_worker_err(&input.name))?;
let items = match fetch_history(&record.socket_path).await {
Ok(items) => items,
Err(_) => {
return Ok(ToolOutput {
summary: format!("worker `{}` is stopped (unreachable)", input.name),
content: None,
});
}
};
let cursor = self.registry.cursor(&input.name).await;
let new_items = if cursor >= items.len() {
&[] as &[serde_json::Value]
} else {
&items[cursor..]
};
let new_text = extract_assistant_text(new_items);
self.registry.set_cursor(&input.name, items.len()).await;
let summary = if new_text.is_empty() {
format!("worker `{}` running; no new assistant text", input.name)
} else {
let lines = new_text.lines().count();
format!(
"worker `{}`: {lines} new line(s) of assistant text",
input.name
)
};
let content = if new_text.is_empty() {
None
} else {
Some(new_text)
};
Ok(ToolOutput { summary, content })
Err(unknown_worker_err(&input.name))
}
}
@@ -298,9 +220,7 @@ pub fn sub_worker_read_output_tool(registry: Arc<SpawnedWorkerRegistry>) -> Tool
// SubWorkerStop
// ---------------------------------------------------------------------------
const STOP_POD_DESCRIPTION: &str = "Terminate a spawned SubWorker and reclaim the delegated scope. The SubWorker \
receives `Shutdown`; its scope entry is released in the machine-wide \
registry so the parent Worker can spawn a new SubWorker over the same paths.";
const STOP_POD_DESCRIPTION: &str = "Cancel and stop a spawned Internal SubWorker session, remove it from the parent's direct-child registry, and reclaim delegated Write scope.";
struct SubWorkerStopTool {
registry: Arc<SpawnedWorkerRegistry>,
@@ -331,34 +251,7 @@ impl Tool for SubWorkerStopTool {
content: None,
});
}
let record = self
.registry
.get(&input.name)
.await
.ok_or_else(|| unknown_worker_err(&input.name))?;
// Best-effort Shutdown. The child's own `ScopeAllocationGuard`
// releases its entry on clean exit; the parent reclaim below is the
// authoritative operation for removing the child record and returning
// delegated Write scope to the spawner.
let _ = connect_and_send(&record.socket_path, &Method::Shutdown).await;
let scope_summary = summarize_scope(&record);
self.registry
.remove(&record.worker_name)
.await
.map_err(|e| {
ToolError::ExecutionFailed(format!("update spawned worker registry: {e}"))
})?;
Ok(ToolOutput {
summary: format!(
"stopped worker `{}`; reclaimed scope: {scope_summary}",
record.worker_name
),
content: None,
})
Err(unknown_worker_err(&input.name))
}
}
@@ -384,29 +277,6 @@ fn unknown_worker_err(name: &str) -> ToolError {
ToolError::InvalidArgument(format!("no spawned worker named `{name}`"))
}
fn summarize_scope(record: &SpawnedWorkerRecord) -> String {
if record.scope_delegated.is_empty() {
return "(none)".into();
}
let parts: Vec<String> = record
.scope_delegated
.iter()
.map(|rule| {
let perm = match rule.permission {
manifest::Permission::Read => "read",
manifest::Permission::Write => "write",
};
let recursive = if rule.recursive {
""
} else {
" [non-recursive]"
};
format!("{perm}:{}{recursive}", rule.target.display())
})
.collect();
parts.join(", ")
}
/// Connect with a timeout, drain the server's connect-time snapshot,
/// write one `Method` line, flush, and close.
///
@@ -453,125 +323,6 @@ where
}
}
/// Failure modes distinguished by `SubWorkerSend`.
#[derive(Debug)]
pub(crate) enum SendRunError {
/// Target SubWorker responded with `Error { AlreadyRunning }` — the
/// caller can retry once the current turn ends.
AlreadyRunning,
/// Target SubWorker explicitly rejected the run after delivery reached the
/// controller.
Rejected { code: ErrorCode, message: String },
/// Transport, protocol, timeout, or unexpected EOF before acceptance
/// evidence was observed.
Io(String),
}
/// Write `Method::Run` to the target and read back events until we see
/// evidence that the controller accepted the run (`UserMessage`,
/// `TurnStart`, or a user-send `InvokeStart`) or rejected it. The connect-time
/// event prelude is drained before sending the method so large Snapshots and
/// large Run payloads cannot block each other on the same socket. Times out
/// per operation so a stuck Worker doesn't hang the tool.
pub(crate) async fn send_run_and_confirm(socket: &Path, input: String) -> Result<(), SendRunError> {
let stream = tokio::time::timeout(SOCKET_OP_TIMEOUT, UnixStream::connect(socket))
.await
.map_err(|_| SendRunError::Io("connect timed out".into()))?
.map_err(|e| SendRunError::Io(format!("connect: {e}")))?;
let (r, w) = stream.into_split();
let mut writer = JsonLineWriter::new(w);
let mut reader = JsonLineReader::new(r);
loop {
let event = tokio::time::timeout(SOCKET_OP_TIMEOUT, reader.next::<Event>())
.await
.map_err(|_| SendRunError::Io("read initial Snapshot timed out".into()))?
.map_err(|e| SendRunError::Io(format!("read initial Snapshot: {e}")))?;
match event {
Some(Event::Snapshot { .. }) => break,
Some(Event::Alert(_)) => continue,
Some(Event::Error {
code: ErrorCode::AlreadyRunning,
..
}) => return Err(SendRunError::AlreadyRunning),
Some(Event::Error { code, message }) => {
return Err(SendRunError::Rejected { code, message });
}
Some(_) => continue,
None => {
return Err(SendRunError::Io(
"connection closed before initial Snapshot".into(),
));
}
}
}
tokio::time::timeout(
SOCKET_OP_TIMEOUT,
writer.write(&Method::Run {
input: vec![protocol::Segment::text(input)],
}),
)
.await
.map_err(|_| SendRunError::Io("write timed out".into()))?
.map_err(|e| SendRunError::Io(format!("write: {e}")))?;
loop {
let event = tokio::time::timeout(SOCKET_OP_TIMEOUT, reader.next::<Event>())
.await
.map_err(|_| SendRunError::Io("read response timed out".into()))?
.map_err(|e| SendRunError::Io(format!("read response: {e}")))?;
match event {
Some(Event::Error {
code: ErrorCode::AlreadyRunning,
..
}) => return Err(SendRunError::AlreadyRunning),
Some(Event::Error { code, message }) => {
return Err(SendRunError::Rejected { code, message });
}
Some(Event::InvokeStart {
kind: InvokeKind::UserSend,
})
| Some(Event::UserMessage { .. })
| Some(Event::TurnStart { .. }) => return Ok(()),
// Other post-Snapshot events can race with the controller's
// response; keep reading until the Run is accepted or rejected.
Some(_) => continue,
None => return Err(SendRunError::Io("connection closed before response".into())),
}
}
}
/// Connect to a Worker's socket and read the connect-time `Event::Snapshot`.
///
/// Workers deliver the session-log mirror as the first non-Alert event on
/// every new connection, so consuming it is sufficient — no explicit
/// `GetHistory` method round trip. Returns the entries as raw JSON
/// values; callers deserialize as `session_store::LogEntry` if they
/// need typed access.
async fn fetch_history(socket: &Path) -> std::io::Result<Vec<serde_json::Value>> {
let stream = tokio::time::timeout(SOCKET_OP_TIMEOUT, UnixStream::connect(socket))
.await
.map_err(|_| std::io::Error::new(std::io::ErrorKind::TimedOut, "connect timed out"))??;
let (r, _w) = stream.into_split();
let mut reader = JsonLineReader::new(r);
loop {
let event = tokio::time::timeout(SOCKET_OP_TIMEOUT, reader.next::<Event>())
.await
.map_err(|_| std::io::Error::new(std::io::ErrorKind::TimedOut, "read timed out"))??;
match event {
Some(Event::Snapshot { entries, .. }) => return Ok(entries),
Some(_) => continue,
None => {
return Err(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"worker closed connection before Snapshot event",
));
}
}
}
}
fn extract_assistant_text(entries: &[serde_json::Value]) -> String {
let mut out = String::new();
for value in entries {
@@ -656,119 +407,6 @@ mod tests {
})
}
fn serve_initial_events_then_run_ack(
listener: UnixListener,
initial_events: Vec<Event>,
ack: Event,
) -> JoinHandle<Option<Method>> {
tokio::spawn(async move {
let (stream, _) = listener.accept().await.ok()?;
let (r, w) = stream.into_split();
let mut reader = JsonLineReader::new(r);
let mut writer = JsonLineWriter::new(w);
for event in initial_events {
writer.write(&event).await.ok()?;
}
let method = reader.next::<Method>().await.ok().flatten()?;
writer.write(&ack).await.ok()?;
Some(method)
})
}
#[tokio::test]
async fn send_run_and_confirm_keeps_connection_open_until_user_message_ack() {
let tmp = TempDir::new().unwrap();
let socket = tmp.path().join("worker.sock");
let listener = UnixListener::bind(&socket).unwrap();
let received = serve_initial_events_then_run_ack(
listener,
vec![
Event::Alert(Alert {
level: AlertLevel::Warn,
source: AlertSource::Worker,
message: "replayed alert".into(),
timestamp_ms: 0,
}),
snapshot(Vec::new()),
],
Event::UserMessage {
segments: vec![protocol::Segment::text("hello")],
},
);
send_run_and_confirm(&socket, "hello".into()).await.unwrap();
let method = received.await.unwrap().expect("expected method");
match method {
Method::Run { input } => {
assert_eq!(protocol::Segment::flatten_to_text(&input), "hello");
}
other => panic!("expected Run, got {other:?}"),
}
}
#[tokio::test]
async fn send_run_and_confirm_drains_alert_and_large_snapshot_before_large_run() {
let tmp = TempDir::new().unwrap();
let socket = tmp.path().join("worker.sock");
let listener = UnixListener::bind(&socket).unwrap();
let large_snapshot_payload = "s".repeat(2 * 1024 * 1024);
let large_run_payload = "r".repeat(2 * 1024 * 1024);
let received = serve_initial_events_then_run_ack(
listener,
vec![
Event::Alert(Alert {
level: AlertLevel::Warn,
source: AlertSource::Worker,
message: "replayed alert".into(),
timestamp_ms: 0,
}),
snapshot(vec![
serde_json::json!({ "payload": large_snapshot_payload }),
]),
],
Event::InvokeStart {
kind: InvokeKind::UserSend,
},
);
send_run_and_confirm(&socket, large_run_payload.clone())
.await
.unwrap();
let method = received.await.unwrap().expect("expected method");
match method {
Method::Run { input } => {
assert_eq!(
protocol::Segment::flatten_to_text(&input),
large_run_payload
);
}
other => panic!("expected Run, got {other:?}"),
}
}
#[tokio::test]
async fn send_run_and_confirm_reports_already_running() {
let tmp = TempDir::new().unwrap();
let socket = tmp.path().join("worker.sock");
let listener = UnixListener::bind(&socket).unwrap();
let received = serve_initial_events_then_run_ack(
listener,
vec![snapshot(Vec::new())],
Event::Error {
code: ErrorCode::AlreadyRunning,
message: "busy".into(),
},
);
let err = send_run_and_confirm(&socket, "hello".into())
.await
.expect_err("expected AlreadyRunning");
assert!(matches!(err, SendRunError::AlreadyRunning));
assert!(matches!(received.await.unwrap(), Some(Method::Run { .. })));
}
#[tokio::test]
async fn connect_and_send_drains_initial_alert_and_snapshot_before_method() {
let tmp = TempDir::new().unwrap();
crates/worker/src/spawn/registry.rs
+52 -38
View File
@@ -1,22 +1,16 @@
//! Shared registry of Workers spawned by this Worker.
//! Parent-owned registry of direct Internal SubWorker sessions.
//!
//! `SubWorkerSpawn` writes here; the worker-comm tools (`SubWorkerSend`,
//! `SubWorkerReadOutput`, `SubWorkerStop`) read and mutate the same instance. Discovery
//! tools consult this registry together with durable Worker state. Runtime
//! write-through still materialises `spawned_workers.json`, but durable state lives
//! in the spawner's Worker metadata.
//! `SubWorkerSpawn` inserts typed `InternalWorkerSessionHandle`s; List/Send/ReadOutput/Stop use
//! the same in-memory authority. Internal children are not persisted, restored, discovered as
//! Runtime Workers, or addressed through sockets. Restore consumes any legacy persisted process
//! child records only to reclaim their delegated scope and clear obsolete metadata.
//!
//! `SubWorkerReadOutput` additionally owns a per-spawned-worker cursor here so
//! two consecutive reads yield only new assistant text. The cursor is
//! an item-index into the child's history; push-only history makes
//! index stable across reads.
//!
//! Cursors intentionally do not persist; a restored registry starts with
//! fresh read positions.
//! `SubWorkerReadOutput` owns a per-child, process-lifetime history cursor so consecutive reads
//! 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::io;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
@@ -25,7 +19,6 @@ use session_store::{
WorkerMetadataStore, WorkerReclaimedChild, WorkerSpawnedChild, WorkerSpawnedScopeRule,
WorkerStoreError,
};
use tokio::net::UnixStream;
use tokio::sync::Mutex;
use tracing::warn;
@@ -36,7 +29,6 @@ use crate::runtime::worker_allocation;
type RegistryStateWriter = Arc<dyn Fn(&[SpawnedWorkerRecord]) -> io::Result<()> + Send + Sync>;
type RegistryReclaimWriter = Arc<dyn Fn(&SpawnedWorkerRecord) -> io::Result<()> + Send + Sync>;
const RESTORE_REACHABILITY_TIMEOUT: Duration = Duration::from_millis(500);
const REGISTRY_CLEANUP_TIMEOUT: Duration = Duration::from_secs(15);
#[derive(Clone)]
@@ -51,7 +43,7 @@ pub struct SpawnedWorkerRegistry {
internal_records: std::sync::Mutex<Vec<InternalSpawnedWorkerRecord>>,
cursors: Mutex<HashMap<String, usize>>,
mutations: Mutex<()>,
runtime_dir: Arc<RuntimeDir>,
runtime_dir: Option<Arc<RuntimeDir>>,
state_writer: Option<RegistryStateWriter>,
reclaim_writer: Option<RegistryReclaimWriter>,
parent_name: Option<String>,
@@ -70,7 +62,7 @@ impl SpawnedWorkerRegistry {
internal_records: std::sync::Mutex::new(Vec::new()),
cursors: Mutex::new(HashMap::new()),
mutations: Mutex::new(()),
runtime_dir,
runtime_dir: Some(runtime_dir),
state_writer: None,
reclaim_writer: None,
parent_name: None,
@@ -78,6 +70,20 @@ impl SpawnedWorkerRegistry {
})
}
pub(crate) fn new_internal(parent_name: String, parent_scope: SharedScope) -> 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: None,
state_writer: None,
reclaim_writer: None,
parent_name: Some(parent_name),
parent_scope: Some(parent_scope),
})
}
/// Build a registry from the spawner's durable Worker state, pruning child
/// records whose socket path is already gone. The surviving list is
/// written through to both `spawned_workers.json` and Worker state so runtime
@@ -113,7 +119,7 @@ impl SpawnedWorkerRegistry {
.map(|m| m.spawned_children.clone())
.unwrap_or_default();
let mut records = Vec::with_capacity(persisted_children.len());
let records = Vec::with_capacity(persisted_children.len());
let mut pruned_records = Vec::new();
for child in &persisted_children {
let record = match record_from_worker_state(child) {
@@ -127,16 +133,11 @@ impl SpawnedWorkerRegistry {
continue;
}
};
if is_reachable(&record.socket_path).await {
records.push(record);
} else {
warn!(
worker = %record.worker_name,
socket = %record.socket_path.display(),
"dropping unreachable persisted spawned-worker record"
);
pruned_records.push(record);
}
warn!(
worker = %record.worker_name,
"reclaiming legacy persisted process Sub-worker during Internal session restore"
);
pruned_records.push(record);
}
runtime_dir.write_spawned_workers(&records).await?;
@@ -183,7 +184,7 @@ impl SpawnedWorkerRegistry {
internal_records: std::sync::Mutex::new(Vec::new()),
cursors: Mutex::new(HashMap::new()),
mutations: Mutex::new(()),
runtime_dir,
runtime_dir: Some(runtime_dir),
state_writer: Some(state_writer),
reclaim_writer: Some(reclaim_writer),
parent_name: Some(worker_name),
@@ -343,7 +344,9 @@ impl SpawnedWorkerRegistry {
}
async fn persist_records(&self, records: &[SpawnedWorkerRecord]) -> io::Result<()> {
self.runtime_dir.write_spawned_workers(records).await?;
if let Some(runtime_dir) = &self.runtime_dir {
runtime_dir.write_spawned_workers(records).await?;
}
if let Some(write_state) = &self.state_writer {
write_state(records)?;
}
@@ -516,13 +519,24 @@ fn record_from_worker_state(
})
}
impl Drop for SpawnedWorkerRegistry {
fn drop(&mut self) {
let Some(parent_scope) = &self.parent_scope else {
return;
};
let Ok(records) = self.internal_records.lock() else {
return;
};
let write_rules = records
.iter()
.flat_map(|record| record.scope_delegated.iter())
.filter(|rule| rule.permission == Permission::Write)
.cloned()
.collect::<Vec<_>>();
let _ = parent_scope.update(|current| current.with_removed_deny_rules(write_rules));
}
}
fn store_error_to_io(error: WorkerStoreError) -> io::Error {
io::Error::other(error)
}
async fn is_reachable(socket: &Path) -> bool {
tokio::time::timeout(RESTORE_REACHABILITY_TIMEOUT, UnixStream::connect(socket))
.await
.map(|result| result.is_ok())
.unwrap_or(false)
}
crates/worker/src/spawn/tool.rs
+298 -311
View File
@@ -1,18 +1,14 @@
//! `SubWorkerSpawn` tool — launch a new SubWorker process as a child of this one.
//! `SubWorkerSpawn` tool — start a parent-owned Internal Worker session.
//!
//! Wires worker-allocation delegation, child manifest-config construction, subprocess
//! launch, and socket handoff into a single `Tool` implementation. When
//! the LLM calls `SubWorkerSpawn`, a fresh SubWorker runtime command is exec'd in its own
//! process group, the worker-allocation is updated atomically, and the child's
//! first turn is kicked off by handing its socket a `Method::Run`.
//! Resolves a child profile, validates filesystem delegation, constructs a normal Worker with the
//! parent's explicit Workspace authority, installs its enabled features, and hands it to the
//! in-process Internal Worker session actor. No Runtime Worker record, OS process, PID, Unix socket,
//! or machine-wide child allocation is created.
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use client::WorkerRuntimeCommand;
use llm_engine::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use manifest::{
CompactionConfigPartial, DelegationScope, EngineManifestConfig, FileUploadLimitsPartial,
@@ -22,25 +18,17 @@ use manifest::{
WorkerManifest, WorkerManifestConfig, WorkerMetaConfig,
};
use serde::Deserialize;
use tokio::net::UnixStream;
use tokio::process::Command;
use tokio::time::sleep;
use crate::ipc::event;
use crate::PromptLoader;
use crate::controller::register_worker_tools;
use crate::internal_worker::{EphemeralSessionStore, spawn_prepared_internal_worker_session};
use crate::prompt::catalog::PromptCatalog;
use crate::runtime::dir::SpawnedWorkerRecord;
use crate::runtime::worker_allocation::{self, LockFileGuard, ScopeLockError};
use crate::spawn::comm_tools::{SendRunError, send_run_and_confirm};
use crate::spawn::registry::SpawnedWorkerRegistry;
use protocol::WorkerEvent;
/// How long we will wait for the spawned SubWorker's socket to become
/// connectable before treating the spawn as failed.
const SOCKET_WAIT_TIMEOUT: Duration = Duration::from_secs(10);
use crate::worker::{Worker, WorkerFilesystemAuthority};
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct SubWorkerSpawnInput {
/// Identifier for the spawned SubWorker. Must be unique machine-wide.
/// Identifier for the spawned Internal SubWorker. Must be unique among this Worker's direct children.
name: String,
/// Profile selector for child role configuration. Omit or use `default`
/// for the effective child default profile, use `inherit` to derive
@@ -219,14 +207,12 @@ fn parse_spawn_profile_selector(raw: Option<&str>) -> Result<SpawnProfileSelecto
/// child SubWorker and record the handoff locally. Constructed by the Worker
/// controller once per Worker lifetime.
pub struct SubWorkerSpawnTool {
/// Spawner's own worker name — becomes the spawned SubWorker's
/// `delegated_from` in the worker-allocation.
/// Spawner's own Worker name, used for direct-child identity collision checks.
spawner_name: String,
/// Path to the spawner's Unix socket. Handed to the child via
/// `--callback` so its `WorkerEvent` callbacks have somewhere to land.
callback_socket: PathBuf,
/// Root of the `$XDG_RUNTIME_DIR/yoi/` tree, used to predict
/// the spawned SubWorker's socket path before the child has bound it.
workspace_context: crate::worker::WorkerWorkspaceContext,
parent_notifies: crate::ipc::notify_buffer::NotifyBuffer,
/// Runtime-owned root used only for bounded Internal Worker tool artifacts such as Bash spill
/// output. It is not an Internal Worker identity or catalog location.
runtime_base: PathBuf,
/// Inherited runtime workspace root for Profile/project/Ticket/workflow/
/// memory context. SubWorkerSpawn `cwd` must not affect this value.
@@ -234,20 +220,8 @@ pub struct SubWorkerSpawnTool {
/// Directory the spawned SubWorker's tools should use when the LLM did not
/// override it. Defaults to the spawner's cwd.
spawner_cwd: PathBuf,
/// Optional typed runtime command injected by tests. Production resolves
/// the runtime command from `std::env::current_exe()` at launch time.
runtime_command: Option<WorkerRuntimeCommand>,
/// Shared registry of spawned children, also used by the
/// worker-comm tools (`SubWorkerSend` / `SubWorkerReadOutput` / `SubWorkerStop`) and by
/// Worker discovery. Writes the list to runtime and durable Worker state on
/// each add.
/// Parent-owned in-memory registry shared by the five SubWorker tools.
registry: Arc<SpawnedWorkerRegistry>,
/// THIS Worker's own parent-callback socket, if any. After a
/// successful spawn we fire `WorkerEvent::ScopeSubDelegated` upward
/// so the grandparent can register the grandchild directly.
/// `None` for top-level Workers — in that case the re-emission is a
/// no-op.
parent_socket: Option<PathBuf>,
/// Spawner's resolved Manifest. `profile = "inherit"` derives the
/// child config from reusable fields here, and selected profiles are
/// merged into the same internal handoff shape before launch.
@@ -266,36 +240,42 @@ pub struct SubWorkerSpawnTool {
/// This is intentionally separate from `spawner_scope`, which authorizes
/// the current Worker's own direct tools.
delegation_scope: DelegationScope,
internal_client_override: Option<Box<dyn llm_engine::llm_client::LlmClient>>,
}
impl SubWorkerSpawnTool {
#[cfg(test)]
fn with_internal_client(mut self, client: Box<dyn llm_engine::llm_client::LlmClient>) -> Self {
self.internal_client_override = Some(client);
self
}
fn new(
spawner_name: String,
callback_socket: PathBuf,
workspace_context: crate::worker::WorkerWorkspaceContext,
parent_notifies: crate::ipc::notify_buffer::NotifyBuffer,
runtime_base: PathBuf,
workspace_root: PathBuf,
spawner_cwd: PathBuf,
registry: Arc<SpawnedWorkerRegistry>,
parent_socket: Option<PathBuf>,
spawner_manifest: WorkerManifest,
available_profiles: AvailableProfiles,
spawner_scope: SharedScope,
delegation_scope: DelegationScope,
runtime_command: Option<WorkerRuntimeCommand>,
) -> Self {
Self {
spawner_name,
callback_socket,
workspace_context,
parent_notifies,
runtime_base,
workspace_root,
spawner_cwd,
runtime_command,
registry,
parent_socket,
spawner_manifest,
available_profiles,
spawner_scope,
delegation_scope,
internal_client_override: None,
}
}
}
@@ -341,173 +321,109 @@ impl Tool for SubWorkerSpawnTool {
)
.map_err(|e| ToolError::InvalidArgument(format!("{e}")))?;
let predicted_socket = self.runtime_base.join(&input.name).join("sock");
let lock_path = worker_allocation::default_allocation_path()
.map_err(|e| ToolError::ExecutionFailed(format!("worker-allocation path: {e}")))?;
let mut child_config: WorkerManifestConfig = serde_json::from_str(&spawn_config_json)
.map_err(|error| {
ToolError::ExecutionFailed(format!("resolve child manifest: {error}"))
})?;
child_config.delegation_scope = ScopeConfig {
allow: scope_allow.clone(),
deny: Vec::new(),
};
let child_manifest =
WorkerManifest::try_from(WorkerManifestConfig::builtin_defaults().merge(child_config))
.map_err(|error| {
ToolError::ExecutionFailed(format!("resolve child manifest: {error}"))
})?;
let store = EphemeralSessionStore::default();
let filesystem_authority =
WorkerFilesystemAuthority::local(self.workspace_root.clone(), child_cwd.clone());
let mut child = Worker::<Box<dyn llm_engine::llm_client::LlmClient>, EphemeralSessionStore>::from_internal_manifest_with_context(
child_manifest,
store.clone(),
PromptLoader::builtins_only(),
self.workspace_context.clone(),
filesystem_authority,
self.internal_client_override
.as_ref()
.map(|client| client.clone_boxed()),
)
.await
.map_err(|error| ToolError::ExecutionFailed(format!("build Internal Worker: {error}")))?;
let child_scope = child.scope_handle();
let child_registry =
SpawnedWorkerRegistry::new_internal(input.name.clone(), child_scope);
register_worker_tools(
&mut child,
self.runtime_base
.join("internal-workers")
.join(&input.name)
.join("bash-output"),
self.runtime_base.clone(),
child_registry,
)
.await
.map_err(|error| {
ToolError::ExecutionFailed(format!("install Internal Worker features: {error}"))
})?;
let child_name = input.name.clone();
let parent_notifies = self.parent_notifies.clone();
let session = spawn_prepared_internal_worker_session(
child,
store,
input.task.clone(),
Some(Arc::new(move |status| {
parent_notifies.push_notify(
format!("SubWorker `{child_name}` turn ended with status {status:?}. Read its output before making completion decisions."),
false,
);
})),
)
.await
.map_err(|error| {
ToolError::ExecutionFailed(format!("start Internal Worker session: {error}"))
})?;
// Reserve the allocation up front. Spawner's pid is a live
// placeholder; the child will rewrite it via `adopt_allocation`.
{
let mut guard = LockFileGuard::open(&lock_path)
.map_err(|e| ToolError::ExecutionFailed(format!("worker-allocation open: {e}")))?;
worker_allocation::delegate_scope(
&mut guard,
&self.spawner_name,
input.name.clone(),
std::process::id(),
predicted_socket.clone(),
scope_allow.clone(),
&self.delegation_scope,
)
.map_err(worker_allocation_err_to_tool)?;
}
// `start_outcome` covers steps that happen before the child is
// observably alive (exec + socket bind). Once its socket is
// listening, the child owns the allocation and we must not roll
// it back — even if later steps (Method::Run delivery, record
// write) fail, the child is running and will release its own
// entry on exit.
let start_outcome = self
.exec_child(
&input.name,
&spawn_config_json,
&predicted_socket,
&child_cwd,
)
.await;
if let Err(e) = start_outcome {
self.release_reservation(&lock_path, &input.name);
return Err(e);
}
// Child is live. Post-start errors propagate but do not roll
// back the scope allocation — the child already owns it.
//
// Mirror that ownership transfer in the spawner's in-memory
// scope: every `Permission::Write` rule in the delegated scope
// is shadowed by a `deny(Write, target)` so subsequent tool
// calls (Edit/Write) on the delegated paths fail with
// `ReadOnly`. Read access is left intact — the registry only
// arbitrates Write, and keeping Read lets the spawner observe
// the child's intermediate output through Read/Glob/Grep.
// Transfer delegated Write authority within this parent-owned process. The machine-wide
// allocation remains owned by the parent Worker; no fake child PID/socket identity is
// introduced.
let revoke_write: Vec<ScopeRule> = scope_allow
.iter()
.filter(|r| r.permission == Permission::Write)
.filter(|rule| rule.permission == Permission::Write)
.cloned()
.collect();
if !revoke_write.is_empty() {
self.spawner_scope
.update(|cur| cur.with_added_deny_rules(revoke_write.clone()))
.map_err(|e| ToolError::ExecutionFailed(format!("revoke spawner scope: {e}")))?;
.update(|current| current.with_added_deny_rules(revoke_write.clone()))
.map_err(|error| {
ToolError::ExecutionFailed(format!("revoke spawner scope: {error}"))
})?;
}
let record = SpawnedWorkerRecord {
let record = crate::spawn::registry::InternalSpawnedWorkerRecord {
worker_name: input.name.clone(),
socket_path: predicted_socket.clone(),
scope_delegated: scope_allow.clone(),
callback_address: self.callback_socket.clone(),
scope_delegated: scope_allow,
session: session.clone(),
};
self.registry.add(record).await.map_err(|e| {
ToolError::ExecutionFailed(format!("write spawned worker registry: {e}"))
})?;
// Notify this Worker's own parent so the grandparent can register
// the new grandchild directly. Fire-and-forget; top-level Workers
// (with no parent) skip the send inside `fire_and_forget`.
event::fire_and_forget(
self.parent_socket.clone(),
WorkerEvent::ScopeSubDelegated {
parent_worker: self.spawner_name.clone(),
sub_worker: input.name.clone(),
sub_socket: predicted_socket.clone(),
scope: scope_allow,
},
);
send_run_and_confirm(&predicted_socket, input.task.clone())
.await
.map_err(|err| spawn_delivery_error(&input.name, err))?;
if let Err(error) = self.registry.add_internal(record) {
let _ = session.stop().await;
if !revoke_write.is_empty() {
let _ = self
.spawner_scope
.update(|current| current.with_removed_deny_rules(revoke_write));
}
return Err(ToolError::ExecutionFailed(format!(
"register Internal Worker session: {error}"
)));
}
Ok(ToolOutput {
summary: format!(
"spawned worker `{}` listening on {}",
input.name,
predicted_socket.display()
),
summary: format!("spawned internal worker `{}`", input.name),
content: None,
})
}
}
impl SubWorkerSpawnTool {
async fn exec_child(
&self,
worker_name: &str,
spawn_config_json: &str,
predicted_socket: &Path,
child_cwd: &Path,
) -> Result<(), ToolError> {
let runtime_command = match &self.runtime_command {
Some(command) => command.clone(),
None => WorkerRuntimeCommand::resolve().map_err(|error| {
ToolError::ExecutionFailed(format!(
"failed to resolve Worker runtime command: {error}"
))
})?,
};
// Pre-create the child's runtime dir so we have a stable place to
// capture its stderr before it has had a chance to bind anything.
// The child's own `RuntimeDir::create` will `create_dir_all` the
// same path again — that's idempotent. On clean exit the child's
// RuntimeDir Drop tears the dir (and this log) down with it.
let worker_runtime_dir = self.runtime_base.join(worker_name);
tokio::fs::create_dir_all(&worker_runtime_dir)
.await
.map_err(|e| {
ToolError::ExecutionFailed(format!(
"create runtime dir {}: {e}",
worker_runtime_dir.display()
))
})?;
let stderr_path = worker_runtime_dir.join("stderr.log");
let stderr_file = std::fs::File::create(&stderr_path).map_err(|e| {
ToolError::ExecutionFailed(format!("open {}: {e}", stderr_path.display()))
})?;
let mut cmd = Command::new(runtime_command.program());
cmd.args(runtime_command.prefix_args())
.arg("--adopt")
.arg("--callback")
.arg(&self.callback_socket)
.arg("--spawn-config-json")
.arg(spawn_config_json)
.arg("--workspace")
.arg(&self.workspace_root)
.current_dir(child_cwd)
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::from(stderr_file))
.process_group(0);
let child = cmd.spawn().map_err(|e| {
ToolError::ExecutionFailed(format!("failed to spawn `{runtime_command}`: {e}"))
})?;
// Default `kill_on_drop = false` keeps the process alive after
// the `Child` is dropped. We intentionally do not `.wait()` —
// when the spawner later exits, init adopts any remaining
// orphans. Lifecycle tracking lives in `spawned_workers.json`.
drop(child);
match wait_for_socket(predicted_socket, SOCKET_WAIT_TIMEOUT).await {
Ok(()) => Ok(()),
Err(e) => Err(annotate_with_stderr(e, &stderr_path).await),
}
}
fn validate_delegation_scope(&self, scope_allow: &[ScopeRule]) -> Result<(), ToolError> {
if self.delegation_scope.is_empty() && !scope_allow.is_empty() {
return Err(ToolError::InvalidArgument(
@@ -529,12 +445,6 @@ impl SubWorkerSpawnTool {
}
Ok(())
}
fn release_reservation(&self, lock_path: &Path, worker_name: &str) {
if let Ok(mut g) = LockFileGuard::open(lock_path) {
let _ = worker_allocation::release_worker(&mut g, worker_name);
}
}
}
fn parse_scope(rules: &[ScopeRuleInput]) -> Result<Vec<ScopeRule>, ToolError> {
@@ -816,143 +726,44 @@ fn manifest_to_reusable_config(manifest: &WorkerManifest) -> WorkerManifestConfi
/// failure message. Capped so a chatty child can't blow up the LLM's
/// tool-result budget — debugging beyond this should read the file
/// directly.
const STDERR_TAIL_BYTES: usize = 4 * 1024;
async fn annotate_with_stderr(err: ToolError, stderr_path: &Path) -> ToolError {
let tail = match tokio::fs::read(stderr_path).await {
Ok(bytes) => {
let start = bytes.len().saturating_sub(STDERR_TAIL_BYTES);
String::from_utf8_lossy(&bytes[start..]).into_owned()
}
Err(_) => return err,
};
let trimmed = tail.trim();
if trimmed.is_empty() {
return err;
}
match err {
ToolError::ExecutionFailed(msg) => ToolError::ExecutionFailed(format!(
"{msg}\n--- child stderr ({}) ---\n{trimmed}",
stderr_path.display()
)),
other => other,
}
}
async fn wait_for_socket(path: &Path, timeout: Duration) -> Result<(), ToolError> {
let deadline = tokio::time::Instant::now() + timeout;
loop {
if path.exists() {
if let Ok(stream) = UnixStream::connect(path).await {
drop(stream);
return Ok(());
}
}
if tokio::time::Instant::now() >= deadline {
return Err(ToolError::ExecutionFailed(format!(
"spawned worker socket did not appear within {timeout:?}: {}",
path.display()
)));
}
sleep(Duration::from_millis(50)).await;
}
}
fn spawn_delivery_error(worker_name: &str, err: SendRunError) -> ToolError {
match err {
SendRunError::AlreadyRunning => ToolError::ExecutionFailed(format!(
"spawned worker `{worker_name}` rejected its initial task as already running; the worker remains registered and can be inspected or stopped"
)),
SendRunError::Rejected { code, message } => ToolError::ExecutionFailed(format!(
"spawned worker `{worker_name}` rejected its initial task with {code:?}: {message}; the worker remains registered and can be inspected or stopped"
)),
SendRunError::Io(msg) => ToolError::ExecutionFailed(format!(
"spawned worker `{worker_name}` did not confirm initial task delivery: {msg}; the worker remains registered and can be inspected or stopped"
)),
}
}
fn worker_allocation_err_to_tool(e: ScopeLockError) -> ToolError {
match e {
ScopeLockError::NotSubset { .. }
| ScopeLockError::WriteConflict { .. }
| ScopeLockError::DuplicateWorkerName(_)
| ScopeLockError::UnknownWorker(_)
| ScopeLockError::InvalidScope { .. }
| ScopeLockError::SegmentConflict { .. } => ToolError::InvalidArgument(e.to_string()),
ScopeLockError::Io(_) => ToolError::ExecutionFailed(e.to_string()),
}
}
/// Factory for the `SubWorkerSpawn` tool.
pub fn sub_worker_spawn_tool(
spawner_name: String,
callback_socket: PathBuf,
workspace_context: crate::worker::WorkerWorkspaceContext,
parent_notifies: crate::ipc::notify_buffer::NotifyBuffer,
runtime_base: PathBuf,
workspace_root: PathBuf,
spawner_cwd: PathBuf,
registry: Arc<SpawnedWorkerRegistry>,
parent_socket: Option<PathBuf>,
spawner_manifest: WorkerManifest,
spawner_scope: SharedScope,
prompts: Arc<PromptCatalog>,
) -> ToolDefinition {
sub_worker_spawn_tool_impl(
spawner_name,
callback_socket,
workspace_context,
parent_notifies,
runtime_base,
workspace_root,
spawner_cwd,
registry,
parent_socket,
spawner_manifest,
spawner_scope,
prompts,
None,
)
}
#[doc(hidden)]
pub fn sub_worker_spawn_tool_with_runtime_command(
spawner_name: String,
callback_socket: PathBuf,
runtime_base: PathBuf,
workspace_root: PathBuf,
spawner_cwd: PathBuf,
registry: Arc<SpawnedWorkerRegistry>,
parent_socket: Option<PathBuf>,
spawner_manifest: WorkerManifest,
spawner_scope: SharedScope,
prompts: Arc<PromptCatalog>,
runtime_command: WorkerRuntimeCommand,
) -> ToolDefinition {
sub_worker_spawn_tool_impl(
spawner_name,
callback_socket,
runtime_base,
workspace_root,
spawner_cwd,
registry,
parent_socket,
spawner_manifest,
spawner_scope,
prompts,
Some(runtime_command),
)
}
fn sub_worker_spawn_tool_impl(
spawner_name: String,
callback_socket: PathBuf,
workspace_context: crate::worker::WorkerWorkspaceContext,
parent_notifies: crate::ipc::notify_buffer::NotifyBuffer,
runtime_base: PathBuf,
workspace_root: PathBuf,
spawner_cwd: PathBuf,
registry: Arc<SpawnedWorkerRegistry>,
parent_socket: Option<PathBuf>,
spawner_manifest: WorkerManifest,
spawner_scope: SharedScope,
prompts: Arc<PromptCatalog>,
runtime_command: Option<WorkerRuntimeCommand>,
) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(SubWorkerSpawnInput);
@@ -966,7 +777,7 @@ fn sub_worker_spawn_tool_impl(
)
.unwrap_or_else(|e| {
format!(
"Spawn a new SubWorker process to split context for a delegated task. Profile description rendering failed: {e}. Available profiles:\n{}",
"Spawn an in-process Internal SubWorker session to split context for a delegated task. Profile description rendering failed: {e}. Available profiles:\n{}",
available_profiles.compact_list()
)
});
@@ -975,18 +786,17 @@ fn sub_worker_spawn_tool_impl(
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(SubWorkerSpawnTool::new(
spawner_name.clone(),
callback_socket.clone(),
workspace_context.clone(),
parent_notifies.clone(),
runtime_base.clone(),
workspace_root.clone(),
spawner_cwd.clone(),
registry.clone(),
parent_socket.clone(),
spawner_manifest.clone(),
available_profiles,
spawner_scope.clone(),
DelegationScope::from_config(&spawner_manifest.delegation_scope)
.expect("resolved Worker manifest has a valid delegation scope"),
runtime_command.clone(),
));
(meta, tool)
})
@@ -995,9 +805,21 @@ fn sub_worker_spawn_tool_impl(
#[cfg(test)]
mod tests {
use super::*;
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use crate::WorkspaceId;
use async_trait::async_trait;
use futures::Stream;
use llm_engine::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent};
use llm_engine::llm_client::{ClientError, LlmClient, Request};
use manifest::{AuthRef, ModelManifest, SchemeKind, WorkerManifest};
use tempfile::TempDir;
use crate::worker::{
WorkspaceClient, WorkspaceClientError, WorkspaceRequest, WorkspaceResponse,
};
fn abs_rule(path: &Path, permission: Permission) -> ScopeRule {
ScopeRule {
target: path.to_path_buf(),
@@ -1006,6 +828,119 @@ mod tests {
}
}
#[tokio::test]
async fn reviewer_profile_spawns_as_workspace_aware_internal_session() {
let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("../..")
.canonicalize()
.unwrap();
let runtime = TempDir::new().unwrap();
let mut manifest = parent_manifest(&workspace_root, None);
manifest.delegation_scope = ScopeConfig {
allow: vec![abs_rule(&workspace_root, Permission::Read)],
deny: Vec::new(),
};
let spawner_scope = SharedScope::new(Scope::from_config(&manifest.scope).unwrap());
let registry = SpawnedWorkerRegistry::new_internal("parent".into(), spawner_scope.clone());
let workspace_context = crate::worker::WorkerWorkspaceContext::with_client(
Some(WorkspaceId::new("workspace-test").unwrap()),
Arc::new(AvailableWorkspaceClient),
);
let calls = Arc::new(AtomicUsize::new(0));
let parent_notifies = crate::ipc::notify_buffer::NotifyBuffer::new();
let tool = SubWorkerSpawnTool::new(
"parent".into(),
workspace_context,
parent_notifies.clone(),
runtime.path().to_path_buf(),
workspace_root.clone(),
workspace_root.clone(),
registry.clone(),
manifest.clone(),
AvailableProfiles::discover(&workspace_root),
spawner_scope,
DelegationScope::from_config(&manifest.delegation_scope).unwrap(),
)
.with_internal_client(Box::new(ScriptedInternalClient {
calls: calls.clone(),
}));
let input = serde_json::json!({
"name": "reviewer-child",
"profile": "builtin:reviewer",
"task": "review immutable commit",
"scope": [{
"target": workspace_root,
"permission": "read",
"recursive": true
}]
});
let output = tool
.execute(
&serde_json::to_string(&input).unwrap(),
llm_engine::tool::ToolExecutionContext::direct(),
)
.await
.expect("spawn project reviewer as Internal Worker");
assert!(output.summary.contains("internal worker `reviewer-child`"));
let record = registry
.get_internal("reviewer-child")
.expect("Internal reviewer registry record");
assert_eq!(
record.session.wait_until_idle().await,
crate::internal_worker::InternalWorkerSessionStatus::Idle
);
assert_eq!(calls.load(Ordering::SeqCst), 1);
assert_eq!(parent_notifies.len(), 1);
assert!(!runtime.path().join("reviewer-child/sock").exists());
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();
assert!(
listed
.content
.unwrap_or_default()
.contains("reviewer-child")
);
let read = (crate::spawn::comm_tools::sub_worker_read_output_tool(registry.clone()))().1;
let first_output = read
.execute(r#"{"name":"reviewer-child"}"#, context.clone())
.await
.unwrap();
assert!(
first_output
.content
.unwrap_or_default()
.contains("reviewed")
);
let second_output = read
.execute(r#"{"name":"reviewer-child"}"#, context.clone())
.await
.unwrap();
assert!(second_output.content.is_none());
let send = (crate::spawn::comm_tools::sub_worker_send_tool(registry.clone()))().1;
send.execute(
r#"{"name":"reviewer-child","message":"review follow-up"}"#,
context.clone(),
)
.await
.unwrap();
assert_eq!(
record.session.wait_until_idle().await,
crate::internal_worker::InternalWorkerSessionStatus::Idle
);
assert_eq!(calls.load(Ordering::SeqCst), 2);
let stop = (crate::spawn::comm_tools::sub_worker_stop_tool(registry.clone()))().1;
stop.execute(r#"{"name":"reviewer-child"}"#, context)
.await
.unwrap();
assert!(registry.get_internal("reviewer-child").is_none());
}
#[test]
fn spawn_worker_input_schema_includes_optional_cwd() {
let schema = serde_json::to_value(schemars::schema_for!(SubWorkerSpawnInput)).unwrap();
@@ -1103,6 +1038,58 @@ mod tests {
);
}
#[derive(Clone)]
struct ScriptedInternalClient {
calls: Arc<AtomicUsize>,
}
#[async_trait]
impl LlmClient for ScriptedInternalClient {
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
async fn stream(
&self,
_request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + Send>>, ClientError>
{
self.calls.fetch_add(1, Ordering::SeqCst);
Ok(Box::pin(futures::stream::iter(vec![
Ok(LlmEvent::text_block_start(0)),
Ok(LlmEvent::text_delta(0, "reviewed")),
Ok(LlmEvent::text_block_stop(0, None)),
Ok(LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
})),
])))
}
}
#[derive(Debug)]
struct AvailableWorkspaceClient;
impl WorkspaceClient for AvailableWorkspaceClient {
fn workspace_id(&self) -> Option<&str> {
Some("workspace-test")
}
fn kind(&self) -> &str {
"test"
}
fn is_available(&self) -> bool {
true
}
fn execute(
&self,
_request: WorkspaceRequest,
) -> Result<WorkspaceResponse, WorkspaceClientError> {
Err(WorkspaceClientError::Unavailable("not invoked".into()))
}
}
fn parent_manifest(root: &Path, deny: Option<&Path>) -> WorkerManifest {
WorkerManifestConfig {
worker: WorkerMetaConfig {
crates/worker/src/worker.rs
+77
View File
@@ -1162,6 +1162,10 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
self.workspace_context.client()
}
pub(crate) fn workspace_context_handle(&self) -> WorkerWorkspaceContext {
self.workspace_context.clone()
}
pub fn workspace_client_handle(&self) -> Arc<dyn WorkspaceClient> {
self.workspace_context.client_handle()
}
@@ -3949,6 +3953,79 @@ where
Ok(worker)
}
/// Build an in-process Internal Worker without machine-wide allocation or durable Worker metadata.
pub(crate) async fn from_internal_manifest_with_context(
manifest: WorkerManifest,
store: St,
loader: PromptLoader,
workspace_context: WorkerWorkspaceContext,
filesystem_authority: WorkerFilesystemAuthority,
client_override: Option<Box<dyn LlmClient>>,
) -> Result<Self, WorkerError> {
let mut common = prepare_worker_common_with_context(
&manifest,
&loader,
true,
workspace_context,
filesystem_authority,
manifest.scope.clone(),
)?;
if let Some(client) = client_override {
common.client = client;
}
let session_id = session_store::new_session_id();
let segment_id = session_store::new_segment_id();
let mut engine = Engine::new(common.client);
apply_worker_manifest(&mut engine, &manifest.engine);
engine.set_cache_key(Some(segment_id.to_string()));
let scope = SharedScope::new(common.scope);
let workdir_session = workdir_session_from_authority(&common.filesystem_authority, &scope);
let mut worker = Self {
manifest,
engine: Some(engine),
store,
worker_metadata_writer: None,
segment_state: SegmentState::new(session_id, segment_id, 0),
filesystem_authority: common.filesystem_authority,
workdir_session,
workspace_context: common.workspace_context,
scope,
delegation_scope: common.delegation_scope,
hook_builder: HookRegistryBuilder::new(),
interceptor_installed: false,
compact_state: None,
usage_tracker: Arc::new(UsageTracker::new()),
metrics_tracker: Arc::new(crate::compact::metrics_tracker::MetricsTracker::new()),
usage_history: Arc::new(Mutex::new(Vec::new())),
tracker: None,
task_feature: TaskFeature::new(),
system_prompt_template: common.system_prompt_template,
feature_instructions: common.feature_instructions,
alerter: None,
event_tx: None,
in_flight: None,
ai_activity_counter: Arc::new(AtomicUsize::new(0)),
pending_notifies: NotifyBuffer::new(),
pending_attachments: Arc::new(Mutex::new(Vec::<SystemItem>::new())),
scope_allocation: None,
callback_socket: None,
runtime_ticket_role: None,
prompts: common.prompts,
inject_resident_summary: true,
extract_in_flight: Arc::new(AtomicBool::new(false)),
consolidation_in_flight: Arc::new(AtomicBool::new(false)),
extract_pointer: Arc::new(Mutex::new(None)),
memory_task: None,
user_segments: Vec::new(),
sink: SegmentLogSink::new(),
history_persistence_wired: false,
log_writer: None,
};
worker.apply_permissions_from_manifest();
worker.apply_prune_from_manifest();
Ok(worker)
}
/// Build a Worker spawned by another Worker (sibling process).
///
/// Behaves like [`Worker::from_manifest`] but claims the scope
crates/worker/tests/controller_test.rs
+3 -5
View File
@@ -483,7 +483,7 @@ permission = "write"
}
#[tokio::test]
async fn sub_worker_feature_requires_delegation_scope() {
async fn sub_worker_feature_exposure_does_not_require_delegation_scope() {
let manifest = r#"
[worker]
name = "worker-management-feature-test"
@@ -507,11 +507,9 @@ permission = "write"
let worker = make_worker_with_pwd_and_manifest(client, manifest).await.0;
let tmp = tempfile::tempdir().unwrap();
let result = WorkerController::spawn(worker, tmp.path()).await;
assert!(result.is_err());
let message = result.err().unwrap().to_string();
assert!(
message.contains("[feature.sub_worker].enabled = true requires non-empty"),
"unexpected error: {message}"
result.is_ok(),
"feature exposure must not imply delegation authority"
);
}
crates/worker/tests/spawn_worker_test.rs
-743
View File
@@ -1,743 +0,0 @@
//! Integration tests for the `SubWorkerSpawn` tool.
//!
//! These tests exercise the tool's worker-allocation delegation, subprocess
//! launch, socket handoff, and `spawned_workers.json` write through an injected
//! typed runtime command. The mock command exits immediately while a
//! test-owned Unix listener pre-binds the predicted socket path, so the tool
//! sees the "child" as live.
use std::path::{Path, PathBuf};
use std::sync::{LazyLock, Mutex};
use client::WorkerRuntimeCommand;
use llm_engine::tool::{ToolError, ToolOutput};
use manifest::{
AuthRef, ModelManifest, Permission, SchemeKind, Scope, ScopeConfig, ScopeRule, SharedScope,
WorkerManifest, WorkerManifestConfig, WorkerMetaConfig,
};
use protocol::stream::{JsonLineReader, JsonLineWriter};
use protocol::{Event, Method};
use serde_json::json;
use std::sync::Arc;
use tempfile::TempDir;
use tokio::net::UnixListener;
use worker::runtime::dir::{RuntimeDir, SpawnedWorkerRecord};
use worker::runtime::worker_allocation::{self, LockFileGuard};
use worker::spawn::registry::SpawnedWorkerRegistry;
use worker::spawn::tool::sub_worker_spawn_tool_with_runtime_command;
/// Serialises tests that mutate `YOI_RUNTIME_DIR` across the
/// thread-pooled test harness.
static ENV_LOCK: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));
struct EnvGuard {
_lock: std::sync::MutexGuard<'static, ()>,
}
impl EnvGuard {
fn acquire() -> Self {
Self {
_lock: ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner()),
}
}
}
/// Set up a tempdir, point `YOI_RUNTIME_DIR` at it (so
/// `workers.json` and per-Worker runtime subdirs both land in the
/// sandbox), and install a live top-level "spawner" allocation so the
/// tool has something to delegate from. Returns the tempdir (keeps it
/// alive for the test's lifetime), runtime base, spawner socket, and
/// the spawner's runtime dir.
async fn setup_spawner(
spawner_name: &str,
allow_root: &Path,
) -> (TempDir, PathBuf, PathBuf, Arc<RuntimeDir>) {
let tmp = TempDir::new().unwrap();
let runtime_base = tmp.path().to_path_buf();
unsafe {
// Outranking env vars must be cleared so `paths::runtime_dir`
// resolves to our sandbox instead of the developer's real one.
std::env::remove_var("YOI_HOME");
std::env::remove_var("XDG_RUNTIME_DIR");
std::env::set_var("YOI_RUNTIME_DIR", &runtime_base);
}
let spawner_rd = RuntimeDir::create(&runtime_base, spawner_name)
.await
.unwrap();
let spawner_socket = spawner_rd.socket_path();
let _guard = worker_allocation::install_top_level(
spawner_name.into(),
std::process::id(),
spawner_socket.clone(),
vec![ScopeRule {
target: allow_root.to_path_buf(),
permission: Permission::Write,
recursive: true,
}],
session_store::new_segment_id(),
)
.unwrap();
// Leak the guard — the spawner allocation needs to outlive the
// tool call. Dropping it would auto-release the allocation, which
// defeats the point of the test.
std::mem::forget(_guard);
(tmp, runtime_base, spawner_socket, Arc::new(spawner_rd))
}
/// Bind a Unix listener at the path the tool will predict for the
/// spawned worker. The tool only needs the socket to accept a connection
/// and receive one `Method::Run` line; the returned `UnixListener` is
/// read from by the caller in a joined task.
async fn bind_mock_worker_socket(
runtime_base: &Path,
worker_name: &str,
) -> (PathBuf, UnixListener) {
let dir = runtime_base.join(worker_name);
tokio::fs::create_dir_all(&dir).await.unwrap();
let socket = dir.join("sock");
let listener = UnixListener::bind(&socket).unwrap();
(socket, listener)
}
/// Launch a tokio task that accepts connections until one carries a
/// `Method` line, then acknowledges it and returns it. `wait_for_socket`
/// inside the tool makes a probe connection that carries no data, so the
/// task must tolerate an empty connection and keep listening.
fn accept_one_method(listener: UnixListener) -> tokio::task::JoinHandle<Option<Method>> {
tokio::spawn(async move {
loop {
let (stream, _) = listener.accept().await.ok()?;
let (reader, writer) = stream.into_split();
let mut r = JsonLineReader::new(reader);
let mut w = JsonLineWriter::new(writer);
if w.write(&Event::Snapshot {
entries: Vec::new(),
greeting: protocol::Greeting {
worker_name: "child".into(),
cwd: "/tmp".into(),
provider: "test".into(),
model: "test".into(),
scope_summary: String::new(),
tools: Vec::new(),
context_window: 200_000,
context_tokens: 0,
},
status: protocol::WorkerStatus::Idle,
in_flight: Default::default(),
})
.await
.is_err()
{
continue;
}
if let Ok(Some(method)) = r.next::<Method>().await {
w.write(&Event::UserMessage {
segments: vec![protocol::Segment::text("accepted")],
})
.await
.ok()?;
return Some(method);
}
}
})
}
fn mock_runtime_command() -> WorkerRuntimeCommand {
WorkerRuntimeCommand::new(which_true(), Vec::new())
}
fn cwd_recording_runtime_command(script_path: &Path, output_path: &Path) -> WorkerRuntimeCommand {
let output = output_path.display();
std::fs::write(
script_path,
format!(
"tmp=\"{output}.tmp\"\npwd > \"$tmp\"\nprintf '%s\\n' \"$@\" >> \"$tmp\"\nmv \"$tmp\" \"{output}\"\n"
),
)
.unwrap();
WorkerRuntimeCommand::new(which_sh(), vec![script_path.as_os_str().to_os_string()])
}
async fn read_recorded_runtime_invocation(output_path: &Path) -> Vec<String> {
for _ in 0..50 {
if let Ok(content) = std::fs::read_to_string(output_path) {
return content.lines().map(str::to_owned).collect();
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
panic!(
"runtime command did not record invocation at {}",
output_path.display()
);
}
/// `/bin/true` only exists on FHS-compliant systems. Resolve it via PATH
/// so the tests work regardless of distro.
fn which_true() -> String {
for dir in std::env::var_os("PATH")
.map(|p| std::env::split_paths(&p).collect::<Vec<_>>())
.unwrap_or_default()
{
let candidate = dir.join("true");
if candidate.is_file() {
return candidate.to_string_lossy().into_owned();
}
}
"/bin/true".into()
}
fn which_sh() -> String {
for dir in std::env::var_os("PATH")
.map(|p| std::env::split_paths(&p).collect::<Vec<_>>())
.unwrap_or_default()
{
let candidate = dir.join("sh");
if candidate.is_file() {
return candidate.to_string_lossy().into_owned();
}
}
"/bin/sh".into()
}
/// Tests don't exercise the model — they intercept the spawned
/// child via a mock socket — but `sub_worker_spawn_tool` needs a value to
/// embed in the overlay TOML. Any well-formed `ModelManifest` works.
fn dummy_model() -> ModelManifest {
ModelManifest {
scheme: Some(SchemeKind::Anthropic),
base_url: None,
model_id: Some("claude-test".into()),
auth: Some(AuthRef::None),
capability: None,
..Default::default()
}
}
fn dummy_manifest(allow_root: &Path) -> WorkerManifest {
dummy_manifest_with_delegation(allow_root, true)
}
fn dummy_manifest_with_delegation(allow_root: &Path, allow_delegation: bool) -> WorkerManifest {
let direct_scope = ScopeConfig {
allow: vec![ScopeRule {
target: allow_root.to_path_buf(),
permission: Permission::Write,
recursive: true,
}],
deny: Vec::new(),
};
let delegation_scope = if allow_delegation {
direct_scope.clone()
} else {
ScopeConfig::default()
};
dummy_manifest_with_scopes(direct_scope, delegation_scope)
}
fn dummy_manifest_with_scopes(
direct_scope: ScopeConfig,
delegation_scope: ScopeConfig,
) -> WorkerManifest {
WorkerManifestConfig {
worker: WorkerMetaConfig {
name: Some("root".into()),
prompt_pack: None,
},
model: dummy_model(),
scope: direct_scope,
delegation_scope,
..Default::default()
}
.try_into()
.unwrap()
}
fn builtin_prompts() -> Arc<worker::PromptCatalog> {
worker::PromptCatalog::builtins_only().unwrap()
}
/// Spawner-side `SharedScope` mirroring the `allow_root` granted by
/// `setup_spawner`. The tool revokes Write rules from this scope on
/// successful spawn — tests can `load()` it to assert the
/// revocation took effect.
fn shared_scope_for(allow_root: &Path) -> SharedScope {
SharedScope::new(Scope::writable(allow_root).unwrap())
}
fn clear_env() {
unsafe {
std::env::remove_var("YOI_RUNTIME_DIR");
}
}
#[tokio::test]
async fn spawn_worker_launches_runtime_in_workspace_and_process_cwd() {
let _env = EnvGuard::acquire();
let allow_root = TempDir::new().unwrap();
let child_cwd = allow_root.path().join("child-cwd");
std::fs::create_dir(&child_cwd).unwrap();
let script = allow_root.path().join("record-pwd.sh");
let output_path = allow_root.path().join("pwd.txt");
let (_tmp, runtime_base, spawner_socket, spawner_rd) =
setup_spawner("root", allow_root.path()).await;
let (_predicted_socket, listener) = bind_mock_worker_socket(&runtime_base, "child-cwd").await;
let received = accept_one_method(listener);
let registry = SpawnedWorkerRegistry::new(spawner_rd);
let def = sub_worker_spawn_tool_with_runtime_command(
"root".into(),
spawner_socket,
runtime_base,
allow_root.path().to_path_buf(),
allow_root.path().to_path_buf(),
registry,
None,
dummy_manifest(allow_root.path()),
shared_scope_for(allow_root.path()),
builtin_prompts(),
cwd_recording_runtime_command(&script, &output_path),
);
let (_meta, tool) = def();
let input = json!({
"name": "child-cwd",
"task": "hello",
"profile": "inherit",
"cwd": child_cwd.to_str().unwrap(),
"scope": [{
"target": allow_root.path().to_str().unwrap(),
"permission": "write"
}]
})
.to_string();
tool.execute(&input, Default::default()).await.unwrap();
assert!(matches!(received.await.unwrap(), Some(Method::Run { .. })));
let invocation = read_recorded_runtime_invocation(&output_path).await;
assert_eq!(invocation[0], child_cwd.to_str().unwrap());
assert!(
invocation
.windows(2)
.any(|pair| pair[0] == "--workspace" && pair[1] == allow_root.path().to_str().unwrap()),
"invocation should carry inherited workspace root: {invocation:?}"
);
assert!(
!invocation.iter().any(|arg| arg == "--tool-cwd"),
"cwd should be process current directory, not a runtime argument: {invocation:?}"
);
clear_env();
}
#[tokio::test]
async fn spawn_worker_omitted_cwd_preserves_spawner_cwd() {
let _env = EnvGuard::acquire();
let allow_root = TempDir::new().unwrap();
let script = allow_root.path().join("record-pwd.sh");
let output_path = allow_root.path().join("pwd.txt");
let (_tmp, runtime_base, spawner_socket, spawner_rd) =
setup_spawner("root", allow_root.path()).await;
let (_predicted_socket, listener) =
bind_mock_worker_socket(&runtime_base, "child-default-cwd").await;
let received = accept_one_method(listener);
let registry = SpawnedWorkerRegistry::new(spawner_rd);
let def = sub_worker_spawn_tool_with_runtime_command(
"root".into(),
spawner_socket,
runtime_base,
allow_root.path().to_path_buf(),
allow_root.path().to_path_buf(),
registry,
None,
dummy_manifest(allow_root.path()),
shared_scope_for(allow_root.path()),
builtin_prompts(),
cwd_recording_runtime_command(&script, &output_path),
);
let (_meta, tool) = def();
let input = json!({
"name": "child-default-cwd",
"task": "hello",
"profile": "inherit",
"scope": [{
"target": allow_root.path().to_str().unwrap(),
"permission": "write"
}]
})
.to_string();
tool.execute(&input, Default::default()).await.unwrap();
assert!(matches!(received.await.unwrap(), Some(Method::Run { .. })));
let invocation = read_recorded_runtime_invocation(&output_path).await;
assert_eq!(invocation[0], allow_root.path().to_str().unwrap());
assert!(
!invocation.iter().any(|arg| arg == "--tool-cwd"),
"omitted cwd should preserve spawner cwd as process cwd: {invocation:?}"
);
clear_env();
}
#[tokio::test]
async fn spawn_worker_delegates_scope_and_sends_run() {
let _env = EnvGuard::acquire();
let allow_root = TempDir::new().unwrap();
let (_tmp, runtime_base, spawner_socket, spawner_rd) =
setup_spawner("root", allow_root.path()).await;
let (_predicted_socket, listener) = bind_mock_worker_socket(&runtime_base, "child").await;
let received = accept_one_method(listener);
let registry = SpawnedWorkerRegistry::new(spawner_rd.clone());
let spawner_scope = shared_scope_for(allow_root.path());
let def = sub_worker_spawn_tool_with_runtime_command(
"root".into(),
spawner_socket.clone(),
runtime_base.clone(),
allow_root.path().to_path_buf(),
allow_root.path().to_path_buf(),
registry,
None,
dummy_manifest(allow_root.path()),
spawner_scope.clone(),
builtin_prompts(),
mock_runtime_command(),
);
let (_meta, tool) = def();
let input = json!({
"name": "child",
"task": "hello",
"profile": "inherit",
"scope": [{
"target": allow_root.path().to_str().unwrap(),
"permission": "write"
}]
})
.to_string();
// Pre-spawn: the spawner can write to the delegated path.
assert!(
spawner_scope
.load()
.is_writable(&allow_root.path().join("a.txt"))
);
let output: ToolOutput = tool.execute(&input, Default::default()).await.unwrap();
assert!(
output.summary.contains("child"),
"summary: {}",
output.summary
);
// Verify the tool delivered Method::Run to the socket.
let method = received.await.unwrap().expect("expected one Method line");
match method {
Method::Run { input } => match input.as_slice() {
[protocol::Segment::Text { content }] => assert_eq!(content, "hello"),
other => panic!("expected single Text segment, got {other:?}"),
},
other => panic!("expected Run, got {other:?}"),
}
// Verify worker_allocation has the child allocation under `root`.
let lock_path = worker_allocation::default_allocation_path().unwrap();
let guard = LockFileGuard::open(&lock_path).unwrap();
let child = guard
.data()
.find("child")
.expect("child allocation missing after spawn");
assert_eq!(child.delegated_from.as_deref(), Some("root"));
drop(guard);
// Verify spawned_workers.json was written.
let spawned_file = spawner_rd.path().join("spawned_workers.json");
let contents = std::fs::read_to_string(&spawned_file).unwrap();
let records: Vec<SpawnedWorkerRecord> = serde_json::from_str(&contents).unwrap();
assert_eq!(records.len(), 1);
assert_eq!(records[0].worker_name, "child");
assert_eq!(records[0].callback_address, spawner_socket);
// Post-spawn: the spawner's runtime scope has been demoted on the
// delegated path. Write is gone, Read remains.
let post = spawner_scope.load();
assert_eq!(
post.permission_at(&allow_root.path().join("a.txt")),
Some(Permission::Read),
"spawner should still be able to read delegated path"
);
clear_env();
}
#[tokio::test]
async fn spawn_worker_requires_explicit_delegation_even_with_direct_scope() {
let _env = EnvGuard::acquire();
let allow_root = TempDir::new().unwrap();
let (_tmp, runtime_base, spawner_socket, spawner_rd) =
setup_spawner("root", allow_root.path()).await;
let manifest = dummy_manifest_with_delegation(allow_root.path(), false);
let direct = Scope::from_config(&manifest.scope).unwrap();
assert!(direct.is_writable(&allow_root.path().join("direct.txt")));
let registry = SpawnedWorkerRegistry::new(spawner_rd.clone());
let def = sub_worker_spawn_tool_with_runtime_command(
"root".into(),
spawner_socket,
runtime_base,
allow_root.path().to_path_buf(),
allow_root.path().to_path_buf(),
registry,
None,
manifest,
shared_scope_for(allow_root.path()),
builtin_prompts(),
mock_runtime_command(),
);
let (_meta, tool) = def();
let input = json!({
"name": "child-no-delegation",
"task": "hello",
"profile": "inherit",
"scope": [{
"target": allow_root.path().to_str().unwrap(),
"permission": "write"
}]
})
.to_string();
let err = tool.execute(&input, Default::default()).await.unwrap_err();
match err {
ToolError::InvalidArgument(message) => {
assert!(message.contains("no delegation scope grant"), "{message}");
assert!(message.contains("direct filesystem scope"), "{message}");
}
other => panic!("expected InvalidArgument, got {other:?}"),
}
clear_env();
}
#[tokio::test]
async fn spawn_worker_rejects_child_non_recursive_scope_under_parent_non_recursive_delegation() {
let _env = EnvGuard::acquire();
let allow_root = TempDir::new().unwrap();
let child = allow_root.path().join("child");
std::fs::create_dir(&child).unwrap();
let (_tmp, runtime_base, spawner_socket, spawner_rd) =
setup_spawner("root", allow_root.path()).await;
let direct_scope = ScopeConfig {
allow: vec![ScopeRule {
target: allow_root.path().to_path_buf(),
permission: Permission::Write,
recursive: true,
}],
deny: Vec::new(),
};
let delegation_scope = ScopeConfig {
allow: vec![ScopeRule {
target: allow_root.path().to_path_buf(),
permission: Permission::Write,
recursive: false,
}],
deny: Vec::new(),
};
let manifest = dummy_manifest_with_scopes(direct_scope, delegation_scope);
let registry = SpawnedWorkerRegistry::new(spawner_rd.clone());
let def = sub_worker_spawn_tool_with_runtime_command(
"root".into(),
spawner_socket,
runtime_base,
allow_root.path().to_path_buf(),
allow_root.path().to_path_buf(),
registry,
None,
manifest,
shared_scope_for(allow_root.path()),
builtin_prompts(),
mock_runtime_command(),
);
let (_meta, tool) = def();
let input = json!({
"name": "child-nonrecursive-overgrant",
"task": "hello",
"profile": "inherit",
"scope": [{
"target": child.to_str().unwrap(),
"permission": "write",
"recursive": false
}]
})
.to_string();
let err = tool.execute(&input, Default::default()).await.unwrap_err();
match err {
ToolError::InvalidArgument(message) => {
assert!(
message.contains("outside this Worker's delegation scope grant"),
"{message}"
);
}
other => panic!("expected InvalidArgument, got {other:?}"),
}
clear_env();
}
#[tokio::test]
async fn spawn_worker_rejects_scope_outside_spawner() {
let _env = EnvGuard::acquire();
let allow_root = TempDir::new().unwrap();
let outside = TempDir::new().unwrap();
let (_tmp, runtime_base, spawner_socket, spawner_rd) =
setup_spawner("root", allow_root.path()).await;
let registry = SpawnedWorkerRegistry::new(spawner_rd);
let spawner_scope = shared_scope_for(allow_root.path());
let def = sub_worker_spawn_tool_with_runtime_command(
"root".into(),
spawner_socket,
runtime_base,
allow_root.path().to_path_buf(),
allow_root.path().to_path_buf(),
registry,
None,
dummy_manifest(allow_root.path()),
spawner_scope.clone(),
builtin_prompts(),
mock_runtime_command(),
);
let (_meta, tool) = def();
// Request write access to a path the spawner doesn't own.
let input = json!({
"name": "child",
"task": "nope",
"profile": "inherit",
"scope": [{
"target": outside.path().to_str().unwrap(),
"permission": "write"
}]
})
.to_string();
let err = tool.execute(&input, Default::default()).await.unwrap_err();
match err {
ToolError::InvalidArgument(msg) => {
assert!(
msg.contains("outside this Worker's delegation scope grant"),
"expected delegation-scope wording: {msg}"
);
}
other => panic!("expected InvalidArgument, got {other:?}"),
}
// The spawner's allocation is unchanged; no "child" appeared.
let lock_path = worker_allocation::default_allocation_path().unwrap();
let guard = LockFileGuard::open(&lock_path).unwrap();
assert!(guard.data().find("child").is_none());
// Failed spawn must not have demoted the spawner's scope either.
assert!(
spawner_scope
.load()
.is_writable(&allow_root.path().join("a.txt"))
);
clear_env();
}
#[tokio::test]
async fn spawn_worker_rolls_back_reservation_when_socket_never_appears() {
let _env = EnvGuard::acquire();
let allow_root = TempDir::new().unwrap();
let (_tmp, runtime_base, spawner_socket, spawner_rd) =
setup_spawner("root", allow_root.path()).await;
// Deliberately do NOT bind a socket at the predicted path. The
// tool's wait_for_socket should time out, triggering rollback.
// `SOCKET_WAIT_TIMEOUT` is 10s in production; we override via a
// tighter env-based lock path and just accept the wait in test.
// To keep the test fast, use a shorter wait by constructing a
// short-lived separate instance.
//
// As the tool's timeout is internal, we accept the 10s wait here —
// marked with `// slow_test`. Keep the rest of the test suite fast
// by running this test alone when iterating.
let registry = SpawnedWorkerRegistry::new(spawner_rd);
let spawner_scope = shared_scope_for(allow_root.path());
let def = sub_worker_spawn_tool_with_runtime_command(
"root".into(),
spawner_socket,
runtime_base,
allow_root.path().to_path_buf(),
allow_root.path().to_path_buf(),
registry,
None,
dummy_manifest(allow_root.path()),
spawner_scope.clone(),
builtin_prompts(),
mock_runtime_command(),
);
let (_meta, tool) = def();
let input = json!({
"name": "ghost",
"task": "will never be delivered",
"profile": "inherit",
"scope": [{
"target": allow_root.path().to_str().unwrap(),
"permission": "write"
}]
})
.to_string();
let err = tool.execute(&input, Default::default()).await.unwrap_err();
match err {
ToolError::ExecutionFailed(msg) => {
assert!(
msg.contains("socket did not appear"),
"expected socket timeout wording: {msg}"
);
}
other => panic!("expected ExecutionFailed, got {other:?}"),
}
// Rollback assertion: the reserved "ghost" allocation is gone.
let lock_path = worker_allocation::default_allocation_path().unwrap();
let guard = LockFileGuard::open(&lock_path).unwrap();
assert!(
guard.data().find("ghost").is_none(),
"allocation was not rolled back after socket wait timed out"
);
// Spawner's runtime scope must also be untouched — revoke is
// performed only after exec_child succeeds.
assert!(
spawner_scope
.load()
.is_writable(&allow_root.path().join("a.txt"))
);
clear_env();
}
crates/worker/tests/worker_comm_tools_test.rs
-730
View File
@@ -1,730 +0,0 @@
//! Integration tests for the worker-comm tools (`SubWorkerSend`,
//! `SubWorkerReadOutput`, `SubWorkerStop`).
//!
//! The real child Worker binary is not started. Instead each test stands
//! up a mock `UnixListener` that speaks the socket protocol directly:
//! it emits the connect-time `Event::Snapshot`, accepts methods such as
//! `Method::Run` / `Method::Shutdown`, and responds with the relevant
//! events when needed. This keeps the tests fast and independent of the
//! LLM layer — the tools are exercised for their wire behaviour alone.
use std::path::{Path, PathBuf};
use std::sync::{Arc, LazyLock, Mutex};
use llm_engine::llm_client::types::{ContentPart, Item, Role};
use llm_engine::tool::ToolOutput;
use manifest::{Permission, Scope, ScopeRule, SharedScope};
use protocol::stream::{JsonLineReader, JsonLineWriter};
use protocol::{ErrorCode, Event, Greeting, Method};
use serde_json::json;
use session_store::FsStore;
use session_store::{CombinedStore, FsWorkerStore, WorkerMetadataStore};
use tempfile::TempDir;
use tokio::net::UnixListener;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use worker::runtime::dir::{RuntimeDir, SpawnedWorkerRecord};
use worker::runtime::worker_allocation::{self, LockFileGuard};
use worker::spawn::comm_tools::{
sub_worker_read_output_tool, sub_worker_send_tool, sub_worker_stop_tool,
};
use worker::spawn::registry::SpawnedWorkerRegistry;
/// Serialises env-mutating tests. The test harness runs tasks across
/// threads, and `YOI_RUNTIME_DIR` is a process-wide resource.
static ENV_LOCK: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));
/// Take `ENV_LOCK` and clear any env vars that would outrank
/// `YOI_RUNTIME_DIR` in `paths::runtime_dir` resolution; restore
/// previous values on drop.
struct EnvGuard {
prev_home: Option<String>,
prev_xdg: Option<String>,
_lock: std::sync::MutexGuard<'static, ()>,
}
impl EnvGuard {
fn acquire() -> Self {
let lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prev_home = std::env::var("YOI_HOME").ok();
let prev_xdg = std::env::var("XDG_RUNTIME_DIR").ok();
unsafe {
std::env::remove_var("YOI_HOME");
std::env::remove_var("XDG_RUNTIME_DIR");
}
Self {
prev_home,
prev_xdg,
_lock: lock,
}
}
}
impl Drop for EnvGuard {
fn drop(&mut self) {
unsafe {
match &self.prev_home {
Some(v) => std::env::set_var("YOI_HOME", v),
None => std::env::remove_var("YOI_HOME"),
}
match &self.prev_xdg {
Some(v) => std::env::set_var("XDG_RUNTIME_DIR", v),
None => std::env::remove_var("XDG_RUNTIME_DIR"),
}
std::env::remove_var("YOI_RUNTIME_DIR");
}
}
}
/// Create a spawner-owned `RuntimeDir` + `SpawnedWorkerRegistry` scoped to
/// a fresh tempdir. The returned `TempDir` must be kept alive by the
/// caller for the duration of the test.
async fn setup_registry() -> (TempDir, Arc<SpawnedWorkerRegistry>, Arc<RuntimeDir>) {
let tmp = TempDir::new().unwrap();
let rd = RuntimeDir::create(tmp.path(), "spawner").await.unwrap();
let rd = Arc::new(rd);
let registry = SpawnedWorkerRegistry::new(rd.clone());
(tmp, registry, rd)
}
/// Register a fake spawned-child record pointing at a given socket
/// path, with a trivial write-scope for `scope_path`. Does not touch
/// workers.json.
async fn register_child(
registry: &SpawnedWorkerRegistry,
name: &str,
socket: &Path,
scope_path: &Path,
) {
let record = SpawnedWorkerRecord {
worker_name: name.into(),
socket_path: socket.to_path_buf(),
scope_delegated: vec![ScopeRule {
target: scope_path.to_path_buf(),
permission: Permission::Write,
recursive: true,
}],
callback_address: "/dev/null".into(),
};
registry.add(record).await.unwrap();
}
/// Bind a Unix listener at a socket path inside the given directory.
async fn bind_mock_socket(dir: &Path, name: &str) -> (PathBuf, UnixListener) {
let socket = dir.join(format!("{name}.sock"));
let listener = UnixListener::bind(&socket).unwrap();
(socket, listener)
}
/// Minimal connect-time snapshot used by mock socket servers.
fn empty_snapshot() -> Event {
Event::Snapshot {
entries: Vec::new(),
greeting: Greeting {
worker_name: "child".into(),
cwd: "/tmp".into(),
provider: "anthropic".into(),
model: "x".into(),
scope_summary: String::new(),
tools: Vec::new(),
context_window: 200_000,
context_tokens: 0,
},
status: protocol::WorkerStatus::Idle,
in_flight: Default::default(),
}
}
/// Accept one connection, send the protocol's connect-time snapshot,
/// and read exactly one `Method` line from it.
/// The reader half is kept open; caller awaits the returned handle.
fn accept_one_method(listener: UnixListener) -> JoinHandle<Option<Method>> {
tokio::spawn(async move {
let (stream, _) = listener.accept().await.ok()?;
let (r, w) = stream.into_split();
let mut reader = JsonLineReader::new(r);
let mut writer = JsonLineWriter::new(w);
writer.write(&empty_snapshot()).await.ok()?;
reader.next::<Method>().await.ok().flatten()
})
}
/// Accept one connection, send the protocol's connect-time snapshot,
/// read one `Method`, then write `response` back. Used by `SubWorkerSend`
/// tests to mock the real controller's `TurnStart` acknowledgement (or
/// its `AlreadyRunning` rejection).
fn accept_method_and_respond(
listener: UnixListener,
response: Event,
) -> JoinHandle<Option<Method>> {
tokio::spawn(async move {
let (stream, _) = listener.accept().await.ok()?;
let (r, w) = stream.into_split();
let mut reader = JsonLineReader::new(r);
let mut writer = JsonLineWriter::new(w);
writer.write(&empty_snapshot()).await.ok()?;
let method = reader.next::<Method>().await.ok().flatten();
if method.is_some() {
let _ = writer.write(&response).await;
}
method
})
}
/// Pretend to be a spawned Worker whose connect-time snapshot carries a
/// fixed set of assistant items. Sends `Event::Snapshot` immediately on
/// every accept — the real Worker does the same, so `SubWorkerReadOutput`'s
/// `fetch_history` just consumes the first non-Alert event.
fn serve_history(listener: UnixListener, items: Vec<Item>) -> JoinHandle<()> {
tokio::spawn(async move {
loop {
let Ok((stream, _)) = listener.accept().await else {
return;
};
let (_r, w) = stream.into_split();
let mut writer = JsonLineWriter::new(w);
let entries: Vec<serde_json::Value> = items
.iter()
.map(|item| {
let entry = session_store::LogEntry::AssistantItem {
ts: 0,
item: session_store::LoggedItem::from(item),
};
serde_json::to_value(&entry).unwrap()
})
.collect();
let event = Event::Snapshot {
entries,
greeting: Greeting {
worker_name: "child".into(),
cwd: "/tmp".into(),
provider: "anthropic".into(),
model: "x".into(),
scope_summary: String::new(),
tools: Vec::new(),
context_window: 200_000,
context_tokens: 0,
},
status: protocol::WorkerStatus::Idle,
in_flight: Default::default(),
};
let _ = writer.write(&event).await;
}
})
}
fn serve_worker_methods(listener: UnixListener) -> mpsc::Receiver<Method> {
let (tx, rx) = mpsc::channel(8);
tokio::spawn(async move {
loop {
let Ok((stream, _)) = listener.accept().await else {
return;
};
let (r, w) = stream.into_split();
let mut reader = JsonLineReader::new(r);
let mut writer = JsonLineWriter::new(w);
if writer.write(&empty_snapshot()).await.is_err() {
continue;
}
let Some(method) = reader.next::<Method>().await.ok().flatten() else {
continue;
};
let is_shutdown = matches!(method, Method::Shutdown);
if matches!(method, Method::Run { .. }) {
let _ = writer.write(&Event::TurnStart { turn: 1 }).await;
}
if tx.send(method).await.is_err() || is_shutdown {
return;
}
}
});
rx
}
fn assistant(text: &str) -> Item {
Item::Message {
id: None,
role: Role::Assistant,
content: vec![ContentPart::Text { text: text.into() }],
status: None,
}
}
// ---------------------------------------------------------------------------
// SubWorkerSend
// ---------------------------------------------------------------------------
#[tokio::test]
async fn send_to_worker_delivers_run_method() {
let (tmp, registry, _rd) = setup_registry().await;
let (socket, listener) = bind_mock_socket(tmp.path(), "child").await;
// Mock the controller's accept path: after reading the method,
// ack with `TurnStart` so `SubWorkerSend`'s confirmation loop succeeds.
let received = accept_method_and_respond(listener, Event::TurnStart { turn: 1 });
register_child(&registry, "child", &socket, tmp.path()).await;
let def = sub_worker_send_tool(registry);
let (_meta, tool) = def();
let input = json!({ "name": "child", "message": "hello there" }).to_string();
let output: ToolOutput = tool.execute(&input, Default::default()).await.unwrap();
assert!(
output.summary.contains("child"),
"summary: {}",
output.summary
);
let method = received.await.unwrap().expect("expected a method");
match method {
Method::Run { input } => match input.as_slice() {
[protocol::Segment::Text { content }] => assert_eq!(content, "hello there"),
other => panic!("expected single Text segment, got {other:?}"),
},
other => panic!("expected Run, got {other:?}"),
}
}
#[tokio::test]
async fn send_to_worker_errors_on_unknown_worker() {
let (_tmp, registry, _rd) = setup_registry().await;
let def = sub_worker_send_tool(registry);
let (_meta, tool) = def();
let input = json!({ "name": "nope", "message": "hi" }).to_string();
let err = tool.execute(&input, Default::default()).await.unwrap_err();
assert!(err.to_string().contains("no spawned worker"), "{err}");
}
#[tokio::test]
async fn send_to_worker_errors_when_worker_already_running() {
let (tmp, registry, _rd) = setup_registry().await;
let (socket, listener) = bind_mock_socket(tmp.path(), "child").await;
// Respond with the same `Error { AlreadyRunning }` that the real
// controller emits when `Method::Run` arrives during RUNNING.
let received = accept_method_and_respond(
listener,
Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn".into(),
},
);
register_child(&registry, "child", &socket, tmp.path()).await;
let def = sub_worker_send_tool(registry);
let (_meta, tool) = def();
let input = json!({ "name": "child", "message": "hi" }).to_string();
let err = tool.execute(&input, Default::default()).await.unwrap_err();
assert!(
err.to_string().contains("already running"),
"expected AlreadyRunning wording: {err}"
);
// Ensure the listener was in fact hit with a Method::Run before the
// rejection path fired — otherwise we'd be asserting on an error
// that came from a connect failure.
let method = received.await.unwrap().expect("expected a method");
assert!(matches!(method, Method::Run { .. }));
}
// ---------------------------------------------------------------------------
// SubWorkerReadOutput
// ---------------------------------------------------------------------------
#[tokio::test]
async fn read_worker_output_returns_new_assistant_text_then_empty_on_second_call() {
let (tmp, registry, _rd) = setup_registry().await;
let (socket, listener) = bind_mock_socket(tmp.path(), "child").await;
register_child(&registry, "child", &socket, tmp.path()).await;
let items = vec![
Item::user_message("hello"),
assistant("hi back"),
assistant("still working"),
];
let _server = serve_history(listener, items);
let def = sub_worker_read_output_tool(registry);
let (_meta, tool) = def();
let input = json!({ "name": "child" }).to_string();
let first: ToolOutput = tool.execute(&input, Default::default()).await.unwrap();
let body = first.content.expect("first read should have content");
assert!(body.contains("hi back"), "body: {body}");
assert!(body.contains("still working"), "body: {body}");
// Cursor now points past all items — second call returns no new text.
let second: ToolOutput = tool.execute(&input, Default::default()).await.unwrap();
assert!(
second.content.is_none(),
"unexpected content: {:?}",
second.content
);
assert!(
second.summary.contains("no new assistant text"),
"summary: {}",
second.summary
);
}
#[tokio::test]
async fn read_worker_output_reports_stopped_on_dead_socket() {
let (tmp, registry, _rd) = setup_registry().await;
// Register a record pointing at a socket that nobody is listening
// on. Connect must fail → tool reports "stopped".
let dead_socket = tmp.path().join("dead.sock");
register_child(&registry, "child", &dead_socket, tmp.path()).await;
let def = sub_worker_read_output_tool(registry);
let (_meta, tool) = def();
let input = json!({ "name": "child" }).to_string();
let output: ToolOutput = tool.execute(&input, Default::default()).await.unwrap();
assert!(output.summary.contains("stopped"), "{}", output.summary);
}
// ---------------------------------------------------------------------------
// SubWorkerStop
// ---------------------------------------------------------------------------
#[tokio::test]
async fn stop_worker_sends_shutdown_and_releases_scope() {
let _env = EnvGuard::acquire();
let tmp = TempDir::new().unwrap();
let store_tmp = TempDir::new().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
);
let rd = Arc::new(RuntimeDir::create(tmp.path(), "spawner").await.unwrap());
let parent_scope = SharedScope::new(
Scope::writable(tmp.path())
.unwrap()
.with_added_deny_rules([ScopeRule {
target: tmp.path().to_path_buf(),
permission: Permission::Write,
recursive: true,
}])
.unwrap(),
);
unsafe {
std::env::set_var("YOI_RUNTIME_DIR", tmp.path());
}
let lock_path = tmp.path().join("workers.json");
// Seed workers.json with a restored top-level `spawner` allocation whose
// scope_deny contains the delegated child path plus the live child
// allocation — mimics a parent resumed after SubWorkerSpawn.
{
let mut g = LockFileGuard::open(&lock_path).unwrap();
let rule = ScopeRule {
target: tmp.path().to_path_buf(),
permission: Permission::Write,
recursive: true,
};
worker_allocation::register_worker_with_deny(
&mut g,
"spawner".into(),
std::process::id(),
"/tmp/spawner.sock".into(),
vec![rule.clone()],
vec![rule.clone()],
session_store::new_segment_id(),
)
.unwrap();
worker_allocation::register_worker(
&mut g,
"child".into(),
std::process::id(),
"/tmp/child.sock".into(),
vec![rule],
session_store::new_segment_id(),
)
.unwrap();
}
let loaded = SpawnedWorkerRegistry::load_from_worker_state_with_reclaim(
rd.clone(),
store.clone(),
"spawner".into(),
Some(parent_scope.clone()),
)
.await
.unwrap();
let registry = loaded.registry;
let (socket, listener) = bind_mock_socket(tmp.path(), "child").await;
let received = accept_one_method(listener);
register_child(&registry, "child", &socket, tmp.path()).await;
let def = sub_worker_stop_tool(registry.clone());
let (_meta, tool) = def();
let input = json!({ "name": "child" }).to_string();
let output: ToolOutput = tool.execute(&input, Default::default()).await.unwrap();
assert!(output.summary.contains("stopped"), "{}", output.summary);
// The child got a Shutdown.
let method = received.await.unwrap().expect("expected shutdown");
assert!(matches!(method, Method::Shutdown));
// Allocation for `child` is gone; `spawner` remains and its restored
// dynamic deny layer has been reclaimed.
{
let g = LockFileGuard::open(&lock_path).unwrap();
assert!(g.data().find("child").is_none(), "child still allocated");
let spawner = g.data().find("spawner").expect("spawner missing");
assert!(spawner.scope_deny.is_empty(), "deny not reclaimed");
}
assert_eq!(
parent_scope
.snapshot()
.permission_at(&tmp.path().join("file.txt")),
Some(Permission::Write)
);
// spawned_workers.json now lists zero children.
let spawned = rd.path().join("spawned_workers.json");
let contents = std::fs::read_to_string(&spawned).unwrap();
let records: Vec<SpawnedWorkerRecord> = serde_json::from_str(&contents).unwrap();
assert!(records.is_empty());
}
#[tokio::test]
async fn stop_worker_succeeds_even_when_child_unreachable() {
let _env = EnvGuard::acquire();
let (tmp, registry, _rd) = setup_registry().await;
unsafe {
std::env::set_var("YOI_RUNTIME_DIR", tmp.path());
}
// No live listener — socket never bound. Registered record points
// at a dead path. SubWorkerStop should still clean up local bookkeeping.
let dead_socket = tmp.path().join("dead.sock");
register_child(&registry, "child", &dead_socket, tmp.path()).await;
let def = sub_worker_stop_tool(registry.clone());
let (_meta, tool) = def();
let input = json!({ "name": "child" }).to_string();
let output: ToolOutput = tool.execute(&input, Default::default()).await.unwrap();
assert!(output.summary.contains("stopped"), "{}", output.summary);
// Registry no longer knows about the child.
assert!(registry.get("child").await.is_none());
}
// ---------------------------------------------------------------------------
// Persistence / restore
// ---------------------------------------------------------------------------
#[tokio::test]
async fn restored_registry_uses_worker_state_without_runtime_file() {
let _env = EnvGuard::acquire();
let runtime_tmp = TempDir::new().unwrap();
let store_tmp = TempDir::new().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
);
unsafe {
std::env::set_var("YOI_RUNTIME_DIR", runtime_tmp.path());
}
let rd = Arc::new(
RuntimeDir::create(runtime_tmp.path(), "spawner")
.await
.unwrap(),
);
let registry = SpawnedWorkerRegistry::load_from_worker_state(
rd.clone(),
store.clone(),
"spawner".to_string(),
)
.await
.unwrap();
let (socket, listener) = bind_mock_socket(runtime_tmp.path(), "child").await;
let mut received = serve_worker_methods(listener);
register_child(&registry, "child", &socket, runtime_tmp.path()).await;
std::fs::remove_file(rd.path().join("spawned_workers.json")).unwrap();
let restored = SpawnedWorkerRegistry::load_from_worker_state(
rd.clone(),
store.clone(),
"spawner".to_string(),
)
.await
.unwrap();
let def = sub_worker_send_tool(restored.clone());
let (_meta, tool) = def();
let input = json!({ "name": "child", "message": "after restart" }).to_string();
tool.execute(&input, Default::default()).await.unwrap();
match received.recv().await.expect("expected Run") {
Method::Run { input } => match input.as_slice() {
[protocol::Segment::Text { content }] => assert_eq!(content, "after restart"),
other => panic!("expected single Text segment, got {other:?}"),
},
other => panic!("expected Run, got {other:?}"),
}
let def = sub_worker_stop_tool(restored.clone());
let (_meta, tool) = def();
tool.execute(&json!({ "name": "child" }).to_string(), Default::default())
.await
.unwrap();
assert!(matches!(
received.recv().await.expect("expected Shutdown"),
Method::Shutdown
));
assert!(restored.get("child").await.is_none());
let metadata = store
.read_by_name("spawner")
.unwrap()
.expect("spawner metadata should remain");
assert!(metadata.spawned_children.is_empty());
assert_eq!(metadata.reclaimed_children.len(), 1);
assert_eq!(metadata.reclaimed_children[0].worker_name, "child");
let runtime_contents = std::fs::read_to_string(rd.path().join("spawned_workers.json")).unwrap();
let runtime_records: Vec<SpawnedWorkerRecord> =
serde_json::from_str(&runtime_contents).unwrap();
assert!(runtime_records.is_empty());
}
#[tokio::test]
async fn load_from_worker_state_prunes_runtime_children_and_reclaims_durable_delegation() {
let runtime_tmp = TempDir::new().unwrap();
let store_tmp = TempDir::new().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
);
let rd = Arc::new(
RuntimeDir::create(runtime_tmp.path(), "spawner")
.await
.unwrap(),
);
let registry = SpawnedWorkerRegistry::load_from_worker_state(
rd.clone(),
store.clone(),
"spawner".to_string(),
)
.await
.unwrap();
let (live_socket, listener) = bind_mock_socket(runtime_tmp.path(), "alive").await;
let _server = serve_worker_methods(listener);
register_child(&registry, "alive", &live_socket, runtime_tmp.path()).await;
register_child(
&registry,
"missing",
&runtime_tmp.path().join("missing.sock"),
runtime_tmp.path(),
)
.await;
let restored = SpawnedWorkerRegistry::load_from_worker_state(
rd.clone(),
store.clone(),
"spawner".to_string(),
)
.await
.unwrap();
assert!(restored.get("alive").await.is_some());
assert!(restored.get("missing").await.is_none());
let metadata = store
.read_by_name("spawner")
.unwrap()
.expect("spawner metadata should be written");
assert_eq!(metadata.spawned_children.len(), 1);
assert_eq!(metadata.spawned_children[0].worker_name, "alive");
assert_eq!(metadata.reclaimed_children.len(), 1);
assert_eq!(metadata.reclaimed_children[0].worker_name, "missing");
}
#[tokio::test]
async fn load_from_worker_state_reclaims_missing_child_scope_and_records_history() {
let _env = EnvGuard::acquire();
let runtime_tmp = TempDir::new().unwrap();
let store_tmp = TempDir::new().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
);
unsafe {
std::env::set_var("YOI_RUNTIME_DIR", runtime_tmp.path());
}
let rd = Arc::new(
RuntimeDir::create(runtime_tmp.path(), "spawner")
.await
.unwrap(),
);
let missing_rule = ScopeRule {
target: runtime_tmp.path().to_path_buf(),
permission: Permission::Write,
recursive: true,
};
{
let mut g = LockFileGuard::open(&runtime_tmp.path().join("workers.json")).unwrap();
worker_allocation::register_worker_with_deny(
&mut g,
"spawner".into(),
std::process::id(),
"/tmp/spawner.sock".into(),
vec![missing_rule.clone()],
vec![missing_rule.clone()],
session_store::new_segment_id(),
)
.unwrap();
}
let parent_scope = SharedScope::new(
Scope::writable(runtime_tmp.path())
.unwrap()
.with_added_deny_rules([missing_rule.clone()])
.unwrap(),
);
let seed =
SpawnedWorkerRegistry::load_from_worker_state(rd.clone(), store.clone(), "spawner".into())
.await
.unwrap();
seed.add(SpawnedWorkerRecord {
worker_name: "missing".into(),
socket_path: runtime_tmp.path().join("missing.sock"),
scope_delegated: vec![missing_rule.clone()],
callback_address: "/dev/null".into(),
})
.await
.unwrap();
let loaded = SpawnedWorkerRegistry::load_from_worker_state_with_reclaim(
rd.clone(),
store.clone(),
"spawner".into(),
Some(parent_scope.clone()),
)
.await
.unwrap();
assert!(loaded.reclaimed_unreachable);
assert!(loaded.registry.get("missing").await.is_none());
assert_eq!(
parent_scope
.snapshot()
.permission_at(&runtime_tmp.path().join("file.txt")),
Some(Permission::Write)
);
let g = LockFileGuard::open(&runtime_tmp.path().join("workers.json")).unwrap();
assert!(g.data().find("missing").is_none());
assert!(g.data().find("spawner").unwrap().scope_deny.is_empty());
let metadata = store
.read_by_name("spawner")
.unwrap()
.expect("spawner metadata should remain");
assert!(metadata.spawned_children.is_empty());
assert_eq!(metadata.reclaimed_children.len(), 1);
assert_eq!(metadata.reclaimed_children[0].worker_name, "missing");
let runtime_contents = std::fs::read_to_string(rd.path().join("spawned_workers.json")).unwrap();
let runtime_records: Vec<SpawnedWorkerRecord> =
serde_json::from_str(&runtime_contents).unwrap();
assert!(runtime_records.is_empty());
}