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/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!(