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

refactor: rename pod crate to worker

Browse Source
This commit is contained in:
Hare
2026-06-26 00:05:57 +09:00
parent 4c677640f4
commit 6c59fe927b
194 changed files with 6637 additions and 6146 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/tests/compact_events_test.rs
+821
View File
@@ -0,0 +1,821 @@
//! Compact lifecycle `Event` broadcasting.
//!
//! Covers three paths:
//! - `try_pre_run_compact` success → `CompactStart + CompactDone`
//! - `try_pre_run_compact` failure → `CompactStart + CompactFailed`
//! - mid-turn `do_compact_and_resume` success → `CompactStart + CompactDone`
//! (driven by `compact_request_threshold` → `PreRequestAction::Yield`)
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use async_trait::async_trait;
use futures::Stream;
use llm_engine::Engine;
use llm_engine::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent};
use llm_engine::llm_client::types::Item;
use llm_engine::llm_client::{ClientError, LlmClient, Request};
use pod_store::{CombinedStore, FsWorkerStore, WorkerMetadataStore};
use protocol::{Event, Method, RunResult};
use session_store::{FsStore, LogEntry, Store};
use tokio::sync::broadcast;
use worker::{Worker, WorkerController};
type TestStore = CombinedStore<FsStore, FsWorkerStore>;
#[derive(Clone)]
struct MockClient {
responses: Arc<Vec<Vec<LlmEvent>>>,
call_count: Arc<AtomicUsize>,
}
impl MockClient {
fn new(responses: Vec<Vec<LlmEvent>>) -> Self {
Self {
responses: Arc::new(responses),
call_count: Arc::new(AtomicUsize::new(0)),
}
}
}
#[async_trait]
impl LlmClient for MockClient {
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>
{
let count = self.call_count.fetch_add(1, Ordering::SeqCst);
if count >= self.responses.len() {
return Err(ClientError::Config("mock client exhausted".into()));
}
let events = self.responses[count].clone();
let stream = futures::stream::iter(events.into_iter().map(Ok));
Ok(Box::pin(stream))
}
}
fn single_text_events(text: &str) -> Vec<LlmEvent> {
vec![
LlmEvent::text_block_start(0),
LlmEvent::text_delta(0, text),
LlmEvent::text_block_stop(0, None),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
/// `single_text_events` + a UsageEvent so the Worker's `usage_history`
/// picks up a measurement, which is how `pre_llm_request` decides
/// whether to yield mid-turn.
fn text_events_with_usage(text: &str, input_tokens: u64) -> Vec<LlmEvent> {
vec![
LlmEvent::text_block_start(0),
LlmEvent::text_delta(0, text),
LlmEvent::text_block_stop(0, None),
LlmEvent::usage(input_tokens, 1),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
fn write_summary_tool_use_events(call_id: &str, text: &str) -> Vec<LlmEvent> {
let input = serde_json::json!({ "text": text }).to_string();
vec![
LlmEvent::tool_use_start(0, call_id, "write_summary"),
LlmEvent::tool_input_delta(0, input),
LlmEvent::tool_use_stop(0),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
// A low compact_threshold guarantees `try_pre_run_compact` will fire
// the first time we check after a run.
const POST_RUN_MANIFEST_TOML: &str = r#"
[worker]
name = "test-worker"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[compaction]
compact_threshold = 1
compact_retained_tokens = 0
[[scope.allow]]
target = "./"
permission = "write"
"#;
// `compact_request_threshold` drives the WorkerInterceptor's mid-turn yield
// path. `compact_threshold` is left unset so the post-run check stays inert.
const MID_TURN_MANIFEST_TOML: &str = r#"
[worker]
name = "test-worker"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[compaction]
compact_request_threshold = 100
compact_retained_tokens = 0
[[scope.allow]]
target = "./"
permission = "write"
"#;
async fn make_worker_with_manifest(
manifest_toml: &str,
client: MockClient,
) -> Worker<MockClient, TestStore> {
let manifest = worker::WorkerManifest::from_toml(manifest_toml).unwrap();
let store_tmp = tempfile::tempdir().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
);
std::mem::forget(store_tmp);
let pwd_tmp = tempfile::tempdir().unwrap();
let pwd = pwd_tmp.path().to_path_buf();
let scope = worker::Scope::writable(&pwd).unwrap();
std::mem::forget(pwd_tmp);
let worker = Engine::new(client);
let mut worker = Worker::new(manifest, worker, store, pwd, scope)
.await
.unwrap();
worker.enable_worker_metadata_write_through().unwrap();
worker
}
async fn make_worker(client: MockClient) -> Worker<MockClient, TestStore> {
make_worker_with_manifest(POST_RUN_MANIFEST_TOML, client).await
}
/// Drain whatever events are already queued on `rx`. Non-blocking.
fn drain(rx: &mut broadcast::Receiver<Event>) -> Vec<Event> {
let mut out = Vec::new();
loop {
match rx.try_recv() {
Ok(ev) => out.push(ev),
Err(_) => break,
}
}
out
}
/// Collect every system-message text that the post-compaction
/// `SegmentStart.history` carries, by reading the sink mirror directly.
fn system_texts_in_sink_session_start(
worker: &worker::Worker<
impl llm_engine::llm_client::client::LlmClient + Clone + 'static,
impl session_store::Store + Clone + 'static,
>,
) -> Vec<String> {
let (entries, _rx) = worker.sink().subscribe_with_snapshot();
for entry in entries.into_iter().rev() {
if let session_store::LogEntry::SegmentStart { history, .. } = entry {
return history
.into_iter()
.filter_map(|logged| {
let item: Item = logged.into();
match item {
Item::Message {
role: llm_engine::Role::System,
content,
..
} => Some(
content
.iter()
.map(|p| p.as_text().to_owned())
.collect::<Vec<_>>()
.join(""),
),
_ => None,
}
})
.collect();
}
}
Vec::new()
}
/// Worker metadata starts with a reserved Session and no Segment, then becomes
/// active once the first SegmentStart is materialized by `run`.
#[tokio::test]
async fn worker_metadata_moves_from_pending_to_active_on_first_run() {
let client = MockClient::new(vec![single_text_events("hi")]);
let mut worker = make_worker(client).await;
let store = worker.store().clone();
let session_id = worker.session_id();
let initial_segment_id = worker.segment_id();
let pending = store
.read_by_name("test-worker")
.unwrap()
.expect("metadata should be initialized at Worker construction");
assert_eq!(pending.worker_name, "test-worker");
let pending_active = pending.active.expect("active session pointer missing");
assert_eq!(pending_active.session_id, session_id);
assert_eq!(pending_active.segment_id, None);
worker.run_text("first").await.unwrap();
let resolved = store
.read_by_name("test-worker")
.unwrap()
.expect("metadata should still exist after first run");
let active = resolved.active.expect("active session pointer missing");
assert_eq!(active.session_id, session_id);
assert_eq!(active.segment_id, Some(initial_segment_id));
}
/// Live auto-fork: when another writer extends the segment behind the
/// Worker's back, the next run's `ensure_segment_head` detects the
/// entry-count drift and branches into a fresh segment **within the same
/// Session**. The source segment is left immutable (no terminal marker
/// written back); the new segment records its parentage forward via
/// `SegmentStart.forked_from`.
#[tokio::test]
async fn concurrent_writer_drift_auto_forks_with_forked_from() {
// No compaction: keep run → run deterministic so each run consumes
// exactly one mock response and ensure_segment_head is the only fork
// trigger.
const NO_COMPACT_MANIFEST_TOML: &str = r#"
[worker]
name = "test-worker"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[[scope.allow]]
target = "./"
permission = "write"
"#;
let client = MockClient::new(vec![
single_text_events("first"),
single_text_events("second"),
]);
let mut worker = make_worker_with_manifest(NO_COMPACT_MANIFEST_TOML, client).await;
worker.run_text("first").await.unwrap();
let store = worker.store().clone();
let session_id = worker.session_id();
let source_segment_id = worker.segment_id();
let source_len_before = store.read_all(session_id, source_segment_id).unwrap().len();
// Simulate a foreign writer appending to the same segment. This bumps
// the on-disk entry count past the Worker's own append tally without
// updating the Worker's `entries_written`.
store
.append(
session_id,
source_segment_id,
&LogEntry::UserInput {
ts: 9999,
segments: vec![protocol::Segment::text("interloper")],
},
)
.unwrap();
// Next run triggers ensure_segment_head, which sees the drift.
worker.run_text("second").await.unwrap();
// The Worker moved to a new segment in the same Session.
let new_segment_id = worker.segment_id();
assert_ne!(new_segment_id, source_segment_id);
assert_eq!(
worker.session_id(),
session_id,
"auto-fork stays in-Session"
);
let metadata = store
.read_by_name("test-worker")
.unwrap()
.expect("metadata should exist after auto-fork");
let active = metadata.active.expect("active session pointer missing");
assert_eq!(active.session_id, session_id);
assert_eq!(active.segment_id, Some(new_segment_id));
// New segment records forked_from pointing at the source.
let new_entries = store.read_all(session_id, new_segment_id).unwrap();
match &new_entries[0] {
LogEntry::SegmentStart {
session_id: seg_session,
forked_from: Some(origin),
..
} => {
assert_eq!(*seg_session, session_id);
assert_eq!(origin.segment_id, source_segment_id);
}
other => panic!("expected SegmentStart with forked_from, got {other:?}"),
}
// Source segment is unchanged except for the foreign append — the
// auto-fork wrote no terminal marker back into it.
let source_after = store.read_all(session_id, source_segment_id).unwrap();
assert_eq!(source_after.len(), source_len_before + 1);
assert!(matches!(
source_after.last(),
Some(LogEntry::UserInput { .. })
));
}
#[tokio::test]
async fn compact_emits_session_start_carrying_summary_and_task_snapshot() {
let client = MockClient::new(vec![
single_text_events("hi"),
write_summary_tool_use_events("call-1", "summary"),
single_text_events("done"),
]);
let mut worker = make_worker(client).await;
let (tx, _rx_keep) = broadcast::channel::<Event>(64);
worker.attach_event_tx(tx);
worker.run_text("first").await.unwrap();
let session_id = worker.session_id();
worker.compact(10_000).await.unwrap();
let compacted_segment_id = worker.segment_id();
let metadata = worker
.store()
.read_by_name("test-worker")
.unwrap()
.expect("metadata should exist after compaction");
let active = metadata.active.expect("active session pointer missing");
assert_eq!(active.session_id, session_id);
assert_eq!(active.segment_id, Some(compacted_segment_id));
let system_texts = system_texts_in_sink_session_start(&worker);
// The post-compaction `SegmentStart.history` carries the new system
// messages introduced by the compactor. Clients re-seed their view
// from this entry alone, so it is the load-bearing payload.
assert!(
system_texts
.iter()
.any(|text| text.starts_with("[Compacted context summary]")),
"summary system message missing from {system_texts:?}"
);
assert!(
system_texts
.iter()
.any(|text| text.starts_with("[Session TaskStore snapshot]")),
"task snapshot system message missing from {system_texts:?}"
);
}
#[tokio::test]
async fn pre_run_compact_success_broadcasts_start_and_done() {
// Responses: (1) first run returns short text, (2) compact worker
// emits write_summary then closes (two LLM calls inside the compact
// worker: one for write_summary, one that the compact loop consumes
// as the final "I'm done" close response).
let client = MockClient::new(vec![
single_text_events("hi"),
write_summary_tool_use_events("call-1", "summary"),
single_text_events("done"),
]);
let mut worker = make_worker(client).await;
let (tx, mut rx) = broadcast::channel::<Event>(64);
worker.attach_event_tx(tx);
worker.run_text("first").await.unwrap();
// Drain run events so only compact events remain in `rx`.
let _ = drain(&mut rx);
worker.try_pre_run_compact().await;
let events = drain(&mut rx);
let kinds: Vec<&str> = events
.iter()
.map(|e| match e {
Event::CompactStart => "start",
Event::CompactDone { .. } => "done",
Event::CompactFailed { .. } => "failed",
_ => "other",
})
.collect();
assert!(
kinds.contains(&"start") && kinds.contains(&"done"),
"expected CompactStart + CompactDone in {kinds:?}"
);
assert!(
!kinds.contains(&"failed"),
"unexpected CompactFailed in {kinds:?}"
);
// CompactDone carries the new session id.
let new_id_in_event = events.iter().find_map(|e| match e {
Event::CompactDone { new_segment_id } => Some(*new_segment_id),
_ => None,
});
assert!(new_id_in_event.is_some(), "CompactDone missing");
assert_eq!(new_id_in_event.unwrap(), worker.segment_id());
}
#[tokio::test]
async fn mid_turn_compact_success_broadcasts_start_and_done() {
// Path: `do_compact_and_resume` via PreRequestAction::Yield.
//
// Sequence of LLM calls the mock will serve:
// [0] first run completes with a UsageEvent(1000 > threshold=100) so
// the next run's pre_llm_request will yield.
// [1] compact worker emits `write_summary` tool call.
// [2] compact worker closes (its final "done" response).
// [3] resume() after compact makes one more LLM call.
let client = MockClient::new(vec![
text_events_with_usage("a", 1000),
write_summary_tool_use_events("call-1", "summary"),
single_text_events("done"),
single_text_events("b"),
]);
let mut worker = make_worker_with_manifest(MID_TURN_MANIFEST_TOML, client).await;
let (tx, mut rx) = broadcast::channel::<Event>(64);
worker.attach_event_tx(tx);
// First run populates usage_history above the request threshold.
worker.run_text("first").await.unwrap();
let _ = drain(&mut rx);
// Second run: pre_llm_request yields immediately, Engine returns
// Yielded, handle_worker_result routes into do_compact_and_resume.
worker.run_text("second").await.unwrap();
let events = drain(&mut rx);
let kinds: Vec<&str> = events
.iter()
.map(|e| match e {
Event::CompactStart => "start",
Event::CompactDone { .. } => "done",
Event::CompactFailed { .. } => "failed",
_ => "other",
})
.collect();
assert!(
kinds.contains(&"start") && kinds.contains(&"done"),
"expected CompactStart + CompactDone in {kinds:?}"
);
assert!(
!kinds.contains(&"failed"),
"unexpected CompactFailed in {kinds:?}"
);
let new_id_in_event = events.iter().find_map(|e| match e {
Event::CompactDone { new_segment_id } => Some(*new_segment_id),
_ => None,
});
assert_eq!(new_id_in_event, Some(worker.segment_id()));
}
/// Regression: `Worker::compact()` must reset the in-memory
/// `extract_pointer` so extract keeps firing on the new compacted
/// session.
///
/// Without the reset, the pointer's `processed_through_history_len`
/// holds the old (typically large) item count, while the new compacted
/// session starts with a much shorter history (`[summary, ...]`).
/// `cumulative_input_tokens_since` would then filter every new
/// usage record out (their `history_len` is below the stale pointer)
/// and extract would never re-fire for the rest of the process.
const EXTRACT_PLUS_COMPACT_MANIFEST: &str = r#"
[worker]
name = "test-worker"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[memory]
extract_threshold = 1
[compaction]
compact_threshold = 1
compact_retained_tokens = 0
[[scope.allow]]
target = "./"
permission = "write"
"#;
fn write_extracted_tool_use_events(call_id: &str) -> Vec<LlmEvent> {
let input = serde_json::json!({
"decisions": [],
"discussions": [],
"attempts": [],
"requests": []
})
.to_string();
vec![
LlmEvent::tool_use_start(0, call_id, "write_extracted"),
LlmEvent::tool_input_delta(0, input),
LlmEvent::tool_use_stop(0),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
#[tokio::test]
async fn compact_resets_extract_pointer_so_extract_can_fire_again() {
// Mock LLM responses, in call order:
// [0] first run with usage(1000) so extract threshold (=1) fires.
// [1] extract worker invokes write_extracted with empty payload.
// [2] extract worker closes after the tool result.
// [3] compact worker invokes write_summary.
// [4] compact worker closes after the tool result.
let client = MockClient::new(vec![
text_events_with_usage("hi", 1000),
write_extracted_tool_use_events("ec1"),
single_text_events("done"),
write_summary_tool_use_events("sc1", "summary"),
single_text_events("done"),
]);
let mut worker = make_worker_with_manifest(EXTRACT_PLUS_COMPACT_MANIFEST, client).await;
worker.run_text("first").await.unwrap();
// extract fires; pointer becomes Some.
worker.try_post_run_extract().await.unwrap();
assert!(
worker.extract_pointer().is_some(),
"extract_pointer should be Some after a successful extract"
);
// Compact runs. Without the fix the in-memory pointer would still
// reference the old session's history_len.
worker.try_pre_run_compact().await;
assert!(
worker.extract_pointer().is_none(),
"extract_pointer must be reset to None after compact (matches cold-restore on the new session)"
);
}
/// `extract_threshold = 0` is treated as "disabled" — without this, a
/// raw `>=` comparison against `tokens_since` would fire extract on
/// every post-run regardless of activity. Mirrors the consolidation
/// zero-threshold convention so users have a single way to opt out
/// without removing the `[memory]` section.
const EXTRACT_THRESHOLD_ZERO_MANIFEST: &str = r#"
[worker]
name = "test-worker"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[memory]
extract_threshold = 0
[[scope.allow]]
target = "./"
permission = "write"
"#;
#[tokio::test]
async fn extract_threshold_zero_is_disabled() {
// Mock provides exactly one response — the first run. If extract
// were treated as "fire on any change" because of `tokens_since >= 0`,
// it would call into the extract worker and exhaust the mock.
let client = MockClient::new(vec![text_events_with_usage("hi", 1000)]);
let mut worker = make_worker_with_manifest(EXTRACT_THRESHOLD_ZERO_MANIFEST, client).await;
worker.run_text("first").await.unwrap();
worker
.try_post_run_extract()
.await
.expect("extract_threshold=0 must skip silently, not fail");
assert!(
worker.extract_pointer().is_none(),
"no extract should have run — pointer must remain None"
);
}
#[tokio::test]
async fn pre_run_compact_failure_broadcasts_start_and_failed() {
// Only the first run has a response. Compaction will run the
// compact worker which immediately exhausts the mock → failure.
let client = MockClient::new(vec![single_text_events("hi")]);
let mut worker = make_worker(client).await;
let (tx, mut rx) = broadcast::channel::<Event>(64);
worker.attach_event_tx(tx);
worker.run_text("first").await.unwrap();
let _ = drain(&mut rx);
// Best-effort: returns Ok(()) even on failure, but emits CompactFailed.
worker.try_pre_run_compact().await;
let events = drain(&mut rx);
let kinds: Vec<&str> = events
.iter()
.map(|e| match e {
Event::CompactStart => "start",
Event::CompactDone { .. } => "done",
Event::CompactFailed { .. } => "failed",
_ => "other",
})
.collect();
assert!(
kinds.contains(&"start") && kinds.contains(&"failed"),
"expected CompactStart + CompactFailed in {kinds:?}"
);
assert!(
!kinds.contains(&"done"),
"unexpected CompactDone in {kinds:?}"
);
}
// ---------------------------------------------------------------------------
// Detached post-run memory jobs (`spawn_post_run_memory_jobs` /
// `wait_for_memory_jobs`). Covers the detach round-trip and the structural
// invariant that the cloned memory-task Worker shares `SegmentState` with the
// source Worker, so that `save_extension` from the background extract does not
// leave the next turn's `save_user_input` looking at a stale session pointer.
const EXTRACT_NO_COMPACT_MANIFEST: &str = r#"
[worker]
name = "test-worker"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[memory]
extract_threshold = 1
[[scope.allow]]
target = "./"
permission = "write"
"#;
#[tokio::test]
async fn extract_large_unprocessed_range_does_not_abort_on_input_occupancy() {
let client = MockClient::new(vec![
text_events_with_usage("recorded", 1000),
write_extracted_tool_use_events("ec-large"),
single_text_events("done"),
]);
let mut worker = make_worker_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await;
let large_request = format!("remember this large slice: {}", "x ".repeat(200_000));
worker.run_text(&large_request).await.unwrap();
worker.try_post_run_extract().await.expect(
"large unprocessed extract ranges must reach the extract worker, not abort locally",
);
assert!(
worker.extract_pointer().is_some(),
"successful extract should advance the pointer even when the input range is large"
);
}
#[tokio::test]
async fn spawn_and_wait_drives_extract_to_completion() {
let client = MockClient::new(vec![
text_events_with_usage("hi", 1000),
write_extracted_tool_use_events("ec1"),
single_text_events("done"),
]);
let mut worker = make_worker_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await;
worker.run_text("first").await.unwrap();
assert!(
worker.extract_pointer().is_none(),
"extract has not run yet — pointer must be None"
);
worker.spawn_post_run_memory_jobs();
worker.wait_for_memory_jobs().await;
assert!(
worker.extract_pointer().is_some(),
"spawn + wait must complete extract; pointer should be set"
);
}
#[tokio::test]
async fn detached_extract_does_not_fork_session_log() {
// Source worker and the cloned memory-task worker share `SegmentState` via
// `Arc<_>`. The detached extract advances the entry tally through
// `save_extension`; the next `run` must see that same tally so
// `ensure_head_or_fork` does not spawn a new session.
let client = MockClient::new(vec![
text_events_with_usage("hi", 1000),
write_extracted_tool_use_events("ec1"),
single_text_events("done"),
text_events_with_usage("ok", 1000),
]);
let mut worker = make_worker_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await;
worker.run_text("first").await.unwrap();
let session_before = worker.segment_id();
worker.spawn_post_run_memory_jobs();
worker.wait_for_memory_jobs().await;
worker.run_text("second").await.unwrap();
let session_after = worker.segment_id();
assert_eq!(
session_before, session_after,
"detached extract's save_extension and the next turn's save_user_input \
must share the entry tally through SegmentState — a fork here means the \
clone carried its own counter"
);
}
#[tokio::test]
async fn controller_compact_method_emits_start_and_done() {
let client = MockClient::new(vec![
text_events_with_usage("hi", 1000),
write_summary_tool_use_events("manual-summary", "manual compact summary"),
single_text_events("done"),
]);
let worker = make_worker_with_manifest(POST_RUN_MANIFEST_TOML, client).await;
let runtime_tmp = tempfile::tempdir().unwrap();
let (handle, _shutdown) = WorkerController::spawn(worker, runtime_tmp.path())
.await
.unwrap();
let mut rx = handle.subscribe();
handle
.send(Method::run_text("seed history"))
.await
.expect("send run");
loop {
match tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for run end")
.expect("event")
{
Event::RunEnd {
result: RunResult::Finished,
} => break,
_ => {}
}
}
handle.send(Method::Compact).await.expect("send compact");
let mut saw_start = false;
loop {
match tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for compact events")
.expect("event")
{
Event::CompactStart => saw_start = true,
Event::CompactDone { .. } => {
break;
}
Event::CompactFailed { error } => panic!("manual compact failed: {error}"),
_ => {}
}
}
assert!(saw_start, "manual compact should emit CompactStart");
let _ = handle.send(Method::Shutdown).await;
}
crates/worker/tests/consolidation_test.rs
+530
View File
@@ -0,0 +1,530 @@
//! consolidation (memory.consolidation) post-run trigger.
//!
//! Covers the gating, lock and cleanup behaviour without exercising the
//! full sub-worker tool loop:
//!
//! - no `[memory]` section → no-op
//! - `[memory]` present but no thresholds → no-op
//! - staging empty → skip
//! - staging below thresholds → skip + lock not acquired
//! - staging above threshold → sub-worker runs, consumed entries removed
//! - existing live lock → skip without error
//!
//! The sub-worker is fed a no-op LLM response (plain text) so it returns
//! immediately. The post-run path then exercises lock acquisition,
//! cleanup, and the empty-payload fast path.
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use async_trait::async_trait;
use futures::Stream;
use llm_engine::Engine;
use llm_engine::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent};
use llm_engine::llm_client::{ClientError, LlmClient, Request};
use memory::WorkspaceLayout;
use memory::extract::{ExtractedPayload, write_staging};
use memory::schema::SourceRef;
use pod_store::{CombinedStore, FsWorkerStore};
use session_store::FsStore;
type TestStore = CombinedStore<FsStore, FsWorkerStore>;
use tokio::sync::broadcast;
use worker::{Event, Worker};
#[derive(Clone)]
struct MockClient {
responses: Arc<Vec<Vec<LlmEvent>>>,
call_count: Arc<AtomicUsize>,
}
impl MockClient {
fn new(responses: Vec<Vec<LlmEvent>>) -> Self {
Self {
responses: Arc::new(responses),
call_count: Arc::new(AtomicUsize::new(0)),
}
}
}
#[async_trait]
impl LlmClient for MockClient {
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>
{
let count = self.call_count.fetch_add(1, Ordering::SeqCst);
if count >= self.responses.len() {
return Err(ClientError::Config("mock client exhausted".into()));
}
let events = self.responses[count].clone();
let stream = futures::stream::iter(events.into_iter().map(Ok));
Ok(Box::pin(stream))
}
}
fn done(text: &str) -> Vec<LlmEvent> {
vec![
LlmEvent::text_block_start(0),
LlmEvent::text_delta(0, text),
LlmEvent::text_block_stop(0, None),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
const NO_MEMORY_TOML: &str = r#"
[worker]
name = "test-worker"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[[scope.allow]]
target = "./"
permission = "write"
"#;
const MEMORY_NO_THRESHOLDS_TOML: &str = r#"
[worker]
name = "test-worker"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[memory]
[[scope.allow]]
target = "./"
permission = "write"
"#;
const FILES_THRESHOLD_TOML: &str = r#"
[worker]
name = "test-worker"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[memory]
consolidation_threshold_files = 2
[[scope.allow]]
target = "./"
permission = "write"
"#;
const ZERO_THRESHOLDS_TOML: &str = r#"
[worker]
name = "test-worker"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[memory]
consolidation_threshold_files = 0
consolidation_threshold_bytes = 0
[[scope.allow]]
target = "./"
permission = "write"
"#;
async fn make_worker_with(
manifest_toml: &str,
pwd: std::path::PathBuf,
client: MockClient,
) -> Worker<MockClient, TestStore> {
let manifest = worker::WorkerManifest::from_toml(manifest_toml).unwrap();
let store_tmp = tempfile::tempdir().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
);
std::mem::forget(store_tmp);
let scope = worker::Scope::writable(&pwd).unwrap();
let worker = Engine::new(client);
Worker::new(manifest, worker, store, pwd, scope)
.await
.unwrap()
}
fn write_n_staging(layout: &WorkspaceLayout, n: usize) -> Vec<uuid::Uuid> {
let mut ids = Vec::new();
for i in 0..n {
let (id, _) = write_staging(
layout,
SourceRef {
segment_id: format!("s-{i}"),
range: [i as u64, i as u64],
},
ExtractedPayload::default(),
)
.unwrap();
ids.push(id);
}
ids
}
fn attach_event_receiver(worker: &mut Worker<MockClient, TestStore>) -> broadcast::Receiver<Event> {
let (tx, rx) = broadcast::channel(16);
worker.attach_event_tx(tx);
rx
}
fn collect_memory_worker_reasons(rx: &mut broadcast::Receiver<Event>) -> Vec<String> {
let mut reasons = Vec::new();
loop {
match rx.try_recv() {
Ok(Event::MemoryWorker(event)) => reasons.push(event.reason),
Ok(_) => {}
Err(broadcast::error::TryRecvError::Empty) => break,
Err(err) => panic!("unexpected broadcast receive error: {err}"),
}
}
reasons
}
fn read_audit_jsonl(layout: &WorkspaceLayout) -> Vec<serde_json::Value> {
let text = std::fs::read_to_string(layout.audit_current_log_path()).unwrap();
text.lines()
.map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
.collect()
}
#[tokio::test]
async fn no_memory_section_is_a_noop() {
let pwd = tempfile::tempdir().unwrap();
let client = MockClient::new(vec![]);
let mut worker = make_worker_with(NO_MEMORY_TOML, pwd.path().to_path_buf(), client).await;
worker
.try_post_run_consolidate()
.await
.expect("missing memory section must skip cleanly");
}
#[tokio::test]
async fn no_thresholds_is_a_noop() {
let pwd = tempfile::tempdir().unwrap();
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
write_n_staging(&layout, 5);
let client = MockClient::new(vec![]);
let mut worker =
make_worker_with(MEMORY_NO_THRESHOLDS_TOML, pwd.path().to_path_buf(), client).await;
worker
.try_post_run_consolidate()
.await
.expect("consolidation disabled when both thresholds are None");
// No staging entries removed.
assert_eq!(memory::consolidate::list_staging_entries(&layout).len(), 5);
}
#[tokio::test]
async fn zero_thresholds_treated_as_disabled() {
// Without the `Some(0) → None` collapse, `total_files >= 0` and
// `total_bytes >= 0` would always evaluate true and consolidation would
// fire on every post-run with any staging activity.
let pwd = tempfile::tempdir().unwrap();
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
write_n_staging(&layout, 5);
let client = MockClient::new(vec![]);
let mut worker = make_worker_with(ZERO_THRESHOLDS_TOML, pwd.path().to_path_buf(), client).await;
worker
.try_post_run_consolidate()
.await
.expect("zero thresholds must collapse to disabled, not fire on every staging entry");
assert_eq!(
memory::consolidate::list_staging_entries(&layout).len(),
5,
"staging must be untouched when both thresholds are zero"
);
let lock_path = layout.staging_dir().join(".consolidation.lock");
assert!(!lock_path.exists(), "no lock should be acquired");
}
#[tokio::test]
async fn empty_staging_skips() {
let pwd = tempfile::tempdir().unwrap();
let client = MockClient::new(vec![]);
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
worker.try_post_run_consolidate().await.unwrap();
// No mock calls expected.
}
#[tokio::test]
async fn empty_staging_skip_is_audit_only() {
let pwd = tempfile::tempdir().unwrap();
let client = MockClient::new(vec![]);
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
let mut rx = attach_event_receiver(&mut worker);
worker.try_post_run_consolidate().await.unwrap();
assert!(collect_memory_worker_reasons(&mut rx).is_empty());
}
#[tokio::test]
async fn invalid_only_staging_is_distinct_from_no_staging() {
let pwd = tempfile::tempdir().unwrap();
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
std::fs::create_dir_all(layout.staging_dir()).unwrap();
let invalid_id = uuid::Uuid::now_v7();
let invalid_path = layout.staging_dir().join(format!("{invalid_id}.json"));
std::fs::write(
&invalid_path,
serde_json::json!({
"source": {
"session_id": "legacy-session",
"range": [0, 1]
},
"requests": []
})
.to_string(),
)
.unwrap();
let client = MockClient::new(vec![]);
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
let mut rx = attach_event_receiver(&mut worker);
worker.try_post_run_consolidate().await.unwrap();
assert!(invalid_path.exists(), "invalid staging is not auto-deleted");
let reasons = collect_memory_worker_reasons(&mut rx);
assert_eq!(reasons, vec!["no_valid_staging_entries invalid=1"]);
let audit = read_audit_jsonl(&layout);
let last = audit.last().unwrap();
assert_eq!(last["reason"], "no_valid_staging_entries invalid=1");
assert_eq!(last["consolidation"]["staging_count"], 0);
assert_eq!(last["consolidation"]["invalid_staging_count"], 1);
}
#[tokio::test]
async fn below_threshold_skip_is_audit_only() {
let pwd = tempfile::tempdir().unwrap();
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
write_n_staging(&layout, 1); // threshold is 2
let client = MockClient::new(vec![]);
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
let mut rx = attach_event_receiver(&mut worker);
worker.try_post_run_consolidate().await.unwrap();
assert!(collect_memory_worker_reasons(&mut rx).is_empty());
let audit = read_audit_jsonl(&layout);
let reason = audit.last().unwrap()["reason"]
.as_str()
.expect("audit reason must be a string");
assert!(reason.starts_with("threshold_not_reached "));
}
#[tokio::test]
async fn completed_event_survives_terminal_empty_drain_skip() {
let pwd = tempfile::tempdir().unwrap();
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
write_n_staging(&layout, 2); // threshold is 2 — fires.
let client = MockClient::new(vec![done("ok")]);
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
let mut rx = attach_event_receiver(&mut worker);
worker.try_post_run_consolidate().await.unwrap();
let reasons = collect_memory_worker_reasons(&mut rx);
assert_eq!(reasons.len(), 2);
assert!(reasons[0].starts_with("staging_threshold_reached files=2 bytes="));
assert_eq!(reasons[1], "completed_no_record_changes");
let audit = read_audit_jsonl(&layout);
assert_eq!(audit.last().unwrap()["reason"], "no_staging_entries");
}
#[tokio::test]
async fn below_threshold_skips_and_does_not_take_lock() {
let pwd = tempfile::tempdir().unwrap();
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
write_n_staging(&layout, 1); // threshold is 2
let client = MockClient::new(vec![]);
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
worker.try_post_run_consolidate().await.unwrap();
// Staging untouched.
assert_eq!(memory::consolidate::list_staging_entries(&layout).len(), 1);
// Lock file must not exist.
let lock_path = layout.staging_dir().join(".consolidation.lock");
assert!(!lock_path.exists(), "lock file should not be created");
}
#[tokio::test]
async fn fires_on_threshold_and_cleans_up_consumed_entries() {
let pwd = tempfile::tempdir().unwrap();
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
write_n_staging(&layout, 2); // threshold is 2 — fires.
// Sub-worker is given a single text-only response. The consolidation prompt
// tells it to call memory tools; the mock skips those, but `Engine::run`
// returns Ok regardless once the LLM closes with a final text.
let client = MockClient::new(vec![done("ok")]);
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
worker.try_post_run_consolidate().await.unwrap();
// Consumed entries removed.
assert!(
memory::consolidate::list_staging_entries(&layout).is_empty(),
"consumed staging entries must be cleaned up"
);
// Lock removed too.
let lock_path = layout.staging_dir().join(".consolidation.lock");
assert!(!lock_path.exists(), "lock file must be removed on success");
}
#[tokio::test]
async fn in_flight_guard_skips_reentry_without_clearing() {
use std::sync::atomic::Ordering;
let pwd = tempfile::tempdir().unwrap();
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
write_n_staging(&layout, 2);
let client = MockClient::new(vec![]);
let mut worker = make_worker_with(
FILES_THRESHOLD_TOML,
pwd.path().to_path_buf(),
client.clone(),
)
.await;
// Pre-set the in-flight flag as if another concurrent caller had
// entered run_consolidate_once. The CAS at the top of
// try_post_run_consolidate must take the early return without
// touching staging or the LLM, and must leave the flag intact for
// the holder to clear.
let in_flight = worker.consolidation_in_flight_handle();
in_flight.store(true, Ordering::Release);
worker.try_post_run_consolidate().await.unwrap();
assert!(
in_flight.load(Ordering::Acquire),
"reentry skip must not clear the in-flight flag — that's the holder's job"
);
assert_eq!(
memory::consolidate::list_staging_entries(&layout).len(),
2,
"staging must remain untouched on reentry skip"
);
assert_eq!(
client.call_count.load(Ordering::SeqCst),
0,
"no LLM calls should fire on reentry skip"
);
// Sanity: when the flag is cleared, the same worker fires normally and
// resets the flag itself (i.e. it isn't accidentally sticky).
in_flight.store(false, Ordering::Release);
let client2 = MockClient::new(vec![done("ok")]);
let mut worker2 =
make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client2).await;
worker2.try_post_run_consolidate().await.unwrap();
assert!(
!worker2
.consolidation_in_flight_handle()
.load(Ordering::Acquire),
"in-flight flag must be cleared after a normal run"
);
}
#[tokio::test]
async fn coalesce_loop_terminates_with_one_iteration_when_snapshot_drains_staging() {
use std::sync::atomic::Ordering;
// Coalesce semantics from `docs/plan/memory.md` §並走防止: a single
// run consumes the snapshot taken at acquire time; the loop
// re-evaluates against any post-snapshot extract additions. With no
// concurrent additions, the second iteration sees an empty staging
// and bails out — exercised here by counting LLM calls.
let pwd = tempfile::tempdir().unwrap();
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
write_n_staging(&layout, 4);
// Provide just one mock response. If the loop wrongly re-enters
// run_consolidate_once after Completed, the second sub-worker run
// would exhaust the mock and surface as an error.
let client = MockClient::new(vec![done("ok")]);
let mut worker = make_worker_with(
FILES_THRESHOLD_TOML,
pwd.path().to_path_buf(),
client.clone(),
)
.await;
worker.try_post_run_consolidate().await.unwrap();
assert_eq!(
client.call_count.load(Ordering::SeqCst),
1,
"Coalesce must terminate once the staging snapshot is drained — got an extra LLM call"
);
assert!(
memory::consolidate::list_staging_entries(&layout).is_empty(),
"staging must be empty after the single iteration"
);
}
#[tokio::test]
async fn live_lock_held_by_other_worker_skips() {
let pwd = tempfile::tempdir().unwrap();
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
write_n_staging(&layout, 3);
// Pre-acquire lock with this test's PID — definitely alive — and
// *don't* release it. The consolidation path must skip without error.
let _live_lock = memory::consolidate::StagingLock::acquire(
&layout,
std::process::id(),
"other-worker",
Vec::new(),
)
.unwrap();
let client = MockClient::new(vec![]);
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
worker
.try_post_run_consolidate()
.await
.expect("InUse lock must surface as graceful skip");
// Staging untouched: lock holder owns the snapshot, not us.
assert_eq!(memory::consolidate::list_staging_entries(&layout).len(), 3);
}
crates/worker/tests/controller_test.rs
+2315
View File
File diff suppressed because it is too large Load Diff
crates/worker/tests/restore_test.rs
+207
View File
@@ -0,0 +1,207 @@
//! Integration tests for `Worker::restore_from_manifest`'s pre-build
//! validation paths.
//!
//! These cases all return before `prepare_worker_common` runs, so they
//! do not need a real LLM client or pod-registry environment — only the
//! session store needs to be present.
use std::sync::{LazyLock, Mutex};
use pod_store::{
CombinedStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerMetadata, WorkerMetadataStore,
};
use session_store::{FsStore, StoreError};
use worker::{Worker, WorkerError};
const MINIMAL_MANIFEST_TOML: &str = r#"
[worker]
name = "restore-test"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[[scope.allow]]
target = "./"
permission = "write"
"#;
/// Serialises tests that mutate runtime-dir env vars, mirroring the
/// pattern used by other integration tests in this crate.
static ENV_LOCK: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));
#[tokio::test]
async fn restore_from_worker_metadata_rejects_missing_metadata() {
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let store_tmp = tempfile::tempdir().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
);
let manifest = worker::WorkerManifest::from_toml(MINIMAL_MANIFEST_TOML).unwrap();
let result = Worker::restore_from_worker_metadata(
"restore-test",
manifest,
store,
worker::PromptLoader::builtins_only(),
)
.await;
match result {
Err(WorkerError::WorkerMetadataMissing { worker_name }) => {
assert_eq!(worker_name, "restore-test")
}
Err(other) => panic!("expected WorkerMetadataMissing, got {other:?}"),
Ok(_) => panic!("expected missing worker metadata to fail"),
}
}
#[tokio::test]
async fn restore_from_worker_metadata_rejects_pending_segment() {
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let store_tmp = tempfile::tempdir().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
);
let manifest = worker::WorkerManifest::from_toml(MINIMAL_MANIFEST_TOML).unwrap();
let session_id = session_store::new_session_id();
store
.write(&WorkerMetadata::new(
"restore-test",
Some(WorkerActiveSegmentRef::pending_segment(session_id)),
))
.unwrap();
let result = Worker::restore_from_worker_metadata(
"restore-test",
manifest,
store,
worker::PromptLoader::builtins_only(),
)
.await;
match result {
Err(WorkerError::WorkerMetadataPending {
worker_name,
session_id: actual,
}) => {
assert_eq!(worker_name, "restore-test");
assert_eq!(actual, session_id);
}
Err(other) => panic!("expected WorkerMetadataPending, got {other:?}"),
Ok(_) => panic!("expected pending worker metadata to fail"),
}
}
#[tokio::test]
async fn restore_from_worker_metadata_resolves_active_pointer_through_session_log() {
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let store_tmp = tempfile::tempdir().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
);
let manifest = worker::WorkerManifest::from_toml(MINIMAL_MANIFEST_TOML).unwrap();
let session_id = session_store::new_session_id();
let segment_id = session_store::new_segment_id();
store
.write(&WorkerMetadata::new(
"restore-test",
Some(WorkerActiveSegmentRef::active_segment(
session_id, segment_id,
)),
))
.unwrap();
let result = Worker::restore_from_worker_metadata(
"restore-test",
manifest,
store,
worker::PromptLoader::builtins_only(),
)
.await;
match result {
Err(WorkerError::Store(StoreError::NotFound(id))) => assert_eq!(id, segment_id),
Err(other) => panic!("expected Store(NotFound) from resolved segment, got {other:?}"),
Ok(_) => panic!("expected unknown resolved segment to fail"),
}
}
#[tokio::test]
async fn restore_from_manifest_rejects_unknown_segment() {
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let store_tmp = tempfile::tempdir().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
);
let manifest = worker::WorkerManifest::from_toml(MINIMAL_MANIFEST_TOML).unwrap();
// A freshly-minted id with no jsonl file at all → store returns
// NotFound, which `Worker::restore_from_manifest` surfaces verbatim
// as `WorkerError::Store`.
let unknown_sid = session_store::new_session_id();
let unknown_seg = session_store::new_segment_id();
let result = Worker::restore_from_manifest(
unknown_sid,
unknown_seg,
manifest,
store,
worker::PromptLoader::builtins_only(),
)
.await;
match result {
Err(WorkerError::Store(StoreError::NotFound(id))) => assert_eq!(id, unknown_seg),
Err(other) => panic!("expected Store(NotFound), got {other:?}"),
Ok(_) => panic!("expected unknown segment to fail"),
}
}
#[tokio::test]
async fn restore_from_manifest_rejects_empty_segment_log() {
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let store_tmp = tempfile::tempdir().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
);
let manifest = worker::WorkerManifest::from_toml(MINIMAL_MANIFEST_TOML).unwrap();
// Pre-create an empty `<sid>/<segid>.jsonl` so `read_all` succeeds
// with no entries. `collect_state` returns `entries_count = 0`,
// which `restore_from_manifest` rejects with `SegmentEmpty` *before*
// it gets as far as building the LLM client.
let sid = session_store::new_session_id();
let segid = session_store::new_segment_id();
let dir = store_tmp.path().join(sid.to_string());
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join(format!("{segid}.jsonl")), b"").unwrap();
let result = Worker::restore_from_manifest(
sid,
segid,
manifest,
store,
worker::PromptLoader::builtins_only(),
)
.await;
match result {
Err(WorkerError::SegmentEmpty { segment_id }) => assert_eq!(segment_id, segid),
Err(other) => panic!("expected SegmentEmpty, got {other:?}"),
Ok(_) => panic!("expected empty segment log to fail"),
}
}
crates/worker/tests/session_metrics_test.rs
+544
View File
@@ -0,0 +1,544 @@
//! End-to-end coverage for the prune-projection metrics path.
//!
//! Drives a Worker with a scripted mock LLM client and a custom tool that
//! returns a long `ToolOutput.content`, then inspects the persisted
//! session log to verify:
//!
//! - `prune.skip { reason: "no_candidates" }` lands when usage estimates are
//! unavailable or the protected-token window covers all tool results.
//! - `prune.fire` lands once enough measured history exceeds the protected-token
//! budget for the projection to actually apply.
//! - The fire metric and the immediately-following `prune.post_request`
//! metric share the same `correlation_id`, so cache_read / cache_write
//! from the LlmUsage that triggered the projection can be joined back
//! to the originating event.
//! - `prune.skip { reason: "below_min_savings" }` lands when candidates
//! exist but their estimated savings are below the configured floor.
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use async_trait::async_trait;
use futures::Stream;
use llm_engine::Engine;
use llm_engine::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent, UsageEvent};
use llm_engine::llm_client::{ClientError, LlmClient, Request};
use llm_engine::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use pod_store::{CombinedStore, FsWorkerStore};
use session_metrics::{DOMAIN, Metric, metrics_from_extensions};
use session_store::{FsStore, LogEntry, SegmentId, SessionId, Store, StoreError, TraceEntry};
use worker::{Worker, WorkerManifest};
type TestStore = CombinedStore<FsStore, FsWorkerStore>;
#[derive(Clone)]
struct MockClient {
responses: Arc<Vec<Vec<LlmEvent>>>,
call_count: Arc<AtomicUsize>,
}
impl MockClient {
fn new(responses: Vec<Vec<LlmEvent>>) -> Self {
Self {
responses: Arc::new(responses),
call_count: Arc::new(AtomicUsize::new(0)),
}
}
}
#[async_trait]
impl LlmClient for MockClient {
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>
{
let count = self.call_count.fetch_add(1, Ordering::SeqCst);
if count >= self.responses.len() {
return Err(ClientError::Config("mock client exhausted".into()));
}
let events = self.responses[count].clone();
let stream = futures::stream::iter(events.into_iter().map(Ok));
Ok(Box::pin(stream))
}
}
/// Tool that returns a fixed `ToolOutput { summary, content: Some(big) }`.
/// `content` is long enough for prune savings to comfortably clear small
/// `min_savings` thresholds.
struct BigContentTool {
summary: &'static str,
content: String,
}
#[async_trait]
impl Tool for BigContentTool {
async fn execute(
&self,
_input: &str,
_ctx: llm_engine::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
Ok(ToolOutput {
summary: self.summary.into(),
content: Some(self.content.clone()),
})
}
}
fn big_content_tool_definition(name: &'static str) -> ToolDefinition {
Arc::new(move || {
let summary = "tool result summary";
let content = "x".repeat(2048);
(
ToolMeta::new(name)
.description("test tool that returns a long content")
.input_schema(serde_json::json!({"type": "object"})),
Arc::new(BigContentTool { summary, content }) as Arc<dyn Tool>,
)
})
}
fn usage_event(input_total: u64, cache_read: u64, cache_write: u64, output: u64) -> LlmEvent {
LlmEvent::Usage(UsageEvent {
input_tokens: Some(input_total),
output_tokens: Some(output),
total_tokens: Some(input_total + output),
cache_read_input_tokens: Some(cache_read),
cache_creation_input_tokens: Some(cache_write),
})
}
/// Tool-call response from the assistant: emits a `tool_use` block then a
/// usage event so usage_history gains a measurement on this turn.
fn tool_use_response(call_id: &str, tool_name: &str) -> Vec<LlmEvent> {
vec![
LlmEvent::tool_use_start(0, call_id, tool_name),
LlmEvent::tool_input_delta(0, "{}"),
LlmEvent::tool_use_stop(0),
usage_event(500, 0, 0, 10),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
/// Plain text response with explicit cache_read/cache_write so that
/// `prune.post_request` can carry meaningful values when this is the
/// LLM call that follows a `prune.fire` event.
fn text_response_with_cache(text: &str, cache_read: u64, cache_write: u64) -> Vec<LlmEvent> {
vec![
LlmEvent::text_block_start(0),
LlmEvent::text_delta(0, text),
LlmEvent::text_block_stop(0, None),
usage_event(800, cache_read, cache_write, 5),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
fn manifest_toml(prune_protected_tokens: u64, prune_min_savings: u64) -> String {
format!(
r#"
[worker]
name = "test-worker"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[compaction]
prune_protected_tokens = {prune_protected_tokens}
prune_min_savings = {prune_min_savings}
[[scope.allow]]
target = "./"
permission = "write"
"#
)
}
async fn make_worker(
manifest_toml: String,
client: MockClient,
tool_name: &'static str,
) -> (
Worker<MockClient, TestStore>,
tempfile::TempDir,
tempfile::TempDir,
) {
let manifest = WorkerManifest::from_toml(&manifest_toml).unwrap();
let store_tmp = tempfile::tempdir().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
);
let pwd_tmp = tempfile::tempdir().unwrap();
let pwd = pwd_tmp.path().to_path_buf();
let scope = worker::Scope::writable(&pwd).unwrap();
let mut worker = Engine::new(client);
worker.register_tool(big_content_tool_definition(tool_name));
let worker = Worker::new(manifest, worker, store, pwd, scope)
.await
.unwrap();
(worker, store_tmp, pwd_tmp)
}
/// Drive Worker through enough runs to exercise both skip-no_candidates and
/// fire branches, then read the session log back and assert the metric
/// stream.
#[tokio::test]
async fn prune_metrics_emit_skip_then_fire_with_post_request_join() {
// Run 1 (request 0): tool_use → triggers tool execution → request 1
// on the second iteration to produce the assistant reply.
// Run 2 (request 2): plain assistant text. Prune evaluation here
// sees user1's tool_result outside the protected-token suffix and
// should fire.
let client = MockClient::new(vec![
tool_use_response("call-1", "big_tool"),
text_response_with_cache("ok", 0, 200),
text_response_with_cache("done", 1234, 50),
]);
let (mut worker, _store_tmp, _pwd_tmp) =
make_worker(manifest_toml(1, 1), client, "big_tool").await;
let session_id = worker.session_id();
let segment_id = worker.segment_id();
// Cloning the store handle to read the session log back after the
// runs complete — the Worker retains its own copy.
let store = worker.store().clone();
worker.run_text("first").await.unwrap();
worker.run_text("second").await.unwrap();
let state = session_store::restore(&store, session_id, segment_id).unwrap();
let metrics = metrics_from_extensions(&state.extensions);
// Run 1 has 2 LLM iterations (tool loop), each evaluates prune with
// only one user-message turn → 2x skip{no_candidates}.
// Run 2 has 1 LLM iteration with enough turns → 1x fire +
// 1x post_request paired by correlation_id.
let names: Vec<&str> = metrics.iter().map(|m| m.name.as_str()).collect();
assert!(
names.contains(&"prune.skip"),
"expected prune.skip in {names:?}"
);
assert!(
names.contains(&"prune.fire"),
"expected prune.fire in {names:?}"
);
assert!(
names.contains(&"prune.post_request"),
"expected prune.post_request in {names:?}"
);
// All skips in run 1 must record reason=no_candidates.
for m in metrics.iter().filter(|m| m.name == "prune.skip") {
assert_eq!(
m.dimensions.get("reason").map(String::as_str),
Some("no_candidates"),
"skip metric should be no_candidates here, got {m:?}"
);
assert!(m.correlation_id.is_none());
}
// The fire metric carries dimensions and correlation_id.
let fire = metrics
.iter()
.find(|m| m.name == "prune.fire")
.expect("prune.fire missing");
assert!(
fire.dimensions.contains_key("candidate_count"),
"fire missing candidate_count: {fire:?}"
);
assert!(
fire.dimensions.contains_key("protected_start_index"),
"fire missing protected_start_index: {fire:?}"
);
assert!(fire.value.is_some(), "fire missing estimated_savings value");
let fire_id = fire
.correlation_id
.as_ref()
.expect("fire metric missing correlation_id");
// Exactly one post_request metric should exist with the same id, and
// its value/dimension should reflect the cache numbers from the
// text_response_with_cache call (cache_read=1234, cache_write=50).
let post = metrics
.iter()
.find(|m| m.name == "prune.post_request")
.expect("prune.post_request missing");
assert_eq!(post.correlation_id.as_ref(), Some(fire_id));
assert_eq!(post.value, Some(1234.0));
assert_eq!(
post.dimensions
.get("cache_write_tokens")
.map(String::as_str),
Some("50")
);
assert!(post.dimensions.contains_key("history_len"));
}
#[tokio::test]
async fn prune_metrics_fire_during_single_long_task_without_multiple_user_turns() {
let client = MockClient::new(vec![
tool_use_response("call-1", "big_tool"),
tool_use_response("call-2", "big_tool"),
tool_use_response("call-3", "big_tool"),
tool_use_response("call-4", "big_tool"),
text_response_with_cache("done", 100, 20),
]);
let (mut worker, _store_tmp, _pwd_tmp) =
make_worker(manifest_toml(1, 1), client, "big_tool").await;
let session_id = worker.session_id();
let segment_id = worker.segment_id();
let store = worker.store().clone();
worker.run_text("one long task").await.unwrap();
let state = session_store::restore(&store, session_id, segment_id).unwrap();
let metrics = metrics_from_extensions(&state.extensions);
let fire_count = metrics.iter().filter(|m| m.name == "prune.fire").count();
assert!(
fire_count > 0,
"single-turn tool loop should produce prune.fire once old heavy ToolResults fall outside the protected-token suffix: {metrics:?}"
);
assert!(
metrics.iter().any(|m| {
m.name == "prune.fire" && m.dimensions.contains_key("protected_start_index")
})
);
}
/// `min_savings` set high enough that candidates exist but the estimated
/// savings always fall short → the second run should record
/// `prune.skip { reason: "below_min_savings" }`.
#[tokio::test]
async fn prune_metrics_record_below_min_savings_skip() {
let client = MockClient::new(vec![
tool_use_response("call-1", "big_tool"),
text_response_with_cache("ok", 0, 100),
text_response_with_cache("done", 0, 0),
]);
let (mut worker, _store_tmp, _pwd_tmp) =
make_worker(manifest_toml(1, 1_000_000), client, "big_tool").await;
let session_id = worker.session_id();
let segment_id = worker.segment_id();
let store = worker.store().clone();
worker.run_text("first").await.unwrap();
worker.run_text("second").await.unwrap();
let state = session_store::restore(&store, session_id, segment_id).unwrap();
let metrics = metrics_from_extensions(&state.extensions);
let below = metrics
.iter()
.find(|m| {
m.name == "prune.skip"
&& m.dimensions.get("reason").map(String::as_str) == Some("below_min_savings")
})
.expect("expected prune.skip with reason=below_min_savings");
assert!(
below.dimensions.contains_key("candidate_count"),
"below_min_savings skip should report candidate_count: {below:?}"
);
assert!(
below.value.is_some(),
"below_min_savings skip should report estimated savings as value: {below:?}"
);
// No prune.fire for this scenario.
assert!(metrics.iter().all(|m| m.name != "prune.fire"));
// No prune.post_request either (no fire to join with).
assert!(metrics.iter().all(|m| m.name != "prune.post_request"));
}
/// `Store` wrapper that delegates to an inner `FsStore` for everything
/// except `LogEntry::Extension { domain: "metrics", .. }` appends, which
/// it rejects with an `Io` error. Lets us drive the `try_record_metric`
/// failure path without affecting any other persistence write.
#[derive(Clone)]
struct MetricFailingStore {
inner: FsStore,
}
impl Store for MetricFailingStore {
fn append(
&self,
session_id: SessionId,
segment_id: SegmentId,
entry: &LogEntry,
) -> Result<(), StoreError> {
if let LogEntry::Extension { domain, .. } = entry {
if domain == DOMAIN {
return Err(StoreError::Io(std::io::Error::other("synthetic failure")));
}
}
self.inner.append(session_id, segment_id, entry)
}
fn read_all(
&self,
session_id: SessionId,
segment_id: SegmentId,
) -> Result<Vec<LogEntry>, StoreError> {
self.inner.read_all(session_id, segment_id)
}
fn list_sessions(&self) -> Result<Vec<SessionId>, StoreError> {
self.inner.list_sessions()
}
fn list_segments(&self, session_id: SessionId) -> Result<Vec<SegmentId>, StoreError> {
self.inner.list_segments(session_id)
}
fn lookup_session_of(&self, segment_id: SegmentId) -> Result<Option<SessionId>, StoreError> {
self.inner.lookup_session_of(segment_id)
}
fn create_segment(
&self,
session_id: SessionId,
segment_id: SegmentId,
entries: &[LogEntry],
) -> Result<(), StoreError> {
self.inner.create_segment(session_id, segment_id, entries)
}
fn exists(&self, session_id: SessionId, segment_id: SegmentId) -> Result<bool, StoreError> {
self.inner.exists(session_id, segment_id)
}
fn read_entry_count(
&self,
session_id: SessionId,
segment_id: SegmentId,
) -> Result<usize, StoreError> {
self.inner.read_entry_count(session_id, segment_id)
}
fn append_trace(
&self,
session_id: SessionId,
segment_id: SegmentId,
entry: &TraceEntry,
) -> Result<(), StoreError> {
self.inner.append_trace(session_id, segment_id, entry)
}
}
/// Metric write failures are non-fatal: the run still completes, the
/// session log carries no metric entries (drops), but a `Warn` alert
/// fires on the alerter so the TUI surface picks it up.
#[tokio::test]
async fn metric_write_failure_emits_warn_alert_and_does_not_abort_run() {
use protocol::{AlertLevel, AlertSource, Event};
use tokio::sync::broadcast;
let manifest_toml = manifest_toml(1, 1);
let manifest = WorkerManifest::from_toml(&manifest_toml).unwrap();
let store_tmp = tempfile::tempdir().unwrap();
let inner = FsStore::new(store_tmp.path()).unwrap();
let store = MetricFailingStore { inner };
let pwd_tmp = tempfile::tempdir().unwrap();
let pwd = pwd_tmp.path().to_path_buf();
let scope = worker::Scope::writable(&pwd).unwrap();
// Even with a tool registered, this run will only emit
// `prune.skip { reason: "no_candidates" }` (one user message,
// protected token budget covers the only user message). That is enough to drive
// the failure path: at least one metric attempts to write.
let client = MockClient::new(vec![text_response_with_cache("hi", 0, 0)]);
let worker = Engine::new(client);
let mut worker = Worker::new(manifest, worker, store.clone(), pwd, scope)
.await
.unwrap();
let (tx, mut rx) = broadcast::channel::<Event>(64);
let alerter = worker::Alerter::new(tx);
worker.attach_alerter(alerter);
let session_id = worker.session_id();
let segment_id = worker.segment_id();
// Run completes successfully despite metric failure.
worker.run_text("hello").await.unwrap();
// No metrics ended up in the log (writes were rejected).
let state = session_store::restore(&store, session_id, segment_id).unwrap();
let metrics = metrics_from_extensions(&state.extensions);
assert!(metrics.is_empty(), "metrics must drop on write failure");
// The alerter saw at least one Warn from AlertSource::Worker.
let mut saw_warn = false;
while let Ok(ev) = rx.try_recv() {
if let Event::Alert(a) = ev {
if a.level == AlertLevel::Warn
&& a.source == AlertSource::Worker
&& a.message.contains("metric")
{
saw_warn = true;
break;
}
}
}
assert!(saw_warn, "expected Warn/Worker alert about metric failure");
}
/// Sessions that have no metrics in the log restore cleanly: the
/// `RestoredState.extensions` simply contains no `metrics` domain
/// payloads, and `metrics_from_extensions` returns an empty Vec.
/// Backward-compatibility check for old logs predating this feature.
#[tokio::test]
async fn old_sessions_without_metrics_replay_cleanly() {
// Manifest without any `[compaction]` section → prune (and therefore
// the prune observer) is never installed, so no metrics get written.
let manifest_toml = r#"
[worker]
name = "test-worker"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[[scope.allow]]
target = "./"
permission = "write"
"#;
let client = MockClient::new(vec![text_response_with_cache("hi", 0, 0)]);
let manifest = WorkerManifest::from_toml(manifest_toml).unwrap();
let store_tmp = tempfile::tempdir().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
);
let pwd_tmp = tempfile::tempdir().unwrap();
let pwd = pwd_tmp.path().to_path_buf();
let scope = worker::Scope::writable(&pwd).unwrap();
let worker = Engine::new(client);
let mut worker = Worker::new(manifest, worker, store.clone(), pwd, scope)
.await
.unwrap();
let session_id = worker.session_id();
let segment_id = worker.segment_id();
worker.run_text("hello").await.unwrap();
let state = session_store::restore(&store, session_id, segment_id).unwrap();
let metrics = metrics_from_extensions(&state.extensions);
assert!(
metrics.is_empty(),
"no metrics should be recorded: {metrics:?}"
);
// And no extension entries at all in the metrics domain.
assert!(state.extensions.iter().all(|(d, _)| d != DOMAIN));
// Smoke check that fold helper is robust on a sentinel Metric value:
let m = Metric::now("smoke");
assert_eq!(m.name, "smoke");
}
crates/worker/tests/spawn_worker_test.rs
+743
View File
@@ -0,0 +1,743 @@
//! Integration tests for the `SpawnWorker` tool.
//!
//! These tests exercise the tool's pod-registry 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::pod_registry::{self, LockFileGuard};
use worker::spawn::registry::SpawnedWorkerRegistry;
use worker::spawn::tool::spawn_worker_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 = pod_registry::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 `spawn_worker_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 = spawn_worker_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 = spawn_worker_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 = spawn_worker_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 pod_registry has the child allocation under `root`.
let lock_path = pod_registry::default_registry_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 = spawn_worker_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 = spawn_worker_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 = spawn_worker_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 = pod_registry::default_registry_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 = spawn_worker_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 = pod_registry::default_registry_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/system_prompt_template_test.rs
+342
View File
@@ -0,0 +1,342 @@
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use async_trait::async_trait;
use futures::Stream;
use llm_engine::Engine;
use llm_engine::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent};
use llm_engine::llm_client::{ClientError, LlmClient, Request};
use pod_store::{CombinedStore, FsWorkerStore};
use session_store::{FsStore, LogEntry, Store};
use worker::{PromptLoader, SystemPromptTemplate, Worker, WorkerError};
type TestStore = CombinedStore<FsStore, FsWorkerStore>;
// ---------------------------------------------------------------------------
// Mock LLM Client
// ---------------------------------------------------------------------------
#[derive(Clone)]
struct MockClient {
responses: Arc<Vec<Vec<LlmEvent>>>,
call_count: Arc<AtomicUsize>,
}
impl MockClient {
fn new(responses: Vec<Vec<LlmEvent>>) -> Self {
Self {
responses: Arc::new(responses),
call_count: Arc::new(AtomicUsize::new(0)),
}
}
}
#[async_trait]
impl LlmClient for MockClient {
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>
{
let count = self.call_count.fetch_add(1, Ordering::SeqCst);
let idx = count.min(self.responses.len() - 1);
let events = self.responses[idx].clone();
let stream = futures::stream::iter(events.into_iter().map(Ok));
Ok(Box::pin(stream))
}
}
fn single_text_events(text: &str) -> Vec<LlmEvent> {
vec![
LlmEvent::text_block_start(0),
LlmEvent::text_delta(0, text),
LlmEvent::text_block_stop(0, None),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
/// Emit a single `write_summary(text=...)` tool call as one LLM response.
fn write_summary_tool_use_events(call_id: &str, text: &str) -> Vec<LlmEvent> {
let input = serde_json::json!({ "text": text }).to_string();
vec![
LlmEvent::tool_use_start(0, call_id, "write_summary"),
LlmEvent::tool_input_delta(0, input),
LlmEvent::tool_use_stop(0),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
const MINIMAL_MANIFEST_TOML: &str = r#"
[worker]
name = "test-worker"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[[scope.allow]]
target = "./"
permission = "write"
"#;
/// Build a Worker with a synthetic instruction template.
///
/// Writes `body` to a temp user-prompts dir under `$user/test`, builds a
/// PromptLoader pointing at it, parses the template, and installs it on
/// a Worker constructed directly via `Worker::new`.
async fn make_worker_with_body(
body: &str,
client: MockClient,
) -> Result<(Worker<MockClient, TestStore>, PathBuf), WorkerError> {
let manifest = worker::WorkerManifest::from_toml(MINIMAL_MANIFEST_TOML).unwrap();
let store_tmp = tempfile::tempdir().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
);
std::mem::forget(store_tmp);
let pwd_tmp = tempfile::tempdir().unwrap();
let pwd = pwd_tmp.path().to_path_buf();
let scope = worker::Scope::writable(&pwd).unwrap();
std::mem::forget(pwd_tmp);
let user_prompts_tmp = tempfile::tempdir().unwrap();
std::fs::write(user_prompts_tmp.path().join("test.md"), body).unwrap();
let loader = PromptLoader::new(Some(user_prompts_tmp.path().to_path_buf()), None);
std::mem::forget(user_prompts_tmp);
let worker = Engine::new(client);
let mut worker = Worker::new(manifest, worker, store, pwd.clone(), scope).await?;
let template = SystemPromptTemplate::parse("$user/test", loader)
.map_err(|source| WorkerError::InvalidSystemPromptTemplate { source })?;
worker.set_system_prompt_template(template);
Ok((worker, pwd))
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[tokio::test]
async fn template_parse_rejects_invalid_syntax() {
let user_prompts_tmp = tempfile::tempdir().unwrap();
std::fs::write(user_prompts_tmp.path().join("broken.md"), "{{ unclosed").unwrap();
let loader = PromptLoader::new(Some(user_prompts_tmp.path().to_path_buf()), None);
let err = SystemPromptTemplate::parse("$user/broken", loader).unwrap_err();
let worker_err: WorkerError = WorkerError::InvalidSystemPromptTemplate { source: err };
assert!(matches!(
worker_err,
WorkerError::InvalidSystemPromptTemplate { .. }
));
}
#[tokio::test]
async fn template_is_not_materialised_before_first_run() {
let client = MockClient::new(vec![single_text_events("ok")]);
let (worker, _pwd) = make_worker_with_body("hello", client).await.unwrap();
// Before first run, worker still has no system prompt.
assert!(worker.engine().get_system_prompt().is_none());
}
#[tokio::test]
async fn materialise_on_first_turn_populates_worker() {
let client = MockClient::new(vec![single_text_events("ok")]);
let (mut worker, pwd) = make_worker_with_body(
"date={{ date }} cwd={{ cwd }} tools={{ tools | join(',') }}",
client,
)
.await
.unwrap();
worker.run_text("hi").await.unwrap();
let rendered = worker
.engine()
.get_system_prompt()
.expect("system prompt materialised")
.to_string();
assert!(rendered.contains("date="));
assert!(rendered.contains("cwd="));
assert!(rendered.contains(&pwd.display().to_string()));
assert!(rendered.starts_with("date="));
// Trailing fixed section must be appended.
assert!(rendered.contains("## Working boundaries"));
}
#[tokio::test]
async fn session_start_state_captures_rendered_prompt() {
let client = MockClient::new(vec![single_text_events("ok")]);
let (mut worker, pwd) = make_worker_with_body("hello cwd={{ cwd }}", client)
.await
.unwrap();
worker.run_text("hi").await.unwrap();
let entries = worker
.store()
.read_all(worker.session_id(), worker.segment_id())
.unwrap();
let first = entries.first().expect("at least one entry");
match first {
LogEntry::SegmentStart { system_prompt, .. } => {
let sp = system_prompt.as_deref().expect("system prompt set");
assert!(sp.starts_with("hello cwd="));
assert!(sp.contains(&pwd.display().to_string()));
assert!(sp.contains("## Working boundaries"));
}
other => panic!("expected SegmentStart as first entry, got {other:?}"),
}
}
#[tokio::test]
async fn render_failure_propagates_as_worker_error() {
let client = MockClient::new(vec![single_text_events("ok")]);
let (mut worker, _pwd) = make_worker_with_body("{{ ghost }}", client).await.unwrap();
let err = worker.run_text("hi").await.unwrap_err();
assert!(matches!(err, WorkerError::SystemPromptRender { .. }));
}
#[tokio::test]
async fn materialise_runs_only_once_across_turns() {
let client = MockClient::new(vec![
single_text_events("first"),
single_text_events("second"),
]);
let (mut worker, _pwd) = make_worker_with_body("fixed prompt {{ cwd }}", client)
.await
.unwrap();
worker.run_text("one").await.unwrap();
let first = worker.engine().get_system_prompt().unwrap().to_string();
worker.run_text("two").await.unwrap();
let second = worker.engine().get_system_prompt().unwrap().to_string();
assert_eq!(first, second);
}
#[tokio::test]
async fn agents_md_is_injected_as_trailing_section_when_present() {
let client = MockClient::new(vec![single_text_events("ok")]);
let (mut worker, pwd) = make_worker_with_body("BODY", client).await.unwrap();
std::fs::write(pwd.join("AGENTS.md"), "# project rules\nbe kind").unwrap();
worker.run_text("hi").await.unwrap();
let rendered = worker.engine().get_system_prompt().unwrap().to_string();
assert!(rendered.starts_with("BODY"));
assert!(rendered.contains("## Project instructions (AGENTS.md)"));
assert!(rendered.contains("# project rules"));
assert!(rendered.contains("be kind"));
}
#[tokio::test]
async fn agents_md_absent_omits_trailing_section() {
let client = MockClient::new(vec![single_text_events("ok")]);
let (mut worker, _pwd) = make_worker_with_body("BODY", client).await.unwrap();
worker.run_text("hi").await.unwrap();
let rendered = worker.engine().get_system_prompt().unwrap().to_string();
assert!(!rendered.contains("## Project instructions"));
assert!(!rendered.contains("AGENTS.md"));
}
#[tokio::test]
async fn agents_md_not_reread_after_compact() {
let client = MockClient::new(vec![
single_text_events("a"), // worker.run_text("first")
single_text_events("b"), // worker.run_text("second")
write_summary_tool_use_events("call-1", "compacted summary"), // compact worker: tool_use
single_text_events("done"), // compact worker: close
single_text_events("c"), // worker.run_text("third")
]);
let (mut worker, pwd) = make_worker_with_body("BODY", client).await.unwrap();
let agents_path = pwd.join("AGENTS.md");
std::fs::write(&agents_path, "original").unwrap();
worker.run_text("first").await.unwrap();
let before = worker.engine().get_system_prompt().unwrap().to_string();
assert!(before.contains("original"));
worker.run_text("second").await.unwrap();
// Mutate the file after the first turn — must not affect the cached
// system prompt either on a subsequent turn or across compaction.
std::fs::write(&agents_path, "mutated").unwrap();
worker.compact(0).await.unwrap();
let after_compact = worker.engine().get_system_prompt().unwrap().to_string();
assert!(after_compact.contains("original"));
assert!(!after_compact.contains("mutated"));
worker.run_text("third").await.unwrap();
let after_third = worker.engine().get_system_prompt().unwrap().to_string();
assert!(after_third.contains("original"));
assert!(!after_third.contains("mutated"));
}
#[tokio::test]
async fn compact_aligns_user_segments_with_retained_history() {
// retained_tokens=0 folds the entire conversation into the summary,
// so retained_items has zero user_messages and self.user_segments
// must be drained to match. A subsequent run() then appends fresh
// segments cleanly without ghost entries from the pre-compaction era.
let client = MockClient::new(vec![
single_text_events("a"),
single_text_events("b"),
write_summary_tool_use_events("call-1", "compacted summary"),
single_text_events("done"),
single_text_events("c"),
]);
let (mut worker, _pwd) = make_worker_with_body("BODY", client).await.unwrap();
worker.run_text("first").await.unwrap();
worker.run_text("second").await.unwrap();
assert_eq!(worker.user_segments().len(), 2);
worker.compact(0).await.unwrap();
assert_eq!(
worker.user_segments().len(),
0,
"compact(0) folds every user_message into the summary, so segments \
must be drained to match retained_items"
);
worker.run_text("third").await.unwrap();
assert_eq!(worker.user_segments().len(), 1);
}
#[tokio::test]
async fn compact_preserves_system_prompt() {
let client = MockClient::new(vec![
single_text_events("a"), // worker.run_text("first")
single_text_events("b"), // worker.run_text("second")
write_summary_tool_use_events("call-1", "compacted summary"), // compact worker: tool_use
single_text_events("done"), // compact worker: close
single_text_events("c"), // worker.run_text("third")
]);
let (mut worker, _pwd) = make_worker_with_body("SP cwd={{ cwd }}", client)
.await
.unwrap();
worker.run_text("first").await.unwrap();
let before = worker.engine().get_system_prompt().unwrap().to_string();
worker.run_text("second").await.unwrap();
worker.compact(0).await.unwrap();
let after = worker.engine().get_system_prompt().unwrap().to_string();
assert_eq!(before, after);
worker.run_text("third").await.unwrap();
assert_eq!(worker.engine().get_system_prompt().unwrap(), after.as_str());
}
crates/worker/tests/worker_comm_tools_test.rs
+728
View File
@@ -0,0 +1,728 @@
//! Integration tests for the worker-comm tools (`SendToWorker`,
//! `ReadWorkerOutput`, `StopWorker`).
//!
//! 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 pod_store::{CombinedStore, FsWorkerStore, WorkerMetadataStore};
use protocol::stream::{JsonLineReader, JsonLineWriter};
use protocol::{ErrorCode, Event, Greeting, Method};
use serde_json::json;
use session_store::FsStore;
use tempfile::TempDir;
use tokio::net::UnixListener;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use worker::runtime::dir::{RuntimeDir, SpawnedWorkerRecord};
use worker::runtime::pod_registry::{self, LockFileGuard};
use worker::spawn::comm_tools::{read_worker_output_tool, send_to_worker_tool, stop_worker_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 `SendToWorker`
/// 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 `ReadWorkerOutput`'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,
}
}
// ---------------------------------------------------------------------------
// SendToWorker
// ---------------------------------------------------------------------------
#[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 `SendToWorker`'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 = send_to_worker_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 = send_to_worker_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 = send_to_worker_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 { .. }));
}
// ---------------------------------------------------------------------------
// ReadWorkerOutput
// ---------------------------------------------------------------------------
#[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 = read_worker_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 = read_worker_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);
}
// ---------------------------------------------------------------------------
// StopWorker
// ---------------------------------------------------------------------------
#[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 SpawnWorker.
{
let mut g = LockFileGuard::open(&lock_path).unwrap();
let rule = ScopeRule {
target: tmp.path().to_path_buf(),
permission: Permission::Write,
recursive: true,
};
pod_registry::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();
pod_registry::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 = stop_worker_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. StopWorker 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 = stop_worker_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 = send_to_worker_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 = stop_worker_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();
pod_registry::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());
}
crates/worker/tests/worker_events_test.rs
+423
View File
@@ -0,0 +1,423 @@
//! Integration tests for the `WorkerEvent` send / receive primitive.
//!
//! These tests drive `worker_events::fire_and_forget` and
//! `worker_events::apply_event_side_effects` directly — the full
//! Controller wiring is exercised by the existing controller /
//! spawn-worker tests, which rely on the same primitives.
use std::path::PathBuf;
use std::sync::{Arc, LazyLock, Mutex};
use std::time::Duration;
use protocol::stream::{JsonLineReader, JsonLineWriter};
use protocol::{Event, Greeting, Method, Permission, ScopeRule, WorkerEvent, WorkerStatus};
use tempfile::TempDir;
use tokio::net::UnixListener;
use worker::ipc::event::{apply_event_side_effects, fire_and_forget, render_event};
use worker::runtime::dir::{RuntimeDir, SpawnedWorkerRecord};
use worker::runtime::pod_registry::{self, LockFileGuard};
use worker::spawn::registry::SpawnedWorkerRegistry;
/// Serialises tests that mutate `YOI_RUNTIME_DIR`.
static ENV_LOCK: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));
/// Take `ENV_LOCK` and clear any env vars that would outrank
/// `YOI_RUNTIME_DIR`; 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");
}
}
}
/// Point `YOI_RUNTIME_DIR` at `dir`. The pod-registry then lives at
/// `<dir>/workers.json` and Worker runtime sub-dirs at `<dir>/{worker_name}/`.
fn set_runtime_dir(dir: &std::path::Path) {
unsafe {
std::env::set_var("YOI_RUNTIME_DIR", dir);
}
}
fn clear_runtime_dir() {
unsafe {
std::env::remove_var("YOI_RUNTIME_DIR");
}
}
/// Minimal connect-time snapshot used by mock parent sockets.
fn empty_snapshot() -> Event {
Event::Snapshot {
entries: Vec::new(),
greeting: Greeting {
worker_name: "parent".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: WorkerStatus::Idle,
in_flight: Default::default(),
}
}
/// Accept a single connection, send the protocol's connect-time snapshot,
/// read one `Method`, and return it.
fn accept_one_method(listener: UnixListener) -> tokio::task::JoinHandle<Option<Method>> {
tokio::spawn(async move {
let (stream, _) = listener.accept().await.ok()?;
let (reader, writer) = stream.into_split();
let mut w = JsonLineWriter::new(writer);
w.write(&empty_snapshot()).await.ok()?;
let mut r = JsonLineReader::new(reader);
r.next::<Method>().await.ok().flatten()
})
}
#[test]
fn render_event_all_variants_mention_worker_name() {
let t1 = render_event(&WorkerEvent::TurnEnded {
worker_name: "alpha".into(),
});
assert!(t1.contains("alpha"), "{t1}");
let t2 = render_event(&WorkerEvent::Errored {
worker_name: "bravo".into(),
message: "boom".into(),
});
assert!(t2.contains("bravo") && t2.contains("boom"), "{t2}");
let t3 = render_event(&WorkerEvent::ShutDown {
worker_name: "charlie".into(),
});
assert!(t3.contains("charlie"), "{t3}");
let t4 = render_event(&WorkerEvent::ScopeSubDelegated {
parent_worker: "delta".into(),
sub_worker: "echo".into(),
sub_socket: "/tmp/sock".into(),
scope: vec![],
});
assert!(t4.contains("delta") && t4.contains("echo"), "{t4}");
}
#[tokio::test]
async fn fire_and_forget_delivers_worker_event_to_listener() {
let dir = TempDir::new().unwrap();
let socket_path = dir.path().join("parent.sock");
let listener = UnixListener::bind(&socket_path).unwrap();
let received = accept_one_method(listener);
fire_and_forget(
Some(socket_path.clone()),
WorkerEvent::TurnEnded {
worker_name: "child".into(),
},
);
let method = tokio::time::timeout(Duration::from_secs(2), received)
.await
.expect("send timed out")
.unwrap()
.expect("no method received");
match method {
Method::WorkerEvent(WorkerEvent::TurnEnded { worker_name }) => {
assert_eq!(worker_name, "child")
}
other => panic!("expected TurnEnded, got {other:?}"),
}
}
#[tokio::test]
async fn fire_and_forget_with_none_socket_is_noop() {
// Nothing binds and nothing listens; the call must not panic and
// must not leak a task that never completes.
fire_and_forget(
None,
WorkerEvent::ShutDown {
worker_name: "x".into(),
},
);
// Yield once so any accidentally-spawned task would surface.
tokio::time::sleep(Duration::from_millis(50)).await;
}
/// Build a registry backed by a fresh runtime dir.
async fn fresh_registry(
runtime_base: &std::path::Path,
worker_name: &str,
) -> Arc<SpawnedWorkerRegistry> {
let rd = RuntimeDir::create(runtime_base, worker_name).await.unwrap();
SpawnedWorkerRegistry::new(Arc::new(rd))
}
#[tokio::test]
async fn apply_shutdown_removes_from_registry_and_tolerates_missing() {
let _env = EnvGuard::acquire();
let scope_dir = TempDir::new().unwrap();
set_runtime_dir(scope_dir.path());
let runtime_base = TempDir::new().unwrap();
let registry = fresh_registry(runtime_base.path(), "parent").await;
// Seed a child record; then ShutDown for it should remove it.
registry
.add(SpawnedWorkerRecord {
worker_name: "child".into(),
socket_path: "/tmp/child.sock".into(),
scope_delegated: vec![],
callback_address: "/tmp/parent.sock".into(),
})
.await
.unwrap();
let event = WorkerEvent::ShutDown {
worker_name: "child".into(),
};
apply_event_side_effects(&event, &registry, "parent", &None).await;
assert!(registry.get("child").await.is_none());
// Second ShutDown for the same (now-missing) child must be a no-op,
// not an error — this is the idempotency guarantee for out-of-order
// delivery.
apply_event_side_effects(&event, &registry, "parent", &None).await;
assert!(registry.get("child").await.is_none());
clear_runtime_dir();
}
#[tokio::test]
async fn apply_scope_sub_delegated_adds_grandchild_then_duplicate_is_noop() {
let _env = EnvGuard::acquire();
let scope_dir = TempDir::new().unwrap();
set_runtime_dir(scope_dir.path());
let runtime_base = TempDir::new().unwrap();
let registry = fresh_registry(runtime_base.path(), "grandparent").await;
// Seed the intermediate child so callback_address lookup succeeds.
registry
.add(SpawnedWorkerRecord {
worker_name: "child".into(),
socket_path: "/tmp/child.sock".into(),
scope_delegated: vec![],
callback_address: "/tmp/grandparent.sock".into(),
})
.await
.unwrap();
let event = WorkerEvent::ScopeSubDelegated {
parent_worker: "child".into(),
sub_worker: "grandchild".into(),
sub_socket: "/tmp/grandchild.sock".into(),
scope: vec![ScopeRule {
target: scope_dir.path().to_path_buf(),
permission: Permission::Write,
recursive: true,
}],
};
apply_event_side_effects(&event, &registry, "grandparent", &None).await;
let gc = registry
.get("grandchild")
.await
.expect("grandchild missing after ScopeSubDelegated");
assert_eq!(gc.socket_path, PathBuf::from("/tmp/grandchild.sock"));
assert_eq!(gc.callback_address, PathBuf::from("/tmp/child.sock"));
// Duplicate delivery must not error and must not overwrite.
apply_event_side_effects(&event, &registry, "grandparent", &None).await;
let gc2 = registry.get("grandchild").await.unwrap();
assert_eq!(gc2.socket_path, PathBuf::from("/tmp/grandchild.sock"));
clear_runtime_dir();
}
#[tokio::test]
async fn apply_scope_sub_delegated_reemits_to_own_parent() {
let _env = EnvGuard::acquire();
let scope_dir = TempDir::new().unwrap();
set_runtime_dir(scope_dir.path());
let runtime_base = TempDir::new().unwrap();
let registry = fresh_registry(runtime_base.path(), "B").await;
// Bind a listener at "A's" socket so we can watch the re-emission
// climb one level up the tree.
let sock_dir = TempDir::new().unwrap();
let a_socket = sock_dir.path().join("A.sock");
let listener = UnixListener::bind(&a_socket).unwrap();
let received = accept_one_method(listener);
// Seed the child record that the event claims spawned the grandchild.
registry
.add(SpawnedWorkerRecord {
worker_name: "C".into(),
socket_path: "/tmp/C.sock".into(),
scope_delegated: vec![],
callback_address: "/tmp/B.sock".into(),
})
.await
.unwrap();
let event = WorkerEvent::ScopeSubDelegated {
parent_worker: "C".into(),
sub_worker: "D".into(),
sub_socket: "/tmp/D.sock".into(),
scope: vec![],
};
// Self is B, and B's parent socket is A's listener.
apply_event_side_effects(&event, &registry, "B", &Some(a_socket.clone())).await;
// A must see the re-emission with parent_worker set to "B" (the
// sender from A's perspective), not "C" (the original sender's
// local view).
let method = tokio::time::timeout(Duration::from_secs(2), received)
.await
.expect("re-emission timed out")
.unwrap()
.expect("no method received on A's socket");
match method {
Method::WorkerEvent(WorkerEvent::ScopeSubDelegated {
parent_worker,
sub_worker,
..
}) => {
assert_eq!(parent_worker, "B");
assert_eq!(sub_worker, "D");
}
other => panic!("expected re-emitted ScopeSubDelegated, got {other:?}"),
}
clear_runtime_dir();
}
#[tokio::test]
async fn apply_turn_ended_and_errored_are_system_noops() {
let _env = EnvGuard::acquire();
let scope_dir = TempDir::new().unwrap();
set_runtime_dir(scope_dir.path());
let runtime_base = TempDir::new().unwrap();
let registry = fresh_registry(runtime_base.path(), "parent").await;
// Seed a child to verify it survives the no-op path.
registry
.add(SpawnedWorkerRecord {
worker_name: "child".into(),
socket_path: "/tmp/child.sock".into(),
scope_delegated: vec![],
callback_address: "/tmp/parent.sock".into(),
})
.await
.unwrap();
apply_event_side_effects(
&WorkerEvent::TurnEnded {
worker_name: "child".into(),
},
&registry,
"parent",
&None,
)
.await;
apply_event_side_effects(
&WorkerEvent::Errored {
worker_name: "child".into(),
message: "x".into(),
},
&registry,
"parent",
&None,
)
.await;
assert!(registry.get("child").await.is_some());
clear_runtime_dir();
}
#[tokio::test]
async fn shutdown_releases_scope_allocation_when_present() {
let _env = EnvGuard::acquire();
let scope_dir = TempDir::new().unwrap();
let lock_path = scope_dir.path().join("workers.json");
set_runtime_dir(scope_dir.path());
// Install a top-level allocation for "kid" so ShutDown has
// something to release.
let guard = pod_registry::install_top_level(
"kid".into(),
std::process::id(),
"/tmp/kid.sock".into(),
vec![],
session_store::new_segment_id(),
)
.unwrap();
std::mem::forget(guard);
let runtime_base = TempDir::new().unwrap();
let registry = fresh_registry(runtime_base.path(), "parent").await;
registry
.add(SpawnedWorkerRecord {
worker_name: "kid".into(),
socket_path: "/tmp/kid.sock".into(),
scope_delegated: vec![],
callback_address: "/tmp/parent.sock".into(),
})
.await
.unwrap();
apply_event_side_effects(
&WorkerEvent::ShutDown {
worker_name: "kid".into(),
},
&registry,
"parent",
&None,
)
.await;
// Allocation is gone from the pod-registry.
let g = LockFileGuard::open(&lock_path).unwrap();
assert!(
g.data().find("kid").is_none(),
"ShutDown should have released the scope allocation"
);
clear_runtime_dir();
}