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

feat: add revisioned worker execution state

Browse Source
This commit is contained in:
Hare
2026-09-06 06:54:29 +09:00
parent 3ed1545c3c
commit e8b9adcde4
32 changed files with 2168 additions and 750 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/tests/compact_events_test.rs
+202 -6
View File
@@ -72,6 +72,41 @@ impl LlmClient for MockClient {
}
}
#[derive(Clone)]
struct BlockingCompactClient {
calls: Arc<AtomicUsize>,
}
impl BlockingCompactClient {
fn new() -> Self {
Self {
calls: Arc::new(AtomicUsize::new(0)),
}
}
}
#[async_trait]
impl LlmClient for BlockingCompactClient {
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 call = self.calls.fetch_add(1, Ordering::SeqCst);
if call == 0 {
Ok(Box::pin(futures::stream::iter(
single_text_events("seed").into_iter().map(Ok),
)))
} else {
Ok(Box::pin(futures::stream::pending()))
}
}
}
fn single_text_events(text: &str) -> Vec<LlmEvent> {
vec![
LlmEvent::text_block_start(0),
@@ -156,10 +191,10 @@ target = "./"
permission = "write"
"#;
async fn make_worker_with_manifest(
manifest_toml: &str,
client: MockClient,
) -> Worker<MockClient, TestStore> {
async fn make_worker_with_manifest<C>(manifest_toml: &str, client: C) -> Worker<C, TestStore>
where
C: LlmClient + Clone + Send + Sync + 'static,
{
let manifest = worker::WorkerManifest::from_toml(manifest_toml).unwrap();
let store_tmp = tempfile::tempdir().unwrap();
@@ -614,12 +649,144 @@ async fn pre_run_compact_failure_broadcasts_start_and_failed() {
);
}
#[tokio::test]
async fn manual_compact_cancel_terminalizes_before_returning_idle() {
let worker =
make_worker_with_manifest(POST_RUN_MANIFEST_TOML, BlockingCompactClient::new()).await;
let runtime_tmp = tempfile::tempdir().unwrap();
let bash_output_dir = runtime_tmp.path().join("bash-output");
let (handle, shutdown_receiver) =
WorkerController::spawn(worker, runtime_tmp.path(), &bash_output_dir)
.await
.unwrap();
let mut rx = handle.subscribe();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"seed history",
))
.await
.expect("send seed run");
loop {
if matches!(
tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for seed run")
.expect("event"),
Event::RunEnd {
result: RunResult::Finished
}
) {
break;
}
}
let compact = protocol::WorkerCommandEnvelope::for_snapshot(1, &handle.shared_state.snapshot());
handle
.send(Method::Compact { command: compact })
.await
.expect("send compact");
loop {
if matches!(
tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for compact start")
.expect("event"),
Event::CompactStart { .. }
) {
break;
}
}
let cancel = protocol::WorkerCommandEnvelope::for_snapshot(2, &handle.shared_state.snapshot());
handle
.send(Method::Cancel { command: cancel })
.await
.expect("send compact cancel");
let mut saw_interrupted = false;
let mut saw_idle = false;
while !(saw_interrupted && saw_idle) {
match tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for compact cancellation")
.expect("event")
{
Event::CompactFailed { lifecycle }
if lifecycle.state == protocol::CompactionLifecycleState::Interrupted =>
{
saw_interrupted = true;
}
Event::WorkerState { snapshot }
if snapshot.catalog_status() == protocol::WorkerStatus::Idle =>
{
assert!(
saw_interrupted,
"Idle must follow durable Interrupted evidence"
);
saw_idle = true;
}
_ => {}
}
}
let compact = protocol::WorkerCommandEnvelope::for_snapshot(3, &handle.shared_state.snapshot());
handle
.send(Method::Compact { command: compact })
.await
.expect("send second compact");
loop {
if matches!(
tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for second compact start")
.expect("event"),
Event::CompactStart { .. }
) {
break;
}
}
let shutdown =
protocol::WorkerCommandEnvelope::for_snapshot(4, &handle.shared_state.snapshot());
handle
.send(Method::Shutdown { command: shutdown })
.await
.expect("send shutdown during compact");
let mut interrupted_before_shutdown = false;
loop {
match tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for shutdown")
.expect("event")
{
Event::CompactFailed { lifecycle }
if lifecycle.state == protocol::CompactionLifecycleState::Interrupted =>
{
interrupted_before_shutdown = true;
}
Event::Shutdown => {
assert!(
interrupted_before_shutdown,
"shutdown must await terminal compaction evidence"
);
break;
}
_ => {}
}
}
tokio::time::timeout(std::time::Duration::from_secs(2), shutdown_receiver)
.await
.expect("controller shutdown timeout")
.expect("shutdown confirmation");
}
#[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"),
single_text_events("follow-up"),
]);
let worker = make_worker_with_manifest(POST_RUN_MANIFEST_TOML, client).await;
let runtime_tmp = tempfile::tempdir().unwrap();
@@ -649,7 +816,11 @@ async fn controller_compact_method_emits_start_and_done() {
}
}
handle.send(Method::Compact).await.expect("send compact");
let command = protocol::WorkerCommandEnvelope::for_snapshot(1, &handle.shared_state.snapshot());
handle
.send(Method::Compact { command })
.await
.expect("send compact");
let mut saw_start = false;
loop {
match tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
@@ -670,5 +841,30 @@ async fn controller_compact_method_emits_start_and_done() {
}
assert!(saw_start, "manual compact should emit CompactStart");
let _ = handle.send(Method::Shutdown).await;
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"run after compact",
))
.await
.expect("send follow-up run");
loop {
match tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for follow-up run")
.expect("event")
{
Event::RunEnd {
result: RunResult::Finished,
} => break,
_ => {}
}
}
assert_eq!(
handle.shared_state.catalog_status(),
protocol::WorkerStatus::Idle,
"successful manual compaction must release the execution fence"
);
let command = protocol::WorkerCommandEnvelope::for_snapshot(2, &handle.shared_state.snapshot());
let _ = handle.send(Method::Shutdown { command }).await;
}
crates/worker/tests/controller_test.rs
+152 -48
View File
@@ -1,5 +1,5 @@
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use agen::Engine;
@@ -25,6 +25,15 @@ use worker::{
type TestStore = CombinedStore<FsStore, FsWorkerStore>;
static NEXT_COMMAND_ID: AtomicU64 = AtomicU64::new(1);
fn worker_command(handle: &WorkerHandle) -> protocol::WorkerCommandEnvelope {
protocol::WorkerCommandEnvelope::for_snapshot(
NEXT_COMMAND_ID.fetch_add(1, Ordering::Relaxed),
&handle.shared_state.snapshot(),
)
}
/// Reconstruct a worker-history-like `Vec<Item>` from the live session
/// log mirror held by the Worker's broadcast sink. Replaces the previous
/// `WorkerSharedState.history()` test helper now that the mirror lives in
@@ -313,7 +322,12 @@ async fn controller_grants_read_scope_for_exact_bash_output_directory() {
}));
assert!(!handle.runtime_dir.path().join("bash-output").exists());
handle.send(Method::Shutdown).await.unwrap();
handle
.send(Method::Shutdown {
command: worker_command(&handle),
})
.await
.unwrap();
shutdown_rx.await.unwrap();
}
@@ -345,7 +359,12 @@ async fn shutdown_closes_bound_workdir_session() {
WorkerController::spawn(worker, runtime_base.path(), &bash_output_dir)
.await
.unwrap();
handle.send(Method::Shutdown).await.unwrap();
handle
.send(Method::Shutdown {
command: worker_command(&handle),
})
.await
.unwrap();
tokio::time::timeout(std::time::Duration::from_secs(5), shutdown_rx)
.await
.expect("controller should shut down")
@@ -459,7 +478,12 @@ async fn controller_projects_workdir_command_events_and_snapshot_state() {
!durable_history.contains("ready") && !durable_history.contains("done"),
"operational command chunks must not be appended to Worker history: {durable_history}"
);
handle.send(Method::Shutdown).await.unwrap();
handle
.send(Method::Shutdown {
command: worker_command(&handle),
})
.await
.unwrap();
}
#[tokio::test]
@@ -530,7 +554,12 @@ async fn controller_refreshes_command_snapshot_after_high_output_provider_lag()
.await
.unwrap();
assert_eq!(output.status, workdir::CommandStatus::Cancelled);
handle.send(Method::Shutdown).await.unwrap();
handle
.send(Method::Shutdown {
command: worker_command(&handle),
})
.await
.unwrap();
}
#[tokio::test]
@@ -571,13 +600,13 @@ async fn controller_startup_failure_closes_bound_workdir_session() {
async fn wait_for_status(handle: &WorkerHandle, status: WorkerStatus) {
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
if handle.shared_state.get_status() == status {
if handle.shared_state.catalog_status() == status {
return;
}
assert!(
tokio::time::Instant::now() < deadline,
"timed out waiting for status {status:?}; current={:?}",
handle.shared_state.get_status()
handle.shared_state.catalog_status()
);
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
@@ -1029,7 +1058,8 @@ async fn run_end_returns_to_idle_without_busy_status() {
Ok(Event::RunEnd { result: protocol::RunResult::Finished }) => {
saw_run_end = true;
}
Ok(Event::Status { status: WorkerStatus::Idle }) if saw_run_end => {
Ok(Event::WorkerState { snapshot })
if saw_run_end && snapshot.catalog_status() == WorkerStatus::Idle => {
saw_idle_status = true;
break;
}
@@ -1046,7 +1076,7 @@ async fn run_end_returns_to_idle_without_busy_status() {
saw_idle_status,
"expected idle status immediately after RunEnd"
);
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Idle);
assert_eq!(handle.shared_state.catalog_status(), WorkerStatus::Idle);
}
#[tokio::test]
@@ -1124,9 +1154,7 @@ async fn snapshot_includes_user_input_for_in_flight_turn() {
loop {
if matches!(
events.recv().await,
Ok(Event::Status {
status: WorkerStatus::Running,
})
Ok(Event::WorkerState { snapshot }) if snapshot.catalog_status() == WorkerStatus::Running
) {
break;
}
@@ -1201,8 +1229,8 @@ async fn attach_snapshot_includes_current_status() {
loop {
let event = reader.next::<Event>().await.unwrap().unwrap();
match event {
Event::Snapshot { status, .. } => {
assert_eq!(status, WorkerStatus::Running);
Event::Snapshot { state, .. } => {
assert_eq!(state.catalog_status(), WorkerStatus::Running);
return;
}
Event::Alert(_) => continue,
@@ -1217,7 +1245,7 @@ async fn shared_state_starts_idle() {
let worker = make_worker(client).await;
let handle = spawn_controller(worker).await;
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Idle);
assert_eq!(handle.shared_state.catalog_status(), WorkerStatus::Idle);
}
#[tokio::test]
@@ -1237,7 +1265,7 @@ async fn run_updates_shared_state_to_idle_after_completion() {
// Wait for the run to complete
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Idle);
assert_eq!(handle.shared_state.catalog_status(), WorkerStatus::Idle);
}
#[tokio::test]
@@ -1360,7 +1388,12 @@ async fn submit_while_running_is_durably_queued() {
assert_eq!(accepted, Some(protocol::SubmissionDisposition::Queued));
let pending_snapshot = pending_snapshot.expect("pending snapshot");
assert_eq!(pending_snapshot.submissions.len(), 1);
handle.send(Method::Pause).await.unwrap();
handle
.send(Method::Pause {
command: worker_command(&handle),
})
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Paused).await;
handle
.send(Method::ContinuePending {
@@ -1382,17 +1415,22 @@ async fn submit_while_running_is_durably_queued() {
.await
.expect("paused ContinuePending rejection");
assert!(rejection.contains("Resume or Cancel"));
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Paused);
assert_eq!(handle.shared_state.catalog_status(), WorkerStatus::Paused);
}
#[tokio::test]
async fn resume_without_pause_returns_error() {
async fn resume_without_pause_returns_invalid_state_acknowledgement() {
let client = MockClient::new(simple_text_events());
let worker = make_worker(client).await;
let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe();
handle.send(Method::Resume).await.unwrap();
handle
.send(Method::Resume {
command: worker_command(&handle),
})
.await
.unwrap();
let mut saw_not_paused = false;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(1);
@@ -1400,7 +1438,10 @@ async fn resume_without_pause_returns_error() {
tokio::select! {
event = rx.recv() => {
match event {
Ok(Event::Error { code, .. }) if code == worker::ErrorCode::NotPaused => {
Ok(Event::CommandAcknowledged { acknowledgement })
if acknowledgement.command == protocol::WorkerCommandKind::Resume
&& acknowledgement.disposition
== protocol::WorkerCommandDisposition::InvalidState => {
saw_not_paused = true;
break;
}
@@ -1412,17 +1453,22 @@ async fn resume_without_pause_returns_error() {
}
}
assert!(saw_not_paused, "should see not_paused error");
assert!(saw_not_paused, "should see invalid-state acknowledgement");
}
#[tokio::test]
async fn cancel_without_run_returns_error() {
async fn cancel_without_run_returns_invalid_state_acknowledgement() {
let client = MockClient::new(simple_text_events());
let worker = make_worker(client).await;
let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe();
handle.send(Method::Cancel).await.unwrap();
handle
.send(Method::Cancel {
command: worker_command(&handle),
})
.await
.unwrap();
let mut saw_not_running = false;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(1);
@@ -1430,7 +1476,10 @@ async fn cancel_without_run_returns_error() {
tokio::select! {
event = rx.recv() => {
match event {
Ok(Event::Error { code, .. }) if code == worker::ErrorCode::NotRunning => {
Ok(Event::CommandAcknowledged { acknowledgement })
if acknowledgement.command == protocol::WorkerCommandKind::Cancel
&& acknowledgement.disposition
== protocol::WorkerCommandDisposition::InvalidState => {
saw_not_running = true;
break;
}
@@ -1442,7 +1491,7 @@ async fn cancel_without_run_returns_error() {
}
}
assert!(saw_not_running, "should see not_running error");
assert!(saw_not_running, "should see invalid-state acknowledgement");
}
#[tokio::test]
@@ -1818,7 +1867,7 @@ async fn notify_while_idle_with_auto_run_false_waits_for_explicit_run() {
}
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Idle);
assert_eq!(handle.shared_state.catalog_status(), WorkerStatus::Idle);
assert!(
client_for_assert.captured_requests().is_empty(),
"weak Notify must not stage RunForNotification while idle"
@@ -1915,7 +1964,7 @@ async fn worker_event_turn_ended_while_idle_auto_starts_turn_and_injects_system_
saw_worker_event_in_mirror,
"Method::WorkerEvent should commit a SystemItem::WorkerEvent entry"
);
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Idle);
assert_eq!(handle.shared_state.catalog_status(), WorkerStatus::Idle);
let requests = client_for_assert.captured_requests();
assert_eq!(
@@ -1978,7 +2027,7 @@ async fn worker_event_scope_sub_delegated_while_idle_stays_control_plane_only()
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert_eq!(
handle.shared_state.get_status(),
handle.shared_state.catalog_status(),
WorkerStatus::Idle,
"control-plane ScopeSubDelegated must not auto-start the parent LLM"
);
@@ -2081,7 +2130,12 @@ async fn weak_notify_while_running_is_deduped_and_survives_until_next_submit() {
.await
.unwrap();
}
handle.send(Method::Cancel).await.unwrap();
handle
.send(Method::Cancel {
command: worker_command(&handle),
})
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Idle).await;
let mut rx = handle.subscribe();
@@ -2478,7 +2532,12 @@ async fn pause_then_resume_transitions_and_preserves_history_consistency() {
"text_delta should arrive before pause"
);
handle.send(Method::Pause).await.unwrap();
handle
.send(Method::Pause {
command: worker_command(&handle),
})
.await
.unwrap();
// The controller emits RunEnd { Paused } when the
// EngineError::Cancelled is translated under pause_requested.
@@ -2494,9 +2553,14 @@ async fn pause_then_resume_transitions_and_preserves_history_consistency() {
);
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Paused);
assert_eq!(handle.shared_state.catalog_status(), WorkerStatus::Paused);
handle.send(Method::Resume).await.unwrap();
handle
.send(Method::Resume {
command: worker_command(&handle),
})
.await
.unwrap();
assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
@@ -2510,7 +2574,7 @@ async fn pause_then_resume_transitions_and_preserves_history_consistency() {
);
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Idle);
assert_eq!(handle.shared_state.catalog_status(), WorkerStatus::Idle);
// History consistency: exactly [user "hello", assistant
// "resumed output"]. No artifacts from the aborted stream
@@ -2610,7 +2674,12 @@ async fn paused_then_run_closes_orphan_tool_use_for_next_request() {
"tool_call_done should arrive before pause"
);
handle.send(Method::Pause).await.unwrap();
handle
.send(Method::Pause {
command: worker_command(&handle),
})
.await
.unwrap();
assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
e,
@@ -2622,7 +2691,7 @@ async fn paused_then_run_closes_orphan_tool_use_for_next_request() {
"expected RunEnd::Paused"
);
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Paused);
assert_eq!(handle.shared_state.catalog_status(), WorkerStatus::Paused);
// New user input while Paused → `Worker::run` observes
// `last_run_interrupted` and runs its interrupt-prep step, which
@@ -2781,7 +2850,12 @@ async fn paused_cancel_abandons_resume_and_next_input_is_fresh_run() {
"tool_call_done should arrive before pause"
);
handle.send(Method::Pause).await.unwrap();
handle
.send(Method::Pause {
command: worker_command(&handle),
})
.await
.unwrap();
assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
e,
@@ -2794,7 +2868,12 @@ async fn paused_cancel_abandons_resume_and_next_input_is_fresh_run() {
);
wait_for_status(&handle, WorkerStatus::Paused).await;
handle.send(Method::Cancel).await.unwrap();
handle
.send(Method::Cancel {
command: worker_command(&handle),
})
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Idle).await;
let (entries_after_cancel, _rx_after_cancel) = handle.sink.subscribe_with_snapshot();
assert!(
@@ -2820,17 +2899,22 @@ async fn paused_cancel_abandons_resume_and_next_input_is_fresh_run() {
"paused cancel must not resume or start another LLM request"
);
handle.send(Method::Resume).await.unwrap();
handle
.send(Method::Resume {
command: worker_command(&handle),
})
.await
.unwrap();
assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
e,
Event::Error {
code: worker::ErrorCode::NotPaused,
..
}
Event::CommandAcknowledged { acknowledgement }
if acknowledgement.command == protocol::WorkerCommandKind::Resume
&& acknowledgement.disposition
== protocol::WorkerCommandDisposition::InvalidState
))
.await,
"resume after paused cancel should be rejected as not paused"
"resume after paused cancel should receive invalid-state acknowledgement"
);
assert_eq!(
client_for_assert.captured_requests().len(),
@@ -2939,7 +3023,12 @@ async fn empty_turn_cancel_rolls_back_submit_entries_and_emits_signal() {
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Running).await;
handle.send(Method::Cancel).await.unwrap();
handle
.send(Method::Cancel {
command: worker_command(&handle),
})
.await
.unwrap();
assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
@@ -2977,7 +3066,12 @@ async fn empty_turn_pause_rolls_back_and_snapshot_does_not_restore_input() {
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Running).await;
handle.send(Method::Pause).await.unwrap();
handle
.send(Method::Pause {
command: worker_command(&handle),
})
.await
.unwrap();
assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
@@ -3034,7 +3128,12 @@ async fn empty_turn_rollback_removes_only_the_most_recent_turn() {
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Running).await;
handle.send(Method::Cancel).await.unwrap();
handle
.send(Method::Cancel {
command: worker_command(&handle),
})
.await
.unwrap();
assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
e,
@@ -3091,7 +3190,12 @@ async fn pause_after_assistant_token_does_not_rollback() {
.await,
"assistant token should be visible before pause"
);
handle.send(Method::Pause).await.unwrap();
handle
.send(Method::Pause {
command: worker_command(&handle),
})
.await
.unwrap();
assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(