From e8b9adcde43028bae4bae1cb5ee067c018e7693a Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 06:54:29 +0900 Subject: [PATCH 1/5] 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 From 282a8d31b584c6cf55bd15804eb9741b22925b13 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 07:56:25 +0900 Subject: [PATCH 2/5] fix: enforce monotonic worker state projection --- crates/protocol/src/lib.rs | 99 +++++ crates/tui/src/app.rs | 99 ++++- crates/worker-runtime/src/worker_backend.rs | 352 ++++++++---------- crates/worker/src/controller.rs | 4 +- crates/worker/src/shared_state.rs | 44 ++- .../src/lib/workspace/console/model.test.ts | 60 +++ .../src/lib/workspace/console/model.ts | 68 +++- 7 files changed, 485 insertions(+), 241 deletions(-) diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index ec96e226..f66be5dc 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -203,6 +203,53 @@ impl WorkerStateSnapshot { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WorkerStateSnapshotApply { + Applied, + Duplicate, + Stale, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct WorkerStateSnapshotConflict { + pub execution_generation: u64, + pub revision: u64, +} + +impl std::fmt::Display for WorkerStateSnapshotConflict { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + formatter, + "conflicting worker state snapshots at generation {} revision {}", + self.execution_generation, self.revision + ) + } +} + +impl std::error::Error for WorkerStateSnapshotConflict {} + +pub fn apply_worker_state_snapshot( + current: &mut WorkerStateSnapshot, + incoming: &WorkerStateSnapshot, +) -> Result { + use std::cmp::Ordering; + + let ordering = (incoming.execution_generation, incoming.revision) + .cmp(&(current.execution_generation, current.revision)); + match ordering { + Ordering::Greater => { + *current = incoming.clone(); + Ok(WorkerStateSnapshotApply::Applied) + } + Ordering::Less => Ok(WorkerStateSnapshotApply::Stale), + Ordering::Equal if incoming == current => Ok(WorkerStateSnapshotApply::Duplicate), + Ordering::Equal => Err(WorkerStateSnapshotConflict { + execution_generation: incoming.execution_generation, + revision: incoming.revision, + }), + } +} + impl From for WorkerStateSnapshot { fn from(status: WorkerStatus) -> Self { let state = match status { @@ -1574,6 +1621,58 @@ pub enum Permission { mod tests { use super::*; + #[test] + fn worker_state_snapshot_apply_is_monotonic_and_detects_conflicts() { + let mut current = WorkerStateSnapshot::initial(4); + let mut newer = current.clone(); + newer.revision = 1; + newer.state = WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running)); + + assert_eq!( + apply_worker_state_snapshot(&mut current, &newer), + Ok(WorkerStateSnapshotApply::Applied) + ); + assert_eq!( + apply_worker_state_snapshot(&mut current, &newer), + Ok(WorkerStateSnapshotApply::Duplicate) + ); + + let stale_revision = WorkerStateSnapshot::initial(4); + assert_eq!( + apply_worker_state_snapshot(&mut current, &stale_revision), + Ok(WorkerStateSnapshotApply::Stale) + ); + let stale_generation = WorkerStateSnapshot { + execution_generation: 3, + revision: u64::MAX, + ..newer.clone() + }; + assert_eq!( + apply_worker_state_snapshot(&mut current, &stale_generation), + Ok(WorkerStateSnapshotApply::Stale) + ); + + let conflicting = WorkerStateSnapshot { + state: WorkerState::Idle, + ..newer.clone() + }; + assert_eq!( + apply_worker_state_snapshot(&mut current, &conflicting), + Err(WorkerStateSnapshotConflict { + execution_generation: 4, + revision: 1, + }) + ); + assert_eq!(current, newer); + + let next_generation = WorkerStateSnapshot::initial(5); + assert_eq!( + apply_worker_state_snapshot(&mut current, &next_generation), + Ok(WorkerStateSnapshotApply::Applied) + ); + assert_eq!(current, next_generation); + } + #[test] fn method_submit_json_roundtrip_and_run_is_rejected() { let json = r#"{"method":"submit","params":{"submission_request_id":"request-1","input":[{"kind":"text","content":"Hello"}]}}"#; diff --git a/crates/tui/src/app.rs b/crates/tui/src/app.rs index a76d32b5..c543f991 100644 --- a/crates/tui/src/app.rs +++ b/crates/tui/src/app.rs @@ -1128,6 +1128,22 @@ impl App { command } + fn apply_worker_state_snapshot(&mut self, snapshot: &WorkerStateSnapshot) { + match protocol::apply_worker_state_snapshot(&mut self.worker_state, snapshot) { + Ok(protocol::WorkerStateSnapshotApply::Applied) => { + self.set_worker_status(self.worker_state.catalog_status()); + } + Ok( + protocol::WorkerStateSnapshotApply::Duplicate + | protocol::WorkerStateSnapshotApply::Stale, + ) => {} + Err(error) => self.handle_error( + ErrorCode::Internal, + format!("worker state stream rejected: {error}"), + ), + } + } + pub fn handle_worker_event(&mut self, event: Event) -> Option { if self.rewind_refresh_fence && event_is_stale_after_rewind(&event) { return None; @@ -1465,8 +1481,7 @@ impl App { self.pending_submissions = session.pending_submissions.clone(); self.restore_snapshot(&session, greeting, in_flight); self.replace_internal_worker_snapshots(internal_workers); - self.worker_state = state.clone(); - self.set_worker_status(state.catalog_status()); + self.apply_worker_state_snapshot(&state); } Event::InternalWorker { worker, @@ -1478,12 +1493,10 @@ impl App { } Event::WorkerState { snapshot } => { self.rewind_refresh_fence = false; - self.worker_state = snapshot.clone(); - self.set_worker_status(snapshot.catalog_status()); + self.apply_worker_state_snapshot(&snapshot); } Event::CommandAcknowledged { acknowledgement } => { - self.worker_state = acknowledgement.state.clone(); - self.set_worker_status(acknowledgement.state.catalog_status()); + self.apply_worker_state_snapshot(&acknowledgement.state); } // Command telemetry is an operational Web Console surface. The // TUI continues to render the final Bash ToolResult from history. @@ -3559,7 +3572,7 @@ mod completion_flow_tests { app.handle_worker_event(Event::Snapshot { greeting: test_greeting(), session: public_session(vec![session_start_value]), - state: WorkerStatus::Running.into(), + state: test_worker_state(WorkerStatus::Running), in_flight: Default::default(), internal_workers: Vec::new(), }); @@ -3570,6 +3583,59 @@ mod completion_flow_tests { assert!(matches!(app.blocks.first(), Some(Block::Greeting(_)))); } + #[test] + fn worker_state_events_and_acknowledgements_share_monotonic_application() { + let mut app = App::new("test".into()); + let running = WorkerStateSnapshot { + execution_generation: 4, + revision: 3, + state: protocol::WorkerState::Busy(protocol::WorkerBusyState::Run( + protocol::WorkerRunState::Running, + )), + last_command_id: 2, + }; + app.handle_worker_event(Event::WorkerState { + snapshot: running.clone(), + }); + app.handle_worker_event(Event::WorkerState { + snapshot: WorkerStateSnapshot { + revision: 2, + state: protocol::WorkerState::Idle, + ..running.clone() + }, + }); + assert_eq!(app.worker_state, running); + + let paused = WorkerStateSnapshot { + revision: 4, + state: protocol::WorkerState::Busy(protocol::WorkerBusyState::Run( + protocol::WorkerRunState::Paused, + )), + last_command_id: 3, + ..running.clone() + }; + app.handle_worker_event(Event::CommandAcknowledged { + acknowledgement: protocol::WorkerCommandAcknowledgement { + command_id: 3, + command: protocol::WorkerCommandKind::Pause, + disposition: protocol::WorkerCommandDisposition::Accepted, + state: paused.clone(), + }, + }); + assert_eq!(app.worker_state, paused); + + app.handle_worker_event(Event::WorkerState { + snapshot: WorkerStateSnapshot { + state: protocol::WorkerState::Idle, + ..paused.clone() + }, + }); + assert_eq!(app.worker_state, paused); + assert!(app.run_error_messages.iter().any(|message| { + message.contains("conflicting worker state snapshots at generation 4 revision 4") + })); + } + #[test] fn snapshot_replaces_live_error_with_one_durable_run_error_block() { let mut app = App::new("test".into()); @@ -3603,7 +3669,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()]), - state: WorkerStatus::Idle.into(), + state: test_worker_state(WorkerStatus::Idle), in_flight: Default::default(), internal_workers: Vec::new(), }); @@ -3667,7 +3733,7 @@ mod completion_flow_tests { pending_submissions: protocol::PendingSubmissionsSnapshot::default(), entries: Vec::new(), }, - state: WorkerStatus::Running.into(), + state: test_worker_state(WorkerStatus::Running), in_flight: InFlightSnapshot { blocks: vec![ InFlightBlock::Thinking { @@ -3994,7 +4060,7 @@ mod completion_flow_tests { pending_submissions: protocol::PendingSubmissionsSnapshot::default(), entries: Vec::new(), }, - state: WorkerStatus::Idle.into(), + state: test_worker_state(WorkerStatus::Idle), in_flight: Default::default(), internal_workers: Vec::new(), }); @@ -4046,7 +4112,7 @@ mod completion_flow_tests { pending_submissions: protocol::PendingSubmissionsSnapshot::default(), entries: Vec::new(), }, - state: WorkerStatus::Idle.into(), + state: test_worker_state(WorkerStatus::Idle), in_flight: Default::default(), internal_workers: vec![InternalWorkerSnapshot { worker: InternalWorkerRef { @@ -4194,6 +4260,13 @@ mod completion_flow_tests { .count() } + fn test_worker_state(status: WorkerStatus) -> WorkerStateSnapshot { + let mut snapshot = WorkerStateSnapshot::from(status); + snapshot.execution_generation = 1; + snapshot.revision = 1; + snapshot + } + fn test_greeting() -> protocol::Greeting { protocol::Greeting { worker_name: "test".into(), @@ -4220,7 +4293,7 @@ mod completion_flow_tests { entries: Vec::new(), }, greeting, - state: WorkerStatus::Idle.into(), + state: test_worker_state(WorkerStatus::Idle), in_flight: Default::default(), internal_workers: Vec::new(), }); @@ -4419,7 +4492,7 @@ mod completion_flow_tests { app.handle_worker_event(Event::Snapshot { greeting: test_greeting(), session: public_session(assistant_item_entries), - state: WorkerStatus::Running.into(), + state: test_worker_state(WorkerStatus::Running), in_flight: Default::default(), internal_workers: Vec::new(), }); diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 6c414874..20b05aed 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -10,7 +10,7 @@ use std::collections::HashMap; use std::future::Future; use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex, RwLock, mpsc}; use std::time::Duration; @@ -38,7 +38,9 @@ use crate::working_directory::{ WorkingDirectoryBinding, WorkingDirectoryDiagnostic, WorkingDirectoryMaterializer, }; use async_trait::async_trait; -use protocol::{Event, Method, Segment, WorkerCommandEnvelope, WorkerStatus}; +#[cfg(test)] +use protocol::WorkerStatus; +use protocol::{Event, Method, Segment, WorkerCommandEnvelope}; static NEXT_INTERNAL_COMMAND_ID: AtomicU64 = AtomicU64::new(1); @@ -1197,7 +1199,6 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { struct RuntimeWorkerExecution { handle: WorkerHandle, shutdown: Arc>>, - busy: Arc, worker_state: Arc>, workspace_client: Option>, } @@ -1296,7 +1297,6 @@ where ) -> Result< ( WorkerHandle, - Arc, Arc>, Option>, ), @@ -1323,7 +1323,6 @@ where .map(|execution| { ( execution.handle.clone(), - execution.busy.clone(), execution.worker_state.clone(), execution.workspace_client.clone(), ) @@ -1434,48 +1433,31 @@ where working_directory: Option, 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_state = match &event { - Event::WorkerState { snapshot } - | Event::Snapshot { state: snapshot, .. } => { - Some(snapshot.clone()) + Ok(mut event) => { + match apply_protocol_worker_state(&bridge_worker_state, &mut event) { + Ok(true) => { + let _ = bridge_context.publish_protocol_event(event); } - Event::CommandAcknowledged { acknowledgement } => { - Some(acknowledgement.state.clone()) + Ok(false) => {} + Err(message) => { + let _ = bridge_context.publish_protocol_event(Event::Error { + code: protocol::ErrorCode::Internal, + message: format!("worker state stream rejected: {message}"), + }); + break; } - _ => 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); } } Err(broadcast::error::RecvError::Lagged(_)) => continue, @@ -1520,7 +1502,6 @@ where RuntimeWorkerExecution { handle, shutdown, - busy, worker_state, workspace_client, }, @@ -1543,32 +1524,28 @@ 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, - Method::Submit { .. } - | Method::SubmitTracked { .. } - | Method::Notify { auto_run: true, .. } - | Method::NotifyTracked { auto_run: true, .. } - | Method::Resume { .. } - ) -} - -fn method_can_start_turn_from_status(method: &Method, status: WorkerStatus) -> bool { - match method { - Method::Resume { .. } => matches!(status, WorkerStatus::Idle | WorkerStatus::Paused), - _ => status == WorkerStatus::Idle, +fn apply_protocol_worker_state( + current: &Arc>, + event: &mut Event, +) -> Result { + let (incoming, replace_stale) = match event { + Event::WorkerState { snapshot } => (snapshot, false), + Event::Snapshot { state, .. } => (state, true), + Event::CommandAcknowledged { acknowledgement } => (&mut acknowledgement.state, true), + _ => return Ok(true), + }; + let mut current = current + .write() + .map_err(|_| "worker state projection lock is poisoned".to_string())?; + match protocol::apply_worker_state_snapshot(&mut current, incoming) { + Ok(protocol::WorkerStateSnapshotApply::Applied) + | Ok(protocol::WorkerStateSnapshotApply::Duplicate) => Ok(true), + Ok(protocol::WorkerStateSnapshotApply::Stale) if replace_stale => { + *incoming = current.clone(); + Ok(true) + } + Ok(protocol::WorkerStateSnapshotApply::Stale) => Ok(false), + Err(error) => Err(error.to_string()), } } @@ -1897,7 +1874,7 @@ where handle: &WorkerExecutionHandle, input: WorkerInput, ) -> WorkerExecutionResult { - let (worker, busy, worker_state, _workspace_client) = match self.get_execution(handle) { + let (worker, worker_state, _workspace_client) = match self.get_execution(handle) { Ok(execution) => execution, Err(mut result) => { result.operation = WorkerExecutionOperation::Input; @@ -1906,15 +1883,10 @@ where }; if input.kind == WorkerInputKind::Notify { - let status = worker.shared_state.catalog_status(); - let claimed_here = status == WorkerStatus::Idle - && busy - .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) - .is_ok(); let notification_request_id = input .submission_request_id .unwrap_or_else(protocol::new_submission_request_id); - let result = self.send_method( + return self.send_method( WorkerExecutionOperation::Input, worker, Method::NotifyTracked { @@ -1926,11 +1898,6 @@ where }, }, ); - if claimed_here && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted - { - busy.store(false, Ordering::SeqCst); - } - return result; } if input.kind == WorkerInputKind::Compact { @@ -1947,26 +1914,12 @@ where ); } - let is_user_submit = input.kind == WorkerInputKind::User; - let status = worker.shared_state.catalog_status(); - let claimed_here = status == WorkerStatus::Idle - && busy - .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) - .is_ok(); - if !is_user_submit && !claimed_here { - return WorkerExecutionResult::busy( - WorkerExecutionOperation::Input, - "Worker is already running", - ); - } - let (method, submission_request_id) = match input.kind { WorkerInputKind::User => { let Some(submission_id) = input .submission_request_id .filter(|submission_id| !submission_id.trim().is_empty()) else { - busy.store(false, Ordering::SeqCst); return WorkerExecutionResult::rejected( WorkerExecutionOperation::Input, "Runtime user input is missing its internal submission id", @@ -1986,7 +1939,7 @@ where ) } WorkerInputKind::Notify => { - unreachable!("Notify input is dispatched before the turn-start busy guard") + unreachable!("Notify input is dispatched before ordinary input mapping") } WorkerInputKind::Compact => unreachable!("compact input is dispatched above"), WorkerInputKind::ListRewindTargets => (Method::ListRewindTargets, None), @@ -1999,7 +1952,7 @@ where }; let waits_for_submission_acceptance = submission_request_id.is_some(); - let result = if waits_for_submission_acceptance { + if waits_for_submission_acceptance { self.send_submit_and_wait_for_acceptance( WorkerExecutionOperation::Input, worker, @@ -2008,11 +1961,7 @@ where ) } else { self.send_method(WorkerExecutionOperation::Input, worker, method) - }; - if claimed_here && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted { - busy.store(false, Ordering::SeqCst); } - result } fn upload_file( @@ -2023,7 +1972,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 })?; @@ -2046,7 +1995,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; @@ -2067,7 +2016,7 @@ where handle: &WorkerExecutionHandle, method: Method, ) -> WorkerExecutionResult { - let (worker, busy, _worker_state, _workspace_client) = match self.get_execution(handle) { + let (worker, _worker_state, _workspace_client) = match self.get_execution(handle) { Ok(execution) => execution, Err(mut result) => { result.operation = WorkerExecutionOperation::ProtocolMethod; @@ -2075,44 +2024,7 @@ where } }; - if let Some(auto_run) = match &method { - Method::Notify { auto_run, .. } | Method::NotifyTracked { auto_run, .. } => { - Some(*auto_run) - } - _ => None, - } { - 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); - if claimed_here && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted - { - busy.store(false, Ordering::SeqCst); - } - return result; - } - - let starts_turn = method_starts_turn(&method); - if starts_turn - && (!method_can_start_turn_from_status(&method, worker.shared_state.catalog_status()) - || busy - .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) - .is_err()) - { - return WorkerExecutionResult::busy( - WorkerExecutionOperation::ProtocolMethod, - "Worker is already running; runtime adapter v0 does not queue protocol methods", - ); - } - - let result = self.send_method(WorkerExecutionOperation::ProtocolMethod, worker, method); - if starts_turn && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted { - busy.store(false, Ordering::SeqCst); - } - result + self.send_method(WorkerExecutionOperation::ProtocolMethod, worker, method) } fn stop_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult { @@ -2193,7 +2105,7 @@ where } fn cancel_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult { - let (worker, _busy, worker_state, _workspace_client) = match self.get_execution(handle) { + let (worker, worker_state, _workspace_client) = match self.get_execution(handle) { Ok(execution) => execution, Err(mut result) => { result.operation = WorkerExecutionOperation::Cancel; @@ -2297,6 +2209,56 @@ mod tests { WorkerCommandEnvelope::for_snapshot(state.last_command_id.saturating_add(1), &state) } + #[test] + fn protocol_bridge_applies_state_and_acknowledgement_monotonically() { + let running = protocol::WorkerStateSnapshot { + execution_generation: 4, + revision: 3, + state: protocol::WorkerState::Busy(protocol::WorkerBusyState::Run( + protocol::WorkerRunState::Running, + )), + last_command_id: 2, + }; + let current = Arc::new(RwLock::new(running.clone())); + let mut stale = Event::WorkerState { + snapshot: protocol::WorkerStateSnapshot { + revision: 2, + state: protocol::WorkerState::Idle, + ..running.clone() + }, + }; + assert!(!apply_protocol_worker_state(¤t, &mut stale).unwrap()); + assert_eq!(*current.read().unwrap(), running); + + let paused = protocol::WorkerStateSnapshot { + revision: 4, + state: protocol::WorkerState::Busy(protocol::WorkerBusyState::Run( + protocol::WorkerRunState::Paused, + )), + last_command_id: 3, + ..running.clone() + }; + let mut acknowledgement = Event::CommandAcknowledged { + acknowledgement: protocol::WorkerCommandAcknowledgement { + command_id: 3, + command: protocol::WorkerCommandKind::Pause, + disposition: protocol::WorkerCommandDisposition::Accepted, + state: paused.clone(), + }, + }; + assert!(apply_protocol_worker_state(¤t, &mut acknowledgement).unwrap()); + assert_eq!(*current.read().unwrap(), paused); + + let mut conflict = Event::WorkerState { + snapshot: protocol::WorkerStateSnapshot { + state: protocol::WorkerState::Idle, + ..paused.clone() + }, + }; + assert!(apply_protocol_worker_state(¤t, &mut conflict).is_err()); + assert_eq!(*current.read().unwrap(), paused); + } + #[test] fn workspace_prompt_projection_notification_advances_shared_cache() { let cache = WorkspacePromptProjectionCache::default(); @@ -2443,44 +2405,6 @@ mod tests { assert_eq!(after_restore_workspace_id.as_deref(), Some("workspace-a")); } - #[test] - 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 { - command: test_command() - }, - WorkerStatus::Paused - )); - assert!(method_can_start_turn_from_status( - &Method::Resume { - command: test_command() - }, - WorkerStatus::Idle - )); - assert!(!method_can_start_turn_from_status( - &Method::Resume { - command: test_command() - }, - WorkerStatus::Running - )); - assert!(!method_can_start_turn_from_status( - &Method::Compact { - command: test_command() - }, - WorkerStatus::Paused - )); - } - #[derive(Clone)] enum MockResponse { Complete(Vec), @@ -2681,11 +2605,38 @@ mod tests { .collect() } + fn wait_for_adapter_command( + backend: &WorkerRuntimeExecutionBackend, + worker_ref: &WorkerRef, + expected_command_id: u64, + ) { + let deadline = std::time::Instant::now() + Duration::from_secs(5); + loop { + let observed = { + let workers = backend.workers.lock().unwrap(); + workers + .get(worker_ref) + .expect("live Worker execution") + .worker_state + .read() + .unwrap() + .last_command_id + }; + if observed >= expected_command_id { + return; + } + assert!( + std::time::Instant::now() < deadline, + "timed out waiting for adapter command {expected_command_id}; last observed={observed}", + ); + std::thread::sleep(Duration::from_millis(10)); + } + } + fn wait_for_adapter_state( backend: &WorkerRuntimeExecutionBackend, worker_ref: &WorkerRef, expected_status: WorkerStatus, - expected_busy: bool, ) { let deadline = std::time::Instant::now() + Duration::from_secs(5); loop { @@ -2693,21 +2644,16 @@ mod tests { 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.catalog_status(), - projected, - execution.busy.load(Ordering::SeqCst), - ) + (execution.handle.shared_state.catalog_status(), projected) }; - if observed == (expected_status, expected_status, expected_busy) { + if observed == (expected_status, expected_status) { return; } assert!( std::time::Instant::now() < deadline, - "timed out waiting for adapter state {expected_status:?}, busy={expected_busy}; last observed controller={:?}, projected={:?}, busy={}", + "timed out waiting for adapter state {expected_status:?}; last observed controller={:?}, projected={:?}", observed.0, observed.1, - observed.2, ); std::thread::sleep(Duration::from_millis(10)); } @@ -3694,22 +3640,18 @@ mod tests { runtime .send_input(&detail.worker_ref, WorkerInput::user("pause and resume")) .expect("start initial turn"); - wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Running, true); + wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Running); - let running_resume = runtime + let running_resume = adapter_command(&backend, &detail.worker_ref); + runtime .send_protocol_method( &detail.worker_ref, Method::Resume { - command: adapter_command(&backend, &detail.worker_ref), + command: running_resume, }, ) - .expect_err("Resume while Running must be rejected"); - assert!( - running_resume - .to_string() - .contains("does not queue protocol methods"), - "unexpected Running Resume error: {running_resume}" - ); + .expect("running Resume is forwarded for controller admission"); + wait_for_adapter_command(&backend, &detail.worker_ref, running_resume.command_id); runtime .send_protocol_method( @@ -3719,7 +3661,7 @@ mod tests { }, ) .expect("pause initial turn"); - wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Paused, false); + wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Paused); runtime .send_protocol_method( @@ -3729,22 +3671,18 @@ mod tests { }, ) .expect("resume paused turn"); - wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Running, true); + wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Running); - let duplicate_resume = runtime + let duplicate_resume = adapter_command(&backend, &detail.worker_ref); + runtime .send_protocol_method( &detail.worker_ref, Method::Resume { - command: adapter_command(&backend, &detail.worker_ref), + command: duplicate_resume, }, ) - .expect_err("duplicate Resume must be rejected"); - assert!( - duplicate_resume - .to_string() - .contains("does not queue protocol methods"), - "unexpected duplicate Resume error: {duplicate_resume}" - ); + .expect("duplicate Resume is forwarded for controller admission"); + wait_for_adapter_command(&backend, &detail.worker_ref, duplicate_resume.command_id); runtime .send_protocol_method( @@ -3754,7 +3692,7 @@ mod tests { }, ) .expect("pause resumed turn"); - wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Paused, false); + wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Paused); runtime .send_protocol_method( &detail.worker_ref, @@ -3763,18 +3701,20 @@ mod tests { }, ) .expect("resume paused turn a second time"); - wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Idle, false); + wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Idle); assert_eq!(call_count.load(Ordering::SeqCst), 3); + let idle_resume = adapter_command(&backend, &detail.worker_ref); runtime .send_protocol_method( &detail.worker_ref, Method::Resume { - command: adapter_command(&backend, &detail.worker_ref), + command: idle_resume, }, ) .expect("Idle Resume preserves controller NotPaused semantics"); - wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Idle, false); + wait_for_adapter_command(&backend, &detail.worker_ref, idle_resume.command_id); + wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Idle); let events = runtime .read_worker_observation_events(&detail.worker_ref, WorkerObservationCursor::zero()) .expect("read protocol events"); diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index f9cc2388..beeb4791 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -3708,7 +3708,7 @@ mod tests { WorkerCommandEnvelope { command_id: 1, expected_execution_generation: 9, - expected_worker_state_revision: 0, + expected_worker_state_revision: 1, }, &shared, ), @@ -3719,7 +3719,7 @@ mod tests { WorkerCommandEnvelope { command_id: 2, expected_execution_generation: 9, - expected_worker_state_revision: 0, + expected_worker_state_revision: 1, }, &shared, ) diff --git a/crates/worker/src/shared_state.rs b/crates/worker/src/shared_state.rs index 58c5700b..573f2bb5 100644 --- a/crates/worker/src/shared_state.rs +++ b/crates/worker/src/shared_state.rs @@ -1,6 +1,6 @@ use std::sync::{ OnceLock, RwLock, - atomic::{AtomicBool, AtomicU64, Ordering}, + atomic::{AtomicBool, Ordering}, }; use protocol::{ @@ -23,7 +23,6 @@ pub struct WorkerSharedState { pub manifest_toml: String, pub greeting: protocol::Greeting, 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 @@ -56,7 +55,6 @@ impl WorkerSharedState { manifest_toml, greeting, state: RwLock::new(WorkerStateSnapshot::initial(execution_generation)), - last_command_id: AtomicU64::new(0), fs_view: OnceLock::new(), flow_transition_enabled: AtomicBool::new(false), } @@ -91,26 +89,27 @@ impl WorkerSharedState { 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() + let mut snapshot = self + .state + .write() + .expect("worker state lock poisoned; refusing command admission"); + if command_id <= snapshot.last_command_id { + return false; + } + snapshot.last_command_id = command_id; + snapshot.revision = snapshot.revision.saturating_add(1); + true } pub fn snapshot(&self) -> WorkerStateSnapshot { - let mut snapshot = self - .state + 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 + .clone() } /// Runtime catalog projection. This must not be used as live command @@ -190,6 +189,23 @@ mod tests { assert_eq!(state.catalog_status(), WorkerStatus::Paused); } + #[test] + fn accepted_command_id_advances_the_snapshot_revision_atomically() { + let state = test_state(); + assert!(state.accept_command_id(9)); + assert_eq!( + state.snapshot(), + WorkerStateSnapshot { + execution_generation: 7, + revision: 1, + last_command_id: 9, + state: WorkerState::Idle, + } + ); + assert!(!state.accept_command_id(9)); + assert_eq!(state.snapshot().revision, 1); + } + #[test] fn status_json_contains_full_snapshot_and_catalog_projection() { let state = test_state(); diff --git a/web/workspace/src/lib/workspace/console/model.test.ts b/web/workspace/src/lib/workspace/console/model.test.ts index c1f1b72a..de1e46b3 100644 --- a/web/workspace/src/lib/workspace/console/model.test.ts +++ b/web/workspace/src/lib/workspace/console/model.test.ts @@ -218,6 +218,66 @@ Deno.test("console routing projects live errors but not completion replies", () ); }); +Deno.test("Worker state events and acknowledgements apply monotonically", () => { + const projector = createConsoleProjector(); + const running: WorkerStateSnapshot = { + execution_generation: 4, + revision: 3, + last_command_id: 2, + state: { kind: "busy", state: { kind: "run", state: "running" } }, + }; + const paused: WorkerStateSnapshot = { + ...running, + revision: 4, + last_command_id: 3, + state: { kind: "busy", state: { kind: "run", state: "paused" } }, + }; + let projection = projector.append([ + { + eventId: "running", + event: { event: "worker_state", data: { snapshot: running } }, + }, + { + eventId: "stale", + event: { + event: "worker_state", + data: { snapshot: { ...running, revision: 2, state: { kind: "idle" } } }, + }, + }, + { + eventId: "pause-ack", + event: { + event: "command_acknowledged", + data: { + acknowledgement: { + command_id: 3, + command: "pause", + disposition: "accepted", + state: paused, + }, + }, + }, + }, + ]); + assertEquals(projection.workerState, paused); + assertEquals(projection.status, "paused"); + + projection = projector.append([{ + eventId: "conflict", + event: { + event: "worker_state", + data: { snapshot: { ...paused, state: { kind: "idle" } } }, + }, + }]); + assertEquals(projection.workerState, paused); + assert( + projection.lines.some((line) => + line.eventId === "conflict:worker-state-conflict" && line.error + ), + "conflicting equal-version snapshots must fail closed", + ); +}); + Deno.test("snapshot replaces a live error with one durable run_errored row", () => { const projector = createConsoleProjector(); let projection = projector.append([ diff --git a/web/workspace/src/lib/workspace/console/model.ts b/web/workspace/src/lib/workspace/console/model.ts index 857e1b30..95578cbd 100644 --- a/web/workspace/src/lib/workspace/console/model.ts +++ b/web/workspace/src/lib/workspace/console/model.ts @@ -10,6 +10,7 @@ import type { InternalWorkerRef, InternalWorkerSnapshot, Segment, + WorkerState, WorkerStateSnapshot, WorkerStatus, } from "$lib/generated/protocol"; @@ -796,6 +797,60 @@ function refreshCompactionActivity( return changed ? { ...projection, lines } : projection; } +function workerStateEqual(left: WorkerState, right: WorkerState): boolean { + if (left.kind !== right.kind) return false; + if (left.kind === "idle" || right.kind === "idle") return true; + return left.state.kind === right.state.kind && + left.state.state === right.state.state; +} + +function workerStateSnapshotEqual( + left: WorkerStateSnapshot, + right: WorkerStateSnapshot, +): boolean { + return left.execution_generation === right.execution_generation && + left.revision === right.revision && + left.last_command_id === right.last_command_id && + workerStateEqual(left.state, right.state); +} + +function applyWorkerStateSnapshot( + projection: ConsoleProjection, + incoming: WorkerStateSnapshot, + eventId: string, +): void { + const current = projection.workerState; + if (!current) { + projection.workerState = incoming; + projection.status = workerStatusFromState(incoming); + return; + } + const generationOrder = incoming.execution_generation - + current.execution_generation; + const revisionOrder = incoming.revision - current.revision; + if (generationOrder > 0 || (generationOrder === 0 && revisionOrder > 0)) { + projection.workerState = incoming; + projection.status = workerStatusFromState(incoming); + return; + } + if (generationOrder < 0 || (generationOrder === 0 && revisionOrder < 0)) { + return; + } + if (!workerStateSnapshotEqual(current, incoming)) { + projection.lines.push( + line( + `${eventId}:worker-state-conflict`, + "error", + "error · internal", + `worker state stream rejected: conflicting snapshots at generation ${incoming.execution_generation} revision ${incoming.revision}`, + undefined, + false, + true, + ), + ); + } +} + export function applyProtocolEvent( projection: ConsoleProjection, envelope: ConsoleEventInput, @@ -917,8 +972,6 @@ export function applyProtocolEvent( ); break; case "snapshot": { - next.workerState = event.data.state; - next.status = workerStatusFromState(event.data.state); next.cwd = event.data.greeting.cwd; const snapshot = snapshotProjectionFromSession( envelope.eventId, @@ -968,6 +1021,7 @@ export function applyProtocolEvent( }; } } + applyWorkerStateSnapshot(next, event.data.state, envelope.eventId); break; } case "internal_worker": { @@ -1016,12 +1070,14 @@ export function applyProtocolEvent( break; } case "worker_state": - next.workerState = event.data.snapshot; - next.status = workerStatusFromState(event.data.snapshot); + applyWorkerStateSnapshot(next, event.data.snapshot, envelope.eventId); break; case "command_acknowledged": - next.workerState = event.data.acknowledgement.state; - next.status = workerStatusFromState(event.data.acknowledgement.state); + applyWorkerStateSnapshot( + next, + event.data.acknowledgement.state, + envelope.eventId, + ); break; case "command": applyCommandEvent(next, envelope.eventId, event.data.event); From 2d1956b653a4fca802a003cac73392787161f6be Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 10:05:54 +0900 Subject: [PATCH 3/5] fix: separate worker state from runtime lifecycle --- crates/protocol/src/lib.rs | 1 + crates/protocol/src/subscription.rs | 6 + crates/tui/src/app.rs | 53 +++-- crates/tui/src/backend_worker_picker.rs | 34 ++- crates/worker-runtime/src/catalog.rs | 6 + crates/worker-runtime/src/http_server.rs | 11 +- crates/worker-runtime/src/runtime.rs | 209 +++++++++++------- crates/worker/src/controller.rs | 138 +++++++++--- crates/worker/src/shared_state.rs | 149 ++++++++++++- crates/workspace-api/src/lib.rs | 5 + crates/workspace-server/src/hosts.rs | 10 + .../src/runtime_subscription_tests.rs | 51 ++++- crates/workspace-server/src/server.rs | 3 + web/workspace/src/lib/generated/protocol.ts | 13 +- .../src/lib/workspace/sidebar/types.ts | 2 + .../src/lib/workspace/sidebar/worker-state.ts | 12 + .../sidebar/worker-subscription.test.ts | 17 ++ .../workspace/sidebar/worker-subscription.ts | 4 +- .../workers/[workerId]/console/+page.svelte | 60 ++--- .../w/[workspaceId]/workers/+page.svelte | 3 +- .../workers/[workerRef]/+page.svelte | 3 +- 21 files changed, 602 insertions(+), 188 deletions(-) create mode 100644 web/workspace/src/lib/workspace/sidebar/worker-state.ts diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index f66be5dc..d54c9232 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -125,6 +125,7 @@ pub enum WorkerCommandDisposition { StaleExecutionGeneration, StaleWorkerStateRevision, StaleCommandId, + Conflict, InvalidState, } diff --git a/crates/protocol/src/subscription.rs b/crates/protocol/src/subscription.rs index 468e241c..bf4a8635 100644 --- a/crates/protocol/src/subscription.rs +++ b/crates/protocol/src/subscription.rs @@ -573,6 +573,11 @@ pub struct SubscriptionWorker { pub resource_key: Option, /// Producer-owned monotonic revision for this Worker subject. pub subject_revision: u64, + /// Latest revisioned foreground state observed from the Worker. This remains + /// absent until an authoritative Worker snapshot/event has been applied. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub worker_state: Option, + /// Runtime catalog lifecycle compatibility projection; not foreground-state authority. pub state: SubscriptionWorkerState, #[serde(default)] pub has_running_internal_workers: bool, @@ -874,6 +879,7 @@ mod tests { runtime_id: None, resource_key: None, subject_revision: 0, + worker_state: None, state: SubscriptionWorkerState::Idle, has_running_internal_workers: false, workspace_id: Some("workspace-1".to_string()), diff --git a/crates/tui/src/app.rs b/crates/tui/src/app.rs index c543f991..9858dd39 100644 --- a/crates/tui/src/app.rs +++ b/crates/tui/src/app.rs @@ -1180,18 +1180,14 @@ impl App { self.assistant_streaming = false; } Event::TurnStart { .. } => { - self.set_worker_status(WorkerStatus::Running); self.run_requests += 1; self.current_tool = None; self.latest_llm_wait_event = None; self.assistant_streaming = false; } - Event::InvokeStart { .. } => { - self.set_worker_status(WorkerStatus::Running); - } + Event::InvokeStart { .. } => {} // UI consumers of per-attempt LlmCall semantics remain out of scope; - // the run-level status starts at InvokeStart and TurnStart counts each - // LLM request within that run. + // authoritative run state comes only from WorkerStateSnapshot. Event::LlmCallStart { .. } | Event::LlmCallEnd { .. } => { self.latest_llm_wait_event = None; } @@ -1398,12 +1394,7 @@ impl App { output_tokens: self.run_output_tokens, }); self.pending_submit_rollback = None; - self.reset_run_state(match result { - RunResult::Paused => WorkerStatus::Paused, - RunResult::Finished | RunResult::LimitReached | RunResult::RolledBack => { - WorkerStatus::Idle - } - }); + self.reset_run_state(); } } Event::CompactStart { .. } => { @@ -1536,7 +1527,7 @@ impl App { }; self.completion = None; self.close_rewind_picker(); - self.reset_run_state(self.worker_status); + self.reset_run_state(); let mut message = if restored_composer { format!( "Rewound session: discarded {} log entries; restored selected input to composer.", @@ -1584,8 +1575,7 @@ impl App { None } - fn reset_run_state(&mut self, status: WorkerStatus) { - self.set_worker_status(status); + fn reset_run_state(&mut self) { self.run_requests = 0; self.run_upload_tokens = 0; self.run_output_tokens = 0; @@ -1615,7 +1605,7 @@ impl App { "Rolled back empty assistant turn; no local submitted input was available to restore." .to_owned() }; - self.reset_run_state(WorkerStatus::Idle); + self.reset_run_state(); self.blocks.push(Block::Alert { level: AlertLevel::Warn, source: AlertSource::Worker, @@ -3583,6 +3573,37 @@ mod completion_flow_tests { assert!(matches!(app.blocks.first(), Some(Block::Greeting(_)))); } + #[test] + fn occurrence_events_do_not_infer_foreground_worker_state() { + let mut app = App::new("test".into()); + app.handle_worker_event(Event::TurnStart { turn: 1 }); + app.handle_worker_event(Event::InvokeStart { + kind: protocol::InvokeKind::UserSend, + }); + app.handle_worker_event(Event::RunEnd { + result: RunResult::Paused, + }); + assert_eq!(app.worker_state.state, protocol::WorkerState::Idle); + assert_eq!(app.worker_status, WorkerStatus::Idle); + + let running = WorkerStateSnapshot { + execution_generation: 1, + revision: 1, + state: protocol::WorkerState::Busy(protocol::WorkerBusyState::Run( + protocol::WorkerRunState::Running, + )), + last_command_id: 0, + }; + app.handle_worker_event(Event::WorkerState { + snapshot: running.clone(), + }); + app.handle_worker_event(Event::RunEnd { + result: RunResult::Finished, + }); + assert_eq!(app.worker_state, running); + assert_eq!(app.worker_status, WorkerStatus::Running); + } + #[test] fn worker_state_events_and_acknowledgements_share_monotonic_application() { let mut app = App::new("test".into()); diff --git a/crates/tui/src/backend_worker_picker.rs b/crates/tui/src/backend_worker_picker.rs index e9a37271..473ca5eb 100644 --- a/crates/tui/src/backend_worker_picker.rs +++ b/crates/tui/src/backend_worker_picker.rs @@ -321,7 +321,7 @@ fn row_line( Span::raw(" "), Span::styled( pad_column(&worker_state(worker), widths.state), - state_style(worker.state.as_str()), + state_style(worker_state_label(worker)), ), Span::raw(" "), Span::styled( @@ -341,8 +341,20 @@ fn worker_name(worker: &BackendWorkerSummary) -> &str { } } +fn worker_state_label(worker: &BackendWorkerSummary) -> &str { + match worker.worker_state.as_ref().map(|state| &state.state) { + Some(protocol::WorkerState::Idle) => "idle", + Some(protocol::WorkerState::Busy(protocol::WorkerBusyState::Run( + protocol::WorkerRunState::Paused, + ))) => "paused", + Some(protocol::WorkerState::Busy(_)) => "running", + None if worker.state == "stopped" => "stopped", + None => "unknown", + } +} + fn worker_state(worker: &BackendWorkerSummary) -> String { - format!("[{}]", worker.state) + format!("[{}]", worker_state_label(worker)) } fn text_width(value: &str) -> usize { @@ -413,7 +425,15 @@ mod tests { identity: "ws".to_string(), workspace_id: Some("ws".to_string()), }, - state: "running".to_string(), + state: "idle".to_string(), + worker_state: Some(protocol::WorkerStateSnapshot { + execution_generation: 1, + revision: 1, + state: protocol::WorkerState::Busy(protocol::WorkerBusyState::Run( + protocol::WorkerRunState::Running, + )), + last_command_id: 0, + }), last_seen_at: None, pinned: false, retention_state: String::new(), @@ -450,6 +470,7 @@ mod tests { worker.display_name = "Coder".to_string(); worker.label = "Coder · T-585".to_string(); worker.state = "stopped".to_string(); + worker.worker_state = None; worker.working_directory = Some( serde_json::from_value(serde_json::json!({ "working_directory_id": "001a06a9f0202000000", @@ -478,12 +499,19 @@ mod tests { short.label = "Coder".to_string(); short.display_name = short.label.clone(); short.state = "idle".to_string(); + short.worker_state = Some(protocol::WorkerStateSnapshot { + execution_generation: 1, + revision: 2, + state: protocol::WorkerState::Idle, + last_command_id: 0, + }); let mut long = worker("runtime-a", "worker-b", None); long.resource_key = "W-100".to_string(); long.label = "Longer worker · T-9".to_string(); long.display_name = long.label.clone(); long.state = "stopped".to_string(); + long.worker_state = None; for worker in [&mut short, &mut long] { worker.working_directory = Some( diff --git a/crates/worker-runtime/src/catalog.rs b/crates/worker-runtime/src/catalog.rs index 3cc27d01..96bb2beb 100644 --- a/crates/worker-runtime/src/catalog.rs +++ b/crates/worker-runtime/src/catalog.rs @@ -307,6 +307,8 @@ pub struct WorkerSummary { pub worker_id: WorkerId, pub status: WorkerStatus, #[serde(default, skip_serializing_if = "Option::is_none")] + pub worker_state: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] pub workspace_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub working_directory: Option, @@ -325,6 +327,8 @@ pub struct WorkerDetail { pub worker_id: WorkerId, pub status: WorkerStatus, #[serde(default, skip_serializing_if = "Option::is_none")] + pub worker_state: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] pub workspace_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub working_directory: Option, @@ -341,6 +345,8 @@ pub struct WorkerDetail { pub struct WorkerLifecycleAck { pub worker_ref: WorkerRef, pub status: WorkerStatus, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub worker_state: Option, } #[cfg(test)] diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index 28cf26aa..c05bdb91 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -3486,7 +3486,16 @@ mod ws_tests { .. }) if delivered_subscription_id == subscription_id && worker.worker_id.as_str() == worker_ref.worker_id.to_string() - && worker.state == protocol::subscription::SubscriptionWorkerState::Running + && worker.state == protocol::subscription::SubscriptionWorkerState::Idle + && matches!( + worker.worker_state, + Some(protocol::WorkerStateSnapshot { + state: protocol::WorkerState::Busy(protocol::WorkerBusyState::Run( + protocol::WorkerRunState::Running + )), + .. + }) + ) )); let unsubscribe_request_id = diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 808788f4..76066336 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -700,6 +700,7 @@ impl Runtime { worker_ref: worker_ref.clone(), worker_id: worker_id.clone(), status: WorkerStatus::Stopped, + worker_state: None, workspace_id: scope.map(|scope| scope.workspace_id.clone()), request: durable_request, run_generation: 1, @@ -787,7 +788,6 @@ impl Runtime { let detail = self.commit_created_worker( &worker_ref, handle, - WorkerStatus::Running, working_directory, dispatch_result, )?; @@ -797,7 +797,6 @@ impl Runtime { self.commit_created_worker( &worker_ref, handle, - WorkerStatus::Idle, working_directory, WorkerExecutionResult::accepted(WorkerExecutionOperation::Spawn), ) @@ -1220,17 +1219,7 @@ impl Runtime { state.ensure_running()?; let worker = state.worker_mut(worker_ref)?; 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 _ = worker.apply_worker_state(snapshot); } let status = worker.status; #[cfg(feature = "ws-server")] @@ -1490,16 +1479,19 @@ impl Runtime { &self, worker_ref: &WorkerRef, handle: WorkerExecutionHandle, - status: WorkerStatus, working_directory: Option, - _result: WorkerExecutionResult, + result: WorkerExecutionResult, ) -> Result { let mut state = self.lock()?; let detail = { let worker = state.worker_mut(worker_ref)?; worker.execution_handle = Some(handle); worker.execution_bound = true; - worker.status = status; + worker.status = WorkerStatus::Idle; + worker.worker_state = None; + if let Some(snapshot) = result.worker_state.as_ref() { + let _ = worker.apply_worker_state(snapshot); + } worker.restore_intent = restore_intent_for_status(worker.status); worker.working_directory = working_directory; worker.detail() @@ -1536,16 +1528,14 @@ impl Runtime { 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); + let applied = worker + .apply_worker_state(&snapshot) + .is_ok_and(|result| matches!(result, protocol::WorkerStateSnapshotApply::Applied)); + if !applied { + return Ok(()); + } state.publish_worker_upsert(worker_ref.worker_id)?; state.persist_runtime_snapshot()?; state.persist_worker(&worker_ref.worker_id)?; @@ -1636,20 +1626,26 @@ impl Runtime { worker_ref: &WorkerRef, reason: Option, ) -> Result { - let current = { + { let state = self.lock()?; state.ensure_running()?; - state.worker(worker_ref)?.status - }; - if matches!(current, WorkerStatus::Idle | WorkerStatus::Stopped) { - return Ok(WorkerLifecycleAck { - worker_ref: worker_ref.clone(), - status: current, - }); + if state.worker(worker_ref)?.status == WorkerStatus::Stopped { + return Ok(WorkerLifecycleAck { + worker_ref: worker_ref.clone(), + status: WorkerStatus::Stopped, + worker_state: None, + }); + } } self.dispatch_lifecycle_to_backend(worker_ref, WorkerExecutionOperation::Cancel)?; let _ = reason; - self.transition_worker_preserving_execution(worker_ref, WorkerStatus::Idle) + let state = self.lock()?; + let worker = state.worker(worker_ref)?; + Ok(WorkerLifecycleAck { + worker_ref: worker_ref.clone(), + status: worker.status, + worker_state: worker.worker_state.clone(), + }) } /// Delete a non-running Worker through a workspace-scoped Runtime authorization context. @@ -1795,12 +1791,13 @@ impl Runtime { ) -> Result { let mut state = self.lock()?; state.ensure_worker_ref(worker_ref)?; - let status_changed = state.project_protocol_event_to_status(worker_ref, &payload); + let worker_state_changed = + state.project_protocol_event_to_worker_state(worker_ref, &payload); let activity_changed = state.project_internal_worker_activity(worker_ref, &payload); - if status_changed || activity_changed { + if worker_state_changed || activity_changed { state.publish_worker_upsert(worker_ref.worker_id)?; } - if status_changed { + if worker_state_changed { state.persist_runtime_snapshot()?; state.persist_worker(&worker_ref.worker_id)?; } @@ -1836,26 +1833,6 @@ impl Runtime { Ok(()) } - fn transition_worker_preserving_execution( - &self, - worker_ref: &WorkerRef, - status: WorkerStatus, - ) -> Result { - let mut state = self.lock()?; - state.ensure_running()?; - let worker = state.worker_mut(worker_ref)?; - worker.status = status; - worker.restore_intent = restore_intent_for_status(status); - let status = worker.status; - state.publish_worker_upsert(worker_ref.worker_id)?; - state.persist_runtime_snapshot()?; - state.persist_worker(&worker_ref.worker_id)?; - Ok(WorkerLifecycleAck { - worker_ref: worker_ref.clone(), - status, - }) - } - fn transition_worker( &self, worker_ref: &WorkerRef, @@ -1867,6 +1844,7 @@ impl Runtime { let worker = state.worker_mut(worker_ref)?; worker.status = status; + worker.worker_state = None; worker.restore_intent = restore_intent_for_status(status); worker.execution_handle = None; worker.internal_workers.clear(); @@ -1877,6 +1855,7 @@ impl Runtime { Ok(WorkerLifecycleAck { worker_ref: worker_ref.clone(), status, + worker_state: None, }) } @@ -2301,6 +2280,7 @@ impl RuntimeState { worker_ref: worker.worker_ref, worker_id: worker.worker_id, status: worker.status, + worker_state: None, workspace_id: worker.workspace_id, request: worker.request, run_generation, @@ -2630,6 +2610,7 @@ impl RuntimeState { .get(&worker.worker_id) .copied() .unwrap_or(0), + worker_state: worker.worker_state.clone(), state: subscription_worker_state(worker.status), has_running_internal_workers: worker .internal_workers @@ -2965,7 +2946,7 @@ impl RuntimeState { Self::update_internal_worker_activity(&mut worker.internal_workers, event) } - fn project_protocol_event_to_status( + fn project_protocol_event_to_worker_state( &mut self, worker_ref: &WorkerRef, event: &protocol::Event, @@ -2973,7 +2954,7 @@ impl RuntimeState { let Some(worker) = self.workers.get_mut(&worker_ref.worker_id) else { return false; }; - let next_status = match event { + let incoming = match event { protocol::Event::WorkerState { snapshot } | protocol::Event::Snapshot { state: snapshot, .. @@ -2983,21 +2964,16 @@ impl RuntimeState { 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, + } => snapshot, + _ => return false, }; - if let Some(next_status) = next_status { - let changed = worker.status != next_status; - worker.status = next_status; - worker.restore_intent = restore_intent_for_status(next_status); - changed - } else { - false + match worker.apply_worker_state(incoming) { + Ok(protocol::WorkerStateSnapshotApply::Applied) => true, + Ok( + protocol::WorkerStateSnapshotApply::Duplicate + | protocol::WorkerStateSnapshotApply::Stale, + ) + | Err(_) => false, } } } @@ -3013,6 +2989,7 @@ struct WorkerRecord { worker_ref: WorkerRef, worker_id: WorkerId, status: WorkerStatus, + worker_state: Option, workspace_id: Option, request: CreateWorkerRequest, run_generation: u64, @@ -3024,6 +3001,19 @@ struct WorkerRecord { } impl WorkerRecord { + fn apply_worker_state( + &mut self, + incoming: &protocol::WorkerStateSnapshot, + ) -> Result { + match self.worker_state.as_mut() { + Some(current) => protocol::apply_worker_state_snapshot(current, incoming), + None => { + self.worker_state = Some(incoming.clone()); + Ok(protocol::WorkerStateSnapshotApply::Applied) + } + } + } + fn belongs_to_workspace(&self, workspace_id: &str) -> bool { self.workspace_id.as_deref() == Some(workspace_id) } @@ -3033,6 +3023,7 @@ impl WorkerRecord { worker_ref: self.worker_ref.clone(), worker_id: self.worker_id, status: self.status, + worker_state: self.worker_state.clone(), workspace_id: self.workspace_id.clone(), working_directory: self.working_directory.clone(), profile: self.request.profile.clone(), @@ -3047,6 +3038,7 @@ impl WorkerRecord { worker_ref: self.worker_ref.clone(), worker_id: self.worker_id, status: self.status, + worker_state: self.worker_state.clone(), workspace_id: self.workspace_id.clone(), working_directory: self.working_directory.clone(), profile: self.request.profile.clone(), @@ -4719,7 +4711,7 @@ mod tests { } #[test] - fn create_worker_uses_started_submission_ack_for_initial_running_status() { + fn create_worker_does_not_infer_state_from_started_submission_ack() { let (runtime, backend) = runtime_and_backend(); backend.set_dispatch_result(WorkerExecutionResult::accepted_submission( WorkerExecutionOperation::Input, @@ -4732,7 +4724,69 @@ mod tests { let detail = runtime.create_worker(request).unwrap(); - assert_eq!(detail.status, WorkerStatus::Running); + assert_eq!(detail.status, WorkerStatus::Idle); + assert_eq!(detail.worker_state, None); + } + + #[test] + fn runtime_applies_only_newer_worker_state_snapshots() { + let (runtime, _) = runtime_and_backend(); + let detail = runtime + .create_worker(task_request("state ordering")) + .unwrap(); + let running = protocol::WorkerStateSnapshot { + execution_generation: 7, + revision: 3, + state: protocol::WorkerState::Busy(protocol::WorkerBusyState::Run( + protocol::WorkerRunState::Running, + )), + last_command_id: 2, + }; + assert!({ + let mut state = runtime.lock().unwrap(); + state.project_protocol_event_to_worker_state( + &detail.worker_ref, + &protocol::Event::WorkerState { + snapshot: running.clone(), + }, + ) + }); + assert_eq!( + runtime + .worker_detail(&detail.worker_ref) + .unwrap() + .worker_state, + Some(running.clone()) + ); + + assert!({ + let mut state = runtime.lock().unwrap(); + !state.project_protocol_event_to_worker_state( + &detail.worker_ref, + &protocol::Event::WorkerState { + snapshot: protocol::WorkerStateSnapshot { + revision: 2, + state: protocol::WorkerState::Idle, + ..running.clone() + }, + }, + ) + }); + assert!({ + let mut state = runtime.lock().unwrap(); + !state.project_protocol_event_to_worker_state( + &detail.worker_ref, + &protocol::Event::WorkerState { + snapshot: protocol::WorkerStateSnapshot { + state: protocol::WorkerState::Idle, + ..running.clone() + }, + }, + ) + }); + let after = runtime.worker_detail(&detail.worker_ref).unwrap(); + assert_eq!(after.status, WorkerStatus::Idle); + assert_eq!(after.worker_state, Some(running)); } #[test] @@ -5027,10 +5081,9 @@ mod tests { assert_eq!(*backend.restore_count.lock().unwrap(), 1); assert_eq!(*backend.run_generations.lock().unwrap(), vec![1, 2]); - assert_eq!( - runtime.worker_detail(&detail.worker_ref).unwrap().status, - WorkerStatus::Running - ); + let restored = runtime.worker_detail(&detail.worker_ref).unwrap(); + assert_eq!(restored.status, WorkerStatus::Idle); + assert_eq!(restored.worker_state, None); } #[test] diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index beeb4791..1655a387 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -17,7 +17,7 @@ use crate::ipc::notify_buffer::NotifyBuffer; use crate::ipc::server::SocketServer; use crate::runtime::dir::RuntimeDir; use crate::segment_log_sink::SegmentLogSink; -use crate::shared_state::WorkerSharedState; +use crate::shared_state::{WorkerCommandAdmission, WorkerSharedState}; use crate::shutdown_after_idle::{ ShutdownAfterIdleRequest, TicketIntakeReadyShutdownHook, is_ticket_intake_role, take_shutdown_request_after_status, @@ -181,21 +181,40 @@ impl WorkerHandle { } } +fn command_admission_disposition( + admission: WorkerCommandAdmission, +) -> Result<(), WorkerCommandDisposition> { + match admission { + WorkerCommandAdmission::Accepted => Ok(()), + WorkerCommandAdmission::Retry | WorkerCommandAdmission::StaleCommandId => { + Err(WorkerCommandDisposition::StaleCommandId) + } + WorkerCommandAdmission::Conflict => Err(WorkerCommandDisposition::Conflict), + WorkerCommandAdmission::ExecutionGenerationMismatch => { + Err(WorkerCommandDisposition::StaleExecutionGeneration) + } + WorkerCommandAdmission::StateRevisionMismatch => { + Err(WorkerCommandDisposition::StaleWorkerStateRevision) + } + } +} + fn validate_command( + envelope: WorkerCommandEnvelope, + kind: WorkerCommandKind, + shared_state: &WorkerSharedState, +) -> Result<(), WorkerCommandDisposition> { + command_admission_disposition(shared_state.admit_command(envelope, kind, true)) +} + +fn validate_shutdown_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); + match shared_state.admit_command(envelope, WorkerCommandKind::Shutdown, false) { + WorkerCommandAdmission::Accepted | WorkerCommandAdmission::Retry => Ok(()), + admission => command_admission_disposition(admission), } - 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( @@ -205,6 +224,7 @@ fn acknowledge_command( command: WorkerCommandKind, disposition: WorkerCommandDisposition, ) { + shared_state.complete_command(command_id, command, disposition); let _ = working_event_tx.send(Event::CommandAcknowledged { acknowledgement: WorkerCommandAcknowledgement { command_id, @@ -1913,7 +1933,9 @@ async fn controller_loop( } } Method::Resume { command } => { - if let Err(disposition) = validate_command(command, &shared_state) { + if let Err(disposition) = + validate_command(command, WorkerCommandKind::Resume, &shared_state) + { acknowledge_command( &working_event_tx, &shared_state, @@ -1953,7 +1975,9 @@ async fn controller_loop( } Method::Cancel { command } => { - if let Err(disposition) = validate_command(command, &shared_state) { + if let Err(disposition) = + validate_command(command, WorkerCommandKind::Cancel, &shared_state) + { acknowledge_command( &working_event_tx, &shared_state, @@ -2017,7 +2041,9 @@ async fn controller_loop( } Method::Pause { command } => { - if let Err(disposition) = validate_command(command, &shared_state) { + if let Err(disposition) = + validate_command(command, WorkerCommandKind::Pause, &shared_state) + { acknowledge_command( &working_event_tx, &shared_state, @@ -2036,7 +2062,9 @@ async fn controller_loop( } Method::Compact { command } => { - if let Err(disposition) = validate_command(command, &shared_state) { + if let Err(disposition) = + validate_command(command, WorkerCommandKind::Compact, &shared_state) + { acknowledge_command( &working_event_tx, &shared_state, @@ -2081,7 +2109,11 @@ async fn controller_loop( method = method_rx.recv() => { match method { Some(Method::Cancel { command }) => { - if let Err(disposition) = validate_command(command, &shared_state) { + if let Err(disposition) = validate_command( + command, + WorkerCommandKind::Cancel, + &shared_state, + ) { acknowledge_command( &working_event_tx, &shared_state, @@ -2101,7 +2133,18 @@ async fn controller_loop( let _ = cancel_tx.send(true); } Some(Method::Shutdown { command }) => { - shared_state.accept_command_id(command.command_id); + if let Err(disposition) = + validate_shutdown_command(command, &shared_state) + { + acknowledge_command( + &working_event_tx, + &shared_state, + command.command_id, + WorkerCommandKind::Shutdown, + disposition, + ); + continue; + } shutdown_after_compaction = true; acknowledge_command( &working_event_tx, @@ -2196,9 +2239,18 @@ async fn controller_loop( }, 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); + // Shutdown ignores the state-revision fence but remains bound to the + // current execution generation and command payload identity. + if let Err(disposition) = validate_shutdown_command(command, &shared_state) { + acknowledge_command( + &working_event_tx, + &shared_state, + command.command_id, + WorkerCommandKind::Shutdown, + disposition, + ); + continue; + } acknowledge_command( &working_event_tx, &shared_state, @@ -2544,7 +2596,9 @@ where method = method_rx.recv(), if input_commit.is_none() => { match method { Some(Method::Cancel { command }) => { - if let Err(disposition) = validate_command(command, shared_state) { + if let Err(disposition) = + validate_command(command, WorkerCommandKind::Cancel, shared_state) + { acknowledge_command( working_event_tx, shared_state, @@ -2583,7 +2637,9 @@ where let _ = cancel_tx.try_send(()); } Some(Method::Pause { command }) => { - if let Err(disposition) = validate_command(command, shared_state) { + if let Err(disposition) = + validate_command(command, WorkerCommandKind::Pause, shared_state) + { acknowledge_command( working_event_tx, shared_state, @@ -2623,7 +2679,16 @@ where let _ = pause_tx.try_send(()); } Some(Method::Shutdown { command }) => { - shared_state.accept_command_id(command.command_id); + if let Err(disposition) = validate_shutdown_command(command, shared_state) { + acknowledge_command( + working_event_tx, + shared_state, + command.command_id, + WorkerCommandKind::Shutdown, + disposition, + ); + continue; + } shutdown_requested = true; set_controller_state( shared_state, @@ -2705,7 +2770,9 @@ where } } Some(Method::Resume { command }) => { - if let Err(disposition) = validate_command(command, shared_state) { + if let Err(disposition) = + validate_command(command, WorkerCommandKind::Resume, shared_state) + { acknowledge_command( working_event_tx, shared_state, @@ -2763,7 +2830,9 @@ where } } Some(Method::Compact { command }) => { - if let Err(disposition) = validate_command(command, shared_state) { + if let Err(disposition) = + validate_command(command, WorkerCommandKind::Compact, shared_state) + { acknowledge_command( working_event_tx, shared_state, @@ -3677,6 +3746,7 @@ mod tests { expected_execution_generation: 8, expected_worker_state_revision: 0, }, + WorkerCommandKind::Pause, &shared, ), Err(WorkerCommandDisposition::StaleExecutionGeneration) @@ -3688,6 +3758,7 @@ mod tests { expected_execution_generation: 9, expected_worker_state_revision: 1, }, + WorkerCommandKind::Pause, &shared, ), Err(WorkerCommandDisposition::StaleWorkerStateRevision) @@ -3699,6 +3770,7 @@ mod tests { expected_execution_generation: 9, expected_worker_state_revision: 0, }, + WorkerCommandKind::Pause, &shared, ) .is_ok() @@ -3708,12 +3780,25 @@ mod tests { WorkerCommandEnvelope { command_id: 1, expected_execution_generation: 9, - expected_worker_state_revision: 1, + expected_worker_state_revision: 0, }, + WorkerCommandKind::Pause, &shared, ), Err(WorkerCommandDisposition::StaleCommandId) ); + assert_eq!( + validate_command( + WorkerCommandEnvelope { + command_id: 1, + expected_execution_generation: 9, + expected_worker_state_revision: 0, + }, + WorkerCommandKind::Cancel, + &shared, + ), + Err(WorkerCommandDisposition::Conflict) + ); assert!( validate_command( WorkerCommandEnvelope { @@ -3721,6 +3806,7 @@ mod tests { expected_execution_generation: 9, expected_worker_state_revision: 1, }, + WorkerCommandKind::Pause, &shared, ) .is_ok() diff --git a/crates/worker/src/shared_state.rs b/crates/worker/src/shared_state.rs index 573f2bb5..810b59f4 100644 --- a/crates/worker/src/shared_state.rs +++ b/crates/worker/src/shared_state.rs @@ -1,17 +1,37 @@ +use std::collections::VecDeque; use std::sync::{ OnceLock, RwLock, atomic::{AtomicBool, Ordering}, }; use protocol::{ - WorkerBusyState, WorkerMaintenanceState, WorkerRunState, WorkerState, WorkerStateSnapshot, - WorkerStatus, + WorkerBusyState, WorkerCommandDisposition, WorkerCommandEnvelope, WorkerCommandKind, + WorkerMaintenanceState, WorkerRunState, WorkerState, WorkerStateSnapshot, WorkerStatus, }; use serde_json::json; use session_store::SegmentId; use crate::fs_view::WorkerFsView; +const COMPLETED_COMMAND_RETENTION: usize = 256; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct AcceptedWorkerCommand { + envelope: WorkerCommandEnvelope, + kind: WorkerCommandKind, + disposition: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum WorkerCommandAdmission { + Accepted, + Retry, + Conflict, + StaleCommandId, + ExecutionGenerationMismatch, + StateRevisionMismatch, +} + /// Shared state between WorkerController and runtime directory. /// /// `WorkerStateSnapshot` is the sole live execution-state authority. Runtime @@ -23,6 +43,7 @@ pub struct WorkerSharedState { pub manifest_toml: String, pub greeting: protocol::Greeting, state: RwLock, + accepted_commands: RwLock>, /// 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 @@ -55,6 +76,7 @@ impl WorkerSharedState { manifest_toml, greeting, state: RwLock::new(WorkerStateSnapshot::initial(execution_generation)), + accepted_commands: RwLock::new(VecDeque::new()), fs_view: OnceLock::new(), flow_transition_enabled: AtomicBool::new(false), } @@ -92,17 +114,99 @@ impl WorkerSharedState { snapshot.clone() } - pub fn accept_command_id(&self, command_id: u64) -> bool { + pub(crate) fn admit_command( + &self, + envelope: WorkerCommandEnvelope, + kind: WorkerCommandKind, + require_state_revision: bool, + ) -> WorkerCommandAdmission { let mut snapshot = self .state .write() .expect("worker state lock poisoned; refusing command admission"); - if command_id <= snapshot.last_command_id { - return false; + let mut accepted = self + .accepted_commands + .write() + .expect("worker command ledger lock poisoned; refusing command admission"); + if let Some(existing) = accepted + .iter() + .find(|accepted| accepted.envelope.command_id == envelope.command_id) + { + return if existing.envelope == envelope && existing.kind == kind { + WorkerCommandAdmission::Retry + } else { + WorkerCommandAdmission::Conflict + }; } - snapshot.last_command_id = command_id; + if envelope.expected_execution_generation != snapshot.execution_generation { + return WorkerCommandAdmission::ExecutionGenerationMismatch; + } + if require_state_revision && envelope.expected_worker_state_revision != snapshot.revision { + return WorkerCommandAdmission::StateRevisionMismatch; + } + if envelope.command_id <= snapshot.last_command_id { + return WorkerCommandAdmission::StaleCommandId; + } + + snapshot.last_command_id = envelope.command_id; snapshot.revision = snapshot.revision.saturating_add(1); - true + accepted.push_back(AcceptedWorkerCommand { + envelope, + kind, + disposition: None, + }); + WorkerCommandAdmission::Accepted + } + + pub(crate) fn complete_command( + &self, + command_id: u64, + kind: WorkerCommandKind, + disposition: WorkerCommandDisposition, + ) { + if !matches!( + disposition, + WorkerCommandDisposition::Accepted | WorkerCommandDisposition::InvalidState + ) { + return; + } + let mut accepted = self + .accepted_commands + .write() + .expect("worker command ledger lock poisoned; refusing command completion"); + if let Some(command) = accepted + .iter_mut() + .find(|command| command.envelope.command_id == command_id && command.kind == kind) + { + command.disposition.get_or_insert(disposition); + } + while accepted + .iter() + .filter(|command| command.disposition.is_some()) + .count() + > COMPLETED_COMMAND_RETENTION + { + let Some(index) = accepted + .iter() + .position(|command| command.disposition.is_some()) + else { + break; + }; + accepted.remove(index); + } + } + + #[cfg(test)] + pub(crate) fn command_result( + &self, + command_id: u64, + ) -> Option> { + self.accepted_commands + .read() + .expect("worker command ledger lock poisoned") + .iter() + .find(|command| command.envelope.command_id == command_id) + .map(|command| command.disposition) } pub fn snapshot(&self) -> WorkerStateSnapshot { @@ -190,9 +294,17 @@ mod tests { } #[test] - fn accepted_command_id_advances_the_snapshot_revision_atomically() { + fn accepted_command_identity_advances_revision_and_detects_reuse_conflicts() { let state = test_state(); - assert!(state.accept_command_id(9)); + let envelope = WorkerCommandEnvelope { + command_id: 9, + expected_execution_generation: 7, + expected_worker_state_revision: 0, + }; + assert_eq!( + state.admit_command(envelope, WorkerCommandKind::Pause, true), + WorkerCommandAdmission::Accepted + ); assert_eq!( state.snapshot(), WorkerStateSnapshot { @@ -202,7 +314,24 @@ mod tests { state: WorkerState::Idle, } ); - assert!(!state.accept_command_id(9)); + assert_eq!(state.command_result(9), Some(None)); + state.complete_command( + 9, + WorkerCommandKind::Pause, + WorkerCommandDisposition::Accepted, + ); + assert_eq!( + state.command_result(9), + Some(Some(WorkerCommandDisposition::Accepted)) + ); + assert_eq!( + state.admit_command(envelope, WorkerCommandKind::Pause, true), + WorkerCommandAdmission::Retry + ); + assert_eq!( + state.admit_command(envelope, WorkerCommandKind::Cancel, true), + WorkerCommandAdmission::Conflict + ); assert_eq!(state.snapshot().revision, 1); } diff --git a/crates/workspace-api/src/lib.rs b/crates/workspace-api/src/lib.rs index a2cf8f87..9bdd7bce 100644 --- a/crates/workspace-api/src/lib.rs +++ b/crates/workspace-api/src/lib.rs @@ -1568,7 +1568,12 @@ pub struct WorkerSummary { #[serde(default)] pub tags: Vec, pub workspace: WorkerWorkspaceSummary, + /// Runtime catalog lifecycle compatibility state. Live foreground state, when + /// available, is carried separately in `worker_state`. pub state: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "typescript", ts(optional))] + pub worker_state: Option, pub last_seen_at: Option, #[serde(default)] pub pinned: bool, diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index f615685d..a6b065b2 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -243,7 +243,10 @@ pub struct WorkerSummary { #[serde(default)] pub tags: Vec, pub workspace: WorkerWorkspaceSummary, + /// Runtime catalog lifecycle compatibility state. pub state: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub worker_state: Option, pub last_seen_at: Option, #[serde(default)] pub pinned: bool, @@ -335,6 +338,7 @@ pub(crate) fn workspace_worker_summary( workspace_id: summary.workspace.workspace_id, }, state: summary.state, + worker_state: summary.worker_state, last_seen_at: summary.last_seen_at, pinned: summary.pinned, retention_state: summary.retention_state, @@ -1998,6 +2002,7 @@ impl EmbeddedWorkerRuntime { workspace_id: summary.workspace_id.clone(), }, state: embedded_worker_status_label(summary.status).to_string(), + worker_state: summary.worker_state.clone(), last_seen_at: None, pinned: false, retention_state: "transient".to_string(), @@ -2037,6 +2042,7 @@ impl EmbeddedWorkerRuntime { workspace_id: detail.workspace_id.clone(), }, state: embedded_worker_status_label(detail.status).to_string(), + worker_state: detail.worker_state.clone(), last_seen_at: None, pinned: false, retention_state: "transient".to_string(), @@ -3340,6 +3346,7 @@ impl RemoteWorkerRuntime { workspace_id: summary.workspace_id.clone(), }, state: embedded_worker_status_label(summary.status).to_string(), + worker_state: summary.worker_state.clone(), last_seen_at: None, pinned: false, retention_state: "transient".to_string(), @@ -3383,6 +3390,7 @@ impl RemoteWorkerRuntime { workspace_id: detail.workspace_id.clone(), }, state: embedded_worker_status_label(detail.status).to_string(), + worker_state: detail.worker_state.clone(), last_seen_at: None, pinned: false, retention_state: "transient".to_string(), @@ -4730,6 +4738,7 @@ pub fn placeholder_worker(host_id: impl Into) -> WorkerSummary { workspace_id: None, }, state: "unsupported".to_string(), + worker_state: None, last_seen_at: None, pinned: false, retention_state: "transient".to_string(), @@ -5251,6 +5260,7 @@ mod tests { workspace_id: None, }, state: "available".to_string(), + worker_state: None, last_seen_at: None, pinned: false, retention_state: "transient".to_string(), diff --git a/crates/workspace-server/src/runtime_subscription_tests.rs b/crates/workspace-server/src/runtime_subscription_tests.rs index 5c67f983..ab5d65ac 100644 --- a/crates/workspace-server/src/runtime_subscription_tests.rs +++ b/crates/workspace-server/src/runtime_subscription_tests.rs @@ -202,7 +202,16 @@ async fn equal_downstream_selectors_share_one_upstream_subscription() { BrokerSubscriptionEvent::Event { payload: SubscriptionEventPayload::WorkerUpserted { ref worker }, .. - } if worker.state == SubscriptionWorkerState::Running + } if worker.state == SubscriptionWorkerState::Idle + && matches!( + worker.worker_state, + Some(protocol::WorkerStateSnapshot { + state: protocol::WorkerState::Busy(protocol::WorkerBusyState::Run( + protocol::WorkerRunState::Running + )), + .. + }) + ) )); } let mut late = broker.subscribe("runtime-test", selector.clone()).unwrap(); @@ -212,7 +221,18 @@ async fn equal_downstream_selectors_share_one_upstream_subscription() { assert!(matches!( snapshot, SubscriptionSnapshot::Workers { workers } - if workers.iter().any(|worker| worker.state == SubscriptionWorkerState::Running) + if workers.iter().any(|worker| { + worker.state == SubscriptionWorkerState::Idle + && matches!( + worker.worker_state, + Some(protocol::WorkerStateSnapshot { + state: protocol::WorkerState::Busy( + protocol::WorkerBusyState::Run(protocol::WorkerRunState::Running) + ), + .. + }) + ) + }) )); drop(late); @@ -337,8 +357,18 @@ async fn embedded_runtime_uses_in_process_subscription_source() { ) .unwrap(); assert!(matches!(next_event(&mut subscription).await, - BrokerSubscriptionEvent::Event { payload: SubscriptionEventPayload::WorkerUpserted { worker }, .. } - if worker.runtime_id.as_deref() == Some("embedded-worker-runtime") && worker.state == SubscriptionWorkerState::Running)); + BrokerSubscriptionEvent::Event { payload: SubscriptionEventPayload::WorkerUpserted { worker }, .. } + if worker.runtime_id.as_deref() == Some("embedded-worker-runtime") + && worker.state == SubscriptionWorkerState::Idle + && matches!( + worker.worker_state, + Some(protocol::WorkerStateSnapshot { + state: protocol::WorkerState::Busy(protocol::WorkerBusyState::Run( + protocol::WorkerRunState::Running + )), + .. + }) + ))); let mut late = broker .subscribe( "embedded-worker-runtime", @@ -351,7 +381,18 @@ async fn embedded_runtime_uses_in_process_subscription_source() { assert!(matches!( snapshot, SubscriptionSnapshot::Workers { workers } - if workers.iter().any(|worker| worker.state == SubscriptionWorkerState::Running) + if workers.iter().any(|worker| { + worker.state == SubscriptionWorkerState::Idle + && matches!( + worker.worker_state, + Some(protocol::WorkerStateSnapshot { + state: protocol::WorkerState::Busy( + protocol::WorkerBusyState::Run(protocol::WorkerRunState::Running) + ), + .. + }) + ) + }) )); runtime diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 39206e5d..9cbeb899 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -15252,6 +15252,7 @@ fn worker_summary_from_registry(record: &WorkerRegistryRecord) -> WorkerSummary singleton_key: None, tags: Vec::new(), state: "missing".to_string(), + worker_state: None, last_seen_at: Some(record.updated_at.clone()), pinned: record.retention_state == "pinned", retention_state: record.retention_state.clone(), @@ -24990,6 +24991,7 @@ mod tests { workspace_id: Some(TEST_WORKSPACE_ID.to_string()), }, state: "idle".to_string(), + worker_state: None, last_seen_at: None, pinned: false, retention_state: "normal".to_string(), @@ -25086,6 +25088,7 @@ mod tests { workspace_id: Some(TEST_WORKSPACE_ID.to_string()), }, state: "idle".to_string(), + worker_state: None, last_seen_at: None, pinned: false, retention_state: "normal".to_string(), diff --git a/web/workspace/src/lib/generated/protocol.ts b/web/workspace/src/lib/generated/protocol.ts index c15fbdca..b3482873 100644 --- a/web/workspace/src/lib/generated/protocol.ts +++ b/web/workspace/src/lib/generated/protocol.ts @@ -19,7 +19,7 @@ command_id: number, expected_execution_generation: number, expected_worker_state 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 WorkerCommandDisposition = "accepted" | "stale_execution_generation" | "stale_worker_state_revision" | "stale_command_id" | "conflict" | "invalid_state"; export type WorkerCommandAcknowledgement = { command_id: number, command: WorkerCommandKind, disposition: WorkerCommandDisposition, /** @@ -233,7 +233,16 @@ resource_key?: string | null, /** * Producer-owned monotonic revision for this Worker subject. */ -subject_revision: number, state: SubscriptionWorkerState, has_running_internal_workers: boolean, workspace_id?: string | null, display_name?: string | null, profile?: string | null, +subject_revision: number, +/** + * Latest revisioned foreground state observed from the Worker. This remains + * absent until an authoritative Worker snapshot/event has been applied. + */ +worker_state?: WorkerStateSnapshot | null, +/** + * Runtime catalog lifecycle compatibility projection; not foreground-state authority. + */ +state: SubscriptionWorkerState, has_running_internal_workers: boolean, workspace_id?: string | null, display_name?: string | null, profile?: string | null, /** * Workspace-facing Repository key. Runtime producers leave this unset and * Workspace Server projections replace `repository_id` with this field. diff --git a/web/workspace/src/lib/workspace/sidebar/types.ts b/web/workspace/src/lib/workspace/sidebar/types.ts index 67d7aa10..856076fd 100644 --- a/web/workspace/src/lib/workspace/sidebar/types.ts +++ b/web/workspace/src/lib/workspace/sidebar/types.ts @@ -19,6 +19,7 @@ import type { Event as PodProtocolEvent, Method as PodProtocolMethod, Segment as PodProtocolSegment, + WorkerStateSnapshot, } from "$lib/generated/protocol"; import type { GitCommitSummary as SharedGitCommitSummary, @@ -99,6 +100,7 @@ export type Worker = { tags: string[]; workspace: { visibility: string; identity: string }; state: string; + worker_state?: WorkerStateSnapshot | null; pinned?: boolean; retention_state?: string; last_seen_at?: string | null; diff --git a/web/workspace/src/lib/workspace/sidebar/worker-state.ts b/web/workspace/src/lib/workspace/sidebar/worker-state.ts new file mode 100644 index 00000000..0f45e9f3 --- /dev/null +++ b/web/workspace/src/lib/workspace/sidebar/worker-state.ts @@ -0,0 +1,12 @@ +import type { WorkerStateSnapshot } from "$lib/generated/protocol"; + +export function liveWorkerState(worker: { + state: string; + worker_state?: WorkerStateSnapshot | null; +}): string { + const state = worker.worker_state?.state; + if (!state) return worker.state === "stopped" ? "stopped" : "unknown"; + if (state.kind === "idle") return "idle"; + if (state.state.kind === "maintenance") return "running"; + return state.state.state === "paused" ? "paused" : "running"; +} diff --git a/web/workspace/src/lib/workspace/sidebar/worker-subscription.test.ts b/web/workspace/src/lib/workspace/sidebar/worker-subscription.test.ts index 279afa0d..84e81598 100644 --- a/web/workspace/src/lib/workspace/sidebar/worker-subscription.test.ts +++ b/web/workspace/src/lib/workspace/sidebar/worker-subscription.test.ts @@ -5,6 +5,7 @@ function assertEquals(actual: unknown, expected: unknown): void { throw new Error(`expected ${JSON.stringify(expected)}, got ${JSON.stringify(actual)}`); } } +import { liveWorkerState } from './worker-state'; import { applyWorkspaceWorkersFrame, createWorkspaceWorkersProjection, @@ -33,6 +34,22 @@ function worker( }; } +Deno.test('Worker list state uses the authoritative live snapshot separately from lifecycle', () => { + const active = worker('runtime-a', 'worker-1', 1); + active.worker_state = { + execution_generation: 4, + revision: 2, + last_command_id: 1, + state: { kind: 'busy', state: { kind: 'run', state: 'paused' } }, + }; + assertEquals(liveWorkerState(active), 'paused'); + + const unavailable = worker('runtime-a', 'worker-2', 1); + assertEquals(liveWorkerState(unavailable), 'unknown'); + unavailable.state = 'stopped'; + assertEquals(liveWorkerState(unavailable), 'stopped'); +}); + Deno.test('workspace Worker snapshot keeps equal local ids from different Runtimes', () => { const projection = createWorkspaceWorkersProjection(); const frame: SubscriptionFrame = { diff --git a/web/workspace/src/lib/workspace/sidebar/worker-subscription.ts b/web/workspace/src/lib/workspace/sidebar/worker-subscription.ts index a1eed656..8c9b7557 100644 --- a/web/workspace/src/lib/workspace/sidebar/worker-subscription.ts +++ b/web/workspace/src/lib/workspace/sidebar/worker-subscription.ts @@ -5,6 +5,7 @@ import { applyWorkspaceWorkersFrame, createWorkspaceWorkersProjection, } from './worker-subscription-model'; +import { liveWorkerState } from './worker-state'; import { compareWorkersForSidebar } from './workers'; import type { Worker } from './types'; @@ -86,7 +87,8 @@ function projectWorker(worker: SubscriptionWorker): SidebarWorker { profile: worker.profile ?? null, tags: [], workspace: { visibility: 'workspace', identity: 'runtime_subscription_worker' }, - state: worker.state, + state: liveWorkerState(worker), + worker_state: worker.worker_state, pinned: false, retention_state: 'transient', implementation: { 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 cf28f294..ed8d1df4 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 @@ -52,11 +52,7 @@ import { pushWorkspaceAlert } from "$lib/workspace/alerts/store"; import { workspaceApiPath } from "$lib/workspace/api/http"; import { workspaceMultiplexer, type WorkspaceMultiplexerSubscription } from "$lib/workspace/multiplexer"; - import type { - Diagnostic, - Worker, - PodProtocolEvent, - } from "$lib/workspace/sidebar/types"; + import type { Diagnostic, Worker } from "$lib/workspace/sidebar/types"; type Props = { data: { @@ -207,7 +203,6 @@ ); let pendingObservationEvents: ConsoleEventInput[] = []; let protocolEventSequence = 0; - let pendingObservedStates: Array = []; let pendingStreamDiagnostics: Diagnostic[] = []; let observationFlushHandle: number | null = null; let nextReloadToken = 0; @@ -249,7 +244,9 @@ const diagnostics = $derived( mergeDiagnostics(worker?.diagnostics ?? [], streamDiagnostics), ); - const workerState = $derived(liveWorkerState ?? worker?.state ?? "loading"); + const workerState = $derived( + liveWorkerState ?? (worker?.state === "stopped" ? "stopped" : "loading"), + ); const workerRunning = $derived(workerState === "running"); const workerPaused = $derived(workerState === "paused"); const composerEditable = $derived(protocolState === "open" && !sending); @@ -343,7 +340,6 @@ observationFlushHandle = null; } pendingObservationEvents = []; - pendingObservedStates = []; pendingStreamDiagnostics = []; } @@ -359,18 +355,15 @@ function flushObservationBatch() { observationFlushHandle = null; const eventBatch = pendingObservationEvents; - const stateBatch = pendingObservedStates; const diagnosticBatch = pendingStreamDiagnostics; pendingObservationEvents = []; - pendingObservedStates = []; pendingStreamDiagnostics = []; if (eventBatch.length > 0) { - const latestState = stateBatch.findLast((state) => state !== null); - if (latestState) { - liveWorkerState = latestState; - } consoleProjection = consoleProjector.append(eventBatch); + liveWorkerState = consoleProjection.status === "shutdown" + ? "shutdown" + : workerStateFromSnapshot(consoleProjection.workerState); advanceEventObservedAtVersion(); } @@ -407,7 +400,6 @@ event: payload, observedAtMs, }); - pendingObservedStates.push(workerStateFromProtocolEvent(payload)); scheduleObservationFlush(); } @@ -925,36 +917,16 @@ handleComposerSubmit(); } - function workerStateFromProtocolEvent( - event: PodProtocolEvent, + function workerStateFromSnapshot( + snapshot: ConsoleProjection["workerState"], ): string | null { - switch (event.event) { - case "snapshot": - 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: - return null; - } + if (!snapshot) return null; + return snapshot.state.kind === "idle" + ? "idle" + : snapshot.state.state.kind === "run" && + snapshot.state.state.state === "paused" + ? "paused" + : "running"; } function connectProtocolTransport( diff --git a/web/workspace/src/routes/w/[workspaceId]/workers/+page.svelte b/web/workspace/src/routes/w/[workspaceId]/workers/+page.svelte index e734def2..e3a0d542 100644 --- a/web/workspace/src/routes/w/[workspaceId]/workers/+page.svelte +++ b/web/workspace/src/routes/w/[workspaceId]/workers/+page.svelte @@ -4,6 +4,7 @@ import { workerHref } from '$lib/workspace/resource-links'; import { formatCurrentWorkdirRevision } from '$lib/workspace/settings/workdir-revision'; import { canOpenWorkerConsole } from '$lib/workspace/sidebar/workers'; + import { liveWorkerState } from '$lib/workspace/sidebar/worker-state'; import type { CleanupWorkerCandidate, RuntimeCleanupExecutionResponse, RuntimeCleanupPlanResponse, Worker } from '$lib/workspace/sidebar/types'; import type { PageProps } from './$types'; @@ -136,7 +137,7 @@ } function workerStatus(worker: Worker): string { - return worker.state; + return liveWorkerState(worker); } function workerProfile(worker: Worker): string { diff --git a/web/workspace/src/routes/w/[workspaceId]/workers/[workerRef]/+page.svelte b/web/workspace/src/routes/w/[workspaceId]/workers/[workerRef]/+page.svelte index ef3d04bb..15d3ffe4 100644 --- a/web/workspace/src/routes/w/[workspaceId]/workers/[workerRef]/+page.svelte +++ b/web/workspace/src/routes/w/[workspaceId]/workers/[workerRef]/+page.svelte @@ -1,5 +1,6 @@ @@ -24,7 +25,7 @@ >Open console
-
Status
{data.worker.state}
+
Status
{liveWorkerState(data.worker)}
Profile
{data.worker.profile}
Internal ID
{data.worker.worker_id}
From 0c48c5dee31877273c36e9b03a4f4306a1fa5ac1 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 10:33:45 +0900 Subject: [PATCH 4/5] fix: fail closed without worker snapshots --- crates/worker-runtime/src/runtime.rs | 44 ++++++++++++++------------- crates/workspace-server/src/server.rs | 28 +++++++++++++++++ 2 files changed, 51 insertions(+), 21 deletions(-) diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 76066336..e44e9f7f 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -1732,27 +1732,9 @@ impl Runtime { return Ok(snapshot); } } - Ok(protocol::Event::Snapshot { - session: protocol::SessionSnapshot { - pending_submissions: protocol::PendingSubmissionsSnapshot::default(), - entries: Vec::new(), - }, - greeting: protocol::Greeting { - worker_name: worker_ref.worker_id.to_string(), - cwd: String::new(), - provider: "worker-runtime".to_string(), - model: "worker-runtime".to_string(), - scope_summary: "runtime worker observation".to_string(), - tools: Vec::new(), - context_window: 0, - context_tokens: 0, - }, - state: protocol::WorkerStateSnapshot::initial(1), - in_flight: protocol::InFlightSnapshot { - blocks: Vec::new(), - commands: Vec::new(), - }, - internal_workers: Vec::new(), + Err(RuntimeError::WorkerExecutionUnavailable { + worker_id: worker_ref.worker_id, + message: "authoritative Worker snapshot is unavailable".to_string(), }) } @@ -4974,6 +4956,26 @@ mod tests { } } + #[cfg(feature = "ws-server")] + #[test] + fn observation_snapshot_fails_closed_when_backend_snapshot_is_unavailable() { + let runtime = runtime_with_backend(); + let detail = runtime + .create_worker(task_request("snapshot unavailable")) + .unwrap(); + + assert!(matches!( + runtime + .worker_observation_snapshot(&detail.worker_ref) + .unwrap_err(), + RuntimeError::WorkerExecutionUnavailable { + worker_id, + message, + } if worker_id == detail.worker_ref.worker_id + && message == "authoritative Worker snapshot is unavailable" + )); + } + struct InputOnlyBackend; impl WorkerExecutionBackend for InputOnlyBackend { diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 9cbeb899..8b70d092 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -18865,6 +18865,34 @@ mod tests { ) } + fn worker_snapshot( + &self, + handle: &worker_runtime::execution::WorkerExecutionHandle, + ) -> Option { + Some(protocol::Event::Snapshot { + session: protocol::SessionSnapshot { + pending_submissions: protocol::PendingSubmissionsSnapshot::default(), + entries: Vec::new(), + }, + greeting: protocol::Greeting { + worker_name: handle.worker_ref().worker_id.to_string(), + cwd: String::new(), + provider: "deterministic-workspace-server-test".to_string(), + model: "deterministic-workspace-server-test".to_string(), + scope_summary: "test execution snapshot".to_string(), + tools: Vec::new(), + context_window: 0, + context_tokens: 0, + }, + state: protocol::WorkerStateSnapshot::initial(1), + in_flight: protocol::InFlightSnapshot { + blocks: Vec::new(), + commands: Vec::new(), + }, + internal_workers: Vec::new(), + }) + } + fn stop_worker( &self, _handle: &worker_runtime::execution::WorkerExecutionHandle, From 88683a8d8fcbad75e9639914f0ab8888f7fd50eb Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 10:41:20 +0900 Subject: [PATCH 5/5] test: provide authoritative websocket snapshot --- crates/worker-runtime/src/http_server.rs | 25 ++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index c05bdb91..a35eaab0 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -3246,6 +3246,31 @@ mod ws_tests { } } + fn worker_snapshot(&self, handle: &WorkerExecutionHandle) -> Option { + Some(protocol::Event::Snapshot { + session: protocol::SessionSnapshot { + pending_submissions: protocol::PendingSubmissionsSnapshot::default(), + entries: Vec::new(), + }, + greeting: protocol::Greeting { + worker_name: handle.worker_ref().worker_id.to_string(), + cwd: String::new(), + provider: "ws-test".to_string(), + model: "ws-test".to_string(), + scope_summary: "WebSocket test execution snapshot".to_string(), + tools: Vec::new(), + context_window: 0, + context_tokens: 0, + }, + state: protocol::WorkerStateSnapshot::initial(1), + in_flight: protocol::InFlightSnapshot { + blocks: Vec::new(), + commands: Vec::new(), + }, + internal_workers: Vec::new(), + }) + } + fn dispatch_method( &self, _handle: &WorkerExecutionHandle,