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/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());
}