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/src/controller.rs
+557 -90
View File
@@ -1,3 +1,4 @@
use std::collections::VecDeque;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::Ordering;
@@ -28,7 +29,9 @@ use protocol::{
AlertLevel, AlertSource, CommandEvent as ProtocolCommandEvent,
CommandSnapshot as ProtocolCommandSnapshot, CommandStatus as ProtocolCommandStatus,
CommandStream as ProtocolCommandStream, CommandStreamSlice as ProtocolCommandStreamSlice,
ErrorCode, Event, Method, RewindTargetId, RunResult, TurnResult, UploadedFileRef, WorkerStatus,
ErrorCode, Event, Method, RewindTargetId, RunResult, TurnResult, UploadedFileRef,
WorkerBusyState, WorkerCommandAcknowledgement, WorkerCommandDisposition, WorkerCommandEnvelope,
WorkerCommandKind, WorkerMaintenanceState, WorkerRunState, WorkerState, WorkerStatus,
};
use workdir::{
CommandEvent as WorkdirCommandEvent, CommandSnapshot as WorkdirCommandSnapshot,
@@ -138,7 +141,7 @@ impl WorkerHandle {
let event = Event::Snapshot {
session,
greeting: self.shared_state.greeting.clone(),
status: self.shared_state.get_status(),
state: self.shared_state.snapshot(),
in_flight,
internal_workers: self.spawned_registry.internal_worker_snapshots(),
};
@@ -178,15 +181,81 @@ impl WorkerHandle {
}
}
fn validate_command(
envelope: WorkerCommandEnvelope,
shared_state: &WorkerSharedState,
) -> Result<(), WorkerCommandDisposition> {
let snapshot = shared_state.snapshot();
if envelope.expected_execution_generation != snapshot.execution_generation {
return Err(WorkerCommandDisposition::StaleExecutionGeneration);
}
if envelope.expected_worker_state_revision != snapshot.revision {
return Err(WorkerCommandDisposition::StaleWorkerStateRevision);
}
if !shared_state.accept_command_id(envelope.command_id) {
return Err(WorkerCommandDisposition::StaleCommandId);
}
Ok(())
}
fn acknowledge_command(
working_event_tx: &broadcast::Sender<Event>,
shared_state: &WorkerSharedState,
command_id: u64,
command: WorkerCommandKind,
disposition: WorkerCommandDisposition,
) {
let _ = working_event_tx.send(Event::CommandAcknowledged {
acknowledgement: WorkerCommandAcknowledgement {
command_id,
command,
disposition,
state: shared_state.snapshot(),
},
});
}
fn reject_invalid_command_state(
working_event_tx: &broadcast::Sender<Event>,
shared_state: &WorkerSharedState,
envelope: WorkerCommandEnvelope,
command: WorkerCommandKind,
) {
acknowledge_command(
working_event_tx,
shared_state,
envelope.command_id,
command,
WorkerCommandDisposition::InvalidState,
);
}
async fn set_controller_state(
shared_state: &Arc<WorkerSharedState>,
runtime_dir: &RuntimeDir,
working_event_tx: &broadcast::Sender<Event>,
state: WorkerState,
) -> protocol::WorkerStateSnapshot {
let snapshot = shared_state.transition(state);
let _ = runtime_dir.write_status(shared_state).await;
let _ = working_event_tx.send(Event::WorkerState {
snapshot: snapshot.clone(),
});
snapshot
}
async fn set_controller_status(
shared_state: &Arc<WorkerSharedState>,
runtime_dir: &RuntimeDir,
working_event_tx: &broadcast::Sender<Event>,
status: WorkerStatus,
) {
shared_state.set_status(status);
let _ = runtime_dir.write_status(shared_state).await;
let _ = working_event_tx.send(Event::Status { status });
let state = match status {
WorkerStatus::Idle | WorkerStatus::Stopped => WorkerState::Idle,
WorkerStatus::Running => WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running)),
WorkerStatus::Paused => WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Paused)),
};
set_controller_state(shared_state, runtime_dir, working_event_tx, state).await;
}
async fn finish_controller_run<C, St>(
@@ -659,12 +728,24 @@ impl WorkerController {
// === 4. Initial runtime files + WorkerSharedState + WorkerHandle +
// SocketServer ===
let manifest_toml = toml::to_string_pretty(worker.manifest()).unwrap_or_default();
worker
.recover_unfinished_compaction()
.await
.map_err(|error| std::io::Error::other(error.to_string()))?;
let greeting = build_greeting(&worker);
let shared_state = Arc::new(WorkerSharedState::new(
let execution_generation = runtime_dir
.path()
.file_name()
.and_then(|name| name.to_str())
.and_then(|name| name.parse::<u64>().ok())
.filter(|generation| *generation > 0)
.unwrap_or(1);
let shared_state = Arc::new(WorkerSharedState::new_with_generation(
worker.manifest().worker.name.clone(),
worker.segment_id(),
manifest_toml.clone(),
greeting,
execution_generation,
));
if let Some(fs_for_view) = fs_for_view {
shared_state.set_fs_view(crate::fs_view::WorkerFsView::new(fs_for_view));
@@ -1432,8 +1513,9 @@ async fn controller_loop<C, St>(
}
};
loop {
// Top-of-iteration: if an event handler staged a run, fire it
let mut deferred_methods = VecDeque::new();
'controller: loop {
// here so the status flip → drive_turn → finish sequence lives
// in one place, regardless of which Method caused it.
if let Some(run) = pending.take() {
@@ -1584,9 +1666,13 @@ async fn controller_loop<C, St>(
continue;
}
let method = match method_rx.recv().await {
Some(m) => m,
None => break,
let method = if let Some(method) = deferred_methods.pop_front() {
method
} else {
match method_rx.recv().await {
Some(method) => method,
None => break,
}
};
match method {
@@ -1784,7 +1870,7 @@ async fn controller_loop<C, St>(
expected_revision,
expected_head_id,
} => {
if shared_state.get_status() != WorkerStatus::Idle {
if shared_state.catalog_status() != WorkerStatus::Idle {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: "ContinuePending requires an idle Worker; Resume or Cancel a paused run first".into(),
@@ -1811,88 +1897,243 @@ async fn controller_loop<C, St>(
}
}
}
Method::Resume => {
if shared_state.get_status() != WorkerStatus::Paused {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::NotPaused,
message: "Worker is not paused".into(),
});
Method::Resume { command } => {
if let Err(disposition) = validate_command(command, &shared_state) {
acknowledge_command(
&working_event_tx,
&shared_state,
command.command_id,
WorkerCommandKind::Resume,
disposition,
);
continue;
}
if !matches!(
shared_state.snapshot().state,
WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Paused))
) {
reject_invalid_command_state(
&working_event_tx,
&shared_state,
command,
WorkerCommandKind::Resume,
);
continue;
}
set_controller_state(
&shared_state,
&runtime_dir,
&working_event_tx,
WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running)),
)
.await;
acknowledge_command(
&working_event_tx,
&shared_state,
command.command_id,
WorkerCommandKind::Resume,
WorkerCommandDisposition::Accepted,
);
pending = Some(PendingRun::Resume);
}
Method::Cancel => match shared_state.get_status() {
WorkerStatus::Paused => match worker.cancel_paused_turn() {
Method::Cancel { command } => {
if let Err(disposition) = validate_command(command, &shared_state) {
acknowledge_command(
&working_event_tx,
&shared_state,
command.command_id,
WorkerCommandKind::Cancel,
disposition,
);
continue;
}
if !matches!(
shared_state.snapshot().state,
WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Paused))
) {
reject_invalid_command_state(
&working_event_tx,
&shared_state,
command,
WorkerCommandKind::Cancel,
);
continue;
}
set_controller_state(
&shared_state,
&runtime_dir,
&working_event_tx,
WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Cancelling)),
)
.await;
acknowledge_command(
&working_event_tx,
&shared_state,
command.command_id,
WorkerCommandKind::Cancel,
WorkerCommandDisposition::Accepted,
);
match worker.cancel_paused_turn() {
Ok(()) => {
worker.clear_in_flight_events();
set_controller_status(
set_controller_state(
&shared_state,
&runtime_dir,
&working_event_tx,
WorkerStatus::Idle,
WorkerState::Idle,
)
.await;
}
Err(error) => {
set_controller_state(
&shared_state,
&runtime_dir,
&working_event_tx,
WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Paused)),
)
.await;
let _ = working_event_tx.send(Event::Error {
code: worker_error_code(&error),
message: error.to_string(),
});
}
},
WorkerStatus::Idle | WorkerStatus::Stopped => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::NotRunning,
message: "Worker is not running".into(),
});
}
WorkerStatus::Running => {
// Running turns receive Cancel through drive_turn; this is
// only reachable across a defensive race window.
let _ = cancel_tx.try_send(());
}
},
Method::Pause => {
// Already paused → idempotent no-op. Otherwise the
// Worker is Idle (Running turns go through `drive_turn`,
// not this outer match), so there is nothing to pause.
if shared_state.get_status() != WorkerStatus::Paused {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::NotRunning,
message: "Worker is not running".into(),
});
}
}
Method::Compact => match shared_state.get_status() {
WorkerStatus::Idle => {
if let Err(error) = worker.manual_compact().await {
let _ = working_event_tx.send(Event::Error {
code: worker_error_code(&error),
message: error.to_string(),
});
}
Method::Pause { command } => {
if let Err(disposition) = validate_command(command, &shared_state) {
acknowledge_command(
&working_event_tx,
&shared_state,
command.command_id,
WorkerCommandKind::Pause,
disposition,
);
} else {
reject_invalid_command_state(
&working_event_tx,
&shared_state,
command,
WorkerCommandKind::Pause,
);
}
WorkerStatus::Paused => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: "Cannot compact while the Worker is paused; resume or start a fresh turn first"
.into(),
});
}
WorkerStatus::Running | WorkerStatus::Stopped => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message:
"Worker is already executing a turn; compact can only run while idle"
.into(),
});
}
},
}
Method::ListRewindTargets => match shared_state.get_status() {
Method::Compact { command } => {
if let Err(disposition) = validate_command(command, &shared_state) {
acknowledge_command(
&working_event_tx,
&shared_state,
command.command_id,
WorkerCommandKind::Compact,
disposition,
);
continue;
}
if !matches!(shared_state.snapshot().state, WorkerState::Idle) {
reject_invalid_command_state(
&working_event_tx,
&shared_state,
command,
WorkerCommandKind::Compact,
);
continue;
}
set_controller_state(
&shared_state,
&runtime_dir,
&working_event_tx,
WorkerState::Busy(WorkerBusyState::Maintenance(
WorkerMaintenanceState::Compacting,
)),
)
.await;
acknowledge_command(
&working_event_tx,
&shared_state,
command.command_id,
WorkerCommandKind::Compact,
WorkerCommandDisposition::Accepted,
);
let (cancel_tx, cancel_rx) = tokio::sync::watch::channel(false);
let mut shutdown_after_compaction = false;
let result = {
let mut compact = Box::pin(worker.manual_compact_with_cancel(cancel_rx));
loop {
tokio::select! {
result = &mut compact => break result,
method = method_rx.recv() => {
match method {
Some(Method::Cancel { command }) => {
if let Err(disposition) = validate_command(command, &shared_state) {
acknowledge_command(
&working_event_tx,
&shared_state,
command.command_id,
WorkerCommandKind::Cancel,
disposition,
);
continue;
}
acknowledge_command(
&working_event_tx,
&shared_state,
command.command_id,
WorkerCommandKind::Cancel,
WorkerCommandDisposition::Accepted,
);
let _ = cancel_tx.send(true);
}
Some(Method::Shutdown { command }) => {
shared_state.accept_command_id(command.command_id);
shutdown_after_compaction = true;
acknowledge_command(
&working_event_tx,
&shared_state,
command.command_id,
WorkerCommandKind::Shutdown,
WorkerCommandDisposition::Accepted,
);
let _ = cancel_tx.send(true);
}
Some(method) => deferred_methods.push_back(method),
None => {
shutdown_after_compaction = true;
let _ = cancel_tx.send(true);
}
}
}
}
}
};
if !matches!(
result,
Err(WorkerError::Store(_))
| Err(WorkerError::WorkerStore(_))
| Err(WorkerError::InvalidState(_))
) {
set_controller_state(
&shared_state,
&runtime_dir,
&working_event_tx,
WorkerState::Idle,
)
.await;
}
if let Err(error) = result {
let _ = working_event_tx.send(Event::Error {
code: worker_error_code(&error),
message: error.to_string(),
});
}
if shutdown_after_compaction {
let _ = working_event_tx.send(Event::Shutdown);
break 'controller;
}
}
Method::ListRewindTargets => match shared_state.catalog_status() {
WorkerStatus::Idle | WorkerStatus::Paused => {
emit_rewind_targets(&worker, &working_event_tx)
}
@@ -1908,7 +2149,7 @@ async fn controller_loop<C, St>(
Method::RewindTo {
target,
expected_head_entries,
} => match shared_state.get_status() {
} => match shared_state.catalog_status() {
WorkerStatus::Idle => {
if apply_rewind(
&mut worker,
@@ -1919,10 +2160,8 @@ async fn controller_loop<C, St>(
.await
{
worker.clear_in_flight_events();
shared_state.set_status(WorkerStatus::Idle);
let _ = working_event_tx.send(Event::Status {
status: WorkerStatus::Idle,
});
let snapshot = shared_state.transition(WorkerState::Idle);
let _ = working_event_tx.send(Event::WorkerState { snapshot });
}
}
WorkerStatus::Paused => {
@@ -1941,7 +2180,17 @@ async fn controller_loop<C, St>(
}
},
Method::Shutdown => {
Method::Shutdown { command } => {
// Shutdown remains unconditional/retryable even when the caller's
// live-state fence is stale.
shared_state.accept_command_id(command.command_id);
acknowledge_command(
&working_event_tx,
&shared_state,
command.command_id,
WorkerCommandKind::Shutdown,
WorkerCommandDisposition::Accepted,
);
let _ = working_event_tx.send(Event::Shutdown);
break;
}
@@ -2023,7 +2272,7 @@ async fn controller_loop<C, St>(
// Auto-kick a turn if the Worker is idle so the
// notification is not stranded. Matches the
// `Method::Notify` idle path.
if shared_state.get_status() == WorkerStatus::Idle {
if shared_state.catalog_status() == WorkerStatus::Idle {
pending = Some(PendingRun::RunForNotification {
invoke_kind: protocol::InvokeKind::WorkerEvent,
notification_request_id: None,
@@ -2270,15 +2519,102 @@ where
}
method = method_rx.recv(), if input_commit.is_none() => {
match method {
Some(Method::Cancel) => {
Some(Method::Cancel { command }) => {
if let Err(disposition) = validate_command(command, shared_state) {
acknowledge_command(
working_event_tx,
shared_state,
command.command_id,
WorkerCommandKind::Cancel,
disposition,
);
continue;
}
if !matches!(
shared_state.snapshot().state,
WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running))
) {
reject_invalid_command_state(
working_event_tx,
shared_state,
command,
WorkerCommandKind::Cancel,
);
continue;
}
set_controller_state(
shared_state,
runtime_dir,
working_event_tx,
WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Cancelling)),
)
.await;
acknowledge_command(
working_event_tx,
shared_state,
command.command_id,
WorkerCommandKind::Cancel,
WorkerCommandDisposition::Accepted,
);
let _ = cancel_tx.try_send(());
}
Some(Method::Pause) => {
Some(Method::Pause { command }) => {
if let Err(disposition) = validate_command(command, shared_state) {
acknowledge_command(
working_event_tx,
shared_state,
command.command_id,
WorkerCommandKind::Pause,
disposition,
);
continue;
}
if !matches!(
shared_state.snapshot().state,
WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running))
) {
reject_invalid_command_state(
working_event_tx,
shared_state,
command,
WorkerCommandKind::Pause,
);
continue;
}
pause_requested = true;
set_controller_state(
shared_state,
runtime_dir,
working_event_tx,
WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Pausing)),
)
.await;
acknowledge_command(
working_event_tx,
shared_state,
command.command_id,
WorkerCommandKind::Pause,
WorkerCommandDisposition::Accepted,
);
let _ = pause_tx.try_send(());
}
Some(Method::Shutdown) => {
Some(Method::Shutdown { command }) => {
shared_state.accept_command_id(command.command_id);
shutdown_requested = true;
set_controller_state(
shared_state,
runtime_dir,
working_event_tx,
WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Cancelling)),
)
.await;
acknowledge_command(
working_event_tx,
shared_state,
command.command_id,
WorkerCommandKind::Shutdown,
WorkerCommandDisposition::Accepted,
);
let _ = cancel_tx.try_send(());
}
Some(Method::Submit {
@@ -2344,7 +2680,25 @@ where
}
}
}
Some(Method::Resume | Method::ContinuePending { .. }) => {
Some(Method::Resume { command }) => {
if let Err(disposition) = validate_command(command, shared_state) {
acknowledge_command(
working_event_tx,
shared_state,
command.command_id,
WorkerCommandKind::Resume,
disposition,
);
} else {
reject_invalid_command_state(
working_event_tx,
shared_state,
command,
WorkerCommandKind::Resume,
);
}
}
Some(Method::ContinuePending { .. }) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn".into(),
@@ -2384,7 +2738,25 @@ where
}
}
}
Some(Method::Compact | Method::ListRewindTargets | Method::RewindTo { .. }) => {
Some(Method::Compact { command }) => {
if let Err(disposition) = validate_command(command, shared_state) {
acknowledge_command(
working_event_tx,
shared_state,
command.command_id,
WorkerCommandKind::Compact,
disposition,
);
} else {
reject_invalid_command_state(
working_event_tx,
shared_state,
command,
WorkerCommandKind::Compact,
);
}
}
Some(Method::ListRewindTargets | Method::RewindTo { .. }) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn; rewind/compact can only run while idle or paused"
@@ -2487,7 +2859,7 @@ where
}
None => {
let _ = cancel_tx.try_send(());
shared_state.set_status(WorkerStatus::Idle);
shared_state.transition(WorkerState::Idle);
return (WorkerStatus::Idle, false, false);
}
}
@@ -2863,7 +3235,7 @@ mod tests {
context_window: 200_000,
context_tokens: 0,
},
status: WorkerStatus::Idle,
state: WorkerStatus::Idle.into(),
in_flight: Default::default(),
internal_workers: Vec::new(),
})
@@ -2919,9 +3291,17 @@ mod tests {
async fn pause_waits_for_run_boundary_and_uses_safe_pause_channel() {
let mut env = make_env().await;
let method_tx = env._method_tx.clone();
env.shared_state
.transition(WorkerState::Busy(WorkerBusyState::Run(
WorkerRunState::Running,
)));
let command = WorkerCommandEnvelope::for_snapshot(1, &env.shared_state.snapshot());
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(10)).await;
method_tx.send(Method::Pause).await.expect("send pause");
method_tx
.send(Method::Pause { command })
.await
.expect("send pause");
});
let worker_future = async {
@@ -3194,8 +3574,13 @@ mod tests {
async fn compact_method_is_rejected_while_running() {
let mut env = make_env().await;
let mut events = env.working_event_tx.subscribe();
env.shared_state
.transition(WorkerState::Busy(WorkerBusyState::Run(
WorkerRunState::Running,
)));
let command = WorkerCommandEnvelope::for_snapshot(1, &env.shared_state.snapshot());
env._method_tx
.send(Method::Compact)
.send(Method::Compact { command })
.await
.expect("send compact");
@@ -3228,11 +3613,93 @@ mod tests {
.expect("event timeout")
.expect("event");
match event {
Event::Error { code, message } => {
assert_eq!(code, ErrorCode::AlreadyRunning);
assert!(message.contains("compact"), "got message: {message}");
Event::CommandAcknowledged { acknowledgement } => {
assert_eq!(acknowledgement.command, WorkerCommandKind::Compact);
assert_eq!(
acknowledgement.disposition,
WorkerCommandDisposition::InvalidState
);
assert!(matches!(
acknowledgement.state.state,
WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running))
));
}
other => panic!("expected compact rejection error, got {other:?}"),
other => panic!("expected compact rejection acknowledgement, got {other:?}"),
}
}
#[test]
fn command_admission_rejects_stale_generation_revision_and_order() {
let shared = WorkerSharedState::new_with_generation(
"worker".into(),
session_store::new_segment_id(),
String::new(),
protocol::Greeting {
worker_name: "worker".into(),
cwd: "/tmp".into(),
provider: "test".into(),
model: "test".into(),
scope_summary: String::new(),
tools: Vec::new(),
context_window: 1,
context_tokens: 0,
},
9,
);
assert_eq!(
validate_command(
WorkerCommandEnvelope {
command_id: 1,
expected_execution_generation: 8,
expected_worker_state_revision: 0,
},
&shared,
),
Err(WorkerCommandDisposition::StaleExecutionGeneration)
);
assert_eq!(
validate_command(
WorkerCommandEnvelope {
command_id: 2,
expected_execution_generation: 9,
expected_worker_state_revision: 1,
},
&shared,
),
Err(WorkerCommandDisposition::StaleWorkerStateRevision)
);
assert!(
validate_command(
WorkerCommandEnvelope {
command_id: 1,
expected_execution_generation: 9,
expected_worker_state_revision: 0,
},
&shared,
)
.is_ok()
);
assert_eq!(
validate_command(
WorkerCommandEnvelope {
command_id: 1,
expected_execution_generation: 9,
expected_worker_state_revision: 0,
},
&shared,
),
Err(WorkerCommandDisposition::StaleCommandId)
);
assert!(
validate_command(
WorkerCommandEnvelope {
command_id: 2,
expected_execution_generation: 9,
expected_worker_state_revision: 0,
},
&shared,
)
.is_ok()
);
}
}
crates/worker/src/discovery.rs
+8 -8
View File
@@ -779,10 +779,10 @@ async fn probe_socket(socket_path: &Path) -> LiveInfo {
loop {
match tokio::time::timeout(PROBE_TIMEOUT, reader.next::<Event>()).await {
Ok(Ok(Some(Event::Snapshot {
status: snapshot_status,
state: snapshot_state,
..
}))) => {
status = Some(snapshot_status);
status = Some(snapshot_state.catalog_status());
break;
}
Ok(Ok(Some(Event::Alert(_)))) => continue,
@@ -1507,7 +1507,7 @@ mod tests {
context_window: 0,
context_tokens: 0,
},
status: WorkerStatus::Idle,
state: WorkerStatus::Idle.into(),
in_flight: Default::default(),
internal_workers: Vec::new(),
})
@@ -1543,7 +1543,7 @@ mod tests {
context_window: 0,
context_tokens: 0,
},
status: WorkerStatus::Idle,
state: WorkerStatus::Idle.into(),
in_flight: Default::default(),
internal_workers: Vec::new(),
})
@@ -1638,7 +1638,7 @@ mod tests {
context_window: 0,
context_tokens: 0,
},
status: WorkerStatus::Idle,
state: WorkerStatus::Idle.into(),
in_flight: Default::default(),
internal_workers: Vec::new(),
})
@@ -1665,7 +1665,7 @@ mod tests {
context_window: 0,
context_tokens: 0,
},
status: WorkerStatus::Idle,
state: WorkerStatus::Idle.into(),
in_flight: Default::default(),
internal_workers: Vec::new(),
})
@@ -1773,7 +1773,7 @@ mod tests {
context_window: 0,
context_tokens: 0,
},
status: WorkerStatus::Paused,
state: WorkerStatus::Paused.into(),
in_flight: Default::default(),
internal_workers: Vec::new(),
})
@@ -1827,7 +1827,7 @@ mod tests {
context_window: 0,
context_tokens: 0,
},
status: WorkerStatus::Idle,
state: WorkerStatus::Idle.into(),
in_flight: Default::default(),
internal_workers: Vec::new(),
})
crates/worker/src/internal_worker.rs
+51 -19
View File
@@ -295,6 +295,38 @@ impl InternalWorkerSessionStatus {
}
}
fn send_internal_worker_state(
event_tx: &broadcast::Sender<Event>,
state_revision: &std::sync::atomic::AtomicU64,
status: InternalWorkerSessionStatus,
) {
let state = match status {
InternalWorkerSessionStatus::Idle
| InternalWorkerSessionStatus::Stopped
| InternalWorkerSessionStatus::Failed => protocol::WorkerState::Idle,
InternalWorkerSessionStatus::Paused => protocol::WorkerState::Busy(
protocol::WorkerBusyState::Run(protocol::WorkerRunState::Paused),
),
InternalWorkerSessionStatus::Running => protocol::WorkerState::Busy(
protocol::WorkerBusyState::Run(protocol::WorkerRunState::Running),
),
InternalWorkerSessionStatus::Stopping => protocol::WorkerState::Busy(
protocol::WorkerBusyState::Run(protocol::WorkerRunState::Cancelling),
),
};
let revision = state_revision
.fetch_add(1, std::sync::atomic::Ordering::AcqRel)
.saturating_add(1);
let _ = event_tx.send(Event::WorkerState {
snapshot: protocol::WorkerStateSnapshot {
execution_generation: 1,
revision,
last_command_id: 0,
state,
},
});
}
fn classify_internal_turn_result(
result: Result<WorkerRunResult, WorkerError>,
) -> (InternalWorkerSessionStatus, Option<String>) {
@@ -351,6 +383,7 @@ pub(crate) struct InternalWorkerSessionSnapshot {
pub(crate) struct InternalWorkerSessionHandle {
command_tx: tokio::sync::mpsc::Sender<InternalWorkerSessionCommand>,
status: Arc<std::sync::atomic::AtomicU8>,
state_revision: Arc<std::sync::atomic::AtomicU64>,
store: EphemeralSessionStore,
session_id: SessionId,
segment_id: SegmentId,
@@ -400,6 +433,10 @@ impl InternalWorkerSessionHandle {
self.in_flight.text_delta(block_id, text.to_owned());
}
fn emit_worker_state(&self, status: InternalWorkerSessionStatus) {
send_internal_worker_state(&self.event_tx, &self.state_revision, status);
}
pub(crate) fn protocol_snapshot(&self) -> InternalWorkerSessionSnapshot {
let (entries, in_flight) = {
let guard = self.in_flight.snapshot_guard();
@@ -473,9 +510,7 @@ impl InternalWorkerSessionHandle {
});
return Err(InternalWorkerSessionError::Unavailable);
}
let _ = self.event_tx.send(Event::Status {
status: WorkerStatus::Running,
});
self.emit_worker_state(InternalWorkerSessionStatus::Running);
Ok(())
}
@@ -767,11 +802,13 @@ pub(crate) async fn prepare_internal_worker_session(
let status = Arc::new(std::sync::atomic::AtomicU8::new(
InternalWorkerSessionStatus::Idle.encode(),
));
let state_revision = Arc::new(std::sync::atomic::AtomicU64::new(0));
let state_changed = Arc::new(tokio::sync::Notify::new());
let last_error = Arc::new(Mutex::new(None));
let handle = InternalWorkerSessionHandle {
command_tx,
status: status.clone(),
state_revision: state_revision.clone(),
store,
session_id,
segment_id,
@@ -807,19 +844,11 @@ pub(crate) async fn prepare_internal_worker_session(
message,
});
}
let protocol_status = match turn_status {
InternalWorkerSessionStatus::Idle => WorkerStatus::Idle,
InternalWorkerSessionStatus::Paused => WorkerStatus::Paused,
InternalWorkerSessionStatus::Stopped
| InternalWorkerSessionStatus::Failed => WorkerStatus::Stopped,
InternalWorkerSessionStatus::Running
| InternalWorkerSessionStatus::Stopping => {
unreachable!("run completion cannot remain active")
}
};
let _ = event_tx.send(Event::Status {
status: protocol_status,
});
send_internal_worker_state(
&event_tx,
&state_revision,
turn_status,
);
if let Some(callback) = &on_turn_end {
callback(turn_status);
}
@@ -861,9 +890,11 @@ pub(crate) async fn prepare_internal_worker_session(
InternalWorkerSessionStatus::Stopped.encode(),
std::sync::atomic::Ordering::Release,
);
let _ = event_tx.send(Event::Status {
status: WorkerStatus::Stopped,
});
send_internal_worker_state(
&event_tx,
&state_revision,
InternalWorkerSessionStatus::Stopped,
);
let _ = event_tx.send(Event::Shutdown);
state_changed.notify_waiters();
if let Some(done) = stop_done {
@@ -1114,6 +1145,7 @@ pub(crate) fn test_internal_worker_session(
status: Arc::new(std::sync::atomic::AtomicU8::new(
InternalWorkerSessionStatus::Idle.encode(),
)),
state_revision: Arc::new(std::sync::atomic::AtomicU64::new(0)),
store,
session_id,
segment_id,
crates/worker/src/runtime/dir.rs
+3 -2
View File
@@ -197,7 +197,6 @@ pub fn default_base() -> Result<PathBuf, io::Error> {
mod tests {
use super::*;
use crate::shared_state::WorkerSharedState;
use protocol::WorkerStatus;
fn test_state() -> WorkerSharedState {
WorkerSharedState::new(
@@ -247,7 +246,9 @@ mod tests {
let rt = RuntimeDir::create(tmp.path(), "my-worker").await.unwrap();
let state = test_state();
state.set_status(WorkerStatus::Running);
state.transition(protocol::WorkerState::Busy(protocol::WorkerBusyState::Run(
protocol::WorkerRunState::Running,
)));
rt.write_status(&state).await.unwrap();
let content = std::fs::read_to_string(rt.path().join("status.json")).unwrap();
crates/worker/src/shared_state.rs
+97 -44
View File
@@ -1,7 +1,12 @@
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{OnceLock, RwLock};
use std::sync::{
OnceLock, RwLock,
atomic::{AtomicBool, AtomicU64, Ordering},
};
use protocol::WorkerStatus;
use protocol::{
WorkerBusyState, WorkerMaintenanceState, WorkerRunState, WorkerState, WorkerStateSnapshot,
WorkerStatus,
};
use serde_json::json;
use session_store::SegmentId;
@@ -9,20 +14,16 @@ use crate::fs_view::WorkerFsView;
/// Shared state between WorkerController and runtime directory.
///
/// Controller updates this in-memory; RuntimeDir writes the status
/// snapshot to disk. Wrapped in `Arc` for sharing.
///
/// History and typed user-segment mirrors used to live here so the
/// IPC layer could answer `Method::GetHistory`. Those reads now go
/// directly through the session-log sink (`Event::Snapshot` +
/// live events), so this struct holds only status, identity,
/// greeting, and filesystem completion lookup hubs.
/// `WorkerStateSnapshot` is the sole live execution-state authority. Runtime
/// catalog status remains a separate lifecycle projection because `Stopped`
/// describes the execution handle rather than a live controller state.
pub struct WorkerSharedState {
pub worker_name: String,
pub segment_id: SegmentId,
pub manifest_toml: String,
pub greeting: protocol::Greeting,
pub status: RwLock<WorkerStatus>,
state: RwLock<WorkerStateSnapshot>,
last_command_id: AtomicU64,
/// Worker-from-the-inside view of the filesystem. Set once in
/// `WorkerController::start` after the local WorkdirSession provider is
/// materialised, and read from the IPC server layer to answer
@@ -38,13 +39,24 @@ impl WorkerSharedState {
segment_id: SegmentId,
manifest_toml: String,
greeting: protocol::Greeting,
) -> Self {
Self::new_with_generation(worker_name, segment_id, manifest_toml, greeting, 1)
}
pub fn new_with_generation(
worker_name: String,
segment_id: SegmentId,
manifest_toml: String,
greeting: protocol::Greeting,
execution_generation: u64,
) -> Self {
Self {
worker_name,
segment_id,
manifest_toml,
greeting,
status: RwLock::new(WorkerStatus::Idle),
state: RwLock::new(WorkerStateSnapshot::initial(execution_generation)),
last_command_id: AtomicU64::new(0),
fs_view: OnceLock::new(),
flow_transition_enabled: AtomicBool::new(false),
}
@@ -70,21 +82,57 @@ impl WorkerSharedState {
self.flow_transition_enabled.load(Ordering::Acquire)
}
pub fn set_status(&self, status: WorkerStatus) {
if let Ok(mut s) = self.status.write() {
*s = status;
pub fn transition(&self, state: WorkerState) -> WorkerStateSnapshot {
let mut snapshot = self
.state
.write()
.expect("worker state lock poisoned; refusing an inferred fallback state");
if snapshot.state != state {
snapshot.revision = snapshot.revision.saturating_add(1);
snapshot.state = state;
}
snapshot.last_command_id = self.last_command_id.load(Ordering::Acquire);
snapshot.clone()
}
pub fn accept_command_id(&self, command_id: u64) -> bool {
self.last_command_id
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| {
(command_id > current).then_some(command_id)
})
.is_ok()
}
pub fn snapshot(&self) -> WorkerStateSnapshot {
let mut snapshot = self
.state
.read()
.expect("worker state lock poisoned; refusing an inferred fallback state")
.clone();
snapshot.last_command_id = self.last_command_id.load(Ordering::Acquire);
snapshot
}
/// Runtime catalog projection. This must not be used as live command
/// admission authority.
pub fn catalog_status(&self) -> WorkerStatus {
match self.snapshot().state {
WorkerState::Idle => WorkerStatus::Idle,
WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Paused)) => WorkerStatus::Paused,
WorkerState::Busy(WorkerBusyState::Run(_))
| WorkerState::Busy(WorkerBusyState::Maintenance(WorkerMaintenanceState::Compacting)) => {
WorkerStatus::Running
}
}
}
pub fn get_status(&self) -> WorkerStatus {
self.status.read().map(|s| *s).unwrap_or(WorkerStatus::Idle)
}
/// Serialize status as JSON.
/// Serialize the runtime-directory lifecycle projection as JSON while
/// retaining the full state snapshot for diagnostics and reconnects.
pub fn status_json(&self) -> String {
let status = self.get_status();
let snapshot = self.snapshot();
json!({
"state": status,
"state": self.catalog_status(),
"worker_state": snapshot,
"segment_id": self.segment_id.to_string(),
"worker_name": self.worker_name,
})
@@ -97,11 +145,12 @@ mod tests {
use super::*;
fn test_state() -> WorkerSharedState {
WorkerSharedState::new(
WorkerSharedState::new_with_generation(
"test-worker".into(),
session_store::new_segment_id(),
"[engine]\nname = \"test-worker\"".into(),
test_greeting(),
7,
)
}
@@ -119,36 +168,40 @@ mod tests {
}
#[test]
fn initial_status_is_idle() {
fn initial_snapshot_is_idle() {
let state = test_state();
assert_eq!(state.get_status(), WorkerStatus::Idle);
assert_eq!(state.snapshot(), WorkerStateSnapshot::initial(7));
assert_eq!(state.catalog_status(), WorkerStatus::Idle);
}
#[test]
fn set_and_get_status() {
fn transitions_increment_revision_only_when_state_changes() {
let state = test_state();
state.set_status(WorkerStatus::Running);
assert_eq!(state.get_status(), WorkerStatus::Running);
state.set_status(WorkerStatus::Paused);
assert_eq!(state.get_status(), WorkerStatus::Paused);
let running = WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running));
let snapshot = state.transition(running.clone());
assert_eq!(snapshot.revision, 1);
assert_eq!(snapshot.state, running);
assert_eq!(state.transition(running).revision, 1);
let paused = WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Paused));
let snapshot = state.transition(paused.clone());
assert_eq!(snapshot.revision, 2);
assert_eq!(snapshot.state, paused);
assert_eq!(state.catalog_status(), WorkerStatus::Paused);
}
#[test]
fn status_json_contains_fields() {
fn status_json_contains_full_snapshot_and_catalog_projection() {
let state = test_state();
let json = state.status_json();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(parsed["state"], "idle");
state.transition(WorkerState::Busy(WorkerBusyState::Maintenance(
WorkerMaintenanceState::Compacting,
)));
let parsed: serde_json::Value = serde_json::from_str(&state.status_json()).unwrap();
assert_eq!(parsed["state"], "running");
assert_eq!(parsed["worker_state"]["execution_generation"], 7);
assert_eq!(parsed["worker_state"]["revision"], 1);
assert_eq!(parsed["worker_state"]["state"]["kind"], "busy");
assert_eq!(parsed["worker_name"], "test-worker");
assert!(parsed["segment_id"].is_string());
}
#[test]
fn status_json_reflects_changes() {
let state = test_state();
state.set_status(WorkerStatus::Running);
let json = state.status_json();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(parsed["state"], "running");
}
}
crates/worker/src/spawn/comm_tools.rs
+9 -3
View File
@@ -97,7 +97,7 @@ mod tests {
context_window: 200_000,
context_tokens: 0,
},
status: WorkerStatus::Idle,
state: WorkerStatus::Idle.into(),
in_flight: Default::default(),
internal_workers: Vec::new(),
}
@@ -137,10 +137,16 @@ mod tests {
],
);
connect_and_send(&socket, &Method::Shutdown).await.unwrap();
let method = Method::Shutdown {
command: protocol::WorkerCommandEnvelope::for_snapshot(
1,
&protocol::WorkerStateSnapshot::initial(1),
),
};
connect_and_send(&socket, &method).await.unwrap();
let method = received.await.unwrap().expect("expected method");
assert!(matches!(method, Method::Shutdown));
assert!(matches!(method, Method::Shutdown { .. }));
}
#[tokio::test]
crates/worker/src/worker.rs
+186 -20
View File
@@ -4571,15 +4571,6 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
Ok(())
}
fn persist_and_send_compact_done(
&mut self,
lifecycle: CompactionLifecycle,
) -> Result<(), WorkerError> {
self.persist_compaction_lifecycle(&lifecycle)?;
self.send_event(Event::CompactDone { lifecycle });
Ok(())
}
fn persist_and_send_compact_failed(
&mut self,
lifecycle: CompactionLifecycle,
@@ -4724,7 +4715,97 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
Ok(rewrite_guard)
}
/// Terminalize and clean up any compaction that was left active by the
/// previous controller generation. This runs before the restored
/// controller publishes its first Idle state.
pub async fn recover_unfinished_compaction(&mut self) -> Result<(), WorkerError> {
let (entries, _) = self.sink.subscribe_with_snapshot();
let latest_payload = entries.iter().rev().find_map(|entry| match entry {
LogEntry::Extension {
domain, payload, ..
} if domain == COMPACTION_EXTENSION_DOMAIN => Some(payload.clone()),
_ => None,
});
let Some(payload) = latest_payload else {
return Ok(());
};
#[derive(serde::Deserialize)]
#[serde(deny_unknown_fields)]
struct CompactionLifecycleWire {
schema_version: u32,
compaction_id: String,
revision: u64,
#[serde(default)]
internal_worker: Option<protocol::InternalWorkerRef>,
state: CompactionLifecycleState,
started_at_ms: u64,
#[serde(default)]
ended_at_ms: Option<u64>,
#[serde(default)]
summary: Option<String>,
#[serde(default)]
error: Option<String>,
#[serde(default)]
new_segment_id: Option<String>,
}
let wire: CompactionLifecycleWire = serde_json::from_value(payload).map_err(|error| {
WorkerError::InvalidState(format!("decode compaction lifecycle: {error}"))
})?;
if !matches!(wire.schema_version, 2 | 3) {
return Err(WorkerError::InvalidState(format!(
"unsupported compaction lifecycle schema version {}",
wire.schema_version
)));
}
let mut lifecycle = CompactionLifecycle {
schema_version: wire.schema_version,
compaction_id: wire.compaction_id,
revision: wire.revision,
internal_worker: wire.internal_worker,
state: wire.state,
started_at_ms: wire.started_at_ms,
ended_at_ms: wire.ended_at_ms,
summary: wire.summary,
error: wire.error,
new_segment_id: wire.new_segment_id,
};
match lifecycle.state {
CompactionLifecycleState::Running => {
lifecycle.schema_version = 3;
lifecycle.revision = lifecycle.revision.saturating_add(1);
lifecycle.state = CompactionLifecycleState::Interrupted;
lifecycle.ended_at_ms = Some(segment_log::now_millis());
lifecycle.error =
Some("worker execution restarted before compaction completed".into());
self.persist_compaction_lifecycle(&lifecycle)?;
self.send_event(Event::CompactFailed {
lifecycle: lifecycle.clone(),
});
self.release_compaction_service(&lifecycle).await;
}
CompactionLifecycleState::Interrupted => {
self.release_compaction_service(&lifecycle).await;
}
CompactionLifecycleState::Done | CompactionLifecycleState::Failed => {}
}
Ok(())
}
pub async fn manual_compact(&mut self) -> Result<ManualCompactResult, WorkerError> {
self.manual_compact_inner(None).await
}
pub async fn manual_compact_with_cancel(
&mut self,
cancel: tokio::sync::watch::Receiver<bool>,
) -> Result<ManualCompactResult, WorkerError> {
self.manual_compact_inner(Some(cancel)).await
}
async fn manual_compact_inner(
&mut self,
mut cancel: Option<tokio::sync::watch::Receiver<bool>>,
) -> Result<ManualCompactResult, WorkerError> {
if self.manifest.compaction.is_none() {
let message =
"manual compact is unavailable because [compaction] is not configured".to_string();
@@ -4764,7 +4845,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
return Ok(ManualCompactResult::Skipped { message });
}
match self.compact(retained).await {
match self.compact_with_cancel(retained, cancel.take()).await {
Ok(new_segment_id) => {
info!(new_segment_id = %new_segment_id, "Manual compaction succeeded");
if let Some(ref state) = state {
@@ -4937,11 +5018,19 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// Runs one parent-owned observable compaction service and returns the new
/// Segment ID. Lifecycle revisions are committed before they are broadcast.
pub async fn compact(&mut self, retained_tokens: u64) -> Result<SegmentId, WorkerError> {
self.compact_with_cancel(retained_tokens, None).await
}
async fn compact_with_cancel(
&mut self,
retained_tokens: u64,
mut cancel: Option<tokio::sync::watch::Receiver<bool>>,
) -> Result<SegmentId, WorkerError> {
let _rewrite_guard = self
.prepare_session_rewrite(SessionRewriteKind::Compact)
.await?;
let mut lifecycle = CompactionLifecycle {
schema_version: 2,
schema_version: 3,
compaction_id: uuid::Uuid::now_v7().to_string(),
revision: 1,
internal_worker: None,
@@ -4953,16 +5042,25 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
new_segment_id: None,
};
self.persist_and_send_compact_start(lifecycle.clone())?;
match self.compact_impl(retained_tokens, &mut lifecycle).await {
Ok((new_segment_id, summary)) => {
lifecycle.revision = lifecycle.revision.saturating_add(1);
lifecycle.state = CompactionLifecycleState::Done;
lifecycle.ended_at_ms = Some(segment_log::now_millis());
lifecycle.summary = Some(summary);
lifecycle.new_segment_id = Some(new_segment_id.to_string());
let terminal = self.persist_and_send_compact_done(lifecycle.clone());
let outcome = if let Some(cancel) = cancel.as_mut() {
tokio::select! {
biased;
changed = cancel.changed() => {
let _ = changed;
Err(WorkerError::CompactCancelled)
}
result = self.compact_impl(retained_tokens, &mut lifecycle) => result,
}
} else {
self.compact_impl(retained_tokens, &mut lifecycle).await
};
match outcome {
Ok((new_segment_id, _summary)) => {
debug_assert_eq!(lifecycle.state, CompactionLifecycleState::Done);
self.send_event(Event::CompactDone {
lifecycle: lifecycle.clone(),
});
self.release_compaction_service(&lifecycle).await;
terminal?;
Ok(new_segment_id)
}
Err(error) => {
@@ -5543,6 +5641,24 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
})?,
});
}
// Commit the terminal lifecycle in the same atomic replacement-segment
// creation as the rewritten history. Restore can therefore never see a
// replacement segment without the Done fact for the compaction that
// created it.
lifecycle.revision = lifecycle.revision.saturating_add(1);
lifecycle.state = CompactionLifecycleState::Done;
lifecycle.ended_at_ms = Some(segment_log::now_millis());
lifecycle.summary = Some(summary_text.clone());
lifecycle.new_segment_id = Some(new_segment_id.to_string());
initial_entries.push(LogEntry::Extension {
ts: segment_log::now_millis(),
domain: COMPACTION_EXTENSION_DOMAIN.to_string(),
payload: serde_json::to_value(&*lifecycle).map_err(|error| {
WorkerError::InvalidState(format!(
"serialize terminal compaction lifecycle: {error}"
))
})?,
});
self.store
.create_segment(old_loc.session_id, new_segment_id, &initial_entries)?;
self.segment_state.set_location(SegmentLocation {
@@ -10165,6 +10281,56 @@ mod build_summary_prompt_tests {
assert_eq!(state.notification_receipts.len(), 1);
}
#[tokio::test]
async fn restore_terminalizes_running_compaction_before_idle_publication() {
let (_dir, mut worker) = rewind_test_worker().await;
let lifecycle = CompactionLifecycle {
schema_version: 3,
compaction_id: "compact-before-restart".into(),
revision: 1,
internal_worker: None,
state: CompactionLifecycleState::Running,
started_at_ms: segment_log::now_millis(),
ended_at_ms: None,
summary: None,
error: None,
new_segment_id: None,
};
worker.persist_compaction_lifecycle(&lifecycle).unwrap();
worker.recover_unfinished_compaction().await.unwrap();
let (entries, _) = worker.sink.subscribe_with_snapshot();
let restored = entries.iter().rev().find_map(|entry| match entry {
LogEntry::Extension {
domain, payload, ..
} if domain == COMPACTION_EXTENSION_DOMAIN => {
serde_json::from_value::<CompactionLifecycle>(payload.clone()).ok()
}
_ => None,
});
let restored = restored.expect("terminal compaction lifecycle");
assert_eq!(restored.state, CompactionLifecycleState::Interrupted);
assert_eq!(restored.revision, 2);
assert!(
restored
.error
.as_deref()
.is_some_and(|error| error.contains("restarted"))
);
let mut future = lifecycle;
future.schema_version = 4;
future.compaction_id = "future-compaction".into();
worker.persist_compaction_lifecycle(&future).unwrap();
let error = worker.recover_unfinished_compaction().await.unwrap_err();
assert!(
error
.to_string()
.contains("unsupported compaction lifecycle schema version 4")
);
}
fn minimal_manifest() -> WorkerManifest {
let toml_str = r#"
[worker]
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!(