From e8b9adcde43028bae4bae1cb5ee067c018e7693a Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 06:54:29 +0900 Subject: [PATCH] feat: add revisioned worker execution state --- crates/client/src/client.rs | 8 +- crates/client/src/transport/in_process.rs | 8 +- crates/client/src/transport/unix_socket.rs | 11 +- crates/client/src/transport/websocket.rs | 8 +- crates/protocol/src/lib.rs | 274 ++++++-- crates/protocol/src/typescript.rs | 14 +- crates/standalone/src/host.rs | 12 +- crates/tui/src/app.rs | 70 +- crates/tui/src/command.rs | 9 +- crates/tui/src/console/mod.rs | 43 +- crates/worker-runtime/src/execution.rs | 42 +- crates/worker-runtime/src/http_server.rs | 40 +- crates/worker-runtime/src/runtime.rs | 216 +++--- crates/worker-runtime/src/worker_backend.rs | 415 ++++++----- crates/worker/src/controller.rs | 647 +++++++++++++++--- crates/worker/src/discovery.rs | 16 +- crates/worker/src/internal_worker.rs | 70 +- crates/worker/src/runtime/dir.rs | 5 +- crates/worker/src/shared_state.rs | 141 ++-- crates/worker/src/spawn/comm_tools.rs | 12 +- crates/worker/src/worker.rs | 206 +++++- crates/worker/tests/compact_events_test.rs | 208 +++++- crates/worker/tests/controller_test.rs | 200 ++++-- crates/workspace-server/src/hosts.rs | 13 +- .../src/runtime_subscription_tests.rs | 22 +- crates/workspace-server/src/server.rs | 20 +- web/workspace/src/lib/generated/protocol.ts | 42 +- .../src/lib/workspace/console/model.test.ts | 37 +- .../src/lib/workspace/console/model.ts | 27 +- .../lib/workspace/console/run-status.test.ts | 7 +- .../src/lib/workspace/console/run-status.ts | 4 +- .../workers/[workerId]/console/+page.svelte | 71 +- 32 files changed, 2168 insertions(+), 750 deletions(-) diff --git a/crates/client/src/client.rs b/crates/client/src/client.rs index b760ce00..bdd73375 100644 --- a/crates/client/src/client.rs +++ b/crates/client/src/client.rs @@ -112,8 +112,8 @@ mod tests { async fn encodes_methods_and_decodes_events_above_transport() { let mut socket = TestSocket::default(); socket.incoming.push_back( - encode_event(&Event::Status { - status: WorkerStatus::Idle, + encode_event(&Event::WorkerState { + snapshot: WorkerStatus::Idle.into(), }) .expect("encode event"), ); @@ -132,9 +132,7 @@ mod tests { )); assert!(matches!( client.next_event().await, - Ok(Some(Event::Status { - status: WorkerStatus::Idle - })) + Ok(Some(Event::WorkerState { .. })) )); } } diff --git a/crates/client/src/transport/in_process.rs b/crates/client/src/transport/in_process.rs index 3111808a..ed50be54 100644 --- a/crates/client/src/transport/in_process.rs +++ b/crates/client/src/transport/in_process.rs @@ -101,8 +101,8 @@ mod tests { )); peer.send( - encode_event(&Event::Status { - status: WorkerStatus::Idle, + encode_event(&Event::WorkerState { + snapshot: WorkerStatus::Idle.into(), }) .expect("encode event"), ) @@ -110,9 +110,7 @@ mod tests { .expect("send event"); assert!(matches!( client.next_event().await, - Ok(Some(Event::Status { - status: WorkerStatus::Idle - })) + Ok(Some(Event::WorkerState { .. })) )); } } diff --git a/crates/client/src/transport/unix_socket.rs b/crates/client/src/transport/unix_socket.rs index 089ed83e..8b57bb2c 100644 --- a/crates/client/src/transport/unix_socket.rs +++ b/crates/client/src/transport/unix_socket.rs @@ -113,8 +113,8 @@ mod tests { let listener = UnixListener::bind(&socket_path).unwrap(); let server = tokio::spawn(async move { let (mut stream, _) = listener.accept().await.unwrap(); - let event = encode_event(&Event::Status { - status: WorkerStatus::Idle, + let event = encode_event(&Event::WorkerState { + snapshot: WorkerStatus::Idle.into(), }) .unwrap(); stream.write_all(event.as_bytes()).await.unwrap(); @@ -126,12 +126,7 @@ mod tests { .await .expect("client should receive event while alive") .expect("transport should succeed"); - assert!(matches!( - event, - Some(Event::Status { - status: WorkerStatus::Idle - }) - )); + assert!(matches!(event, Some(Event::WorkerState { .. }))); server.await.unwrap(); } diff --git a/crates/client/src/transport/websocket.rs b/crates/client/src/transport/websocket.rs index b4c1ed84..883af12e 100644 --- a/crates/client/src/transport/websocket.rs +++ b/crates/client/src/transport/websocket.rs @@ -116,8 +116,8 @@ mod tests { Message::Text(ref text) if matches!(decode_method(text), Ok(Method::Submit { .. })) )); - let event = encode_event(&Event::Status { - status: WorkerStatus::Idle, + let event = encode_event(&Event::WorkerState { + snapshot: WorkerStatus::Idle.into(), }) .unwrap(); socket.send(Message::Text(event.into())).await.unwrap(); @@ -134,9 +134,7 @@ mod tests { .expect("send method"); assert!(matches!( client.next_event().await, - Ok(Some(Event::Status { - status: WorkerStatus::Idle - })) + Ok(Some(Event::WorkerState { .. })) )); server.await.unwrap(); } diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index fd5190c7..ec96e226 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -85,6 +85,142 @@ impl AuthenticatedInputSource { } } +/// Immutable identity and revision fence for one state-changing Worker command. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +pub struct WorkerCommandEnvelope { + /// Caller-owned sequence. A controller accepts command ids in strictly + /// increasing order for one execution generation. + pub command_id: u64, + pub expected_execution_generation: u64, + pub expected_worker_state_revision: u64, +} + +impl WorkerCommandEnvelope { + pub fn for_snapshot(command_id: u64, snapshot: &WorkerStateSnapshot) -> Self { + Self { + command_id, + expected_execution_generation: snapshot.execution_generation, + expected_worker_state_revision: snapshot.revision, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(rename_all = "snake_case")] +pub enum WorkerCommandKind { + Resume, + Cancel, + Pause, + Compact, + Shutdown, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(rename_all = "snake_case")] +pub enum WorkerCommandDisposition { + Accepted, + StaleExecutionGeneration, + StaleWorkerStateRevision, + StaleCommandId, + InvalidState, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +pub struct WorkerCommandAcknowledgement { + pub command_id: u64, + pub command: WorkerCommandKind, + pub disposition: WorkerCommandDisposition, + /// The complete authoritative state observed after command admission. + pub state: WorkerStateSnapshot, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(tag = "kind", content = "state", rename_all = "snake_case")] +pub enum WorkerState { + Idle, + Busy(WorkerBusyState), +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(tag = "kind", content = "state", rename_all = "snake_case")] +pub enum WorkerBusyState { + Run(WorkerRunState), + Maintenance(WorkerMaintenanceState), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(rename_all = "snake_case")] +pub enum WorkerRunState { + Running, + Pausing, + Paused, + Cancelling, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(rename_all = "snake_case")] +pub enum WorkerMaintenanceState { + Compacting, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +pub struct WorkerStateSnapshot { + pub execution_generation: u64, + pub revision: u64, + /// Highest lifecycle command id observed by this controller generation. + pub last_command_id: u64, + pub state: WorkerState, +} + +impl WorkerStateSnapshot { + pub fn initial(execution_generation: u64) -> Self { + Self { + execution_generation, + revision: 0, + last_command_id: 0, + state: WorkerState::Idle, + } + } + + /// Compatibility projection for Runtime catalog lifecycle. This value is + /// never command-admission authority and cannot produce `Stopped`. + pub fn catalog_status(&self) -> WorkerStatus { + match self.state { + WorkerState::Idle => WorkerStatus::Idle, + WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Paused)) => WorkerStatus::Paused, + WorkerState::Busy(WorkerBusyState::Run(_)) + | WorkerState::Busy(WorkerBusyState::Maintenance(_)) => WorkerStatus::Running, + } + } +} + +impl From for WorkerStateSnapshot { + fn from(status: WorkerStatus) -> Self { + 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)), + }; + Self { + execution_generation: 1, + revision: 0, + last_command_id: 0, + state, + } + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "typescript", derive(ts_rs::TS))] #[serde(tag = "method", content = "params", rename_all = "snake_case")] @@ -149,20 +285,28 @@ pub enum Method { expected_revision: u64, expected_head_id: String, }, - Resume, - Cancel, + Resume { + command: WorkerCommandEnvelope, + }, + Cancel { + command: WorkerCommandEnvelope, + }, /// Stop the in-flight turn and transition to `Paused`. /// /// Unlike `Cancel` (which discards and returns to `Idle`), a paused /// Worker can resume the interrupted work via `Resume`, or accept a /// fresh `Submit` (orphan `tool_use` items are closed with a /// synthetic tool result before the new user message is appended). - Pause, + Pause { + command: WorkerCommandEnvelope, + }, /// Request an explicit compaction while the Worker is otherwise idle. /// /// This is a typed control method: clients must not send `compact` as a /// `Method::Submit` user message. - Compact, + Compact { + command: WorkerCommandEnvelope, + }, /// Ask the Worker to list valid rewind targets from its authoritative session log. ListRewindTargets, /// Truncate the current session back to the selected rewind target and @@ -171,7 +315,9 @@ pub enum Method { target: RewindTargetId, expected_head_entries: usize, }, - Shutdown, + Shutdown { + command: WorkerCommandEnvelope, + }, /// Request a list of completion candidates from the Worker. /// /// Reply is sent on the same socket as `Event::Completions` (not @@ -938,8 +1084,9 @@ pub enum Event { Snapshot { session: SessionSnapshot, greeting: Greeting, - #[serde(default)] - status: WorkerStatus, + /// Full revisioned live execution state. `Stopped` remains Runtime + /// catalog authority and is deliberately not represented here. + state: WorkerStateSnapshot, /// Unfinished model output that has already streamed in the current /// run but is not yet represented by committed snapshot entries. #[serde(default, skip_serializing_if = "InFlightSnapshot::is_empty")] @@ -976,8 +1123,11 @@ pub enum Event { }, /// Current Worker controller status. Broadcast on every controller-level /// transition and included in `History` snapshots for late attach. - Status { - status: WorkerStatus, + WorkerState { + snapshot: WorkerStateSnapshot, + }, + CommandAcknowledged { + acknowledgement: WorkerCommandAcknowledgement, }, /// Bounded, provider-owned command telemetry for the live Console. This is /// intentionally not a history entry and is reconstructed from @@ -1612,28 +1762,39 @@ mod tests { } #[test] - fn method_without_params() { - let json = r#"{"method":"resume"}"#; - let method: Method = serde_json::from_str(json).unwrap(); - assert!(matches!(method, Method::Resume)); + fn lifecycle_method_without_command_fails_closed() { + let error = serde_json::from_str::(r#"{"method":"resume"}"#).unwrap_err(); + assert!(error.to_string().contains("params")); } #[test] - fn method_pause_roundtrip() { - let json = r#"{"method":"pause"}"#; - let method: Method = serde_json::from_str(json).unwrap(); - assert!(matches!(method, Method::Pause)); - let serialized = serde_json::to_string(&method).unwrap(); - assert_eq!(serialized, json); - } - - #[test] - fn method_compact_roundtrip() { - let json = r#"{"method":"compact"}"#; - let method: Method = serde_json::from_str(json).unwrap(); - assert!(matches!(method, Method::Compact)); - let serialized = serde_json::to_string(&method).unwrap(); - assert_eq!(serialized, json); + fn lifecycle_methods_roundtrip_with_fences() { + for method in [ + Method::Pause { + command: WorkerCommandEnvelope { + command_id: 11, + expected_execution_generation: 4, + expected_worker_state_revision: 8, + }, + }, + Method::Compact { + command: WorkerCommandEnvelope { + command_id: 12, + expected_execution_generation: 4, + expected_worker_state_revision: 9, + }, + }, + ] { + let json = serde_json::to_string(&method).unwrap(); + let decoded: Method = serde_json::from_str(&json).unwrap(); + match decoded { + Method::Pause { command } | Method::Compact { command } => { + assert_eq!(command.expected_execution_generation, 4); + assert!(command.command_id >= 11); + } + other => panic!("unexpected lifecycle method: {other:?}"), + } + } } #[test] @@ -1902,7 +2063,7 @@ mod tests { context_window: 200_000, context_tokens: 42_000, }, - status: WorkerStatus::Paused, + state: WorkerStatus::Paused.into(), in_flight: InFlightSnapshot::default(), internal_workers: Vec::new(), }; @@ -1919,12 +2080,13 @@ mod tests { assert_eq!(parsed["data"]["greeting"]["tools"][0], "Read"); assert_eq!(parsed["data"]["greeting"]["context_window"], 200_000); assert_eq!(parsed["data"]["greeting"]["context_tokens"], 42_000); - assert_eq!(parsed["data"]["status"], "paused"); + assert_eq!(parsed["data"]["state"]["state"]["kind"], "busy"); + assert_eq!(parsed["data"]["state"]["state"]["state"]["state"], "paused"); } #[test] fn event_snapshot_in_flight_roundtrip_and_default() { - let inbound = r#"{"event":"snapshot","data":{"session":{"entries":[]},"greeting":{"worker_name":"test","cwd":"/tmp","provider":"p","model":"m","scope_summary":"s","tools":[]},"status":"running"}}"#; + let inbound = r#"{"event":"snapshot","data":{"session":{"entries":[]},"greeting":{"worker_name":"test","cwd":"/tmp","provider":"p","model":"m","scope_summary":"s","tools":[]},"state":{"execution_generation":1,"revision":1,"last_command_id":0,"state":{"kind":"busy","state":{"kind":"run","state":"running"}}}}}"#; let decoded: Event = serde_json::from_str(inbound).unwrap(); match decoded { Event::Snapshot { in_flight, .. } => assert!(in_flight.is_empty()), @@ -1946,7 +2108,7 @@ mod tests { context_window: 0, context_tokens: 0, }, - status: WorkerStatus::Running, + state: WorkerStatus::Running.into(), in_flight: InFlightSnapshot { blocks: vec![ InFlightBlock::Text { @@ -2034,20 +2196,32 @@ mod tests { } #[test] - fn event_status_format() { - let event = Event::Status { - status: WorkerStatus::Running, + fn event_worker_state_format() { + let event = Event::WorkerState { + snapshot: WorkerStateSnapshot { + execution_generation: 7, + revision: 3, + last_command_id: 9, + state: WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running)), + }, }; let json = serde_json::to_string(&event).unwrap(); let parsed: serde_json::Value = serde_json::from_str(&json).unwrap(); - assert_eq!(parsed["event"], "status"); - assert_eq!(parsed["data"]["status"], "running"); + assert_eq!(parsed["event"], "worker_state"); + assert_eq!(parsed["data"]["snapshot"]["execution_generation"], 7); + assert_eq!(parsed["data"]["snapshot"]["revision"], 3); + assert_eq!(parsed["data"]["snapshot"]["state"]["kind"], "busy"); let decoded: Event = serde_json::from_str(&json).unwrap(); assert!(matches!( decoded, - Event::Status { - status: WorkerStatus::Running + Event::WorkerState { + snapshot: WorkerStateSnapshot { + execution_generation: 7, + revision: 3, + state: WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running)), + .. + } } )); } @@ -2088,19 +2262,10 @@ mod tests { } #[test] - fn event_snapshot_without_status_defaults_to_idle() { + fn event_snapshot_without_worker_state_fails_closed() { let json = r#"{"event":"snapshot","data":{"session":{"entries":[]},"greeting":{"worker_name":"test","cwd":"/tmp","provider":"anthropic","model":"claude","scope_summary":"","tools":[]}}}"#; - let decoded: Event = serde_json::from_str(json).unwrap(); - match decoded { - Event::Snapshot { - status, greeting, .. - } => { - assert_eq!(status, WorkerStatus::Idle); - assert_eq!(greeting.context_window, 0); - assert_eq!(greeting.context_tokens, 0); - } - other => panic!("expected Snapshot, got {other:?}"), - } + let error = serde_json::from_str::(json).unwrap_err(); + assert!(error.to_string().contains("state")); } #[test] @@ -2513,7 +2678,12 @@ mod tests { "scope_summary": "scope", "tools": [] }, - "status": "idle" + "state": { + "execution_generation": 1, + "revision": 0, + "last_command_id": 0, + "state": { "kind": "idle" } + } } })) .unwrap(); diff --git a/crates/protocol/src/typescript.rs b/crates/protocol/src/typescript.rs index bf63c8d6..ec5588f9 100644 --- a/crates/protocol/src/typescript.rs +++ b/crates/protocol/src/typescript.rs @@ -12,7 +12,10 @@ use crate::{ RewindSummary, RewindTarget, RewindTargetId, RunResult, ScopeRule, Segment, SessionContentPart, SessionEntryProvenance, SessionMessageRole, SessionSnapshot, SessionSnapshotEntry, SessionSnapshotEntryData, SessionToolAttachment, SubmissionDisposition, ToolResultDisposition, - TurnResult, UploadedFileAvailability, UploadedFileRef, WorkerEvent, WorkerStatus, + TurnResult, UploadedFileAvailability, UploadedFileRef, WorkerBusyState, + WorkerCommandAcknowledgement, WorkerCommandDisposition, WorkerCommandEnvelope, + WorkerCommandKind, WorkerEvent, WorkerMaintenanceState, WorkerRunState, WorkerState, + WorkerStateSnapshot, WorkerStatus, subscription::{ EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame, SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest, @@ -46,6 +49,15 @@ pub fn generated_protocol_types() -> String { push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); diff --git a/crates/standalone/src/host.rs b/crates/standalone/src/host.rs index c4c5df73..1131bccd 100644 --- a/crates/standalone/src/host.rs +++ b/crates/standalone/src/host.rs @@ -318,7 +318,11 @@ impl StandaloneHost { } pub async fn shutdown(mut self) -> Result<(), StandaloneShutdownError> { - let _ = self.handle.send(Method::Shutdown).await; + let command = protocol::WorkerCommandEnvelope::for_snapshot( + u64::MAX, + &self.handle.shared_state.snapshot(), + ); + let _ = self.handle.send(Method::Shutdown { command }).await; let Some(shutdown) = self.shutdown.take() else { self.retain_lease(); return Err(StandaloneShutdownError::ConfirmationLost); @@ -500,7 +504,11 @@ fn active_pointer( } async fn stop_started_worker(started: BootstrappedWorker) { - let _ = started.handle.send(Method::Shutdown).await; + let command = protocol::WorkerCommandEnvelope::for_snapshot( + u64::MAX, + &started.handle.shared_state.snapshot(), + ); + let _ = started.handle.send(Method::Shutdown { command }).await; let _ = tokio::time::timeout(Duration::from_secs(2), started.shutdown).await; } diff --git a/crates/tui/src/app.rs b/crates/tui/src/app.rs index d23430c2..a76d32b5 100644 --- a/crates/tui/src/app.rs +++ b/crates/tui/src/app.rs @@ -5,7 +5,7 @@ use std::time::{Duration, Instant}; use protocol::{ AlertLevel, AlertSource, CompletionEntry, CompletionKind, ErrorCode, Event, InFlightBlock, InFlightSnapshot, InFlightToolCallState, InternalWorkerRef, InternalWorkerSnapshot, Method, - RewindTarget, RunResult, Segment, WorkerStatus, + RewindTarget, RunResult, Segment, WorkerCommandEnvelope, WorkerStateSnapshot, WorkerStatus, }; use crate::block::{ @@ -225,8 +225,10 @@ pub struct WorkerViewTab { pub struct App { pub worker_name: String, pub connected: bool, - /// Last controller status reported by the Worker. Drives the status line - /// and Ctrl-key routing; do not infer this solely from replayed history. + /// Latest authoritative revisioned live execution state. + pub worker_state: WorkerStateSnapshot, + next_command_id: u64, + /// Derived Runtime-catalog compatibility projection used by existing UI. pub worker_status: WorkerStatus, /// True while the Worker is in `WorkerStatus::Running`. pub running: bool, @@ -337,6 +339,8 @@ impl App { Self { worker_name, connected: false, + worker_state: WorkerStateSnapshot::initial(1), + next_command_id: 1, worker_status: WorkerStatus::Idle, running: false, paused: false, @@ -745,7 +749,8 @@ impl App { if self.paused { self.input_history.cancel_browse(); self.input.clear(); - return Some(Method::Resume); + let command = self.next_command_envelope(); + return Some(Method::Resume { command }); } return None; } @@ -1114,6 +1119,15 @@ impl App { } } + pub fn next_command_envelope(&mut self) -> WorkerCommandEnvelope { + let command_id = self + .next_command_id + .max(self.worker_state.last_command_id.saturating_add(1)); + let command = WorkerCommandEnvelope::for_snapshot(command_id, &self.worker_state); + self.next_command_id = command_id.saturating_add(1); + command + } + pub fn handle_worker_event(&mut self, event: Event) -> Option { if self.rewind_refresh_fence && event_is_stale_after_rewind(&event) { return None; @@ -1443,7 +1457,7 @@ impl App { Event::Snapshot { session, greeting, - status, + state, in_flight, internal_workers, } => { @@ -1451,7 +1465,8 @@ impl App { self.pending_submissions = session.pending_submissions.clone(); self.restore_snapshot(&session, greeting, in_flight); self.replace_internal_worker_snapshots(internal_workers); - self.set_worker_status(status); + self.worker_state = state.clone(); + self.set_worker_status(state.catalog_status()); } Event::InternalWorker { worker, @@ -1461,9 +1476,14 @@ impl App { Event::InternalWorkerRemoved { worker, revision } => { self.remove_internal_worker(worker, revision) } - Event::Status { status } => { + Event::WorkerState { snapshot } => { self.rewind_refresh_fence = false; - self.set_worker_status(status); + self.worker_state = snapshot.clone(); + self.set_worker_status(snapshot.catalog_status()); + } + Event::CommandAcknowledged { acknowledgement } => { + self.worker_state = acknowledgement.state.clone(); + self.set_worker_status(acknowledgement.state.catalog_status()); } // Command telemetry is an operational Web Console surface. The // TUI continues to render the final Bash ToolResult from history. @@ -2026,12 +2046,18 @@ impl App { self.input_mode = CommandInputMode::Composer; self.command_completion_selected = None; } - if let Some(Method::ListRewindTargets) = result.method.as_ref() { + let mut method = result.method; + if let Some(Method::Compact { .. }) = method { + method = Some(Method::Compact { + command: self.next_command_envelope(), + }); + } + if let Some(Method::ListRewindTargets) = method.as_ref() { self.completion = None; self.rewind_picker = None; self.rewind_request_pending = true; } - result.method + method } fn push_command_diagnostic(&mut self, message: impl Into) { @@ -2761,8 +2787,8 @@ mod rewind_refresh_tests { }); assert!(!blocks_contain(&app, "stale tail after rewind")); - app.handle_worker_event(Event::Status { - status: WorkerStatus::Idle, + app.handle_worker_event(Event::WorkerState { + snapshot: WorkerStatus::Idle.into(), }); app.handle_worker_event(Event::TextDelta { text: "new live tail after status".into(), @@ -3478,7 +3504,7 @@ mod completion_flow_tests { let mut app = App::new("test".into()); app.set_worker_status(WorkerStatus::Paused); - assert!(matches!(app.submit_input(), Some(Method::Resume))); + assert!(matches!(app.submit_input(), Some(Method::Resume { .. }))); assert_eq!(app.queued_input_count(), 0); } @@ -3533,7 +3559,7 @@ mod completion_flow_tests { app.handle_worker_event(Event::Snapshot { greeting: test_greeting(), session: public_session(vec![session_start_value]), - status: WorkerStatus::Running, + state: WorkerStatus::Running.into(), in_flight: Default::default(), internal_workers: Vec::new(), }); @@ -3551,8 +3577,8 @@ mod completion_flow_tests { code: ErrorCode::ProviderError, message: "provider unavailable".into(), }); - app.handle_worker_event(Event::Status { - status: WorkerStatus::Idle, + app.handle_worker_event(Event::WorkerState { + snapshot: WorkerStatus::Idle.into(), }); let live_errors = app @@ -3577,7 +3603,7 @@ mod completion_flow_tests { app.handle_worker_event(Event::Snapshot { greeting: test_greeting(), session: public_session(vec![serde_json::to_value(run_errored).unwrap()]), - status: WorkerStatus::Idle, + state: WorkerStatus::Idle.into(), in_flight: Default::default(), internal_workers: Vec::new(), }); @@ -3641,7 +3667,7 @@ mod completion_flow_tests { pending_submissions: protocol::PendingSubmissionsSnapshot::default(), entries: Vec::new(), }, - status: WorkerStatus::Running, + state: WorkerStatus::Running.into(), in_flight: InFlightSnapshot { blocks: vec![ InFlightBlock::Thinking { @@ -3968,7 +3994,7 @@ mod completion_flow_tests { pending_submissions: protocol::PendingSubmissionsSnapshot::default(), entries: Vec::new(), }, - status: WorkerStatus::Idle, + state: WorkerStatus::Idle.into(), in_flight: Default::default(), internal_workers: Vec::new(), }); @@ -4020,7 +4046,7 @@ mod completion_flow_tests { pending_submissions: protocol::PendingSubmissionsSnapshot::default(), entries: Vec::new(), }, - status: WorkerStatus::Idle, + state: WorkerStatus::Idle.into(), in_flight: Default::default(), internal_workers: vec![InternalWorkerSnapshot { worker: InternalWorkerRef { @@ -4194,7 +4220,7 @@ mod completion_flow_tests { entries: Vec::new(), }, greeting, - status: WorkerStatus::Idle, + state: WorkerStatus::Idle.into(), in_flight: Default::default(), internal_workers: Vec::new(), }); @@ -4393,7 +4419,7 @@ mod completion_flow_tests { app.handle_worker_event(Event::Snapshot { greeting: test_greeting(), session: public_session(assistant_item_entries), - status: WorkerStatus::Running, + state: WorkerStatus::Running.into(), in_flight: Default::default(), internal_workers: Vec::new(), }); diff --git a/crates/tui/src/command.rs b/crates/tui/src/command.rs index 44004a30..ba24299e 100644 --- a/crates/tui/src/command.rs +++ b/crates/tui/src/command.rs @@ -409,7 +409,12 @@ fn compact_command(invocation: CommandInvocation<'_>) -> CommandExecution { let _ = invocation.environment; let _ = invocation.args.raw(); CommandExecution { - method: Some(Method::Compact), + method: Some(Method::Compact { + command: protocol::WorkerCommandEnvelope::for_snapshot( + 0, + &protocol::WorkerStateSnapshot::initial(1), + ), + }), diagnostics: vec![CommandDiagnostic::new("compact requested")], exit_command_mode: true, clear_input: true, @@ -483,7 +488,7 @@ mod tests { fn compact_command_returns_compact_method_not_run() { let registry = CommandRegistry::builtins(); let result = registry.dispatch("compact", &env()); - assert!(matches!(result.method, Some(Method::Compact))); + assert!(matches!(result.method, Some(Method::Compact { .. }))); assert!(result.exit_command_mode); assert!(result.clear_input); assert!(result.diagnostics[0].message.contains("compact requested")); diff --git a/crates/tui/src/console/mod.rs b/crates/tui/src/console/mod.rs index a7ff703e..af041066 100644 --- a/crates/tui/src/console/mod.rs +++ b/crates/tui/src/console/mod.rs @@ -572,7 +572,7 @@ async fn run_e2e_rewind_fixture( pending_submissions: protocol::PendingSubmissionsSnapshot::default(), entries: Vec::new(), }, - status: WorkerStatus::Idle, + state: WorkerStatus::Idle.into(), greeting: Greeting { worker_name: worker_name.clone(), cwd: workspace_root.display().to_string(), @@ -1438,13 +1438,15 @@ fn handle_cancel_or_shutdown(app: &mut App) -> Option { WorkerStatus::Running | WorkerStatus::Paused ) { app.shutdown_confirm = None; - return Some(Method::Cancel); + let command = app.next_command_envelope(); + return Some(Method::Cancel { command }); } if let Some(pressed_at) = app.shutdown_confirm && pressed_at.elapsed() < CONFIRM_TIMEOUT { app.shutdown_confirm = None; - return Some(Method::Shutdown); + let command = app.next_command_envelope(); + return Some(Method::Shutdown { command }); } app.shutdown_confirm = Some(std::time::Instant::now()); app.flash_actionbar_notice( @@ -1460,7 +1462,8 @@ fn handle_cancel_or_shutdown(app: &mut App) -> Option { /// Idle / Paused → 2-tap to quit the TUI (the Worker keeps running). fn handle_pause_or_quit(app: &mut App) -> Option { if app.worker_status == WorkerStatus::Running { - return Some(Method::Pause); + let command = app.next_command_envelope(); + return Some(Method::Pause { command }); } if let Some(t) = app.quit_confirm && t.elapsed() < CONFIRM_TIMEOUT @@ -2090,7 +2093,7 @@ mod tests { &mut app, KeyEvent::new(KeyCode::Char('c'), KeyModifiers::CONTROL), ), - Some(Method::Pause) + Some(Method::Pause { .. }) )); assert_eq!(app.queued_input_count(), 1); @@ -2100,7 +2103,7 @@ mod tests { &mut app, KeyEvent::new(KeyCode::Char('x'), KeyModifiers::CONTROL), ), - Some(Method::Cancel) + Some(Method::Cancel { .. }) )); assert_eq!(app.queued_input_count(), 1); } @@ -2114,7 +2117,7 @@ mod tests { &mut app, KeyEvent::new(KeyCode::Char('x'), KeyModifiers::CONTROL), ); - assert!(matches!(cancel, Some(Method::Cancel))); + assert!(matches!(cancel, Some(Method::Cancel { .. }))); } #[test] @@ -2136,7 +2139,7 @@ mod tests { assert!(matches!( handle_key(&mut app, ctrl_x()), - Some(Method::Shutdown) + Some(Method::Shutdown { .. }) )); assert!(app.shutdown_confirm.is_none()); } @@ -2466,7 +2469,7 @@ mod tests { } let method = handle_key(&mut app, KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); - assert!(matches!(method, Some(protocol::Method::Compact))); + assert!(matches!(method, Some(protocol::Method::Compact { .. }))); assert!(!app.is_command_mode()); assert_eq!(input_text(&app), ""); assert_eq!(app.queued_input_count(), 0); @@ -2573,7 +2576,7 @@ mod tests { pending_submissions: protocol::PendingSubmissionsSnapshot::default(), entries: vec![], }, - status: WorkerStatus::Idle, + state: WorkerStatus::Idle.into(), in_flight: Default::default(), internal_workers: Vec::new(), }); @@ -2606,7 +2609,7 @@ mod tests { pending_submissions: protocol::PendingSubmissionsSnapshot::default(), entries: vec![], }, - status: WorkerStatus::Idle, + state: WorkerStatus::Idle.into(), in_flight: Default::default(), internal_workers: Vec::new(), }); @@ -2743,8 +2746,8 @@ mod tests { kind: protocol::InternalWorkerKind::SubWorker, }, revision: 1, - event: Box::new(Event::Status { - status: WorkerStatus::Running, + event: Box::new(Event::WorkerState { + snapshot: WorkerStatus::Running.into(), }), }); enter_command_mode(&mut app); @@ -2859,8 +2862,8 @@ mod tests { kind: protocol::InternalWorkerKind::SubWorker, }, revision: 1, - event: Box::new(Event::Status { - status: WorkerStatus::Running, + event: Box::new(Event::WorkerState { + snapshot: WorkerStatus::Running.into(), }), }); @@ -2885,8 +2888,8 @@ mod tests { kind: protocol::InternalWorkerKind::SubWorker, }, revision: 1, - event: Box::new(Event::Status { - status: WorkerStatus::Running, + event: Box::new(Event::WorkerState { + snapshot: WorkerStatus::Running.into(), }), }); handle_key(&mut app, key(KeyCode::Tab)); @@ -2902,7 +2905,7 @@ mod tests { ); assert!(first.is_none()); - assert!(matches!(second, Some(Method::Shutdown))); + assert!(matches!(second, Some(Method::Shutdown { .. }))); assert_eq!(app.worker_status, WorkerStatus::Idle); } @@ -2924,8 +2927,8 @@ mod tests { kind: protocol::InternalWorkerKind::SubWorker, }, revision: 1, - event: Box::new(Event::Status { - status: WorkerStatus::Running, + event: Box::new(Event::WorkerState { + snapshot: WorkerStatus::Running.into(), }), }); diff --git a/crates/worker-runtime/src/execution.rs b/crates/worker-runtime/src/execution.rs index c85e9f29..3c75dc2a 100644 --- a/crates/worker-runtime/src/execution.rs +++ b/crates/worker-runtime/src/execution.rs @@ -15,18 +15,6 @@ use std::fmt; use std::sync::Arc; use workdir::WorkdirSessionHandle; -/// Current execution-side run state for a Worker. -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum WorkerExecutionRunState { - #[default] - Stopped, - Idle, - Busy, - Rejected, - Errored, -} - /// Execution operation that produced a result. #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] @@ -55,7 +43,8 @@ pub struct WorkerSubmissionAck { pub struct WorkerExecutionResult { pub operation: WorkerExecutionOperation, pub outcome: WorkerExecutionOutcome, - pub run_state: WorkerExecutionRunState, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub worker_state: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub message: Option, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -74,22 +63,23 @@ pub enum WorkerExecutionOutcome { } impl WorkerExecutionResult { - pub fn accepted( - operation: WorkerExecutionOperation, - run_state: WorkerExecutionRunState, - ) -> Self { + pub fn accepted(operation: WorkerExecutionOperation) -> Self { Self { operation, outcome: WorkerExecutionOutcome::Accepted, - run_state, + worker_state: None, message: None, submission: None, } } + pub fn with_worker_state(mut self, worker_state: protocol::WorkerStateSnapshot) -> Self { + self.worker_state = Some(worker_state); + self + } + pub fn accepted_submission( operation: WorkerExecutionOperation, - run_state: WorkerExecutionRunState, submission_request_id: impl Into, submission_id: impl Into, disposition: protocol::SubmissionDisposition, @@ -97,7 +87,7 @@ impl WorkerExecutionResult { Self { operation, outcome: WorkerExecutionOutcome::Accepted, - run_state, + worker_state: None, message: None, submission: Some(WorkerSubmissionAck { submission_request_id: submission_request_id.into(), @@ -111,7 +101,7 @@ impl WorkerExecutionResult { Self { operation, outcome: WorkerExecutionOutcome::Busy, - run_state: WorkerExecutionRunState::Busy, + worker_state: None, message: Some(message.into()), submission: None, } @@ -121,7 +111,7 @@ impl WorkerExecutionResult { Self { operation, outcome: WorkerExecutionOutcome::Rejected, - run_state: WorkerExecutionRunState::Stopped, + worker_state: None, message: Some(message.into()), submission: None, } @@ -131,7 +121,7 @@ impl WorkerExecutionResult { Self { operation, outcome: WorkerExecutionOutcome::Errored, - run_state: WorkerExecutionRunState::Errored, + worker_state: None, message: Some(message.into()), submission: None, } @@ -141,7 +131,7 @@ impl WorkerExecutionResult { Self { operation, outcome: WorkerExecutionOutcome::Unsupported, - run_state: WorkerExecutionRunState::Stopped, + worker_state: None, message: Some(message.into()), submission: None, } @@ -280,7 +270,6 @@ pub struct WorkerExecutionRestoreRequest { pub enum WorkerExecutionSpawnResult { Connected { handle: WorkerExecutionHandle, - run_state: WorkerExecutionRunState, working_directory: Option, }, Rejected(WorkerExecutionResult), @@ -290,12 +279,10 @@ pub enum WorkerExecutionSpawnResult { impl WorkerExecutionSpawnResult { pub fn connected( handle: WorkerExecutionHandle, - run_state: WorkerExecutionRunState, working_directory: Option, ) -> Self { Self::Connected { handle, - run_state, working_directory, } } @@ -623,7 +610,6 @@ mod tests { fn submission_ack_survives_json_round_trip() { let result = WorkerExecutionResult::accepted_submission( WorkerExecutionOperation::Input, - WorkerExecutionRunState::Busy, "request-1", "submission-1", protocol::SubmissionDisposition::Started, diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index 6bcc5adf..7762ac86 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -2206,8 +2206,8 @@ mod tests { }; use crate::execution::{ WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, - WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionRunState, - WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, + WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionSpawnRequest, + WorkerExecutionSpawnResult, }; use crate::management::RuntimeOptions; use axum::body::to_bytes; @@ -2979,7 +2979,6 @@ mod tests { fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { WorkerExecutionSpawnResult::Connected { handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), - run_state: WorkerExecutionRunState::Idle, working_directory: request .working_directory .as_ref() @@ -2993,7 +2992,6 @@ mod tests { ) -> WorkerExecutionSpawnResult { WorkerExecutionSpawnResult::Connected { handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), - run_state: WorkerExecutionRunState::Idle, working_directory: request.previous_working_directory, } } @@ -3006,24 +3004,17 @@ mod tests { if let Some(submission_id) = input.submission_request_id { WorkerExecutionResult::accepted_submission( WorkerExecutionOperation::Input, - WorkerExecutionRunState::Idle, submission_id.clone(), submission_id, protocol::SubmissionDisposition::Started, ) } else { - WorkerExecutionResult::accepted( - WorkerExecutionOperation::Input, - WorkerExecutionRunState::Idle, - ) + WorkerExecutionResult::accepted(WorkerExecutionOperation::Input) } } fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { - WorkerExecutionResult::accepted( - WorkerExecutionOperation::Stop, - WorkerExecutionRunState::Stopped, - ) + WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop) } } @@ -3295,8 +3286,7 @@ mod ws_tests { }; use crate::execution::{ WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, - WorkerExecutionResult, WorkerExecutionRunState, WorkerExecutionSpawnRequest, - WorkerExecutionSpawnResult, + WorkerExecutionResult, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, }; use crate::management::RuntimeOptions; use futures::{SinkExt, StreamExt}; @@ -3316,7 +3306,6 @@ mod ws_tests { fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { WorkerExecutionSpawnResult::Connected { handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), - run_state: WorkerExecutionRunState::Idle, working_directory: request .working_directory .as_ref() @@ -3332,16 +3321,12 @@ mod ws_tests { if let Some(submission_id) = input.submission_request_id { WorkerExecutionResult::accepted_submission( WorkerExecutionOperation::Input, - WorkerExecutionRunState::Idle, submission_id.clone(), submission_id, protocol::SubmissionDisposition::Started, ) } else { - WorkerExecutionResult::accepted( - WorkerExecutionOperation::Input, - WorkerExecutionRunState::Idle, - ) + WorkerExecutionResult::accepted(WorkerExecutionOperation::Input) } } @@ -3350,10 +3335,7 @@ mod ws_tests { _handle: &WorkerExecutionHandle, _method: protocol::Method, ) -> WorkerExecutionResult { - WorkerExecutionResult::accepted( - WorkerExecutionOperation::ProtocolMethod, - WorkerExecutionRunState::Idle, - ) + WorkerExecutionResult::accepted(WorkerExecutionOperation::ProtocolMethod) } } @@ -3564,16 +3546,16 @@ mod ws_tests { runtime .observe_worker_event( &other.worker_ref, - protocol::Event::Status { - status: protocol::WorkerStatus::Running, + protocol::Event::WorkerState { + snapshot: protocol::WorkerStatus::Running.into(), }, ) .unwrap(); runtime .observe_worker_event( &worker_ref, - protocol::Event::Status { - status: protocol::WorkerStatus::Running, + protocol::Event::WorkerState { + snapshot: protocol::WorkerStatus::Running.into(), }, ) .unwrap(); diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 9701aae1..808788f4 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -13,8 +13,8 @@ use crate::error::RuntimeError; use crate::execution::WorkerExecutionRestoreRequest; use crate::execution::{ WorkerExecutionBackend, WorkerExecutionBackendRef, WorkerExecutionHandle, - WorkerExecutionOperation, WorkerExecutionResult, WorkerExecutionRunState, - WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, + WorkerExecutionOperation, WorkerExecutionResult, WorkerExecutionSpawnRequest, + WorkerExecutionSpawnResult, }; #[cfg(feature = "fs-store")] use crate::fs_store::{ @@ -725,12 +725,11 @@ impl Runtime { }; let spawn_result = backend.spawn_worker(spawn_request); - let (handle, run_state, working_directory) = match spawn_result { + let (handle, working_directory) = match spawn_result { WorkerExecutionSpawnResult::Connected { handle, - run_state, working_directory, - } => (handle, run_state, working_directory), + } => (handle, working_directory), WorkerExecutionSpawnResult::Rejected(result) | WorkerExecutionSpawnResult::Errored(result) => { self.rollback_failed_create(&worker_ref)?; @@ -785,11 +784,10 @@ impl Runtime { result, }); } - let initial_run_state = dispatch_result.run_state; let detail = self.commit_created_worker( &worker_ref, handle, - initial_run_state, + WorkerStatus::Running, working_directory, dispatch_result, )?; @@ -799,9 +797,9 @@ impl Runtime { self.commit_created_worker( &worker_ref, handle, - run_state, + WorkerStatus::Idle, working_directory, - WorkerExecutionResult::accepted(WorkerExecutionOperation::Spawn, run_state), + WorkerExecutionResult::accepted(WorkerExecutionOperation::Spawn), ) } } @@ -1086,13 +1084,12 @@ impl Runtime { match backend.restore_worker(request) { WorkerExecutionSpawnResult::Connected { handle, - run_state, working_directory, } => { self.commit_restored_worker_execution( worker_ref, handle, - run_state, + WorkerStatus::Idle, working_directory, )?; self.worker_detail(worker_ref) @@ -1222,7 +1219,19 @@ impl Runtime { let mut state = self.lock()?; state.ensure_running()?; let worker = state.worker_mut(worker_ref)?; - worker.status = worker_status_from_run_state(dispatch_result.run_state); + if let Some(snapshot) = dispatch_result.worker_state.as_ref() { + worker.status = match snapshot.catalog_status() { + protocol::WorkerStatus::Idle => WorkerStatus::Idle, + protocol::WorkerStatus::Running => WorkerStatus::Running, + protocol::WorkerStatus::Paused => WorkerStatus::Paused, + protocol::WorkerStatus::Stopped => WorkerStatus::Stopped, + }; + } else if matches!( + submission.as_ref().map(|ack| ack.disposition), + Some(protocol::SubmissionDisposition::Started) + ) { + worker.status = WorkerStatus::Running; + } let status = worker.status; #[cfg(feature = "ws-server")] if let Some(payload) = input_protocol_event(&input) { @@ -1431,7 +1440,7 @@ impl Runtime { let entries = self.worker_completions(worker_ref, kind, &prefix)?; return Ok(vec![Event::Completions { kind, entries }]); } - if matches!(&method, Method::Shutdown) { + if matches!(&method, Method::Shutdown { .. }) { self.stop_worker(worker_ref, Some("worker protocol shutdown".to_string()))?; return Ok(Vec::new()); } @@ -1481,7 +1490,7 @@ impl Runtime { &self, worker_ref: &WorkerRef, handle: WorkerExecutionHandle, - run_state: WorkerExecutionRunState, + status: WorkerStatus, working_directory: Option, _result: WorkerExecutionResult, ) -> Result { @@ -1490,7 +1499,7 @@ impl Runtime { let worker = state.worker_mut(worker_ref)?; worker.execution_handle = Some(handle); worker.execution_bound = true; - worker.status = worker_status_from_run_state(run_state); + worker.status = status; worker.restore_intent = restore_intent_for_status(worker.status); worker.working_directory = working_directory; worker.detail() @@ -1518,16 +1527,28 @@ impl Runtime { worker_ref: &WorkerRef, result: WorkerExecutionResult, ) -> Result<(), RuntimeError> { - let mut state = self.lock()?; - if result.is_accepted() { - let status = worker_status_from_run_state(result.run_state); - let worker = state.worker_mut(worker_ref)?; - worker.status = status; - worker.restore_intent = restore_intent_for_status(status); - state.publish_worker_upsert(worker_ref.worker_id)?; - state.persist_runtime_snapshot()?; - state.persist_worker(&worker_ref.worker_id)?; + // Accepted dispatch without a state snapshot is transport evidence only; + // the revisioned protocol stream remains live authority. Test/detached + // backends may return an exact full snapshot as their acknowledgement. + if !result.is_accepted() { + return Ok(()); } + let Some(snapshot) = result.worker_state else { + return Ok(()); + }; + let status = match snapshot.catalog_status() { + protocol::WorkerStatus::Idle => WorkerStatus::Idle, + protocol::WorkerStatus::Running => WorkerStatus::Running, + protocol::WorkerStatus::Paused => WorkerStatus::Paused, + protocol::WorkerStatus::Stopped => WorkerStatus::Stopped, + }; + let mut state = self.lock()?; + let worker = state.worker_mut(worker_ref)?; + worker.status = status; + worker.restore_intent = restore_intent_for_status(status); + state.publish_worker_upsert(worker_ref.worker_id)?; + state.persist_runtime_snapshot()?; + state.persist_worker(&worker_ref.worker_id)?; Ok(()) } @@ -1730,7 +1751,7 @@ impl Runtime { context_window: 0, context_tokens: 0, }, - status: protocol::WorkerStatus::Idle, + state: protocol::WorkerStateSnapshot::initial(1), in_flight: protocol::InFlightSnapshot { blocks: Vec::new(), commands: Vec::new(), @@ -1968,12 +1989,11 @@ impl Runtime { match backend.restore_worker(request) { WorkerExecutionSpawnResult::Connected { handle, - run_state, working_directory, } => self.commit_restored_worker_execution( &candidate.worker_ref, handle, - run_state, + WorkerStatus::Idle, working_directory, )?, WorkerExecutionSpawnResult::Rejected(result) @@ -1990,7 +2010,7 @@ impl Runtime { &self, worker_ref: &WorkerRef, handle: WorkerExecutionHandle, - run_state: WorkerExecutionRunState, + status: WorkerStatus, working_directory: Option, ) -> Result<(), RuntimeError> { let mut state = self.lock()?; @@ -1999,7 +2019,7 @@ impl Runtime { let worker = state.worker_mut(worker_ref)?; worker.execution_handle = Some(handle); worker.execution_bound = true; - worker.status = worker_status_from_run_state(run_state); + worker.status = status; worker.restore_intent = restore_intent_for_status(worker.status); worker.working_directory = working_directory; } @@ -2867,7 +2887,7 @@ impl RuntimeState { ) { match event { protocol::Event::Snapshot { - status, + state, internal_workers, .. } => { @@ -2875,7 +2895,7 @@ impl RuntimeState { statuses.insert( worker.session_id.clone(), InternalWorkerActivity { - status: *status, + status: state.catalog_status(), parent_session_id: worker.parent_session_id.clone(), }, ); @@ -2888,26 +2908,17 @@ impl RuntimeState { event, .. } => Self::project_internal_worker_event(statuses, nested_worker, event), - protocol::Event::Status { status } => { - statuses.insert( - worker.session_id.clone(), - InternalWorkerActivity { - status: *status, - parent_session_id: worker.parent_session_id.clone(), + protocol::Event::WorkerState { snapshot } + | protocol::Event::CommandAcknowledged { + acknowledgement: + protocol::WorkerCommandAcknowledgement { + state: snapshot, .. }, - ); - } - protocol::Event::RunEnd { result } => { - let status = match result { - protocol::RunResult::Paused => protocol::WorkerStatus::Paused, - protocol::RunResult::Finished - | protocol::RunResult::LimitReached - | protocol::RunResult::RolledBack => protocol::WorkerStatus::Idle, - }; + } => { statuses.insert( worker.session_id.clone(), InternalWorkerActivity { - status, + status: snapshot.catalog_status(), parent_session_id: worker.parent_session_id.clone(), }, ); @@ -2963,28 +2974,21 @@ impl RuntimeState { return false; }; let next_status = match event { - protocol::Event::Status { - status: protocol::WorkerStatus::Running, - } => Some(WorkerStatus::Running), - protocol::Event::Status { - status: protocol::WorkerStatus::Idle, - } => Some(WorkerStatus::Idle), - protocol::Event::Status { - status: protocol::WorkerStatus::Paused, - } => Some(WorkerStatus::Paused), - protocol::Event::Snapshot { status, .. } => match status { - protocol::WorkerStatus::Running => Some(WorkerStatus::Running), - protocol::WorkerStatus::Idle => Some(WorkerStatus::Idle), - protocol::WorkerStatus::Paused => Some(WorkerStatus::Paused), - protocol::WorkerStatus::Stopped => Some(WorkerStatus::Stopped), - }, - protocol::Event::RunEnd { result } => match result { - protocol::RunResult::Finished | protocol::RunResult::RolledBack => { - Some(WorkerStatus::Idle) - } - protocol::RunResult::Paused => Some(WorkerStatus::Paused), - protocol::RunResult::LimitReached => Some(WorkerStatus::Idle), - }, + protocol::Event::WorkerState { snapshot } + | protocol::Event::Snapshot { + state: snapshot, .. + } + | protocol::Event::CommandAcknowledged { + acknowledgement: + protocol::WorkerCommandAcknowledgement { + state: snapshot, .. + }, + } => Some(match snapshot.catalog_status() { + protocol::WorkerStatus::Idle => WorkerStatus::Idle, + protocol::WorkerStatus::Running => WorkerStatus::Running, + protocol::WorkerStatus::Paused => WorkerStatus::Paused, + protocol::WorkerStatus::Stopped => WorkerStatus::Stopped, + }), _ => None, }; if let Some(next_status) = next_status { @@ -3081,16 +3085,6 @@ fn restore_intent_for_status(status: WorkerStatus) -> WorkerRestoreIntent { } } -fn worker_status_from_run_state(run_state: WorkerExecutionRunState) -> WorkerStatus { - match run_state { - WorkerExecutionRunState::Idle => WorkerStatus::Idle, - WorkerExecutionRunState::Busy => WorkerStatus::Running, - WorkerExecutionRunState::Stopped - | WorkerExecutionRunState::Rejected - | WorkerExecutionRunState::Errored => WorkerStatus::Stopped, - } -} - fn repository_resource_error(error: BackendResourceError) -> RuntimeError { let (code, message) = match error { BackendResourceError::Expired => ( @@ -3304,7 +3298,7 @@ mod tests { }; use crate::execution::{ WorkerExecutionBackend, WorkerExecutionContext, WorkerExecutionHandle, - WorkerExecutionRestoreRequest, WorkerExecutionRunState, + WorkerExecutionRestoreRequest, }; use crate::working_directory::WorkingDirectoryDiagnostic; use async_trait::async_trait; @@ -3313,6 +3307,14 @@ mod tests { use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; + fn test_command() -> protocol::WorkerCommandEnvelope { + protocol::WorkerCommandEnvelope { + command_id: 1, + expected_execution_generation: 1, + expected_worker_state_revision: 0, + } + } + #[test] fn repository_resource_failures_keep_typed_credential_diagnostics() { let cases = [ @@ -3359,7 +3361,9 @@ mod tests { protocol::Event::InternalWorker { worker, revision: 1, - event: Box::new(protocol::Event::Status { status }), + event: Box::new(protocol::Event::WorkerState { + snapshot: status.into(), + }), } } @@ -3452,7 +3456,7 @@ mod tests { context_window: 0, context_tokens: 0, }, - status: protocol::WorkerStatus::Idle, + state: protocol::WorkerStatus::Idle.into(), in_flight: protocol::InFlightSnapshot::default(), internal_workers: Vec::new(), }; @@ -3967,7 +3971,6 @@ mod tests { .insert(request.worker_ref.worker_id.clone(), request.context); WorkerExecutionSpawnResult::Connected { handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), - run_state: WorkerExecutionRunState::Idle, working_directory: request .working_directory .as_ref() @@ -3997,7 +4000,6 @@ mod tests { .insert(request.worker_ref.worker_id.clone(), request.context); WorkerExecutionSpawnResult::Connected { handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), - run_state: WorkerExecutionRunState::Idle, working_directory: request .working_directory .as_ref() @@ -4020,7 +4022,6 @@ mod tests { .unwrap_or_else(|| { WorkerExecutionResult::accepted_submission( WorkerExecutionOperation::Input, - WorkerExecutionRunState::Idle, "request-test", "test-submission", protocol::SubmissionDisposition::Started, @@ -4038,17 +4039,11 @@ mod tests { } fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { - WorkerExecutionResult::accepted( - WorkerExecutionOperation::Stop, - WorkerExecutionRunState::Stopped, - ) + WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop) } fn cancel_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { - WorkerExecutionResult::accepted( - WorkerExecutionOperation::Cancel, - WorkerExecutionRunState::Stopped, - ) + WorkerExecutionResult::accepted(WorkerExecutionOperation::Cancel) } #[cfg(feature = "ws-server")] @@ -4374,7 +4369,9 @@ mod tests { .send_protocol_method_scoped( &scope("workspace-a", "server-a"), &workspace_b.worker_ref, - Method::Shutdown, + Method::Shutdown { + command: test_command(), + }, ) .unwrap_err(); assert!(matches!( @@ -4722,11 +4719,10 @@ mod tests { } #[test] - fn create_worker_uses_committed_input_ack_run_state() { + fn create_worker_uses_started_submission_ack_for_initial_running_status() { let (runtime, backend) = runtime_and_backend(); backend.set_dispatch_result(WorkerExecutionResult::accepted_submission( WorkerExecutionOperation::Input, - WorkerExecutionRunState::Idle, "request-test", "test-submission", protocol::SubmissionDisposition::Started, @@ -4736,7 +4732,7 @@ mod tests { let detail = runtime.create_worker(request).unwrap(); - assert_eq!(detail.status, WorkerStatus::Idle); + assert_eq!(detail.status, WorkerStatus::Running); } #[test] @@ -4745,7 +4741,6 @@ mod tests { backend.preserve_commit_ack_submission_id(); backend.set_dispatch_result(WorkerExecutionResult::accepted_submission( WorkerExecutionOperation::Input, - WorkerExecutionRunState::Busy, "request-test", "forged-submission", protocol::SubmissionDisposition::Started, @@ -4770,7 +4765,6 @@ mod tests { let (runtime, backend) = runtime_and_backend(); backend.set_dispatch_result(WorkerExecutionResult::accepted( WorkerExecutionOperation::Input, - WorkerExecutionRunState::Busy, )); let mut request = task_request("missing initial input commit ack"); request.initial_input = Some(WorkerInput::user("start the ticket")); @@ -4898,7 +4892,7 @@ mod tests { context_window: 128, context_tokens: 64, }, - status: protocol::WorkerStatus::Running, + state: protocol::WorkerStatus::Running.into(), in_flight: protocol::InFlightSnapshot { blocks: Vec::new(), commands: Vec::new(), @@ -4914,13 +4908,13 @@ mod tests { protocol::Event::Snapshot { session, greeting, - status, + state, .. } => { assert_eq!(session.entries.len(), 1); assert_eq!(session.entries[0].entry_id, "restored-log-entry"); assert_eq!(greeting.worker_name, "live-worker"); - assert_eq!(status, protocol::WorkerStatus::Running); + assert_eq!(state.catalog_status(), protocol::WorkerStatus::Running); } other => panic!("expected snapshot, got {other:?}"), } @@ -4936,7 +4930,6 @@ mod tests { fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { WorkerExecutionSpawnResult::Connected { handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), - run_state: WorkerExecutionRunState::Idle, working_directory: request .working_directory .as_ref() @@ -4951,7 +4944,6 @@ mod tests { ) -> WorkerExecutionResult { WorkerExecutionResult::accepted_submission( WorkerExecutionOperation::Input, - WorkerExecutionRunState::Idle, "request-test", input .submission_request_id @@ -4993,7 +4985,12 @@ mod tests { .unwrap(); runtime - .send_protocol_method(&detail.worker_ref, Method::Shutdown) + .send_protocol_method( + &detail.worker_ref, + Method::Shutdown { + command: test_command(), + }, + ) .unwrap(); assert_eq!( @@ -5009,7 +5006,12 @@ mod tests { .create_worker(task_request("restore explicitly")) .unwrap(); runtime - .send_protocol_method(&detail.worker_ref, Method::Shutdown) + .send_protocol_method( + &detail.worker_ref, + Method::Shutdown { + command: test_command(), + }, + ) .unwrap(); assert!(matches!( @@ -5027,7 +5029,7 @@ mod tests { assert_eq!(*backend.run_generations.lock().unwrap(), vec![1, 2]); assert_eq!( runtime.worker_detail(&detail.worker_ref).unwrap().status, - WorkerStatus::Idle + WorkerStatus::Running ); } diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 5225dd8b..6c414874 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -10,8 +10,8 @@ use std::collections::HashMap; use std::future::Future; use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, Mutex, mpsc}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, RwLock, mpsc}; use std::time::Duration; use crate::auth::{ @@ -25,8 +25,8 @@ use crate::catalog::{ }; use crate::execution::{ WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, - WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionRunState, - WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, + WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionSpawnRequest, + WorkerExecutionSpawnResult, }; use crate::identity::WorkerRef; use crate::interaction::{WorkerInput, WorkerInputKind}; @@ -38,7 +38,26 @@ use crate::working_directory::{ WorkingDirectoryBinding, WorkingDirectoryDiagnostic, WorkingDirectoryMaterializer, }; use async_trait::async_trait; -use protocol::{ErrorCode, Event, Method, Segment, WorkerStatus}; +use protocol::{Event, Method, Segment, WorkerCommandEnvelope, WorkerStatus}; + +static NEXT_INTERNAL_COMMAND_ID: AtomicU64 = AtomicU64::new(1); + +fn next_internal_command( + state: &RwLock, +) -> Result { + let snapshot = state + .read() + .map_err(|_| "worker state lock is poisoned".to_string())? + .clone(); + let floor = snapshot.last_command_id.saturating_add(1); + let command_id = NEXT_INTERNAL_COMMAND_ID + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + Some(current.max(floor).saturating_add(1)) + }) + .unwrap_or(floor) + .max(floor); + Ok(WorkerCommandEnvelope::for_snapshot(command_id, &snapshot)) +} use session_store::{CombinedStore, WorkerAggregateStore, WorkerSessionStore}; #[cfg(test)] use session_store::{FsStore, FsWorkerStore}; @@ -172,7 +191,7 @@ impl WorkerObservationProvider for RuntimeGrantedWorkerObservationProvider { }, display_name: grant.worker_id.clone(), relation: "granted_peer".to_string(), - status: format!("{:?}", state.get_status()).to_lowercase(), + status: format!("{:?}", state.catalog_status()).to_lowercase(), }); } subjects.sort_by(|left, right| left.subject.cmp(&right.subject)); @@ -1174,10 +1193,12 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { } } +#[derive(Clone)] struct RuntimeWorkerExecution { handle: WorkerHandle, shutdown: Arc>>, busy: Arc, + worker_state: Arc>, workspace_client: Option>, } @@ -1276,6 +1297,7 @@ where ( WorkerHandle, Arc, + Arc>, Option>, ), WorkerExecutionResult, @@ -1302,6 +1324,7 @@ where ( execution.handle.clone(), execution.busy.clone(), + execution.worker_state.clone(), execution.workspace_client.clone(), ) }) @@ -1318,7 +1341,6 @@ where operation: WorkerExecutionOperation, worker: WorkerHandle, method: Method, - accepted_run_state: WorkerExecutionRunState, ) -> WorkerExecutionResult { self.run_on_adapter_runtime(async move { worker @@ -1326,7 +1348,7 @@ where .await .map_err(|err| format!("failed to send Worker method: {err}")) }) - .map(|_| WorkerExecutionResult::accepted(operation, accepted_run_state)) + .map(|_| WorkerExecutionResult::accepted(operation)) .unwrap_or_else(|message| WorkerExecutionResult::errored(operation, message)) } @@ -1336,7 +1358,6 @@ where worker: WorkerHandle, method: Method, submission_request_id: String, - accepted_run_state: WorkerExecutionRunState, ) -> WorkerExecutionResult { let request_id = submission_request_id.clone(); self.run_on_adapter_runtime(async move { @@ -1395,7 +1416,6 @@ where .map(|(submission_id, disposition)| { WorkerExecutionResult::accepted_submission( operation, - accepted_run_state, submission_request_id, submission_id, disposition, @@ -1415,38 +1435,45 @@ where workspace_client: Option>, ) -> WorkerExecutionSpawnResult { let busy = Arc::new(AtomicBool::new(false)); + let worker_state = Arc::new(RwLock::new(handle.shared_state.snapshot())); #[cfg(feature = "ws-server")] { let streams = subscribe_worker_protocol_session(&handle); let mut events = streams.events; let mut entry_events = streams.log_entries; let bridge_busy = busy.clone(); + let bridge_worker_state = worker_state.clone(); if let Err(message) = self.spawn_on_adapter_runtime(async move { loop { tokio::select! { event = events.recv() => { match event { Ok(event) => { - let next_busy = match &event { - Event::InvokeStart { .. } - | Event::Status { - status: WorkerStatus::Running, - } => Some(true), - Event::RunEnd { .. } - | Event::Error { - code: ErrorCode::NotPaused, - .. + let next_state = match &event { + Event::WorkerState { snapshot } + | Event::Snapshot { state: snapshot, .. } => { + Some(snapshot.clone()) } - | Event::Status { - status: - WorkerStatus::Idle - | WorkerStatus::Paused - | WorkerStatus::Stopped, + Event::CommandAcknowledged { acknowledgement } => { + Some(acknowledgement.state.clone()) } - | Event::Shutdown => Some(false), _ => None, }; + let next_busy = next_state + .as_ref() + .map(worker_state_is_executing) + .or_else(|| matches!(event, Event::Shutdown).then_some(false)); let _ = bridge_context.publish_protocol_event(event); + if let Some(next_state) = next_state { + if let Ok(mut current) = bridge_worker_state.write() { + if next_state.execution_generation > current.execution_generation + || (next_state.execution_generation == current.execution_generation + && next_state.revision >= current.revision) + { + *current = next_state; + } + } + } if let Some(next_busy) = next_busy { bridge_busy.store(next_busy, Ordering::SeqCst); } @@ -1494,13 +1521,13 @@ where handle, shutdown, busy, + worker_state, workspace_client, }, ); WorkerExecutionSpawnResult::Connected { handle: WorkerExecutionHandle::new(worker_ref, self.backend_id()), - run_state: WorkerExecutionRunState::Idle, working_directory: working_directory.map(|binding| binding.status()), } } @@ -1516,6 +1543,17 @@ impl Drop for WorkerRuntimeExecutionBackend { } } +fn worker_state_is_executing(snapshot: &protocol::WorkerStateSnapshot) -> bool { + matches!( + snapshot.state, + protocol::WorkerState::Busy(protocol::WorkerBusyState::Run( + protocol::WorkerRunState::Running + | protocol::WorkerRunState::Pausing + | protocol::WorkerRunState::Cancelling + )) | protocol::WorkerState::Busy(protocol::WorkerBusyState::Maintenance(_)) + ) +} + fn method_starts_turn(method: &Method) -> bool { matches!( method, @@ -1523,41 +1561,17 @@ fn method_starts_turn(method: &Method) -> bool { | Method::SubmitTracked { .. } | Method::Notify { auto_run: true, .. } | Method::NotifyTracked { auto_run: true, .. } - | Method::Resume - | Method::Compact + | Method::Resume { .. } ) } fn method_can_start_turn_from_status(method: &Method, status: WorkerStatus) -> bool { match method { - Method::Resume => matches!(status, WorkerStatus::Idle | WorkerStatus::Paused), + Method::Resume { .. } => matches!(status, WorkerStatus::Idle | WorkerStatus::Paused), _ => status == WorkerStatus::Idle, } } -fn accepted_notify_run_state(status: WorkerStatus, auto_run: bool) -> WorkerExecutionRunState { - match status { - WorkerStatus::Running => WorkerExecutionRunState::Busy, - WorkerStatus::Idle if auto_run => WorkerExecutionRunState::Busy, - WorkerStatus::Idle | WorkerStatus::Paused | WorkerStatus::Stopped => { - WorkerExecutionRunState::Idle - } - } -} - -fn accepted_run_state_for_method(method: &Method) -> WorkerExecutionRunState { - match method { - Method::Submit { .. } - | Method::SubmitTracked { .. } - | Method::Notify { auto_run: true, .. } - | Method::NotifyTracked { auto_run: true, .. } - | Method::Resume - | Method::Compact => WorkerExecutionRunState::Busy, - Method::Shutdown => WorkerExecutionRunState::Stopped, - _ => WorkerExecutionRunState::Idle, - } -} - impl WorkerExecutionBackend for WorkerRuntimeExecutionBackend where F: RuntimeWorkerFactory, @@ -1883,7 +1897,7 @@ where handle: &WorkerExecutionHandle, input: WorkerInput, ) -> WorkerExecutionResult { - let (worker, busy, _workspace_client) = match self.get_execution(handle) { + let (worker, busy, worker_state, _workspace_client) = match self.get_execution(handle) { Ok(execution) => execution, Err(mut result) => { result.operation = WorkerExecutionOperation::Input; @@ -1892,8 +1906,7 @@ where }; if input.kind == WorkerInputKind::Notify { - let status = worker.shared_state.get_status(); - let accepted_run_state = accepted_notify_run_state(status, true); + let status = worker.shared_state.catalog_status(); let claimed_here = status == WorkerStatus::Idle && busy .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) @@ -1912,7 +1925,6 @@ where operation_id: notification_request_id, }, }, - accepted_run_state, ); if claimed_here && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted { @@ -1921,8 +1933,22 @@ where return result; } + if input.kind == WorkerInputKind::Compact { + let command = match next_internal_command(&worker_state) { + Ok(command) => command, + Err(error) => { + return WorkerExecutionResult::errored(WorkerExecutionOperation::Input, error); + } + }; + return self.send_method( + WorkerExecutionOperation::Input, + worker, + Method::Compact { command }, + ); + } + let is_user_submit = input.kind == WorkerInputKind::User; - let status = worker.shared_state.get_status(); + let status = worker.shared_state.catalog_status(); let claimed_here = status == WorkerStatus::Idle && busy .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) @@ -1962,7 +1988,7 @@ where WorkerInputKind::Notify => { unreachable!("Notify input is dispatched before the turn-start busy guard") } - WorkerInputKind::Compact => (Method::Compact, None), + WorkerInputKind::Compact => unreachable!("compact input is dispatched above"), WorkerInputKind::ListRewindTargets => (Method::ListRewindTargets, None), WorkerInputKind::RegisterPeer => ( Method::RegisterPeer { @@ -1971,15 +1997,6 @@ where None, ), }; - let accepted_run_state = match method { - Method::Submit { .. } - | Method::SubmitTracked { .. } - | Method::Notify { .. } - | Method::NotifyTracked { .. } - | Method::Compact => WorkerExecutionRunState::Busy, - _ => WorkerExecutionRunState::Idle, - }; - let accepted_is_idle = accepted_run_state == WorkerExecutionRunState::Idle; let waits_for_submission_acceptance = submission_request_id.is_some(); let result = if waits_for_submission_acceptance { @@ -1988,20 +2005,11 @@ where worker, method, submission_request_id.expect("Submit must have a submission request id"), - accepted_run_state, ) } else { - self.send_method( - WorkerExecutionOperation::Input, - worker, - method, - accepted_run_state, - ) + self.send_method(WorkerExecutionOperation::Input, worker, method) }; - if accepted_is_idle - || (claimed_here - && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted) - { + if claimed_here && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted { busy.store(false, Ordering::SeqCst); } result @@ -2015,7 +2023,7 @@ where content: &[u8], context: Option<&session_store::UploadedFileUploadContext>, ) -> Result { - let (worker, _, _) = self.get_execution(handle).map_err(|mut result| { + let (worker, _, _, _) = self.get_execution(handle).map_err(|mut result| { result.operation = WorkerExecutionOperation::UploadFile; result })?; @@ -2038,7 +2046,7 @@ where handle: &WorkerExecutionHandle, artifact_id: &str, ) -> WorkerExecutionResult { - let (worker, _, _) = match self.get_execution(handle) { + let (worker, _, _, _) = match self.get_execution(handle) { Ok(execution) => execution, Err(mut result) => { result.operation = WorkerExecutionOperation::DeleteUploadedFile; @@ -2046,10 +2054,7 @@ where } }; match worker.delete_uploaded_file(artifact_id) { - Ok(_) => WorkerExecutionResult::accepted( - WorkerExecutionOperation::DeleteUploadedFile, - WorkerExecutionRunState::Idle, - ), + Ok(_) => WorkerExecutionResult::accepted(WorkerExecutionOperation::DeleteUploadedFile), Err(error) => WorkerExecutionResult::rejected( WorkerExecutionOperation::DeleteUploadedFile, format!("uploaded_file_delete_rejected: {error}"), @@ -2062,7 +2067,7 @@ where handle: &WorkerExecutionHandle, method: Method, ) -> WorkerExecutionResult { - let (worker, busy, _workspace_client) = match self.get_execution(handle) { + let (worker, busy, _worker_state, _workspace_client) = match self.get_execution(handle) { Ok(execution) => execution, Err(mut result) => { result.operation = WorkerExecutionOperation::ProtocolMethod; @@ -2076,19 +2081,13 @@ where } _ => None, } { - let status = worker.shared_state.get_status(); - let accepted_run_state = accepted_notify_run_state(status, auto_run); + let status = worker.shared_state.catalog_status(); let claimed_here = status == WorkerStatus::Idle && auto_run && busy .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) .is_ok(); - let result = self.send_method( - WorkerExecutionOperation::ProtocolMethod, - worker, - method, - accepted_run_state, - ); + let result = self.send_method(WorkerExecutionOperation::ProtocolMethod, worker, method); if claimed_here && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted { busy.store(false, Ordering::SeqCst); @@ -2098,7 +2097,7 @@ where let starts_turn = method_starts_turn(&method); if starts_turn - && (!method_can_start_turn_from_status(&method, worker.shared_state.get_status()) + && (!method_can_start_turn_from_status(&method, worker.shared_state.catalog_status()) || busy .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) .is_err()) @@ -2109,17 +2108,8 @@ where ); } - let accepted_run_state = accepted_run_state_for_method(&method); - let accepted_is_idle = accepted_run_state == WorkerExecutionRunState::Idle; - let result = self.send_method( - WorkerExecutionOperation::ProtocolMethod, - worker, - method, - accepted_run_state, - ); - if (starts_turn && accepted_is_idle) - || (starts_turn && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted) - { + let result = self.send_method(WorkerExecutionOperation::ProtocolMethod, worker, method); + if starts_turn && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted { busy.store(false, Ordering::SeqCst); } result @@ -2137,7 +2127,7 @@ where ); } let execution = match self.workers.lock() { - Ok(mut workers) => workers.remove(handle.worker_ref()), + Ok(workers) => workers.get(handle.worker_ref()).cloned(), Err(_) => { return WorkerExecutionResult::errored( WorkerExecutionOperation::Stop, @@ -2153,48 +2143,73 @@ where }; let artifact_cleanup = execution.handle.clone(); let shutdown = execution.shutdown.clone(); + let command = match next_internal_command(&execution.worker_state) { + Ok(command) => command, + Err(error) => { + return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, error); + } + }; let result = self.send_method( WorkerExecutionOperation::Stop, - execution.handle, - Method::Shutdown, - WorkerExecutionRunState::Stopped, + execution.handle.clone(), + Method::Shutdown { command }, ); if result.outcome != crate::execution::WorkerExecutionOutcome::Accepted { return result; } - match self.run_on_adapter_runtime(async move { - let receiver = shutdown.lock().await.take(); - if let Some(receiver) = receiver { - receiver - .await - .map_err(|_| "Worker shutdown completion channel closed".to_string())?; + let shutdown_wait = self.run_on_adapter_runtime(async move { + let mut guard = shutdown.lock().await; + let Some(mut receiver) = guard.take() else { + return Ok(()); + }; + match tokio::time::timeout(Duration::from_secs(5), &mut receiver).await { + Ok(Ok(())) => Ok(()), + Ok(Err(_)) => Err("Worker shutdown completion channel closed".to_string()), + Err(_) => { + *guard = Some(receiver); + Err("Worker shutdown confirmation timed out; stop remains retryable".into()) + } } - Ok(()) - }) { - Ok(()) => match artifact_cleanup.delete_uncommitted_uploaded_files() { - Ok(_) => result, - Err(error) => WorkerExecutionResult::errored( - WorkerExecutionOperation::Stop, - format!("uploaded_file_cleanup_failed: {error}"), - ), - }, - Err(message) => WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message), + }); + if let Err(message) = shutdown_wait { + return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message); + } + if let Err(error) = artifact_cleanup.delete_uncommitted_uploaded_files() { + return WorkerExecutionResult::errored( + WorkerExecutionOperation::Stop, + format!("uploaded_file_cleanup_failed: {error}"), + ); + } + match self.workers.lock() { + Ok(mut workers) => { + workers.remove(handle.worker_ref()); + result + } + Err(_) => WorkerExecutionResult::errored( + WorkerExecutionOperation::Stop, + "worker adapter registry lock is poisoned after shutdown", + ), } } fn cancel_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult { - let (worker, _busy, _workspace_client) = match self.get_execution(handle) { + let (worker, _busy, worker_state, _workspace_client) = match self.get_execution(handle) { Ok(execution) => execution, Err(mut result) => { result.operation = WorkerExecutionOperation::Cancel; return result; } }; + let command = match next_internal_command(&worker_state) { + Ok(command) => command, + Err(error) => { + return WorkerExecutionResult::errored(WorkerExecutionOperation::Cancel, error); + } + }; self.send_method( WorkerExecutionOperation::Cancel, worker, - Method::Cancel, - WorkerExecutionRunState::Idle, + Method::Cancel { command }, ) } @@ -2259,6 +2274,29 @@ mod tests { use manifest::{Scope, WorkerManifest}; use session_store::{LogEntry, WorkerMetadataStore}; + fn test_command() -> WorkerCommandEnvelope { + WorkerCommandEnvelope { + command_id: 1, + expected_execution_generation: 1, + expected_worker_state_revision: 0, + } + } + + fn adapter_command( + backend: &WorkerRuntimeExecutionBackend, + worker_ref: &WorkerRef, + ) -> WorkerCommandEnvelope { + let workers = backend.workers.lock().unwrap(); + let state = workers + .get(worker_ref) + .expect("worker execution") + .worker_state + .read() + .unwrap() + .clone(); + WorkerCommandEnvelope::for_snapshot(state.last_command_id.saturating_add(1), &state) + } + #[test] fn workspace_prompt_projection_notification_advances_shared_cache() { let cache = WorkspacePromptProjectionCache::default(); @@ -2406,41 +2444,39 @@ mod tests { } #[test] - fn notify_run_state_allows_running_worker_inbox_delivery() { - assert_eq!( - accepted_notify_run_state(WorkerStatus::Running, true), - WorkerExecutionRunState::Busy - ); - assert_eq!( - accepted_notify_run_state(WorkerStatus::Idle, true), - WorkerExecutionRunState::Busy - ); - assert_eq!( - accepted_notify_run_state(WorkerStatus::Idle, false), - WorkerExecutionRunState::Idle - ); - assert_eq!( - accepted_notify_run_state(WorkerStatus::Paused, true), - WorkerExecutionRunState::Idle - ); + fn compact_is_maintenance_not_a_turn_start() { + assert!(!method_starts_turn(&Method::Compact { + command: test_command(), + })); + assert!(method_starts_turn(&Method::Resume { + command: test_command(), + })); } #[test] fn resume_turn_claim_accepts_paused_and_idle_but_not_running_status() { assert!(method_can_start_turn_from_status( - &Method::Resume, + &Method::Resume { + command: test_command() + }, WorkerStatus::Paused )); assert!(method_can_start_turn_from_status( - &Method::Resume, + &Method::Resume { + command: test_command() + }, WorkerStatus::Idle )); assert!(!method_can_start_turn_from_status( - &Method::Resume, + &Method::Resume { + command: test_command() + }, WorkerStatus::Running )); assert!(!method_can_start_turn_from_status( - &Method::Compact, + &Method::Compact { + command: test_command() + }, WorkerStatus::Paused )); } @@ -2656,19 +2692,22 @@ mod tests { let observed = { let workers = backend.workers.lock().unwrap(); let execution = workers.get(worker_ref).expect("live Worker execution"); + let projected = execution.worker_state.read().unwrap().catalog_status(); ( - execution.handle.shared_state.get_status(), + execution.handle.shared_state.catalog_status(), + projected, execution.busy.load(Ordering::SeqCst), ) }; - if observed == (expected_status, expected_busy) { + if observed == (expected_status, expected_status, expected_busy) { return; } assert!( std::time::Instant::now() < deadline, - "timed out waiting for adapter state {expected_status:?}, busy={expected_busy}; last observed status={:?}, busy={}", + "timed out waiting for adapter state {expected_status:?}, busy={expected_busy}; last observed controller={:?}, projected={:?}, busy={}", observed.0, observed.1, + observed.2, ); std::thread::sleep(Duration::from_millis(10)); } @@ -3169,13 +3208,19 @@ mod tests { .expect("in-process restore must not bind the overlong Unix socket path"); assert_eq!( - controller.handle.shared_state.get_status(), + controller.handle.shared_state.catalog_status(), WorkerStatus::Idle ); assert!(!socket_path.exists()); assert!(run_dir.join("worker.out.log").is_file()); assert!(run_dir.join("worker.err.log").is_file()); - controller.handle.send(Method::Shutdown).await.unwrap(); + controller + .handle + .send(Method::Shutdown { + command: test_command(), + }) + .await + .unwrap(); if let Some(receiver) = controller.shutdown.lock().await.take() { receiver.await.unwrap(); } @@ -3289,7 +3334,9 @@ mod tests { backend .run_on_adapter_runtime(async move { handle - .send(Method::Shutdown) + .send(Method::Shutdown { + command: test_command(), + }) .await .map_err(|error| error.to_string())?; if let Some(receiver) = shutdown.lock().await.take() { @@ -3614,6 +3661,7 @@ mod tests { #[test] #[cfg(feature = "ws-server")] + #[serial_test::serial(worker_allocation)] fn adapter_resumes_paused_turn_once_and_preserves_idle_not_paused_error() { let hanging_events = || simple_text_events().into_iter().take(2).collect::>(); let client = MockClient::sequential(vec![ @@ -3649,7 +3697,12 @@ mod tests { wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Running, true); let running_resume = runtime - .send_protocol_method(&detail.worker_ref, Method::Resume) + .send_protocol_method( + &detail.worker_ref, + Method::Resume { + command: adapter_command(&backend, &detail.worker_ref), + }, + ) .expect_err("Resume while Running must be rejected"); assert!( running_resume @@ -3659,17 +3712,32 @@ mod tests { ); runtime - .send_protocol_method(&detail.worker_ref, Method::Pause) + .send_protocol_method( + &detail.worker_ref, + Method::Pause { + command: adapter_command(&backend, &detail.worker_ref), + }, + ) .expect("pause initial turn"); wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Paused, false); runtime - .send_protocol_method(&detail.worker_ref, Method::Resume) + .send_protocol_method( + &detail.worker_ref, + Method::Resume { + command: adapter_command(&backend, &detail.worker_ref), + }, + ) .expect("resume paused turn"); wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Running, true); let duplicate_resume = runtime - .send_protocol_method(&detail.worker_ref, Method::Resume) + .send_protocol_method( + &detail.worker_ref, + Method::Resume { + command: adapter_command(&backend, &detail.worker_ref), + }, + ) .expect_err("duplicate Resume must be rejected"); assert!( duplicate_resume @@ -3679,17 +3747,32 @@ mod tests { ); runtime - .send_protocol_method(&detail.worker_ref, Method::Pause) + .send_protocol_method( + &detail.worker_ref, + Method::Pause { + command: adapter_command(&backend, &detail.worker_ref), + }, + ) .expect("pause resumed turn"); wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Paused, false); runtime - .send_protocol_method(&detail.worker_ref, Method::Resume) + .send_protocol_method( + &detail.worker_ref, + Method::Resume { + command: adapter_command(&backend, &detail.worker_ref), + }, + ) .expect("resume paused turn a second time"); wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Idle, false); assert_eq!(call_count.load(Ordering::SeqCst), 3); runtime - .send_protocol_method(&detail.worker_ref, Method::Resume) + .send_protocol_method( + &detail.worker_ref, + Method::Resume { + command: adapter_command(&backend, &detail.worker_ref), + }, + ) .expect("Idle Resume preserves controller NotPaused semantics"); wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Idle, false); let events = runtime @@ -3698,10 +3781,10 @@ mod tests { assert!(events.iter().any(|event| { matches!( &event.payload, - Event::Error { - code: protocol::ErrorCode::NotPaused, - .. - } + Event::CommandAcknowledged { acknowledgement } + if acknowledgement.command == protocol::WorkerCommandKind::Resume + && acknowledgement.disposition + == protocol::WorkerCommandDisposition::InvalidState ) })); assert_eq!(call_count.load(Ordering::SeqCst), 3); diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 4a493a0b..98bf1be1 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -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, + 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, + 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, + runtime_dir: &RuntimeDir, + working_event_tx: &broadcast::Sender, + 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, runtime_dir: &RuntimeDir, working_event_tx: &broadcast::Sender, 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( @@ -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::().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( } }; - 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( 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( 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( } } } - 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( 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( .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( } }, - 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( // 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() + ); + } } diff --git a/crates/worker/src/discovery.rs b/crates/worker/src/discovery.rs index 755180f4..844d8c52 100644 --- a/crates/worker/src/discovery.rs +++ b/crates/worker/src/discovery.rs @@ -779,10 +779,10 @@ async fn probe_socket(socket_path: &Path) -> LiveInfo { loop { match tokio::time::timeout(PROBE_TIMEOUT, reader.next::()).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(), }) diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index ec7ef586..63c5b85b 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -295,6 +295,38 @@ impl InternalWorkerSessionStatus { } } +fn send_internal_worker_state( + event_tx: &broadcast::Sender, + 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, ) -> (InternalWorkerSessionStatus, Option) { @@ -351,6 +383,7 @@ pub(crate) struct InternalWorkerSessionSnapshot { pub(crate) struct InternalWorkerSessionHandle { command_tx: tokio::sync::mpsc::Sender, status: Arc, + state_revision: Arc, 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, diff --git a/crates/worker/src/runtime/dir.rs b/crates/worker/src/runtime/dir.rs index ad462f39..37d8d99e 100644 --- a/crates/worker/src/runtime/dir.rs +++ b/crates/worker/src/runtime/dir.rs @@ -197,7 +197,6 @@ pub fn default_base() -> Result { 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(); diff --git a/crates/worker/src/shared_state.rs b/crates/worker/src/shared_state.rs index 20563691..58c5700b 100644 --- a/crates/worker/src/shared_state.rs +++ b/crates/worker/src/shared_state.rs @@ -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, + state: RwLock, + 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"); - } } diff --git a/crates/worker/src/spawn/comm_tools.rs b/crates/worker/src/spawn/comm_tools.rs index 0389f288..28efb9a6 100644 --- a/crates/worker/src/spawn/comm_tools.rs +++ b/crates/worker/src/spawn/comm_tools.rs @@ -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] diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 7b051970..93ec40ac 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -4571,15 +4571,6 @@ impl Worker { 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 Worker { 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, + state: CompactionLifecycleState, + started_at_ms: u64, + #[serde(default)] + ended_at_ms: Option, + #[serde(default)] + summary: Option, + #[serde(default)] + error: Option, + #[serde(default)] + new_segment_id: Option, + } + 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 { + self.manual_compact_inner(None).await + } + + pub async fn manual_compact_with_cancel( + &mut self, + cancel: tokio::sync::watch::Receiver, + ) -> Result { + self.manual_compact_inner(Some(cancel)).await + } + + async fn manual_compact_inner( + &mut self, + mut cancel: Option>, + ) -> Result { if self.manifest.compaction.is_none() { let message = "manual compact is unavailable because [compaction] is not configured".to_string(); @@ -4764,7 +4845,7 @@ impl Worker { 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 Worker { /// 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 { + self.compact_with_cancel(retained_tokens, None).await + } + + async fn compact_with_cancel( + &mut self, + retained_tokens: u64, + mut cancel: Option>, + ) -> Result { 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 Worker { 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 Worker { })?, }); } + // 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::(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] diff --git a/crates/worker/tests/compact_events_test.rs b/crates/worker/tests/compact_events_test.rs index 8265f9e2..1521b909 100644 --- a/crates/worker/tests/compact_events_test.rs +++ b/crates/worker/tests/compact_events_test.rs @@ -72,6 +72,41 @@ impl LlmClient for MockClient { } } +#[derive(Clone)] +struct BlockingCompactClient { + calls: Arc, +} + +impl BlockingCompactClient { + fn new() -> Self { + Self { + calls: Arc::new(AtomicUsize::new(0)), + } + } +} + +#[async_trait] +impl LlmClient for BlockingCompactClient { + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + async fn stream( + &self, + _request: Request, + ) -> Result> + 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 { 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 { +async fn make_worker_with_manifest(manifest_toml: &str, client: C) -> Worker +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; } diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index 319f773d..35bf5193 100644 --- a/crates/worker/tests/controller_test.rs +++ b/crates/worker/tests/controller_test.rs @@ -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; +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` 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::().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!( diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index 4a90e192..f615685d 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -36,8 +36,6 @@ use worker_runtime::config_bundle::{ ConfigBundleMetadata, ConfigBundleProvenance, ConfigProfileDescriptor, }; use worker_runtime::error::RuntimeError as EmbeddedRuntimeError; -#[cfg(test)] -use worker_runtime::execution::WorkerExecutionRunState; use worker_runtime::fs_store::FsRuntimeStoreOptions; use worker_runtime::http_server::{ RUNTIME_PING_PERMISSION, RUNTIME_WORKSPACE_SCOPE_HEADER, @@ -5170,7 +5168,6 @@ mod tests { request.worker_ref, self.backend_id(), ), - run_state: WorkerExecutionRunState::Idle, working_directory: request .working_directory .as_ref() @@ -5199,8 +5196,8 @@ mod tests { let content = input.content; std::thread::spawn(move || { std::thread::sleep(std::time::Duration::from_millis(10)); - let _ = context.publish_protocol_event(protocol::Event::Status { - status: protocol::WorkerStatus::Running, + let _ = context.publish_protocol_event(protocol::Event::WorkerState { + snapshot: protocol::WorkerStatus::Running.into(), }); let _ = context.publish_protocol_event(protocol::Event::TextDone { text: format!("echo: {content}"), @@ -5208,14 +5205,13 @@ mod tests { let _ = context.publish_protocol_event(protocol::Event::RunEnd { result: protocol::RunResult::Finished, }); - let _ = context.publish_protocol_event(protocol::Event::Status { - status: protocol::WorkerStatus::Idle, + let _ = context.publish_protocol_event(protocol::Event::WorkerState { + snapshot: protocol::WorkerStatus::Idle.into(), }); }); if let Some(submission_request_id) = submission_request_id { worker_runtime::execution::WorkerExecutionResult::accepted_submission( worker_runtime::execution::WorkerExecutionOperation::Input, - WorkerExecutionRunState::Busy, submission_request_id, uuid::Uuid::now_v7().to_string(), protocol::SubmissionDisposition::Started, @@ -5223,7 +5219,6 @@ mod tests { } else { worker_runtime::execution::WorkerExecutionResult::accepted( worker_runtime::execution::WorkerExecutionOperation::Input, - WorkerExecutionRunState::Busy, ) } } diff --git a/crates/workspace-server/src/runtime_subscription_tests.rs b/crates/workspace-server/src/runtime_subscription_tests.rs index 6c14c56a..5c67f983 100644 --- a/crates/workspace-server/src/runtime_subscription_tests.rs +++ b/crates/workspace-server/src/runtime_subscription_tests.rs @@ -6,7 +6,7 @@ use worker_runtime::catalog::{ }; use worker_runtime::execution::{ WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, WorkerExecutionResult, - WorkerExecutionRunState, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, + WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, }; use worker_runtime::identity::WorkerId; use worker_runtime::profile_archive::{ProfileSourceArchiveRef, ProfileSourceGraphSummary}; @@ -22,7 +22,6 @@ impl WorkerExecutionBackend for TestExecutionBackend { fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { WorkerExecutionSpawnResult::connected( WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), - WorkerExecutionRunState::Idle, None, ) } @@ -35,24 +34,17 @@ impl WorkerExecutionBackend for TestExecutionBackend { if let Some(submission_request_id) = input.submission_request_id { WorkerExecutionResult::accepted_submission( WorkerExecutionOperation::Input, - WorkerExecutionRunState::Busy, submission_request_id, uuid::Uuid::now_v7().to_string(), protocol::SubmissionDisposition::Started, ) } else { - WorkerExecutionResult::accepted( - WorkerExecutionOperation::Input, - WorkerExecutionRunState::Busy, - ) + WorkerExecutionResult::accepted(WorkerExecutionOperation::Input) } } fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { - WorkerExecutionResult::accepted( - WorkerExecutionOperation::Stop, - WorkerExecutionRunState::Stopped, - ) + WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop) } } @@ -199,8 +191,8 @@ async fn equal_downstream_selectors_share_one_upstream_subscription() { runtime .observe_worker_event( &worker.worker_ref, - protocol::Event::Status { - status: protocol::WorkerStatus::Running, + protocol::Event::WorkerState { + snapshot: protocol::WorkerStatus::Running.into(), }, ) .unwrap(); @@ -339,8 +331,8 @@ async fn embedded_runtime_uses_in_process_subscription_source() { runtime .observe_worker_event( &worker.worker_ref, - protocol::Event::Status { - status: protocol::WorkerStatus::Running, + protocol::Event::WorkerState { + snapshot: protocol::WorkerStatus::Running.into(), }, ) .unwrap(); diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 3fab8457..b55f3d2a 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -18559,7 +18559,6 @@ mod tests { request.worker_ref, self.backend_id(), ), - run_state: worker_runtime::execution::WorkerExecutionRunState::Idle, working_directory, } } @@ -18575,7 +18574,6 @@ mod tests { .push((handle.worker_ref().clone(), method)); worker_runtime::execution::WorkerExecutionResult::accepted( worker_runtime::execution::WorkerExecutionOperation::ProtocolMethod, - worker_runtime::execution::WorkerExecutionRunState::Idle, ) } @@ -18585,7 +18583,6 @@ mod tests { ) -> worker_runtime::execution::WorkerExecutionResult { worker_runtime::execution::WorkerExecutionResult::accepted( worker_runtime::execution::WorkerExecutionOperation::Stop, - worker_runtime::execution::WorkerExecutionRunState::Stopped, ) } @@ -18595,7 +18592,6 @@ mod tests { ) -> worker_runtime::execution::WorkerExecutionResult { worker_runtime::execution::WorkerExecutionResult::accepted( worker_runtime::execution::WorkerExecutionOperation::Cancel, - worker_runtime::execution::WorkerExecutionRunState::Stopped, ) } @@ -18632,16 +18628,16 @@ mod tests { if let Some(submission_request_id) = submission_request_id { worker_runtime::execution::WorkerExecutionResult::accepted_submission( worker_runtime::execution::WorkerExecutionOperation::Input, - worker_runtime::execution::WorkerExecutionRunState::Idle, submission_request_id, uuid::Uuid::now_v7().to_string(), protocol::SubmissionDisposition::Started, ) + .with_worker_state(protocol::WorkerStateSnapshot::initial(1)) } else { worker_runtime::execution::WorkerExecutionResult::accepted( worker_runtime::execution::WorkerExecutionOperation::Input, - worker_runtime::execution::WorkerExecutionRunState::Idle, ) + .with_worker_state(protocol::WorkerStateSnapshot::initial(1)) } } } @@ -27497,7 +27493,13 @@ mod tests { protocol::subscription::SubscriptionFramePayload::WorkerProtocol( protocol::subscription::SubscriptionWorkerProtocolMethod { subscription_id: second_protocol_subscription_id, - method: protocol::Method::Resume, + method: protocol::Method::Resume { + command: protocol::WorkerCommandEnvelope { + command_id: 1, + expected_execution_generation: 1, + expected_worker_state_revision: 0, + }, + }, }, ), ); @@ -27514,7 +27516,7 @@ mod tests { .iter() .any(|(worker_ref, method)| { worker_ref.worker_id.to_string() == worker_id - && matches!(method, protocol::Method::Resume) + && matches!(method, protocol::Method::Resume { .. }) }) { break; @@ -27527,7 +27529,7 @@ mod tests { let protocol_methods = execution_backend.protocol_methods(); assert!(protocol_methods.iter().any(|(worker_ref, method)| { worker_ref.worker_id.to_string() == worker_id - && matches!(method, protocol::Method::Resume) + && matches!(method, protocol::Method::Resume { .. }) })); server.abort(); } diff --git a/web/workspace/src/lib/generated/protocol.ts b/web/workspace/src/lib/generated/protocol.ts index 6055d3e7..c15fbdca 100644 --- a/web/workspace/src/lib/generated/protocol.ts +++ b/web/workspace/src/lib/generated/protocol.ts @@ -10,6 +10,37 @@ export type CompletionKind = "file"; export type WorkerStatus = "idle" | "running" | "paused" | "stopped"; +export type WorkerCommandEnvelope = { +/** + * Caller-owned sequence. A controller accepts command ids in strictly + * increasing order for one execution generation. + */ +command_id: number, expected_execution_generation: number, expected_worker_state_revision: number, }; + +export type WorkerCommandKind = "resume" | "cancel" | "pause" | "compact" | "shutdown"; + +export type WorkerCommandDisposition = "accepted" | "stale_execution_generation" | "stale_worker_state_revision" | "stale_command_id" | "invalid_state"; + +export type WorkerCommandAcknowledgement = { command_id: number, command: WorkerCommandKind, disposition: WorkerCommandDisposition, +/** + * The complete authoritative state observed after command admission. + */ +state: WorkerStateSnapshot, }; + +export type WorkerRunState = "running" | "pausing" | "paused" | "cancelling"; + +export type WorkerMaintenanceState = "compacting"; + +export type WorkerBusyState = { "kind": "run", "state": WorkerRunState } | { "kind": "maintenance", "state": WorkerMaintenanceState }; + +export type WorkerState = { "kind": "idle" } | { "kind": "busy", "state": WorkerBusyState }; + +export type WorkerStateSnapshot = { execution_generation: number, revision: number, +/** + * Highest lifecycle command id observed by this controller generation. + */ +last_command_id: number, state: WorkerState, }; + export type TurnResult = "finished" | "paused"; export type InvokeKind = "user_send" | "notify" | "worker_event" | "system_reminder" | "wakeup"; @@ -231,7 +262,7 @@ export type SubscriptionFramePayload = { "frame": "request", "message": Subscrip export type SubscriptionFrame = { protocol_version: number, } & ({ "frame": "request", "message": SubscriptionRequest } | { "frame": "response", "message": SubscriptionResponse } | { "frame": "event", "message": SubscriptionEvent } | { "frame": "worker_protocol", "message": SubscriptionWorkerProtocolMethod }); -export type Method = { "method": "submit", "params": { submission_request_id: string, input: Array, } } | { "method": "notify", "params": { notification_request_id: string, message: string, auto_run?: boolean, } } | { "method": "worker_event", "params": WorkerEvent } | { "method": "list_pending_submissions" } | { "method": "cancel_pending_submission", "params": { submission_id: string, expected_revision: number, } } | { "method": "clear_pending_submissions", "params": { expected_revision: number, } } | { "method": "continue_pending", "params": { expected_revision: number, expected_head_id: string, } } | { "method": "resume" } | { "method": "cancel" } | { "method": "pause" } | { "method": "compact" } | { "method": "list_rewind_targets" } | { "method": "rewind_to", "params": { target: RewindTargetId, expected_head_entries: number, } } | { "method": "shutdown" } | { "method": "list_completions", "params": { kind: CompletionKind, prefix: string, } } | { "method": "list_workers" } | { "method": "restore_worker", "params": { name: string, } } | { "method": "register_peer", "params": { name: string, } }; +export type Method = { "method": "submit", "params": { submission_request_id: string, input: Array, } } | { "method": "notify", "params": { notification_request_id: string, message: string, auto_run?: boolean, } } | { "method": "worker_event", "params": WorkerEvent } | { "method": "list_pending_submissions" } | { "method": "cancel_pending_submission", "params": { submission_id: string, expected_revision: number, } } | { "method": "clear_pending_submissions", "params": { expected_revision: number, } } | { "method": "continue_pending", "params": { expected_revision: number, expected_head_id: string, } } | { "method": "resume", "params": { command: WorkerCommandEnvelope, } } | { "method": "cancel", "params": { command: WorkerCommandEnvelope, } } | { "method": "pause", "params": { command: WorkerCommandEnvelope, } } | { "method": "compact", "params": { command: WorkerCommandEnvelope, } } | { "method": "list_rewind_targets" } | { "method": "rewind_to", "params": { target: RewindTargetId, expected_head_entries: number, } } | { "method": "shutdown", "params": { command: WorkerCommandEnvelope, } } | { "method": "list_completions", "params": { kind: CompletionKind, prefix: string, } } | { "method": "list_workers" } | { "method": "restore_worker", "params": { name: string, } } | { "method": "register_peer", "params": { name: string, } }; export type Event = { "event": "submission_accepted", "data": { submission_request_id: string, submission_id: string, disposition: SubmissionDisposition, } } | { "event": "submission_rejected", "data": { submission_request_id: string, message: string, } } | { "event": "pending_submissions_changed", "data": { pending: PendingSubmissionsSnapshot, } } | { "event": "user_message", "data": { segments: Array, } } | { "event": "system_item", "data": { item: unknown, } } | { "event": "invoke_start", "data": { kind: InvokeKind, } } | { "event": "turn_start", "data": { turn: number, } } | { "event": "turn_end", "data": { turn: number, result: TurnResult, } } | { "event": "llm_call_start", "data": { llm_call: number, } } | { "event": "llm_call_end", "data": { llm_call: number, } } | { "event": "llm_retry", "data": { llm_call: number, /** @@ -247,7 +278,12 @@ summary: string, * Full tool output. Absent when the tool chose to return * summary-only, or when the result was pruned. */ -output?: string | null, disposition?: ToolResultDisposition | null, is_error: boolean, } } | { "event": "usage", "data": { input_tokens: number | null, output_tokens: number | null, cache_read_input_tokens?: number | null, } } | { "event": "run_end", "data": { result: RunResult, } } | { "event": "error", "data": { code: ErrorCode, message: string, } } | { "event": "snapshot", "data": { session: SessionSnapshot, greeting: Greeting, status: WorkerStatus, +output?: string | null, disposition?: ToolResultDisposition | null, is_error: boolean, } } | { "event": "usage", "data": { input_tokens: number | null, output_tokens: number | null, cache_read_input_tokens?: number | null, } } | { "event": "run_end", "data": { result: RunResult, } } | { "event": "error", "data": { code: ErrorCode, message: string, } } | { "event": "snapshot", "data": { session: SessionSnapshot, greeting: Greeting, +/** + * Full revisioned live execution state. `Stopped` remains Runtime + * catalog authority and is deliberately not represented here. + */ +state: WorkerStateSnapshot, /** * Unfinished model output that has already streamed in the current * run but is not yet represented by committed snapshot entries. @@ -257,4 +293,4 @@ in_flight?: InFlightSnapshot, * Parent-owned Internal Worker sessions visible to this client. * Service-private Internal Workers are deliberately excluded. */ -internal_workers?: Array, } } | { "event": "internal_worker", "data": { worker: InternalWorkerRef, revision: number, event: Event, } } | { "event": "internal_worker_removed", "data": { worker: InternalWorkerRef, revision: number, } } | { "event": "segment_rotated", "data": { session: SessionSnapshot, } } | { "event": "status", "data": { status: WorkerStatus, } } | { "event": "command", "data": { event: CommandEvent, } } | { "event": "completions", "data": { kind: CompletionKind, entries: Array, } } | { "event": "rewind_targets", "data": { head_entries: number, targets: Array, } } | { "event": "rewind_applied", "data": { session: SessionSnapshot, input: Array, summary: RewindSummary, } } | { "event": "workers_listed", "data": { workers: unknown, } } | { "event": "worker_restored", "data": { result: unknown, } } | { "event": "peer_registered", "data": { result: unknown, } } | { "event": "alert", "data": Alert } | { "event": "memory_worker", "data": MemoryWorkerEvent } | { "event": "compact_start", "data": { lifecycle: CompactionLifecycle, } } | { "event": "compact_done", "data": { lifecycle: CompactionLifecycle, } } | { "event": "compact_failed", "data": { lifecycle: CompactionLifecycle, } } | { "event": "shutdown" }; +internal_workers?: Array, } } | { "event": "internal_worker", "data": { worker: InternalWorkerRef, revision: number, event: Event, } } | { "event": "internal_worker_removed", "data": { worker: InternalWorkerRef, revision: number, } } | { "event": "segment_rotated", "data": { session: SessionSnapshot, } } | { "event": "worker_state", "data": { snapshot: WorkerStateSnapshot, } } | { "event": "command_acknowledged", "data": { acknowledgement: WorkerCommandAcknowledgement, } } | { "event": "command", "data": { event: CommandEvent, } } | { "event": "completions", "data": { kind: CompletionKind, entries: Array, } } | { "event": "rewind_targets", "data": { head_entries: number, targets: Array, } } | { "event": "rewind_applied", "data": { session: SessionSnapshot, input: Array, summary: RewindSummary, } } | { "event": "workers_listed", "data": { workers: unknown, } } | { "event": "worker_restored", "data": { result: unknown, } } | { "event": "peer_registered", "data": { result: unknown, } } | { "event": "alert", "data": Alert } | { "event": "memory_worker", "data": MemoryWorkerEvent } | { "event": "compact_start", "data": { lifecycle: CompactionLifecycle, } } | { "event": "compact_done", "data": { lifecycle: CompactionLifecycle, } } | { "event": "compact_failed", "data": { lifecycle: CompactionLifecycle, } } | { "event": "shutdown" }; diff --git a/web/workspace/src/lib/workspace/console/model.test.ts b/web/workspace/src/lib/workspace/console/model.test.ts index 60b98633..c1f1b72a 100644 --- a/web/workspace/src/lib/workspace/console/model.test.ts +++ b/web/workspace/src/lib/workspace/console/model.test.ts @@ -1,4 +1,4 @@ -import type { Event } from "$lib/generated/protocol"; +import type { Event, WorkerStateSnapshot, WorkerStatus } from "$lib/generated/protocol"; import { type ConsoleEventInput, type ConsoleLine, @@ -19,6 +19,23 @@ declare const Deno: { test(name: string, fn: () => void): void; }; +function workerState(status: WorkerStatus): WorkerStateSnapshot { + return { + execution_generation: 1, + revision: status === "idle" ? 0 : 1, + last_command_id: 0, + state: status === "idle" + ? { kind: "idle" } + : { + kind: "busy", + state: { + kind: "run", + state: status === "paused" ? "paused" : "running", + }, + }, + }; +} + function assert(condition: unknown, message: string): asserts condition { if (!condition) { throw new Error(message); @@ -131,7 +148,7 @@ function snapshotEvent(cwd: string, entries: unknown[] = []): Event { context_window: 100, context_tokens: 20, }, - status: "idle", + state: workerState("idle"), in_flight: { blocks: [] }, }, }; @@ -213,7 +230,7 @@ Deno.test("snapshot replaces a live error with one durable run_errored row", () }, { eventId: "idle-after-error", - event: { event: "status", data: { status: "idle" } } satisfies Event, + event: { event: "worker_state", data: { snapshot: workerState("idle") } } satisfies Event, }, ]); @@ -653,7 +670,7 @@ Deno.test("projectConsole streams distinct Bash stdout and stderr through termin Deno.test("snapshot restores bounded in-flight Bash command output", () => { const snapshot = snapshotEvent("/repo"); if (snapshot.event !== "snapshot") throw new Error("snapshot fixture expected"); - snapshot.data.status = "running"; + snapshot.data.state = workerState("running"); snapshot.data.in_flight = { blocks: [{ kind: "tool_call", @@ -1403,7 +1420,7 @@ Deno.test("projectConsole hides lifecycle events and renders system items", () = const projection = projectConsole([ { eventId: "30", - event: { event: "status", data: { status: "running" } } satisfies Event, + event: { event: "worker_state", data: { snapshot: workerState("running") } } satisfies Event, }, { eventId: "31", @@ -1527,7 +1544,7 @@ Deno.test("projectConsole renders snapshot entries and in-flight output", () => context_window: 100, context_tokens: 20, }, - status: "running", + state: workerState("running"), in_flight: { blocks: [ { kind: "text", text: "partial" }, @@ -1578,7 +1595,7 @@ Deno.test("projectConsole restores system items from snapshot entries", () => { context_window: 100, context_tokens: 20, }, - status: "idle", + state: workerState("idle"), }, } satisfies Event, }]); @@ -1922,7 +1939,7 @@ Deno.test("console Worker views expose only direct Internal Workers", () => { kind: "sub_worker", }, revision: 1, - event: { event: "status", data: { status: "running" } }, + event: { event: "worker_state", data: { snapshot: workerState("running") } }, }, }, }, @@ -1941,7 +1958,7 @@ Deno.test("console Worker views expose only direct Internal Workers", () => { kind: "sub_worker", }, revision: 1, - event: { event: "status", data: { status: "idle" } }, + event: { event: "worker_state", data: { snapshot: workerState("idle") } }, }, }, }]); @@ -2033,7 +2050,7 @@ Deno.test("parent snapshot authoritatively replaces Internal Worker projections" kind: "sub_worker", }, revision: 1, - event: { event: "status", data: { status: "running" } }, + event: { event: "worker_state", data: { snapshot: workerState("running") } }, }, }, }]); diff --git a/web/workspace/src/lib/workspace/console/model.ts b/web/workspace/src/lib/workspace/console/model.ts index 4464c89a..857e1b30 100644 --- a/web/workspace/src/lib/workspace/console/model.ts +++ b/web/workspace/src/lib/workspace/console/model.ts @@ -10,6 +10,8 @@ import type { InternalWorkerRef, InternalWorkerSnapshot, Segment, + WorkerStateSnapshot, + WorkerStatus, } from "$lib/generated/protocol"; import { stringify as stringifyYaml } from "yaml"; import { workspaceRoute } from "$lib/workspace/api/http"; @@ -169,6 +171,7 @@ export type ConsoleProjection = { tasks: ConsoleTask[]; taskNextId: number; status: string | null; + workerState: WorkerStateSnapshot | null; usage: string | null; runActivity: RunActivityStats; cwd: string | null; @@ -251,12 +254,22 @@ export function isConsoleProjectionEvent(event: ProtocolEvent): boolean { return event.event !== "completions"; } +function workerStatusFromState(snapshot: WorkerStateSnapshot): WorkerStatus { + if (snapshot.state.kind === "idle") return "idle"; + if ( + snapshot.state.state.kind === "run" && + snapshot.state.state.state === "paused" + ) return "paused"; + return "running"; +} + export function emptyConsoleProjection(): ConsoleProjection { return { lines: [], tasks: [], taskNextId: 1, status: null, + workerState: null, usage: null, runActivity: emptyRunActivityStats(), cwd: null, @@ -793,6 +806,7 @@ export function applyProtocolEvent( tasks: [...projection.tasks], taskNextId: projection.taskNextId, status: projection.status, + workerState: projection.workerState, usage: projection.usage, runActivity: applyRunActivityEvent( projection.runActivity, @@ -903,7 +917,8 @@ export function applyProtocolEvent( ); break; case "snapshot": { - next.status = event.data.status; + next.workerState = event.data.state; + next.status = workerStatusFromState(event.data.state); next.cwd = event.data.greeting.cwd; const snapshot = snapshotProjectionFromSession( envelope.eventId, @@ -1000,8 +1015,13 @@ export function applyProtocolEvent( if (existingIndex >= 0) next.internalWorkers.splice(existingIndex, 1); break; } - case "status": - next.status = event.data.status; + case "worker_state": + next.workerState = event.data.snapshot; + next.status = workerStatusFromState(event.data.snapshot); + break; + case "command_acknowledged": + next.workerState = event.data.acknowledgement.state; + next.status = workerStatusFromState(event.data.acknowledgement.state); break; case "command": applyCommandEvent(next, envelope.eventId, event.data.event); @@ -1939,6 +1959,7 @@ function snapshotProjectionFromSession( tasks: [], taskNextId: 1, status: null, + workerState: null, usage: null, runActivity: emptyRunActivityStats(), cwd, diff --git a/web/workspace/src/lib/workspace/console/run-status.test.ts b/web/workspace/src/lib/workspace/console/run-status.test.ts index 575ce733..0063a242 100644 --- a/web/workspace/src/lib/workspace/console/run-status.test.ts +++ b/web/workspace/src/lib/workspace/console/run-status.test.ts @@ -75,7 +75,12 @@ Deno.test("new invoke and running snapshot reset run activity", () => { data: { entries: [], greeting: { text: "", profile: "" }, - status: "idle", + state: { + execution_generation: 1, + revision: 0, + last_command_id: 0, + state: { kind: "idle" }, + }, in_flight: {}, internal_workers: [], }, diff --git a/web/workspace/src/lib/workspace/console/run-status.ts b/web/workspace/src/lib/workspace/console/run-status.ts index aabb45c2..37495dc1 100644 --- a/web/workspace/src/lib/workspace/console/run-status.ts +++ b/web/workspace/src/lib/workspace/console/run-status.ts @@ -25,7 +25,9 @@ export function applyRunActivityEvent( case "invoke_start": return { ...emptyRunActivityStats(), startedAtMs: observedAtMs }; case "snapshot": - return event.data.status === "running" + return event.data.state.state.kind === "busy" && + !(event.data.state.state.state.kind === "run" && + event.data.state.state.state.state === "paused") ? { ...emptyRunActivityStats(), startedAtMs: observedAtMs } : emptyRunActivityStats(); case "turn_start": diff --git a/web/workspace/src/routes/w/[workspaceId]/runtimes/[runtimeId]/workers/[workerId]/console/+page.svelte b/web/workspace/src/routes/w/[workspaceId]/runtimes/[runtimeId]/workers/[workerId]/console/+page.svelte index c7852975..cf28f294 100644 --- a/web/workspace/src/routes/w/[workspaceId]/runtimes/[runtimeId]/workers/[workerId]/console/+page.svelte +++ b/web/workspace/src/routes/w/[workspaceId]/runtimes/[runtimeId]/workers/[workerId]/console/+page.svelte @@ -541,9 +541,42 @@ } } + let nextWorkerCommandId = 1; + + function lifecycleMethod( + command: "pause" | "cancel" | "resume" | "compact", + ): ProtocolMethod | null { + const state = consoleProjection.workerState; + if (!state) { + sendError = "Worker state snapshot is not available; reconnect before sending control."; + return null; + } + const commandId = Math.max( + nextWorkerCommandId, + state.last_command_id + 1, + ); + nextWorkerCommandId = commandId + 1; + const envelope = { + command_id: commandId, + expected_execution_generation: state.execution_generation, + expected_worker_state_revision: state.revision, + }; + switch (command) { + case "pause": + return { method: "pause", params: { command: envelope } }; + case "cancel": + return { method: "cancel", params: { command: envelope } }; + case "resume": + return { method: "resume", params: { command: envelope } }; + case "compact": + return { method: "compact", params: { command: envelope } }; + } + } + function sendWorkerControl(command: "pause" | "cancel" | "resume") { const label = command[0].toUpperCase() + command.slice(1); - sendControl({ method: command }, label); + const method = lifecycleMethod(command); + if (method) sendControl(method, label); } function isEditableTarget(target: EventTarget | null): boolean { @@ -627,8 +660,11 @@ auto_run: true, }, }; - case "compact": - return { method: "compact" }; + case "compact": { + const method = lifecycleMethod("compact"); + if (!method) throw new Error("Worker state snapshot is not available"); + return method; + } case "list_rewind_targets": return { method: "list_rewind_targets" }; case "register_peer": @@ -691,7 +727,7 @@ function handleComposerSubmit() { if (workerRunning) { - sendControl({ method: "cancel" }, "Stop"); + sendWorkerControl("cancel"); return; } void submitDraft(composerInputElement?.snapshot() ?? draft); @@ -894,8 +930,26 @@ ): string | null { switch (event.event) { case "snapshot": - case "status": - return event.data.status; + return event.data.state.state.kind === "idle" + ? "idle" + : event.data.state.state.state.kind === "run" && + event.data.state.state.state.state === "paused" + ? "paused" + : "running"; + case "worker_state": + return event.data.snapshot.state.kind === "idle" + ? "idle" + : event.data.snapshot.state.state.kind === "run" && + event.data.snapshot.state.state.state === "paused" + ? "paused" + : "running"; + case "command_acknowledged": + return event.data.acknowledgement.state.state.kind === "idle" + ? "idle" + : event.data.acknowledgement.state.state.state.kind === "run" && + event.data.acknowledgement.state.state.state.state === "paused" + ? "paused" + : "running"; case "shutdown": return "shutdown"; default: @@ -1620,7 +1674,10 @@ type="button" class="secondary-button" disabled={protocolState !== "open"} - onclick={() => sendControl({ method: "compact" }, "Compact")} + onclick={() => { + const method = lifecycleMethod("compact"); + if (method) sendControl(method, "Compact"); + }} > Compact