feat: add revisioned worker execution state

This commit is contained in:
2026-09-06 06:54:29 +09:00
parent 3ed1545c3c
commit e8b9adcde4
32 changed files with 2168 additions and 750 deletions
+3 -5
View File
@@ -112,8 +112,8 @@ mod tests {
async fn encodes_methods_and_decodes_events_above_transport() { async fn encodes_methods_and_decodes_events_above_transport() {
let mut socket = TestSocket::default(); let mut socket = TestSocket::default();
socket.incoming.push_back( socket.incoming.push_back(
encode_event(&Event::Status { encode_event(&Event::WorkerState {
status: WorkerStatus::Idle, snapshot: WorkerStatus::Idle.into(),
}) })
.expect("encode event"), .expect("encode event"),
); );
@@ -132,9 +132,7 @@ mod tests {
)); ));
assert!(matches!( assert!(matches!(
client.next_event().await, client.next_event().await,
Ok(Some(Event::Status { Ok(Some(Event::WorkerState { .. }))
status: WorkerStatus::Idle
}))
)); ));
} }
} }
+3 -5
View File
@@ -101,8 +101,8 @@ mod tests {
)); ));
peer.send( peer.send(
encode_event(&Event::Status { encode_event(&Event::WorkerState {
status: WorkerStatus::Idle, snapshot: WorkerStatus::Idle.into(),
}) })
.expect("encode event"), .expect("encode event"),
) )
@@ -110,9 +110,7 @@ mod tests {
.expect("send event"); .expect("send event");
assert!(matches!( assert!(matches!(
client.next_event().await, client.next_event().await,
Ok(Some(Event::Status { Ok(Some(Event::WorkerState { .. }))
status: WorkerStatus::Idle
}))
)); ));
} }
} }
+3 -8
View File
@@ -113,8 +113,8 @@ mod tests {
let listener = UnixListener::bind(&socket_path).unwrap(); let listener = UnixListener::bind(&socket_path).unwrap();
let server = tokio::spawn(async move { let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap(); let (mut stream, _) = listener.accept().await.unwrap();
let event = encode_event(&Event::Status { let event = encode_event(&Event::WorkerState {
status: WorkerStatus::Idle, snapshot: WorkerStatus::Idle.into(),
}) })
.unwrap(); .unwrap();
stream.write_all(event.as_bytes()).await.unwrap(); stream.write_all(event.as_bytes()).await.unwrap();
@@ -126,12 +126,7 @@ mod tests {
.await .await
.expect("client should receive event while alive") .expect("client should receive event while alive")
.expect("transport should succeed"); .expect("transport should succeed");
assert!(matches!( assert!(matches!(event, Some(Event::WorkerState { .. })));
event,
Some(Event::Status {
status: WorkerStatus::Idle
})
));
server.await.unwrap(); server.await.unwrap();
} }
+3 -5
View File
@@ -116,8 +116,8 @@ mod tests {
Message::Text(ref text) Message::Text(ref text)
if matches!(decode_method(text), Ok(Method::Submit { .. })) if matches!(decode_method(text), Ok(Method::Submit { .. }))
)); ));
let event = encode_event(&Event::Status { let event = encode_event(&Event::WorkerState {
status: WorkerStatus::Idle, snapshot: WorkerStatus::Idle.into(),
}) })
.unwrap(); .unwrap();
socket.send(Message::Text(event.into())).await.unwrap(); socket.send(Message::Text(event.into())).await.unwrap();
@@ -134,9 +134,7 @@ mod tests {
.expect("send method"); .expect("send method");
assert!(matches!( assert!(matches!(
client.next_event().await, client.next_event().await,
Ok(Some(Event::Status { Ok(Some(Event::WorkerState { .. }))
status: WorkerStatus::Idle
}))
)); ));
server.await.unwrap(); server.await.unwrap();
} }
+222 -52
View File
@@ -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<WorkerStatus> 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)] #[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] #[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(tag = "method", content = "params", rename_all = "snake_case")] #[serde(tag = "method", content = "params", rename_all = "snake_case")]
@@ -149,20 +285,28 @@ pub enum Method {
expected_revision: u64, expected_revision: u64,
expected_head_id: String, expected_head_id: String,
}, },
Resume, Resume {
Cancel, command: WorkerCommandEnvelope,
},
Cancel {
command: WorkerCommandEnvelope,
},
/// Stop the in-flight turn and transition to `Paused`. /// Stop the in-flight turn and transition to `Paused`.
/// ///
/// Unlike `Cancel` (which discards and returns to `Idle`), a paused /// Unlike `Cancel` (which discards and returns to `Idle`), a paused
/// Worker can resume the interrupted work via `Resume`, or accept a /// Worker can resume the interrupted work via `Resume`, or accept a
/// fresh `Submit` (orphan `tool_use` items are closed with a /// fresh `Submit` (orphan `tool_use` items are closed with a
/// synthetic tool result before the new user message is appended). /// synthetic tool result before the new user message is appended).
Pause, Pause {
command: WorkerCommandEnvelope,
},
/// Request an explicit compaction while the Worker is otherwise idle. /// Request an explicit compaction while the Worker is otherwise idle.
/// ///
/// This is a typed control method: clients must not send `compact` as a /// This is a typed control method: clients must not send `compact` as a
/// `Method::Submit` user message. /// `Method::Submit` user message.
Compact, Compact {
command: WorkerCommandEnvelope,
},
/// Ask the Worker to list valid rewind targets from its authoritative session log. /// Ask the Worker to list valid rewind targets from its authoritative session log.
ListRewindTargets, ListRewindTargets,
/// Truncate the current session back to the selected rewind target and /// Truncate the current session back to the selected rewind target and
@@ -171,7 +315,9 @@ pub enum Method {
target: RewindTargetId, target: RewindTargetId,
expected_head_entries: usize, expected_head_entries: usize,
}, },
Shutdown, Shutdown {
command: WorkerCommandEnvelope,
},
/// Request a list of completion candidates from the Worker. /// Request a list of completion candidates from the Worker.
/// ///
/// Reply is sent on the same socket as `Event::Completions` (not /// Reply is sent on the same socket as `Event::Completions` (not
@@ -938,8 +1084,9 @@ pub enum Event {
Snapshot { Snapshot {
session: SessionSnapshot, session: SessionSnapshot,
greeting: Greeting, greeting: Greeting,
#[serde(default)] /// Full revisioned live execution state. `Stopped` remains Runtime
status: WorkerStatus, /// catalog authority and is deliberately not represented here.
state: WorkerStateSnapshot,
/// Unfinished model output that has already streamed in the current /// Unfinished model output that has already streamed in the current
/// run but is not yet represented by committed snapshot entries. /// run but is not yet represented by committed snapshot entries.
#[serde(default, skip_serializing_if = "InFlightSnapshot::is_empty")] #[serde(default, skip_serializing_if = "InFlightSnapshot::is_empty")]
@@ -976,8 +1123,11 @@ pub enum Event {
}, },
/// Current Worker controller status. Broadcast on every controller-level /// Current Worker controller status. Broadcast on every controller-level
/// transition and included in `History` snapshots for late attach. /// transition and included in `History` snapshots for late attach.
Status { WorkerState {
status: WorkerStatus, snapshot: WorkerStateSnapshot,
},
CommandAcknowledged {
acknowledgement: WorkerCommandAcknowledgement,
}, },
/// Bounded, provider-owned command telemetry for the live Console. This is /// Bounded, provider-owned command telemetry for the live Console. This is
/// intentionally not a history entry and is reconstructed from /// intentionally not a history entry and is reconstructed from
@@ -1612,28 +1762,39 @@ mod tests {
} }
#[test] #[test]
fn method_without_params() { fn lifecycle_method_without_command_fails_closed() {
let json = r#"{"method":"resume"}"#; let error = serde_json::from_str::<Method>(r#"{"method":"resume"}"#).unwrap_err();
let method: Method = serde_json::from_str(json).unwrap(); assert!(error.to_string().contains("params"));
assert!(matches!(method, Method::Resume));
} }
#[test] #[test]
fn method_pause_roundtrip() { fn lifecycle_methods_roundtrip_with_fences() {
let json = r#"{"method":"pause"}"#; for method in [
let method: Method = serde_json::from_str(json).unwrap(); Method::Pause {
assert!(matches!(method, Method::Pause)); command: WorkerCommandEnvelope {
let serialized = serde_json::to_string(&method).unwrap(); command_id: 11,
assert_eq!(serialized, json); expected_execution_generation: 4,
} expected_worker_state_revision: 8,
},
#[test] },
fn method_compact_roundtrip() { Method::Compact {
let json = r#"{"method":"compact"}"#; command: WorkerCommandEnvelope {
let method: Method = serde_json::from_str(json).unwrap(); command_id: 12,
assert!(matches!(method, Method::Compact)); expected_execution_generation: 4,
let serialized = serde_json::to_string(&method).unwrap(); expected_worker_state_revision: 9,
assert_eq!(serialized, json); },
},
] {
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] #[test]
@@ -1902,7 +2063,7 @@ mod tests {
context_window: 200_000, context_window: 200_000,
context_tokens: 42_000, context_tokens: 42_000,
}, },
status: WorkerStatus::Paused, state: WorkerStatus::Paused.into(),
in_flight: InFlightSnapshot::default(), in_flight: InFlightSnapshot::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}; };
@@ -1919,12 +2080,13 @@ mod tests {
assert_eq!(parsed["data"]["greeting"]["tools"][0], "Read"); assert_eq!(parsed["data"]["greeting"]["tools"][0], "Read");
assert_eq!(parsed["data"]["greeting"]["context_window"], 200_000); assert_eq!(parsed["data"]["greeting"]["context_window"], 200_000);
assert_eq!(parsed["data"]["greeting"]["context_tokens"], 42_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] #[test]
fn event_snapshot_in_flight_roundtrip_and_default() { 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(); let decoded: Event = serde_json::from_str(inbound).unwrap();
match decoded { match decoded {
Event::Snapshot { in_flight, .. } => assert!(in_flight.is_empty()), Event::Snapshot { in_flight, .. } => assert!(in_flight.is_empty()),
@@ -1946,7 +2108,7 @@ mod tests {
context_window: 0, context_window: 0,
context_tokens: 0, context_tokens: 0,
}, },
status: WorkerStatus::Running, state: WorkerStatus::Running.into(),
in_flight: InFlightSnapshot { in_flight: InFlightSnapshot {
blocks: vec![ blocks: vec![
InFlightBlock::Text { InFlightBlock::Text {
@@ -2034,20 +2196,32 @@ mod tests {
} }
#[test] #[test]
fn event_status_format() { fn event_worker_state_format() {
let event = Event::Status { let event = Event::WorkerState {
status: WorkerStatus::Running, 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 json = serde_json::to_string(&event).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap(); let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(parsed["event"], "status"); assert_eq!(parsed["event"], "worker_state");
assert_eq!(parsed["data"]["status"], "running"); 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(); let decoded: Event = serde_json::from_str(&json).unwrap();
assert!(matches!( assert!(matches!(
decoded, decoded,
Event::Status { Event::WorkerState {
status: WorkerStatus::Running snapshot: WorkerStateSnapshot {
execution_generation: 7,
revision: 3,
state: WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running)),
..
}
} }
)); ));
} }
@@ -2088,19 +2262,10 @@ mod tests {
} }
#[test] #[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 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(); let error = serde_json::from_str::<Event>(json).unwrap_err();
match decoded { assert!(error.to_string().contains("state"));
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:?}"),
}
} }
#[test] #[test]
@@ -2513,7 +2678,12 @@ mod tests {
"scope_summary": "scope", "scope_summary": "scope",
"tools": [] "tools": []
}, },
"status": "idle" "state": {
"execution_generation": 1,
"revision": 0,
"last_command_id": 0,
"state": { "kind": "idle" }
}
} }
})) }))
.unwrap(); .unwrap();
+13 -1
View File
@@ -12,7 +12,10 @@ use crate::{
RewindSummary, RewindTarget, RewindTargetId, RunResult, ScopeRule, Segment, SessionContentPart, RewindSummary, RewindTarget, RewindTargetId, RunResult, ScopeRule, Segment, SessionContentPart,
SessionEntryProvenance, SessionMessageRole, SessionSnapshot, SessionSnapshotEntry, SessionEntryProvenance, SessionMessageRole, SessionSnapshot, SessionSnapshotEntry,
SessionSnapshotEntryData, SessionToolAttachment, SubmissionDisposition, ToolResultDisposition, SessionSnapshotEntryData, SessionToolAttachment, SubmissionDisposition, ToolResultDisposition,
TurnResult, UploadedFileAvailability, UploadedFileRef, WorkerEvent, WorkerStatus, TurnResult, UploadedFileAvailability, UploadedFileRef, WorkerBusyState,
WorkerCommandAcknowledgement, WorkerCommandDisposition, WorkerCommandEnvelope,
WorkerCommandKind, WorkerEvent, WorkerMaintenanceState, WorkerRunState, WorkerState,
WorkerStateSnapshot, WorkerStatus,
subscription::{ subscription::{
EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame, EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest, SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest,
@@ -46,6 +49,15 @@ pub fn generated_protocol_types() -> String {
push_decl::<AlertSource>(&cfg, &mut output); push_decl::<AlertSource>(&cfg, &mut output);
push_decl::<CompletionKind>(&cfg, &mut output); push_decl::<CompletionKind>(&cfg, &mut output);
push_decl::<WorkerStatus>(&cfg, &mut output); push_decl::<WorkerStatus>(&cfg, &mut output);
push_decl::<WorkerCommandEnvelope>(&cfg, &mut output);
push_decl::<WorkerCommandKind>(&cfg, &mut output);
push_decl::<WorkerCommandDisposition>(&cfg, &mut output);
push_decl::<WorkerCommandAcknowledgement>(&cfg, &mut output);
push_decl::<WorkerRunState>(&cfg, &mut output);
push_decl::<WorkerMaintenanceState>(&cfg, &mut output);
push_decl::<WorkerBusyState>(&cfg, &mut output);
push_decl::<WorkerState>(&cfg, &mut output);
push_decl::<WorkerStateSnapshot>(&cfg, &mut output);
push_decl::<TurnResult>(&cfg, &mut output); push_decl::<TurnResult>(&cfg, &mut output);
push_decl::<InvokeKind>(&cfg, &mut output); push_decl::<InvokeKind>(&cfg, &mut output);
push_decl::<RunResult>(&cfg, &mut output); push_decl::<RunResult>(&cfg, &mut output);
+10 -2
View File
@@ -318,7 +318,11 @@ impl StandaloneHost {
} }
pub async fn shutdown(mut self) -> Result<(), StandaloneShutdownError> { 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 { let Some(shutdown) = self.shutdown.take() else {
self.retain_lease(); self.retain_lease();
return Err(StandaloneShutdownError::ConfirmationLost); return Err(StandaloneShutdownError::ConfirmationLost);
@@ -500,7 +504,11 @@ fn active_pointer(
} }
async fn stop_started_worker(started: BootstrappedWorker) { 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; let _ = tokio::time::timeout(Duration::from_secs(2), started.shutdown).await;
} }
+48 -22
View File
@@ -5,7 +5,7 @@ use std::time::{Duration, Instant};
use protocol::{ use protocol::{
AlertLevel, AlertSource, CompletionEntry, CompletionKind, ErrorCode, Event, InFlightBlock, AlertLevel, AlertSource, CompletionEntry, CompletionKind, ErrorCode, Event, InFlightBlock,
InFlightSnapshot, InFlightToolCallState, InternalWorkerRef, InternalWorkerSnapshot, Method, InFlightSnapshot, InFlightToolCallState, InternalWorkerRef, InternalWorkerSnapshot, Method,
RewindTarget, RunResult, Segment, WorkerStatus, RewindTarget, RunResult, Segment, WorkerCommandEnvelope, WorkerStateSnapshot, WorkerStatus,
}; };
use crate::block::{ use crate::block::{
@@ -225,8 +225,10 @@ pub struct WorkerViewTab {
pub struct App { pub struct App {
pub worker_name: String, pub worker_name: String,
pub connected: bool, pub connected: bool,
/// Last controller status reported by the Worker. Drives the status line /// Latest authoritative revisioned live execution state.
/// and Ctrl-key routing; do not infer this solely from replayed history. pub worker_state: WorkerStateSnapshot,
next_command_id: u64,
/// Derived Runtime-catalog compatibility projection used by existing UI.
pub worker_status: WorkerStatus, pub worker_status: WorkerStatus,
/// True while the Worker is in `WorkerStatus::Running`. /// True while the Worker is in `WorkerStatus::Running`.
pub running: bool, pub running: bool,
@@ -337,6 +339,8 @@ impl App {
Self { Self {
worker_name, worker_name,
connected: false, connected: false,
worker_state: WorkerStateSnapshot::initial(1),
next_command_id: 1,
worker_status: WorkerStatus::Idle, worker_status: WorkerStatus::Idle,
running: false, running: false,
paused: false, paused: false,
@@ -745,7 +749,8 @@ impl App {
if self.paused { if self.paused {
self.input_history.cancel_browse(); self.input_history.cancel_browse();
self.input.clear(); self.input.clear();
return Some(Method::Resume); let command = self.next_command_envelope();
return Some(Method::Resume { command });
} }
return None; 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<Method> { pub fn handle_worker_event(&mut self, event: Event) -> Option<Method> {
if self.rewind_refresh_fence && event_is_stale_after_rewind(&event) { if self.rewind_refresh_fence && event_is_stale_after_rewind(&event) {
return None; return None;
@@ -1443,7 +1457,7 @@ impl App {
Event::Snapshot { Event::Snapshot {
session, session,
greeting, greeting,
status, state,
in_flight, in_flight,
internal_workers, internal_workers,
} => { } => {
@@ -1451,7 +1465,8 @@ impl App {
self.pending_submissions = session.pending_submissions.clone(); self.pending_submissions = session.pending_submissions.clone();
self.restore_snapshot(&session, greeting, in_flight); self.restore_snapshot(&session, greeting, in_flight);
self.replace_internal_worker_snapshots(internal_workers); 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 { Event::InternalWorker {
worker, worker,
@@ -1461,9 +1476,14 @@ impl App {
Event::InternalWorkerRemoved { worker, revision } => { Event::InternalWorkerRemoved { worker, revision } => {
self.remove_internal_worker(worker, revision) self.remove_internal_worker(worker, revision)
} }
Event::Status { status } => { Event::WorkerState { snapshot } => {
self.rewind_refresh_fence = false; 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 // Command telemetry is an operational Web Console surface. The
// TUI continues to render the final Bash ToolResult from history. // TUI continues to render the final Bash ToolResult from history.
@@ -2026,12 +2046,18 @@ impl App {
self.input_mode = CommandInputMode::Composer; self.input_mode = CommandInputMode::Composer;
self.command_completion_selected = None; 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.completion = None;
self.rewind_picker = None; self.rewind_picker = None;
self.rewind_request_pending = true; self.rewind_request_pending = true;
} }
result.method method
} }
fn push_command_diagnostic(&mut self, message: impl Into<String>) { fn push_command_diagnostic(&mut self, message: impl Into<String>) {
@@ -2761,8 +2787,8 @@ mod rewind_refresh_tests {
}); });
assert!(!blocks_contain(&app, "stale tail after rewind")); assert!(!blocks_contain(&app, "stale tail after rewind"));
app.handle_worker_event(Event::Status { app.handle_worker_event(Event::WorkerState {
status: WorkerStatus::Idle, snapshot: WorkerStatus::Idle.into(),
}); });
app.handle_worker_event(Event::TextDelta { app.handle_worker_event(Event::TextDelta {
text: "new live tail after status".into(), text: "new live tail after status".into(),
@@ -3478,7 +3504,7 @@ mod completion_flow_tests {
let mut app = App::new("test".into()); let mut app = App::new("test".into());
app.set_worker_status(WorkerStatus::Paused); 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); assert_eq!(app.queued_input_count(), 0);
} }
@@ -3533,7 +3559,7 @@ mod completion_flow_tests {
app.handle_worker_event(Event::Snapshot { app.handle_worker_event(Event::Snapshot {
greeting: test_greeting(), greeting: test_greeting(),
session: public_session(vec![session_start_value]), session: public_session(vec![session_start_value]),
status: WorkerStatus::Running, state: WorkerStatus::Running.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}); });
@@ -3551,8 +3577,8 @@ mod completion_flow_tests {
code: ErrorCode::ProviderError, code: ErrorCode::ProviderError,
message: "provider unavailable".into(), message: "provider unavailable".into(),
}); });
app.handle_worker_event(Event::Status { app.handle_worker_event(Event::WorkerState {
status: WorkerStatus::Idle, snapshot: WorkerStatus::Idle.into(),
}); });
let live_errors = app let live_errors = app
@@ -3577,7 +3603,7 @@ mod completion_flow_tests {
app.handle_worker_event(Event::Snapshot { app.handle_worker_event(Event::Snapshot {
greeting: test_greeting(), greeting: test_greeting(),
session: public_session(vec![serde_json::to_value(run_errored).unwrap()]), session: public_session(vec![serde_json::to_value(run_errored).unwrap()]),
status: WorkerStatus::Idle, state: WorkerStatus::Idle.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}); });
@@ -3641,7 +3667,7 @@ mod completion_flow_tests {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(), pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(), entries: Vec::new(),
}, },
status: WorkerStatus::Running, state: WorkerStatus::Running.into(),
in_flight: InFlightSnapshot { in_flight: InFlightSnapshot {
blocks: vec![ blocks: vec![
InFlightBlock::Thinking { InFlightBlock::Thinking {
@@ -3968,7 +3994,7 @@ mod completion_flow_tests {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(), pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(), entries: Vec::new(),
}, },
status: WorkerStatus::Idle, state: WorkerStatus::Idle.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}); });
@@ -4020,7 +4046,7 @@ mod completion_flow_tests {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(), pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(), entries: Vec::new(),
}, },
status: WorkerStatus::Idle, state: WorkerStatus::Idle.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: vec![InternalWorkerSnapshot { internal_workers: vec![InternalWorkerSnapshot {
worker: InternalWorkerRef { worker: InternalWorkerRef {
@@ -4194,7 +4220,7 @@ mod completion_flow_tests {
entries: Vec::new(), entries: Vec::new(),
}, },
greeting, greeting,
status: WorkerStatus::Idle, state: WorkerStatus::Idle.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}); });
@@ -4393,7 +4419,7 @@ mod completion_flow_tests {
app.handle_worker_event(Event::Snapshot { app.handle_worker_event(Event::Snapshot {
greeting: test_greeting(), greeting: test_greeting(),
session: public_session(assistant_item_entries), session: public_session(assistant_item_entries),
status: WorkerStatus::Running, state: WorkerStatus::Running.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}); });
+7 -2
View File
@@ -409,7 +409,12 @@ fn compact_command(invocation: CommandInvocation<'_>) -> CommandExecution {
let _ = invocation.environment; let _ = invocation.environment;
let _ = invocation.args.raw(); let _ = invocation.args.raw();
CommandExecution { 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")], diagnostics: vec![CommandDiagnostic::new("compact requested")],
exit_command_mode: true, exit_command_mode: true,
clear_input: true, clear_input: true,
@@ -483,7 +488,7 @@ mod tests {
fn compact_command_returns_compact_method_not_run() { fn compact_command_returns_compact_method_not_run() {
let registry = CommandRegistry::builtins(); let registry = CommandRegistry::builtins();
let result = registry.dispatch("compact", &env()); 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.exit_command_mode);
assert!(result.clear_input); assert!(result.clear_input);
assert!(result.diagnostics[0].message.contains("compact requested")); assert!(result.diagnostics[0].message.contains("compact requested"));
+23 -20
View File
@@ -572,7 +572,7 @@ async fn run_e2e_rewind_fixture(
pending_submissions: protocol::PendingSubmissionsSnapshot::default(), pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(), entries: Vec::new(),
}, },
status: WorkerStatus::Idle, state: WorkerStatus::Idle.into(),
greeting: Greeting { greeting: Greeting {
worker_name: worker_name.clone(), worker_name: worker_name.clone(),
cwd: workspace_root.display().to_string(), cwd: workspace_root.display().to_string(),
@@ -1438,13 +1438,15 @@ fn handle_cancel_or_shutdown(app: &mut App) -> Option<Method> {
WorkerStatus::Running | WorkerStatus::Paused WorkerStatus::Running | WorkerStatus::Paused
) { ) {
app.shutdown_confirm = None; 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 if let Some(pressed_at) = app.shutdown_confirm
&& pressed_at.elapsed() < CONFIRM_TIMEOUT && pressed_at.elapsed() < CONFIRM_TIMEOUT
{ {
app.shutdown_confirm = None; 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.shutdown_confirm = Some(std::time::Instant::now());
app.flash_actionbar_notice( app.flash_actionbar_notice(
@@ -1460,7 +1462,8 @@ fn handle_cancel_or_shutdown(app: &mut App) -> Option<Method> {
/// Idle / Paused → 2-tap to quit the TUI (the Worker keeps running). /// Idle / Paused → 2-tap to quit the TUI (the Worker keeps running).
fn handle_pause_or_quit(app: &mut App) -> Option<Method> { fn handle_pause_or_quit(app: &mut App) -> Option<Method> {
if app.worker_status == WorkerStatus::Running { 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 if let Some(t) = app.quit_confirm
&& t.elapsed() < CONFIRM_TIMEOUT && t.elapsed() < CONFIRM_TIMEOUT
@@ -2090,7 +2093,7 @@ mod tests {
&mut app, &mut app,
KeyEvent::new(KeyCode::Char('c'), KeyModifiers::CONTROL), KeyEvent::new(KeyCode::Char('c'), KeyModifiers::CONTROL),
), ),
Some(Method::Pause) Some(Method::Pause { .. })
)); ));
assert_eq!(app.queued_input_count(), 1); assert_eq!(app.queued_input_count(), 1);
@@ -2100,7 +2103,7 @@ mod tests {
&mut app, &mut app,
KeyEvent::new(KeyCode::Char('x'), KeyModifiers::CONTROL), KeyEvent::new(KeyCode::Char('x'), KeyModifiers::CONTROL),
), ),
Some(Method::Cancel) Some(Method::Cancel { .. })
)); ));
assert_eq!(app.queued_input_count(), 1); assert_eq!(app.queued_input_count(), 1);
} }
@@ -2114,7 +2117,7 @@ mod tests {
&mut app, &mut app,
KeyEvent::new(KeyCode::Char('x'), KeyModifiers::CONTROL), KeyEvent::new(KeyCode::Char('x'), KeyModifiers::CONTROL),
); );
assert!(matches!(cancel, Some(Method::Cancel))); assert!(matches!(cancel, Some(Method::Cancel { .. })));
} }
#[test] #[test]
@@ -2136,7 +2139,7 @@ mod tests {
assert!(matches!( assert!(matches!(
handle_key(&mut app, ctrl_x()), handle_key(&mut app, ctrl_x()),
Some(Method::Shutdown) Some(Method::Shutdown { .. })
)); ));
assert!(app.shutdown_confirm.is_none()); assert!(app.shutdown_confirm.is_none());
} }
@@ -2466,7 +2469,7 @@ mod tests {
} }
let method = handle_key(&mut app, KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); 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!(!app.is_command_mode());
assert_eq!(input_text(&app), ""); assert_eq!(input_text(&app), "");
assert_eq!(app.queued_input_count(), 0); assert_eq!(app.queued_input_count(), 0);
@@ -2573,7 +2576,7 @@ mod tests {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(), pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: vec![], entries: vec![],
}, },
status: WorkerStatus::Idle, state: WorkerStatus::Idle.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}); });
@@ -2606,7 +2609,7 @@ mod tests {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(), pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: vec![], entries: vec![],
}, },
status: WorkerStatus::Idle, state: WorkerStatus::Idle.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}); });
@@ -2743,8 +2746,8 @@ mod tests {
kind: protocol::InternalWorkerKind::SubWorker, kind: protocol::InternalWorkerKind::SubWorker,
}, },
revision: 1, revision: 1,
event: Box::new(Event::Status { event: Box::new(Event::WorkerState {
status: WorkerStatus::Running, snapshot: WorkerStatus::Running.into(),
}), }),
}); });
enter_command_mode(&mut app); enter_command_mode(&mut app);
@@ -2859,8 +2862,8 @@ mod tests {
kind: protocol::InternalWorkerKind::SubWorker, kind: protocol::InternalWorkerKind::SubWorker,
}, },
revision: 1, revision: 1,
event: Box::new(Event::Status { event: Box::new(Event::WorkerState {
status: WorkerStatus::Running, snapshot: WorkerStatus::Running.into(),
}), }),
}); });
@@ -2885,8 +2888,8 @@ mod tests {
kind: protocol::InternalWorkerKind::SubWorker, kind: protocol::InternalWorkerKind::SubWorker,
}, },
revision: 1, revision: 1,
event: Box::new(Event::Status { event: Box::new(Event::WorkerState {
status: WorkerStatus::Running, snapshot: WorkerStatus::Running.into(),
}), }),
}); });
handle_key(&mut app, key(KeyCode::Tab)); handle_key(&mut app, key(KeyCode::Tab));
@@ -2902,7 +2905,7 @@ mod tests {
); );
assert!(first.is_none()); assert!(first.is_none());
assert!(matches!(second, Some(Method::Shutdown))); assert!(matches!(second, Some(Method::Shutdown { .. })));
assert_eq!(app.worker_status, WorkerStatus::Idle); assert_eq!(app.worker_status, WorkerStatus::Idle);
} }
@@ -2924,8 +2927,8 @@ mod tests {
kind: protocol::InternalWorkerKind::SubWorker, kind: protocol::InternalWorkerKind::SubWorker,
}, },
revision: 1, revision: 1,
event: Box::new(Event::Status { event: Box::new(Event::WorkerState {
status: WorkerStatus::Running, snapshot: WorkerStatus::Running.into(),
}), }),
}); });
+14 -28
View File
@@ -15,18 +15,6 @@ use std::fmt;
use std::sync::Arc; use std::sync::Arc;
use workdir::WorkdirSessionHandle; 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. /// Execution operation that produced a result.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")] #[serde(rename_all = "snake_case")]
@@ -55,7 +43,8 @@ pub struct WorkerSubmissionAck {
pub struct WorkerExecutionResult { pub struct WorkerExecutionResult {
pub operation: WorkerExecutionOperation, pub operation: WorkerExecutionOperation,
pub outcome: WorkerExecutionOutcome, pub outcome: WorkerExecutionOutcome,
pub run_state: WorkerExecutionRunState, #[serde(default, skip_serializing_if = "Option::is_none")]
pub worker_state: Option<protocol::WorkerStateSnapshot>,
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub message: Option<String>, pub message: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
@@ -74,22 +63,23 @@ pub enum WorkerExecutionOutcome {
} }
impl WorkerExecutionResult { impl WorkerExecutionResult {
pub fn accepted( pub fn accepted(operation: WorkerExecutionOperation) -> Self {
operation: WorkerExecutionOperation,
run_state: WorkerExecutionRunState,
) -> Self {
Self { Self {
operation, operation,
outcome: WorkerExecutionOutcome::Accepted, outcome: WorkerExecutionOutcome::Accepted,
run_state, worker_state: None,
message: None, message: None,
submission: 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( pub fn accepted_submission(
operation: WorkerExecutionOperation, operation: WorkerExecutionOperation,
run_state: WorkerExecutionRunState,
submission_request_id: impl Into<String>, submission_request_id: impl Into<String>,
submission_id: impl Into<String>, submission_id: impl Into<String>,
disposition: protocol::SubmissionDisposition, disposition: protocol::SubmissionDisposition,
@@ -97,7 +87,7 @@ impl WorkerExecutionResult {
Self { Self {
operation, operation,
outcome: WorkerExecutionOutcome::Accepted, outcome: WorkerExecutionOutcome::Accepted,
run_state, worker_state: None,
message: None, message: None,
submission: Some(WorkerSubmissionAck { submission: Some(WorkerSubmissionAck {
submission_request_id: submission_request_id.into(), submission_request_id: submission_request_id.into(),
@@ -111,7 +101,7 @@ impl WorkerExecutionResult {
Self { Self {
operation, operation,
outcome: WorkerExecutionOutcome::Busy, outcome: WorkerExecutionOutcome::Busy,
run_state: WorkerExecutionRunState::Busy, worker_state: None,
message: Some(message.into()), message: Some(message.into()),
submission: None, submission: None,
} }
@@ -121,7 +111,7 @@ impl WorkerExecutionResult {
Self { Self {
operation, operation,
outcome: WorkerExecutionOutcome::Rejected, outcome: WorkerExecutionOutcome::Rejected,
run_state: WorkerExecutionRunState::Stopped, worker_state: None,
message: Some(message.into()), message: Some(message.into()),
submission: None, submission: None,
} }
@@ -131,7 +121,7 @@ impl WorkerExecutionResult {
Self { Self {
operation, operation,
outcome: WorkerExecutionOutcome::Errored, outcome: WorkerExecutionOutcome::Errored,
run_state: WorkerExecutionRunState::Errored, worker_state: None,
message: Some(message.into()), message: Some(message.into()),
submission: None, submission: None,
} }
@@ -141,7 +131,7 @@ impl WorkerExecutionResult {
Self { Self {
operation, operation,
outcome: WorkerExecutionOutcome::Unsupported, outcome: WorkerExecutionOutcome::Unsupported,
run_state: WorkerExecutionRunState::Stopped, worker_state: None,
message: Some(message.into()), message: Some(message.into()),
submission: None, submission: None,
} }
@@ -280,7 +270,6 @@ pub struct WorkerExecutionRestoreRequest {
pub enum WorkerExecutionSpawnResult { pub enum WorkerExecutionSpawnResult {
Connected { Connected {
handle: WorkerExecutionHandle, handle: WorkerExecutionHandle,
run_state: WorkerExecutionRunState,
working_directory: Option<WorkingDirectoryStatus>, working_directory: Option<WorkingDirectoryStatus>,
}, },
Rejected(WorkerExecutionResult), Rejected(WorkerExecutionResult),
@@ -290,12 +279,10 @@ pub enum WorkerExecutionSpawnResult {
impl WorkerExecutionSpawnResult { impl WorkerExecutionSpawnResult {
pub fn connected( pub fn connected(
handle: WorkerExecutionHandle, handle: WorkerExecutionHandle,
run_state: WorkerExecutionRunState,
working_directory: Option<WorkingDirectoryStatus>, working_directory: Option<WorkingDirectoryStatus>,
) -> Self { ) -> Self {
Self::Connected { Self::Connected {
handle, handle,
run_state,
working_directory, working_directory,
} }
} }
@@ -623,7 +610,6 @@ mod tests {
fn submission_ack_survives_json_round_trip() { fn submission_ack_survives_json_round_trip() {
let result = WorkerExecutionResult::accepted_submission( let result = WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input, WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
"request-1", "request-1",
"submission-1", "submission-1",
protocol::SubmissionDisposition::Started, protocol::SubmissionDisposition::Started,
+11 -29
View File
@@ -2206,8 +2206,8 @@ mod tests {
}; };
use crate::execution::{ use crate::execution::{
WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation,
WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionRunState, WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionSpawnRequest,
WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, WorkerExecutionSpawnResult,
}; };
use crate::management::RuntimeOptions; use crate::management::RuntimeOptions;
use axum::body::to_bytes; use axum::body::to_bytes;
@@ -2979,7 +2979,6 @@ mod tests {
fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult {
WorkerExecutionSpawnResult::Connected { WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: request working_directory: request
.working_directory .working_directory
.as_ref() .as_ref()
@@ -2993,7 +2992,6 @@ mod tests {
) -> WorkerExecutionSpawnResult { ) -> WorkerExecutionSpawnResult {
WorkerExecutionSpawnResult::Connected { WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: request.previous_working_directory, working_directory: request.previous_working_directory,
} }
} }
@@ -3006,24 +3004,17 @@ mod tests {
if let Some(submission_id) = input.submission_request_id { if let Some(submission_id) = input.submission_request_id {
WorkerExecutionResult::accepted_submission( WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input, WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
submission_id.clone(), submission_id.clone(),
submission_id, submission_id,
protocol::SubmissionDisposition::Started, protocol::SubmissionDisposition::Started,
) )
} else { } else {
WorkerExecutionResult::accepted( WorkerExecutionResult::accepted(WorkerExecutionOperation::Input)
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
)
} }
} }
fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
WorkerExecutionResult::accepted( WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop)
WorkerExecutionOperation::Stop,
WorkerExecutionRunState::Stopped,
)
} }
} }
@@ -3295,8 +3286,7 @@ mod ws_tests {
}; };
use crate::execution::{ use crate::execution::{
WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation,
WorkerExecutionResult, WorkerExecutionRunState, WorkerExecutionSpawnRequest, WorkerExecutionResult, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult,
WorkerExecutionSpawnResult,
}; };
use crate::management::RuntimeOptions; use crate::management::RuntimeOptions;
use futures::{SinkExt, StreamExt}; use futures::{SinkExt, StreamExt};
@@ -3316,7 +3306,6 @@ mod ws_tests {
fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult {
WorkerExecutionSpawnResult::Connected { WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: request working_directory: request
.working_directory .working_directory
.as_ref() .as_ref()
@@ -3332,16 +3321,12 @@ mod ws_tests {
if let Some(submission_id) = input.submission_request_id { if let Some(submission_id) = input.submission_request_id {
WorkerExecutionResult::accepted_submission( WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input, WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
submission_id.clone(), submission_id.clone(),
submission_id, submission_id,
protocol::SubmissionDisposition::Started, protocol::SubmissionDisposition::Started,
) )
} else { } else {
WorkerExecutionResult::accepted( WorkerExecutionResult::accepted(WorkerExecutionOperation::Input)
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
)
} }
} }
@@ -3350,10 +3335,7 @@ mod ws_tests {
_handle: &WorkerExecutionHandle, _handle: &WorkerExecutionHandle,
_method: protocol::Method, _method: protocol::Method,
) -> WorkerExecutionResult { ) -> WorkerExecutionResult {
WorkerExecutionResult::accepted( WorkerExecutionResult::accepted(WorkerExecutionOperation::ProtocolMethod)
WorkerExecutionOperation::ProtocolMethod,
WorkerExecutionRunState::Idle,
)
} }
} }
@@ -3564,16 +3546,16 @@ mod ws_tests {
runtime runtime
.observe_worker_event( .observe_worker_event(
&other.worker_ref, &other.worker_ref,
protocol::Event::Status { protocol::Event::WorkerState {
status: protocol::WorkerStatus::Running, snapshot: protocol::WorkerStatus::Running.into(),
}, },
) )
.unwrap(); .unwrap();
runtime runtime
.observe_worker_event( .observe_worker_event(
&worker_ref, &worker_ref,
protocol::Event::Status { protocol::Event::WorkerState {
status: protocol::WorkerStatus::Running, snapshot: protocol::WorkerStatus::Running.into(),
}, },
) )
.unwrap(); .unwrap();
+109 -107
View File
@@ -13,8 +13,8 @@ use crate::error::RuntimeError;
use crate::execution::WorkerExecutionRestoreRequest; use crate::execution::WorkerExecutionRestoreRequest;
use crate::execution::{ use crate::execution::{
WorkerExecutionBackend, WorkerExecutionBackendRef, WorkerExecutionHandle, WorkerExecutionBackend, WorkerExecutionBackendRef, WorkerExecutionHandle,
WorkerExecutionOperation, WorkerExecutionResult, WorkerExecutionRunState, WorkerExecutionOperation, WorkerExecutionResult, WorkerExecutionSpawnRequest,
WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, WorkerExecutionSpawnResult,
}; };
#[cfg(feature = "fs-store")] #[cfg(feature = "fs-store")]
use crate::fs_store::{ use crate::fs_store::{
@@ -725,12 +725,11 @@ impl Runtime {
}; };
let spawn_result = backend.spawn_worker(spawn_request); 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 { WorkerExecutionSpawnResult::Connected {
handle, handle,
run_state,
working_directory, working_directory,
} => (handle, run_state, working_directory), } => (handle, working_directory),
WorkerExecutionSpawnResult::Rejected(result) WorkerExecutionSpawnResult::Rejected(result)
| WorkerExecutionSpawnResult::Errored(result) => { | WorkerExecutionSpawnResult::Errored(result) => {
self.rollback_failed_create(&worker_ref)?; self.rollback_failed_create(&worker_ref)?;
@@ -785,11 +784,10 @@ impl Runtime {
result, result,
}); });
} }
let initial_run_state = dispatch_result.run_state;
let detail = self.commit_created_worker( let detail = self.commit_created_worker(
&worker_ref, &worker_ref,
handle, handle,
initial_run_state, WorkerStatus::Running,
working_directory, working_directory,
dispatch_result, dispatch_result,
)?; )?;
@@ -799,9 +797,9 @@ impl Runtime {
self.commit_created_worker( self.commit_created_worker(
&worker_ref, &worker_ref,
handle, handle,
run_state, WorkerStatus::Idle,
working_directory, working_directory,
WorkerExecutionResult::accepted(WorkerExecutionOperation::Spawn, run_state), WorkerExecutionResult::accepted(WorkerExecutionOperation::Spawn),
) )
} }
} }
@@ -1086,13 +1084,12 @@ impl Runtime {
match backend.restore_worker(request) { match backend.restore_worker(request) {
WorkerExecutionSpawnResult::Connected { WorkerExecutionSpawnResult::Connected {
handle, handle,
run_state,
working_directory, working_directory,
} => { } => {
self.commit_restored_worker_execution( self.commit_restored_worker_execution(
worker_ref, worker_ref,
handle, handle,
run_state, WorkerStatus::Idle,
working_directory, working_directory,
)?; )?;
self.worker_detail(worker_ref) self.worker_detail(worker_ref)
@@ -1222,7 +1219,19 @@ impl Runtime {
let mut state = self.lock()?; let mut state = self.lock()?;
state.ensure_running()?; state.ensure_running()?;
let worker = state.worker_mut(worker_ref)?; 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; let status = worker.status;
#[cfg(feature = "ws-server")] #[cfg(feature = "ws-server")]
if let Some(payload) = input_protocol_event(&input) { if let Some(payload) = input_protocol_event(&input) {
@@ -1431,7 +1440,7 @@ impl Runtime {
let entries = self.worker_completions(worker_ref, kind, &prefix)?; let entries = self.worker_completions(worker_ref, kind, &prefix)?;
return Ok(vec![Event::Completions { kind, entries }]); 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()))?; self.stop_worker(worker_ref, Some("worker protocol shutdown".to_string()))?;
return Ok(Vec::new()); return Ok(Vec::new());
} }
@@ -1481,7 +1490,7 @@ impl Runtime {
&self, &self,
worker_ref: &WorkerRef, worker_ref: &WorkerRef,
handle: WorkerExecutionHandle, handle: WorkerExecutionHandle,
run_state: WorkerExecutionRunState, status: WorkerStatus,
working_directory: Option<CatalogWorkingDirectoryStatus>, working_directory: Option<CatalogWorkingDirectoryStatus>,
_result: WorkerExecutionResult, _result: WorkerExecutionResult,
) -> Result<WorkerDetail, RuntimeError> { ) -> Result<WorkerDetail, RuntimeError> {
@@ -1490,7 +1499,7 @@ impl Runtime {
let worker = state.worker_mut(worker_ref)?; let worker = state.worker_mut(worker_ref)?;
worker.execution_handle = Some(handle); worker.execution_handle = Some(handle);
worker.execution_bound = true; 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.restore_intent = restore_intent_for_status(worker.status);
worker.working_directory = working_directory; worker.working_directory = working_directory;
worker.detail() worker.detail()
@@ -1518,16 +1527,28 @@ impl Runtime {
worker_ref: &WorkerRef, worker_ref: &WorkerRef,
result: WorkerExecutionResult, result: WorkerExecutionResult,
) -> Result<(), RuntimeError> { ) -> Result<(), RuntimeError> {
let mut state = self.lock()?; // Accepted dispatch without a state snapshot is transport evidence only;
if result.is_accepted() { // the revisioned protocol stream remains live authority. Test/detached
let status = worker_status_from_run_state(result.run_state); // backends may return an exact full snapshot as their acknowledgement.
let worker = state.worker_mut(worker_ref)?; if !result.is_accepted() {
worker.status = status; return Ok(());
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)?;
} }
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(()) Ok(())
} }
@@ -1730,7 +1751,7 @@ impl Runtime {
context_window: 0, context_window: 0,
context_tokens: 0, context_tokens: 0,
}, },
status: protocol::WorkerStatus::Idle, state: protocol::WorkerStateSnapshot::initial(1),
in_flight: protocol::InFlightSnapshot { in_flight: protocol::InFlightSnapshot {
blocks: Vec::new(), blocks: Vec::new(),
commands: Vec::new(), commands: Vec::new(),
@@ -1968,12 +1989,11 @@ impl Runtime {
match backend.restore_worker(request) { match backend.restore_worker(request) {
WorkerExecutionSpawnResult::Connected { WorkerExecutionSpawnResult::Connected {
handle, handle,
run_state,
working_directory, working_directory,
} => self.commit_restored_worker_execution( } => self.commit_restored_worker_execution(
&candidate.worker_ref, &candidate.worker_ref,
handle, handle,
run_state, WorkerStatus::Idle,
working_directory, working_directory,
)?, )?,
WorkerExecutionSpawnResult::Rejected(result) WorkerExecutionSpawnResult::Rejected(result)
@@ -1990,7 +2010,7 @@ impl Runtime {
&self, &self,
worker_ref: &WorkerRef, worker_ref: &WorkerRef,
handle: WorkerExecutionHandle, handle: WorkerExecutionHandle,
run_state: WorkerExecutionRunState, status: WorkerStatus,
working_directory: Option<CatalogWorkingDirectoryStatus>, working_directory: Option<CatalogWorkingDirectoryStatus>,
) -> Result<(), RuntimeError> { ) -> Result<(), RuntimeError> {
let mut state = self.lock()?; let mut state = self.lock()?;
@@ -1999,7 +2019,7 @@ impl Runtime {
let worker = state.worker_mut(worker_ref)?; let worker = state.worker_mut(worker_ref)?;
worker.execution_handle = Some(handle); worker.execution_handle = Some(handle);
worker.execution_bound = true; 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.restore_intent = restore_intent_for_status(worker.status);
worker.working_directory = working_directory; worker.working_directory = working_directory;
} }
@@ -2867,7 +2887,7 @@ impl RuntimeState {
) { ) {
match event { match event {
protocol::Event::Snapshot { protocol::Event::Snapshot {
status, state,
internal_workers, internal_workers,
.. ..
} => { } => {
@@ -2875,7 +2895,7 @@ impl RuntimeState {
statuses.insert( statuses.insert(
worker.session_id.clone(), worker.session_id.clone(),
InternalWorkerActivity { InternalWorkerActivity {
status: *status, status: state.catalog_status(),
parent_session_id: worker.parent_session_id.clone(), parent_session_id: worker.parent_session_id.clone(),
}, },
); );
@@ -2888,26 +2908,17 @@ impl RuntimeState {
event, event,
.. ..
} => Self::project_internal_worker_event(statuses, nested_worker, event), } => Self::project_internal_worker_event(statuses, nested_worker, event),
protocol::Event::Status { status } => { protocol::Event::WorkerState { snapshot }
statuses.insert( | protocol::Event::CommandAcknowledged {
worker.session_id.clone(), acknowledgement:
InternalWorkerActivity { protocol::WorkerCommandAcknowledgement {
status: *status, state: snapshot, ..
parent_session_id: worker.parent_session_id.clone(),
}, },
); } => {
}
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( statuses.insert(
worker.session_id.clone(), worker.session_id.clone(),
InternalWorkerActivity { InternalWorkerActivity {
status, status: snapshot.catalog_status(),
parent_session_id: worker.parent_session_id.clone(), parent_session_id: worker.parent_session_id.clone(),
}, },
); );
@@ -2963,28 +2974,21 @@ impl RuntimeState {
return false; return false;
}; };
let next_status = match event { let next_status = match event {
protocol::Event::Status { protocol::Event::WorkerState { snapshot }
status: protocol::WorkerStatus::Running, | protocol::Event::Snapshot {
} => Some(WorkerStatus::Running), state: snapshot, ..
protocol::Event::Status { }
status: protocol::WorkerStatus::Idle, | protocol::Event::CommandAcknowledged {
} => Some(WorkerStatus::Idle), acknowledgement:
protocol::Event::Status { protocol::WorkerCommandAcknowledgement {
status: protocol::WorkerStatus::Paused, state: snapshot, ..
} => Some(WorkerStatus::Paused), },
protocol::Event::Snapshot { status, .. } => match status { } => Some(match snapshot.catalog_status() {
protocol::WorkerStatus::Running => Some(WorkerStatus::Running), protocol::WorkerStatus::Idle => WorkerStatus::Idle,
protocol::WorkerStatus::Idle => Some(WorkerStatus::Idle), protocol::WorkerStatus::Running => WorkerStatus::Running,
protocol::WorkerStatus::Paused => Some(WorkerStatus::Paused), protocol::WorkerStatus::Paused => WorkerStatus::Paused,
protocol::WorkerStatus::Stopped => Some(WorkerStatus::Stopped), protocol::WorkerStatus::Stopped => 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),
},
_ => None, _ => None,
}; };
if let Some(next_status) = next_status { 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 { fn repository_resource_error(error: BackendResourceError) -> RuntimeError {
let (code, message) = match error { let (code, message) = match error {
BackendResourceError::Expired => ( BackendResourceError::Expired => (
@@ -3304,7 +3298,7 @@ mod tests {
}; };
use crate::execution::{ use crate::execution::{
WorkerExecutionBackend, WorkerExecutionContext, WorkerExecutionHandle, WorkerExecutionBackend, WorkerExecutionContext, WorkerExecutionHandle,
WorkerExecutionRestoreRequest, WorkerExecutionRunState, WorkerExecutionRestoreRequest,
}; };
use crate::working_directory::WorkingDirectoryDiagnostic; use crate::working_directory::WorkingDirectoryDiagnostic;
use async_trait::async_trait; use async_trait::async_trait;
@@ -3313,6 +3307,14 @@ mod tests {
use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex}; 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] #[test]
fn repository_resource_failures_keep_typed_credential_diagnostics() { fn repository_resource_failures_keep_typed_credential_diagnostics() {
let cases = [ let cases = [
@@ -3359,7 +3361,9 @@ mod tests {
protocol::Event::InternalWorker { protocol::Event::InternalWorker {
worker, worker,
revision: 1, 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_window: 0,
context_tokens: 0, context_tokens: 0,
}, },
status: protocol::WorkerStatus::Idle, state: protocol::WorkerStatus::Idle.into(),
in_flight: protocol::InFlightSnapshot::default(), in_flight: protocol::InFlightSnapshot::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}; };
@@ -3967,7 +3971,6 @@ mod tests {
.insert(request.worker_ref.worker_id.clone(), request.context); .insert(request.worker_ref.worker_id.clone(), request.context);
WorkerExecutionSpawnResult::Connected { WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: request working_directory: request
.working_directory .working_directory
.as_ref() .as_ref()
@@ -3997,7 +4000,6 @@ mod tests {
.insert(request.worker_ref.worker_id.clone(), request.context); .insert(request.worker_ref.worker_id.clone(), request.context);
WorkerExecutionSpawnResult::Connected { WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: request working_directory: request
.working_directory .working_directory
.as_ref() .as_ref()
@@ -4020,7 +4022,6 @@ mod tests {
.unwrap_or_else(|| { .unwrap_or_else(|| {
WorkerExecutionResult::accepted_submission( WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input, WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
"request-test", "request-test",
"test-submission", "test-submission",
protocol::SubmissionDisposition::Started, protocol::SubmissionDisposition::Started,
@@ -4038,17 +4039,11 @@ mod tests {
} }
fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
WorkerExecutionResult::accepted( WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop)
WorkerExecutionOperation::Stop,
WorkerExecutionRunState::Stopped,
)
} }
fn cancel_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { fn cancel_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
WorkerExecutionResult::accepted( WorkerExecutionResult::accepted(WorkerExecutionOperation::Cancel)
WorkerExecutionOperation::Cancel,
WorkerExecutionRunState::Stopped,
)
} }
#[cfg(feature = "ws-server")] #[cfg(feature = "ws-server")]
@@ -4374,7 +4369,9 @@ mod tests {
.send_protocol_method_scoped( .send_protocol_method_scoped(
&scope("workspace-a", "server-a"), &scope("workspace-a", "server-a"),
&workspace_b.worker_ref, &workspace_b.worker_ref,
Method::Shutdown, Method::Shutdown {
command: test_command(),
},
) )
.unwrap_err(); .unwrap_err();
assert!(matches!( assert!(matches!(
@@ -4722,11 +4719,10 @@ mod tests {
} }
#[test] #[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(); let (runtime, backend) = runtime_and_backend();
backend.set_dispatch_result(WorkerExecutionResult::accepted_submission( backend.set_dispatch_result(WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input, WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
"request-test", "request-test",
"test-submission", "test-submission",
protocol::SubmissionDisposition::Started, protocol::SubmissionDisposition::Started,
@@ -4736,7 +4732,7 @@ mod tests {
let detail = runtime.create_worker(request).unwrap(); let detail = runtime.create_worker(request).unwrap();
assert_eq!(detail.status, WorkerStatus::Idle); assert_eq!(detail.status, WorkerStatus::Running);
} }
#[test] #[test]
@@ -4745,7 +4741,6 @@ mod tests {
backend.preserve_commit_ack_submission_id(); backend.preserve_commit_ack_submission_id();
backend.set_dispatch_result(WorkerExecutionResult::accepted_submission( backend.set_dispatch_result(WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input, WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
"request-test", "request-test",
"forged-submission", "forged-submission",
protocol::SubmissionDisposition::Started, protocol::SubmissionDisposition::Started,
@@ -4770,7 +4765,6 @@ mod tests {
let (runtime, backend) = runtime_and_backend(); let (runtime, backend) = runtime_and_backend();
backend.set_dispatch_result(WorkerExecutionResult::accepted( backend.set_dispatch_result(WorkerExecutionResult::accepted(
WorkerExecutionOperation::Input, WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
)); ));
let mut request = task_request("missing initial input commit ack"); let mut request = task_request("missing initial input commit ack");
request.initial_input = Some(WorkerInput::user("start the ticket")); request.initial_input = Some(WorkerInput::user("start the ticket"));
@@ -4898,7 +4892,7 @@ mod tests {
context_window: 128, context_window: 128,
context_tokens: 64, context_tokens: 64,
}, },
status: protocol::WorkerStatus::Running, state: protocol::WorkerStatus::Running.into(),
in_flight: protocol::InFlightSnapshot { in_flight: protocol::InFlightSnapshot {
blocks: Vec::new(), blocks: Vec::new(),
commands: Vec::new(), commands: Vec::new(),
@@ -4914,13 +4908,13 @@ mod tests {
protocol::Event::Snapshot { protocol::Event::Snapshot {
session, session,
greeting, greeting,
status, state,
.. ..
} => { } => {
assert_eq!(session.entries.len(), 1); assert_eq!(session.entries.len(), 1);
assert_eq!(session.entries[0].entry_id, "restored-log-entry"); assert_eq!(session.entries[0].entry_id, "restored-log-entry");
assert_eq!(greeting.worker_name, "live-worker"); 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:?}"), other => panic!("expected snapshot, got {other:?}"),
} }
@@ -4936,7 +4930,6 @@ mod tests {
fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult {
WorkerExecutionSpawnResult::Connected { WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: request working_directory: request
.working_directory .working_directory
.as_ref() .as_ref()
@@ -4951,7 +4944,6 @@ mod tests {
) -> WorkerExecutionResult { ) -> WorkerExecutionResult {
WorkerExecutionResult::accepted_submission( WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input, WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
"request-test", "request-test",
input input
.submission_request_id .submission_request_id
@@ -4993,7 +4985,12 @@ mod tests {
.unwrap(); .unwrap();
runtime runtime
.send_protocol_method(&detail.worker_ref, Method::Shutdown) .send_protocol_method(
&detail.worker_ref,
Method::Shutdown {
command: test_command(),
},
)
.unwrap(); .unwrap();
assert_eq!( assert_eq!(
@@ -5009,7 +5006,12 @@ mod tests {
.create_worker(task_request("restore explicitly")) .create_worker(task_request("restore explicitly"))
.unwrap(); .unwrap();
runtime runtime
.send_protocol_method(&detail.worker_ref, Method::Shutdown) .send_protocol_method(
&detail.worker_ref,
Method::Shutdown {
command: test_command(),
},
)
.unwrap(); .unwrap();
assert!(matches!( assert!(matches!(
@@ -5027,7 +5029,7 @@ mod tests {
assert_eq!(*backend.run_generations.lock().unwrap(), vec![1, 2]); assert_eq!(*backend.run_generations.lock().unwrap(), vec![1, 2]);
assert_eq!( assert_eq!(
runtime.worker_detail(&detail.worker_ref).unwrap().status, runtime.worker_detail(&detail.worker_ref).unwrap().status,
WorkerStatus::Idle WorkerStatus::Running
); );
} }
+249 -166
View File
@@ -10,8 +10,8 @@
use std::collections::HashMap; use std::collections::HashMap;
use std::future::Future; use std::future::Future;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, mpsc}; use std::sync::{Arc, Mutex, RwLock, mpsc};
use std::time::Duration; use std::time::Duration;
use crate::auth::{ use crate::auth::{
@@ -25,8 +25,8 @@ use crate::catalog::{
}; };
use crate::execution::{ use crate::execution::{
WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation,
WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionRunState, WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionSpawnRequest,
WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, WorkerExecutionSpawnResult,
}; };
use crate::identity::WorkerRef; use crate::identity::WorkerRef;
use crate::interaction::{WorkerInput, WorkerInputKind}; use crate::interaction::{WorkerInput, WorkerInputKind};
@@ -38,7 +38,26 @@ use crate::working_directory::{
WorkingDirectoryBinding, WorkingDirectoryDiagnostic, WorkingDirectoryMaterializer, WorkingDirectoryBinding, WorkingDirectoryDiagnostic, WorkingDirectoryMaterializer,
}; };
use async_trait::async_trait; 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<protocol::WorkerStateSnapshot>,
) -> Result<WorkerCommandEnvelope, String> {
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}; use session_store::{CombinedStore, WorkerAggregateStore, WorkerSessionStore};
#[cfg(test)] #[cfg(test)]
use session_store::{FsStore, FsWorkerStore}; use session_store::{FsStore, FsWorkerStore};
@@ -172,7 +191,7 @@ impl WorkerObservationProvider for RuntimeGrantedWorkerObservationProvider {
}, },
display_name: grant.worker_id.clone(), display_name: grant.worker_id.clone(),
relation: "granted_peer".to_string(), 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)); subjects.sort_by(|left, right| left.subject.cmp(&right.subject));
@@ -1174,10 +1193,12 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
} }
} }
#[derive(Clone)]
struct RuntimeWorkerExecution { struct RuntimeWorkerExecution {
handle: WorkerHandle, handle: WorkerHandle,
shutdown: Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>, shutdown: Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>,
busy: Arc<AtomicBool>, busy: Arc<AtomicBool>,
worker_state: Arc<RwLock<protocol::WorkerStateSnapshot>>,
workspace_client: Option<Arc<dyn WorkspaceClient>>, workspace_client: Option<Arc<dyn WorkspaceClient>>,
} }
@@ -1276,6 +1297,7 @@ where
( (
WorkerHandle, WorkerHandle,
Arc<AtomicBool>, Arc<AtomicBool>,
Arc<RwLock<protocol::WorkerStateSnapshot>>,
Option<Arc<dyn WorkspaceClient>>, Option<Arc<dyn WorkspaceClient>>,
), ),
WorkerExecutionResult, WorkerExecutionResult,
@@ -1302,6 +1324,7 @@ where
( (
execution.handle.clone(), execution.handle.clone(),
execution.busy.clone(), execution.busy.clone(),
execution.worker_state.clone(),
execution.workspace_client.clone(), execution.workspace_client.clone(),
) )
}) })
@@ -1318,7 +1341,6 @@ where
operation: WorkerExecutionOperation, operation: WorkerExecutionOperation,
worker: WorkerHandle, worker: WorkerHandle,
method: Method, method: Method,
accepted_run_state: WorkerExecutionRunState,
) -> WorkerExecutionResult { ) -> WorkerExecutionResult {
self.run_on_adapter_runtime(async move { self.run_on_adapter_runtime(async move {
worker worker
@@ -1326,7 +1348,7 @@ where
.await .await
.map_err(|err| format!("failed to send Worker method: {err}")) .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)) .unwrap_or_else(|message| WorkerExecutionResult::errored(operation, message))
} }
@@ -1336,7 +1358,6 @@ where
worker: WorkerHandle, worker: WorkerHandle,
method: Method, method: Method,
submission_request_id: String, submission_request_id: String,
accepted_run_state: WorkerExecutionRunState,
) -> WorkerExecutionResult { ) -> WorkerExecutionResult {
let request_id = submission_request_id.clone(); let request_id = submission_request_id.clone();
self.run_on_adapter_runtime(async move { self.run_on_adapter_runtime(async move {
@@ -1395,7 +1416,6 @@ where
.map(|(submission_id, disposition)| { .map(|(submission_id, disposition)| {
WorkerExecutionResult::accepted_submission( WorkerExecutionResult::accepted_submission(
operation, operation,
accepted_run_state,
submission_request_id, submission_request_id,
submission_id, submission_id,
disposition, disposition,
@@ -1415,38 +1435,45 @@ where
workspace_client: Option<Arc<dyn WorkspaceClient>>, workspace_client: Option<Arc<dyn WorkspaceClient>>,
) -> WorkerExecutionSpawnResult { ) -> WorkerExecutionSpawnResult {
let busy = Arc::new(AtomicBool::new(false)); let busy = Arc::new(AtomicBool::new(false));
let worker_state = Arc::new(RwLock::new(handle.shared_state.snapshot()));
#[cfg(feature = "ws-server")] #[cfg(feature = "ws-server")]
{ {
let streams = subscribe_worker_protocol_session(&handle); let streams = subscribe_worker_protocol_session(&handle);
let mut events = streams.events; let mut events = streams.events;
let mut entry_events = streams.log_entries; let mut entry_events = streams.log_entries;
let bridge_busy = busy.clone(); let bridge_busy = busy.clone();
let bridge_worker_state = worker_state.clone();
if let Err(message) = self.spawn_on_adapter_runtime(async move { if let Err(message) = self.spawn_on_adapter_runtime(async move {
loop { loop {
tokio::select! { tokio::select! {
event = events.recv() => { event = events.recv() => {
match event { match event {
Ok(event) => { Ok(event) => {
let next_busy = match &event { let next_state = match &event {
Event::InvokeStart { .. } Event::WorkerState { snapshot }
| Event::Status { | Event::Snapshot { state: snapshot, .. } => {
status: WorkerStatus::Running, Some(snapshot.clone())
} => Some(true),
Event::RunEnd { .. }
| Event::Error {
code: ErrorCode::NotPaused,
..
} }
| Event::Status { Event::CommandAcknowledged { acknowledgement } => {
status: Some(acknowledgement.state.clone())
WorkerStatus::Idle
| WorkerStatus::Paused
| WorkerStatus::Stopped,
} }
| Event::Shutdown => Some(false),
_ => None, _ => 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); 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 { if let Some(next_busy) = next_busy {
bridge_busy.store(next_busy, Ordering::SeqCst); bridge_busy.store(next_busy, Ordering::SeqCst);
} }
@@ -1494,13 +1521,13 @@ where
handle, handle,
shutdown, shutdown,
busy, busy,
worker_state,
workspace_client, workspace_client,
}, },
); );
WorkerExecutionSpawnResult::Connected { WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(worker_ref, self.backend_id()), handle: WorkerExecutionHandle::new(worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: working_directory.map(|binding| binding.status()), working_directory: working_directory.map(|binding| binding.status()),
} }
} }
@@ -1516,6 +1543,17 @@ impl<F> Drop for WorkerRuntimeExecutionBackend<F> {
} }
} }
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 { fn method_starts_turn(method: &Method) -> bool {
matches!( matches!(
method, method,
@@ -1523,41 +1561,17 @@ fn method_starts_turn(method: &Method) -> bool {
| Method::SubmitTracked { .. } | Method::SubmitTracked { .. }
| Method::Notify { auto_run: true, .. } | Method::Notify { auto_run: true, .. }
| Method::NotifyTracked { auto_run: true, .. } | Method::NotifyTracked { auto_run: true, .. }
| Method::Resume | Method::Resume { .. }
| Method::Compact
) )
} }
fn method_can_start_turn_from_status(method: &Method, status: WorkerStatus) -> bool { fn method_can_start_turn_from_status(method: &Method, status: WorkerStatus) -> bool {
match method { match method {
Method::Resume => matches!(status, WorkerStatus::Idle | WorkerStatus::Paused), Method::Resume { .. } => matches!(status, WorkerStatus::Idle | WorkerStatus::Paused),
_ => status == WorkerStatus::Idle, _ => 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<F> WorkerExecutionBackend for WorkerRuntimeExecutionBackend<F> impl<F> WorkerExecutionBackend for WorkerRuntimeExecutionBackend<F>
where where
F: RuntimeWorkerFactory, F: RuntimeWorkerFactory,
@@ -1883,7 +1897,7 @@ where
handle: &WorkerExecutionHandle, handle: &WorkerExecutionHandle,
input: WorkerInput, input: WorkerInput,
) -> WorkerExecutionResult { ) -> 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, Ok(execution) => execution,
Err(mut result) => { Err(mut result) => {
result.operation = WorkerExecutionOperation::Input; result.operation = WorkerExecutionOperation::Input;
@@ -1892,8 +1906,7 @@ where
}; };
if input.kind == WorkerInputKind::Notify { if input.kind == WorkerInputKind::Notify {
let status = worker.shared_state.get_status(); let status = worker.shared_state.catalog_status();
let accepted_run_state = accepted_notify_run_state(status, true);
let claimed_here = status == WorkerStatus::Idle let claimed_here = status == WorkerStatus::Idle
&& busy && busy
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
@@ -1912,7 +1925,6 @@ where
operation_id: notification_request_id, operation_id: notification_request_id,
}, },
}, },
accepted_run_state,
); );
if claimed_here && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted if claimed_here && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted
{ {
@@ -1921,8 +1933,22 @@ where
return result; 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 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 let claimed_here = status == WorkerStatus::Idle
&& busy && busy
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
@@ -1962,7 +1988,7 @@ where
WorkerInputKind::Notify => { WorkerInputKind::Notify => {
unreachable!("Notify input is dispatched before the turn-start busy guard") 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::ListRewindTargets => (Method::ListRewindTargets, None),
WorkerInputKind::RegisterPeer => ( WorkerInputKind::RegisterPeer => (
Method::RegisterPeer { Method::RegisterPeer {
@@ -1971,15 +1997,6 @@ where
None, 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 waits_for_submission_acceptance = submission_request_id.is_some();
let result = if waits_for_submission_acceptance { let result = if waits_for_submission_acceptance {
@@ -1988,20 +2005,11 @@ where
worker, worker,
method, method,
submission_request_id.expect("Submit must have a submission request id"), submission_request_id.expect("Submit must have a submission request id"),
accepted_run_state,
) )
} else { } else {
self.send_method( self.send_method(WorkerExecutionOperation::Input, worker, method)
WorkerExecutionOperation::Input,
worker,
method,
accepted_run_state,
)
}; };
if accepted_is_idle if claimed_here && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted {
|| (claimed_here
&& result.outcome != crate::execution::WorkerExecutionOutcome::Accepted)
{
busy.store(false, Ordering::SeqCst); busy.store(false, Ordering::SeqCst);
} }
result result
@@ -2015,7 +2023,7 @@ where
content: &[u8], content: &[u8],
context: Option<&session_store::UploadedFileUploadContext>, context: Option<&session_store::UploadedFileUploadContext>,
) -> Result<protocol::UploadedFileRef, WorkerExecutionResult> { ) -> Result<protocol::UploadedFileRef, WorkerExecutionResult> {
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.operation = WorkerExecutionOperation::UploadFile;
result result
})?; })?;
@@ -2038,7 +2046,7 @@ where
handle: &WorkerExecutionHandle, handle: &WorkerExecutionHandle,
artifact_id: &str, artifact_id: &str,
) -> WorkerExecutionResult { ) -> WorkerExecutionResult {
let (worker, _, _) = match self.get_execution(handle) { let (worker, _, _, _) = match self.get_execution(handle) {
Ok(execution) => execution, Ok(execution) => execution,
Err(mut result) => { Err(mut result) => {
result.operation = WorkerExecutionOperation::DeleteUploadedFile; result.operation = WorkerExecutionOperation::DeleteUploadedFile;
@@ -2046,10 +2054,7 @@ where
} }
}; };
match worker.delete_uploaded_file(artifact_id) { match worker.delete_uploaded_file(artifact_id) {
Ok(_) => WorkerExecutionResult::accepted( Ok(_) => WorkerExecutionResult::accepted(WorkerExecutionOperation::DeleteUploadedFile),
WorkerExecutionOperation::DeleteUploadedFile,
WorkerExecutionRunState::Idle,
),
Err(error) => WorkerExecutionResult::rejected( Err(error) => WorkerExecutionResult::rejected(
WorkerExecutionOperation::DeleteUploadedFile, WorkerExecutionOperation::DeleteUploadedFile,
format!("uploaded_file_delete_rejected: {error}"), format!("uploaded_file_delete_rejected: {error}"),
@@ -2062,7 +2067,7 @@ where
handle: &WorkerExecutionHandle, handle: &WorkerExecutionHandle,
method: Method, method: Method,
) -> WorkerExecutionResult { ) -> 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, Ok(execution) => execution,
Err(mut result) => { Err(mut result) => {
result.operation = WorkerExecutionOperation::ProtocolMethod; result.operation = WorkerExecutionOperation::ProtocolMethod;
@@ -2076,19 +2081,13 @@ where
} }
_ => None, _ => None,
} { } {
let status = worker.shared_state.get_status(); let status = worker.shared_state.catalog_status();
let accepted_run_state = accepted_notify_run_state(status, auto_run);
let claimed_here = status == WorkerStatus::Idle let claimed_here = status == WorkerStatus::Idle
&& auto_run && auto_run
&& busy && busy
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_ok(); .is_ok();
let result = self.send_method( let result = self.send_method(WorkerExecutionOperation::ProtocolMethod, worker, method);
WorkerExecutionOperation::ProtocolMethod,
worker,
method,
accepted_run_state,
);
if claimed_here && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted if claimed_here && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted
{ {
busy.store(false, Ordering::SeqCst); busy.store(false, Ordering::SeqCst);
@@ -2098,7 +2097,7 @@ where
let starts_turn = method_starts_turn(&method); let starts_turn = method_starts_turn(&method);
if starts_turn 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 || busy
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_err()) .is_err())
@@ -2109,17 +2108,8 @@ where
); );
} }
let accepted_run_state = accepted_run_state_for_method(&method); let result = self.send_method(WorkerExecutionOperation::ProtocolMethod, worker, method);
let accepted_is_idle = accepted_run_state == WorkerExecutionRunState::Idle; if starts_turn && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted {
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)
{
busy.store(false, Ordering::SeqCst); busy.store(false, Ordering::SeqCst);
} }
result result
@@ -2137,7 +2127,7 @@ where
); );
} }
let execution = match self.workers.lock() { let execution = match self.workers.lock() {
Ok(mut workers) => workers.remove(handle.worker_ref()), Ok(workers) => workers.get(handle.worker_ref()).cloned(),
Err(_) => { Err(_) => {
return WorkerExecutionResult::errored( return WorkerExecutionResult::errored(
WorkerExecutionOperation::Stop, WorkerExecutionOperation::Stop,
@@ -2153,48 +2143,73 @@ where
}; };
let artifact_cleanup = execution.handle.clone(); let artifact_cleanup = execution.handle.clone();
let shutdown = execution.shutdown.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( let result = self.send_method(
WorkerExecutionOperation::Stop, WorkerExecutionOperation::Stop,
execution.handle, execution.handle.clone(),
Method::Shutdown, Method::Shutdown { command },
WorkerExecutionRunState::Stopped,
); );
if result.outcome != crate::execution::WorkerExecutionOutcome::Accepted { if result.outcome != crate::execution::WorkerExecutionOutcome::Accepted {
return result; return result;
} }
match self.run_on_adapter_runtime(async move { let shutdown_wait = self.run_on_adapter_runtime(async move {
let receiver = shutdown.lock().await.take(); let mut guard = shutdown.lock().await;
if let Some(receiver) = receiver { let Some(mut receiver) = guard.take() else {
receiver return Ok(());
.await };
.map_err(|_| "Worker shutdown completion channel closed".to_string())?; 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(()) });
}) { if let Err(message) = shutdown_wait {
Ok(()) => match artifact_cleanup.delete_uncommitted_uploaded_files() { return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message);
Ok(_) => result, }
Err(error) => WorkerExecutionResult::errored( if let Err(error) = artifact_cleanup.delete_uncommitted_uploaded_files() {
WorkerExecutionOperation::Stop, return WorkerExecutionResult::errored(
format!("uploaded_file_cleanup_failed: {error}"), WorkerExecutionOperation::Stop,
), format!("uploaded_file_cleanup_failed: {error}"),
}, );
Err(message) => WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message), }
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 { 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, Ok(execution) => execution,
Err(mut result) => { Err(mut result) => {
result.operation = WorkerExecutionOperation::Cancel; result.operation = WorkerExecutionOperation::Cancel;
return result; return result;
} }
}; };
let command = match next_internal_command(&worker_state) {
Ok(command) => command,
Err(error) => {
return WorkerExecutionResult::errored(WorkerExecutionOperation::Cancel, error);
}
};
self.send_method( self.send_method(
WorkerExecutionOperation::Cancel, WorkerExecutionOperation::Cancel,
worker, worker,
Method::Cancel, Method::Cancel { command },
WorkerExecutionRunState::Idle,
) )
} }
@@ -2259,6 +2274,29 @@ mod tests {
use manifest::{Scope, WorkerManifest}; use manifest::{Scope, WorkerManifest};
use session_store::{LogEntry, WorkerMetadataStore}; 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<MockFactory>,
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] #[test]
fn workspace_prompt_projection_notification_advances_shared_cache() { fn workspace_prompt_projection_notification_advances_shared_cache() {
let cache = WorkspacePromptProjectionCache::default(); let cache = WorkspacePromptProjectionCache::default();
@@ -2406,41 +2444,39 @@ mod tests {
} }
#[test] #[test]
fn notify_run_state_allows_running_worker_inbox_delivery() { fn compact_is_maintenance_not_a_turn_start() {
assert_eq!( assert!(!method_starts_turn(&Method::Compact {
accepted_notify_run_state(WorkerStatus::Running, true), command: test_command(),
WorkerExecutionRunState::Busy }));
); assert!(method_starts_turn(&Method::Resume {
assert_eq!( command: test_command(),
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
);
} }
#[test] #[test]
fn resume_turn_claim_accepts_paused_and_idle_but_not_running_status() { fn resume_turn_claim_accepts_paused_and_idle_but_not_running_status() {
assert!(method_can_start_turn_from_status( assert!(method_can_start_turn_from_status(
&Method::Resume, &Method::Resume {
command: test_command()
},
WorkerStatus::Paused WorkerStatus::Paused
)); ));
assert!(method_can_start_turn_from_status( assert!(method_can_start_turn_from_status(
&Method::Resume, &Method::Resume {
command: test_command()
},
WorkerStatus::Idle WorkerStatus::Idle
)); ));
assert!(!method_can_start_turn_from_status( assert!(!method_can_start_turn_from_status(
&Method::Resume, &Method::Resume {
command: test_command()
},
WorkerStatus::Running WorkerStatus::Running
)); ));
assert!(!method_can_start_turn_from_status( assert!(!method_can_start_turn_from_status(
&Method::Compact, &Method::Compact {
command: test_command()
},
WorkerStatus::Paused WorkerStatus::Paused
)); ));
} }
@@ -2656,19 +2692,22 @@ mod tests {
let observed = { let observed = {
let workers = backend.workers.lock().unwrap(); let workers = backend.workers.lock().unwrap();
let execution = workers.get(worker_ref).expect("live Worker execution"); 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), execution.busy.load(Ordering::SeqCst),
) )
}; };
if observed == (expected_status, expected_busy) { if observed == (expected_status, expected_status, expected_busy) {
return; return;
} }
assert!( assert!(
std::time::Instant::now() < deadline, 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.0,
observed.1, observed.1,
observed.2,
); );
std::thread::sleep(Duration::from_millis(10)); 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"); .expect("in-process restore must not bind the overlong Unix socket path");
assert_eq!( assert_eq!(
controller.handle.shared_state.get_status(), controller.handle.shared_state.catalog_status(),
WorkerStatus::Idle WorkerStatus::Idle
); );
assert!(!socket_path.exists()); assert!(!socket_path.exists());
assert!(run_dir.join("worker.out.log").is_file()); assert!(run_dir.join("worker.out.log").is_file());
assert!(run_dir.join("worker.err.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() { if let Some(receiver) = controller.shutdown.lock().await.take() {
receiver.await.unwrap(); receiver.await.unwrap();
} }
@@ -3289,7 +3334,9 @@ mod tests {
backend backend
.run_on_adapter_runtime(async move { .run_on_adapter_runtime(async move {
handle handle
.send(Method::Shutdown) .send(Method::Shutdown {
command: test_command(),
})
.await .await
.map_err(|error| error.to_string())?; .map_err(|error| error.to_string())?;
if let Some(receiver) = shutdown.lock().await.take() { if let Some(receiver) = shutdown.lock().await.take() {
@@ -3614,6 +3661,7 @@ mod tests {
#[test] #[test]
#[cfg(feature = "ws-server")] #[cfg(feature = "ws-server")]
#[serial_test::serial(worker_allocation)]
fn adapter_resumes_paused_turn_once_and_preserves_idle_not_paused_error() { fn adapter_resumes_paused_turn_once_and_preserves_idle_not_paused_error() {
let hanging_events = || simple_text_events().into_iter().take(2).collect::<Vec<_>>(); let hanging_events = || simple_text_events().into_iter().take(2).collect::<Vec<_>>();
let client = MockClient::sequential(vec![ let client = MockClient::sequential(vec![
@@ -3649,7 +3697,12 @@ mod tests {
wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Running, true); wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Running, true);
let running_resume = runtime 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"); .expect_err("Resume while Running must be rejected");
assert!( assert!(
running_resume running_resume
@@ -3659,17 +3712,32 @@ mod tests {
); );
runtime 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"); .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, false);
runtime 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"); .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, true);
let duplicate_resume = runtime 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"); .expect_err("duplicate Resume must be rejected");
assert!( assert!(
duplicate_resume duplicate_resume
@@ -3679,17 +3747,32 @@ mod tests {
); );
runtime 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"); .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, false);
runtime 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"); .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, false);
assert_eq!(call_count.load(Ordering::SeqCst), 3); assert_eq!(call_count.load(Ordering::SeqCst), 3);
runtime 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"); .expect("Idle Resume preserves controller NotPaused semantics");
wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Idle, false); wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Idle, false);
let events = runtime let events = runtime
@@ -3698,10 +3781,10 @@ mod tests {
assert!(events.iter().any(|event| { assert!(events.iter().any(|event| {
matches!( matches!(
&event.payload, &event.payload,
Event::Error { Event::CommandAcknowledged { acknowledgement }
code: protocol::ErrorCode::NotPaused, if acknowledgement.command == protocol::WorkerCommandKind::Resume
.. && acknowledgement.disposition
} == protocol::WorkerCommandDisposition::InvalidState
) )
})); }));
assert_eq!(call_count.load(Ordering::SeqCst), 3); assert_eq!(call_count.load(Ordering::SeqCst), 3);
+557 -90
View File
@@ -1,3 +1,4 @@
use std::collections::VecDeque;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use std::sync::Arc; use std::sync::Arc;
use std::sync::atomic::Ordering; use std::sync::atomic::Ordering;
@@ -28,7 +29,9 @@ use protocol::{
AlertLevel, AlertSource, CommandEvent as ProtocolCommandEvent, AlertLevel, AlertSource, CommandEvent as ProtocolCommandEvent,
CommandSnapshot as ProtocolCommandSnapshot, CommandStatus as ProtocolCommandStatus, CommandSnapshot as ProtocolCommandSnapshot, CommandStatus as ProtocolCommandStatus,
CommandStream as ProtocolCommandStream, CommandStreamSlice as ProtocolCommandStreamSlice, 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::{ use workdir::{
CommandEvent as WorkdirCommandEvent, CommandSnapshot as WorkdirCommandSnapshot, CommandEvent as WorkdirCommandEvent, CommandSnapshot as WorkdirCommandSnapshot,
@@ -138,7 +141,7 @@ impl WorkerHandle {
let event = Event::Snapshot { let event = Event::Snapshot {
session, session,
greeting: self.shared_state.greeting.clone(), greeting: self.shared_state.greeting.clone(),
status: self.shared_state.get_status(), state: self.shared_state.snapshot(),
in_flight, in_flight,
internal_workers: self.spawned_registry.internal_worker_snapshots(), internal_workers: self.spawned_registry.internal_worker_snapshots(),
}; };
@@ -178,15 +181,81 @@ impl WorkerHandle {
} }
} }
fn validate_command(
envelope: WorkerCommandEnvelope,
shared_state: &WorkerSharedState,
) -> Result<(), WorkerCommandDisposition> {
let snapshot = shared_state.snapshot();
if envelope.expected_execution_generation != snapshot.execution_generation {
return Err(WorkerCommandDisposition::StaleExecutionGeneration);
}
if envelope.expected_worker_state_revision != snapshot.revision {
return Err(WorkerCommandDisposition::StaleWorkerStateRevision);
}
if !shared_state.accept_command_id(envelope.command_id) {
return Err(WorkerCommandDisposition::StaleCommandId);
}
Ok(())
}
fn acknowledge_command(
working_event_tx: &broadcast::Sender<Event>,
shared_state: &WorkerSharedState,
command_id: u64,
command: WorkerCommandKind,
disposition: WorkerCommandDisposition,
) {
let _ = working_event_tx.send(Event::CommandAcknowledged {
acknowledgement: WorkerCommandAcknowledgement {
command_id,
command,
disposition,
state: shared_state.snapshot(),
},
});
}
fn reject_invalid_command_state(
working_event_tx: &broadcast::Sender<Event>,
shared_state: &WorkerSharedState,
envelope: WorkerCommandEnvelope,
command: WorkerCommandKind,
) {
acknowledge_command(
working_event_tx,
shared_state,
envelope.command_id,
command,
WorkerCommandDisposition::InvalidState,
);
}
async fn set_controller_state(
shared_state: &Arc<WorkerSharedState>,
runtime_dir: &RuntimeDir,
working_event_tx: &broadcast::Sender<Event>,
state: WorkerState,
) -> protocol::WorkerStateSnapshot {
let snapshot = shared_state.transition(state);
let _ = runtime_dir.write_status(shared_state).await;
let _ = working_event_tx.send(Event::WorkerState {
snapshot: snapshot.clone(),
});
snapshot
}
async fn set_controller_status( async fn set_controller_status(
shared_state: &Arc<WorkerSharedState>, shared_state: &Arc<WorkerSharedState>,
runtime_dir: &RuntimeDir, runtime_dir: &RuntimeDir,
working_event_tx: &broadcast::Sender<Event>, working_event_tx: &broadcast::Sender<Event>,
status: WorkerStatus, status: WorkerStatus,
) { ) {
shared_state.set_status(status); let state = match status {
let _ = runtime_dir.write_status(shared_state).await; WorkerStatus::Idle | WorkerStatus::Stopped => WorkerState::Idle,
let _ = working_event_tx.send(Event::Status { status }); WorkerStatus::Running => WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running)),
WorkerStatus::Paused => WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Paused)),
};
set_controller_state(shared_state, runtime_dir, working_event_tx, state).await;
} }
async fn finish_controller_run<C, St>( async fn finish_controller_run<C, St>(
@@ -659,12 +728,24 @@ impl WorkerController {
// === 4. Initial runtime files + WorkerSharedState + WorkerHandle + // === 4. Initial runtime files + WorkerSharedState + WorkerHandle +
// SocketServer === // SocketServer ===
let manifest_toml = toml::to_string_pretty(worker.manifest()).unwrap_or_default(); 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 greeting = build_greeting(&worker);
let shared_state = Arc::new(WorkerSharedState::new( let execution_generation = runtime_dir
.path()
.file_name()
.and_then(|name| name.to_str())
.and_then(|name| name.parse::<u64>().ok())
.filter(|generation| *generation > 0)
.unwrap_or(1);
let shared_state = Arc::new(WorkerSharedState::new_with_generation(
worker.manifest().worker.name.clone(), worker.manifest().worker.name.clone(),
worker.segment_id(), worker.segment_id(),
manifest_toml.clone(), manifest_toml.clone(),
greeting, greeting,
execution_generation,
)); ));
if let Some(fs_for_view) = fs_for_view { if let Some(fs_for_view) = fs_for_view {
shared_state.set_fs_view(crate::fs_view::WorkerFsView::new(fs_for_view)); shared_state.set_fs_view(crate::fs_view::WorkerFsView::new(fs_for_view));
@@ -1432,8 +1513,9 @@ async fn controller_loop<C, St>(
} }
}; };
loop { let mut deferred_methods = VecDeque::new();
// Top-of-iteration: if an event handler staged a run, fire it
'controller: loop {
// here so the status flip → drive_turn → finish sequence lives // here so the status flip → drive_turn → finish sequence lives
// in one place, regardless of which Method caused it. // in one place, regardless of which Method caused it.
if let Some(run) = pending.take() { if let Some(run) = pending.take() {
@@ -1584,9 +1666,13 @@ async fn controller_loop<C, St>(
continue; continue;
} }
let method = match method_rx.recv().await { let method = if let Some(method) = deferred_methods.pop_front() {
Some(m) => m, method
None => break, } else {
match method_rx.recv().await {
Some(method) => method,
None => break,
}
}; };
match method { match method {
@@ -1784,7 +1870,7 @@ async fn controller_loop<C, St>(
expected_revision, expected_revision,
expected_head_id, expected_head_id,
} => { } => {
if shared_state.get_status() != WorkerStatus::Idle { if shared_state.catalog_status() != WorkerStatus::Idle {
let _ = working_event_tx.send(Event::Error { let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest, code: ErrorCode::InvalidRequest,
message: "ContinuePending requires an idle Worker; Resume or Cancel a paused run first".into(), message: "ContinuePending requires an idle Worker; Resume or Cancel a paused run first".into(),
@@ -1811,88 +1897,243 @@ async fn controller_loop<C, St>(
} }
} }
} }
Method::Resume => { Method::Resume { command } => {
if shared_state.get_status() != WorkerStatus::Paused { if let Err(disposition) = validate_command(command, &shared_state) {
let _ = working_event_tx.send(Event::Error { acknowledge_command(
code: ErrorCode::NotPaused, &working_event_tx,
message: "Worker is not paused".into(), &shared_state,
}); command.command_id,
WorkerCommandKind::Resume,
disposition,
);
continue; 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); pending = Some(PendingRun::Resume);
} }
Method::Cancel => match shared_state.get_status() { Method::Cancel { command } => {
WorkerStatus::Paused => match worker.cancel_paused_turn() { 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(()) => { Ok(()) => {
worker.clear_in_flight_events(); worker.clear_in_flight_events();
set_controller_status( set_controller_state(
&shared_state, &shared_state,
&runtime_dir, &runtime_dir,
&working_event_tx, &working_event_tx,
WorkerStatus::Idle, WorkerState::Idle,
) )
.await; .await;
} }
Err(error) => { 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 { let _ = working_event_tx.send(Event::Error {
code: worker_error_code(&error), code: worker_error_code(&error),
message: error.to_string(), 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() { Method::Pause { command } => {
WorkerStatus::Idle => { if let Err(disposition) = validate_command(command, &shared_state) {
if let Err(error) = worker.manual_compact().await { acknowledge_command(
let _ = working_event_tx.send(Event::Error { &working_event_tx,
code: worker_error_code(&error), &shared_state,
message: error.to_string(), 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 => { WorkerStatus::Idle | WorkerStatus::Paused => {
emit_rewind_targets(&worker, &working_event_tx) emit_rewind_targets(&worker, &working_event_tx)
} }
@@ -1908,7 +2149,7 @@ async fn controller_loop<C, St>(
Method::RewindTo { Method::RewindTo {
target, target,
expected_head_entries, expected_head_entries,
} => match shared_state.get_status() { } => match shared_state.catalog_status() {
WorkerStatus::Idle => { WorkerStatus::Idle => {
if apply_rewind( if apply_rewind(
&mut worker, &mut worker,
@@ -1919,10 +2160,8 @@ async fn controller_loop<C, St>(
.await .await
{ {
worker.clear_in_flight_events(); worker.clear_in_flight_events();
shared_state.set_status(WorkerStatus::Idle); let snapshot = shared_state.transition(WorkerState::Idle);
let _ = working_event_tx.send(Event::Status { let _ = working_event_tx.send(Event::WorkerState { snapshot });
status: WorkerStatus::Idle,
});
} }
} }
WorkerStatus::Paused => { WorkerStatus::Paused => {
@@ -1941,7 +2180,17 @@ async fn controller_loop<C, St>(
} }
}, },
Method::Shutdown => { Method::Shutdown { command } => {
// Shutdown remains unconditional/retryable even when the caller's
// live-state fence is stale.
shared_state.accept_command_id(command.command_id);
acknowledge_command(
&working_event_tx,
&shared_state,
command.command_id,
WorkerCommandKind::Shutdown,
WorkerCommandDisposition::Accepted,
);
let _ = working_event_tx.send(Event::Shutdown); let _ = working_event_tx.send(Event::Shutdown);
break; break;
} }
@@ -2023,7 +2272,7 @@ async fn controller_loop<C, St>(
// Auto-kick a turn if the Worker is idle so the // Auto-kick a turn if the Worker is idle so the
// notification is not stranded. Matches the // notification is not stranded. Matches the
// `Method::Notify` idle path. // `Method::Notify` idle path.
if shared_state.get_status() == WorkerStatus::Idle { if shared_state.catalog_status() == WorkerStatus::Idle {
pending = Some(PendingRun::RunForNotification { pending = Some(PendingRun::RunForNotification {
invoke_kind: protocol::InvokeKind::WorkerEvent, invoke_kind: protocol::InvokeKind::WorkerEvent,
notification_request_id: None, notification_request_id: None,
@@ -2270,15 +2519,102 @@ where
} }
method = method_rx.recv(), if input_commit.is_none() => { method = method_rx.recv(), if input_commit.is_none() => {
match method { 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(()); 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; 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(()); let _ = pause_tx.try_send(());
} }
Some(Method::Shutdown) => { Some(Method::Shutdown { command }) => {
shared_state.accept_command_id(command.command_id);
shutdown_requested = true; 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(()); let _ = cancel_tx.try_send(());
} }
Some(Method::Submit { 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 { let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning, code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn".into(), 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 { let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning, code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn; rewind/compact can only run while idle or paused" message: "Worker is already executing a turn; rewind/compact can only run while idle or paused"
@@ -2487,7 +2859,7 @@ where
} }
None => { None => {
let _ = cancel_tx.try_send(()); let _ = cancel_tx.try_send(());
shared_state.set_status(WorkerStatus::Idle); shared_state.transition(WorkerState::Idle);
return (WorkerStatus::Idle, false, false); return (WorkerStatus::Idle, false, false);
} }
} }
@@ -2863,7 +3235,7 @@ mod tests {
context_window: 200_000, context_window: 200_000,
context_tokens: 0, context_tokens: 0,
}, },
status: WorkerStatus::Idle, state: WorkerStatus::Idle.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}) })
@@ -2919,9 +3291,17 @@ mod tests {
async fn pause_waits_for_run_boundary_and_uses_safe_pause_channel() { async fn pause_waits_for_run_boundary_and_uses_safe_pause_channel() {
let mut env = make_env().await; let mut env = make_env().await;
let method_tx = env._method_tx.clone(); 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::spawn(async move {
tokio::time::sleep(Duration::from_millis(10)).await; 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 { let worker_future = async {
@@ -3194,8 +3574,13 @@ mod tests {
async fn compact_method_is_rejected_while_running() { async fn compact_method_is_rejected_while_running() {
let mut env = make_env().await; let mut env = make_env().await;
let mut events = env.working_event_tx.subscribe(); 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 env._method_tx
.send(Method::Compact) .send(Method::Compact { command })
.await .await
.expect("send compact"); .expect("send compact");
@@ -3228,11 +3613,93 @@ mod tests {
.expect("event timeout") .expect("event timeout")
.expect("event"); .expect("event");
match event { match event {
Event::Error { code, message } => { Event::CommandAcknowledged { acknowledgement } => {
assert_eq!(code, ErrorCode::AlreadyRunning); assert_eq!(acknowledgement.command, WorkerCommandKind::Compact);
assert!(message.contains("compact"), "got message: {message}"); 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()
);
}
} }
+8 -8
View File
@@ -779,10 +779,10 @@ async fn probe_socket(socket_path: &Path) -> LiveInfo {
loop { loop {
match tokio::time::timeout(PROBE_TIMEOUT, reader.next::<Event>()).await { match tokio::time::timeout(PROBE_TIMEOUT, reader.next::<Event>()).await {
Ok(Ok(Some(Event::Snapshot { Ok(Ok(Some(Event::Snapshot {
status: snapshot_status, state: snapshot_state,
.. ..
}))) => { }))) => {
status = Some(snapshot_status); status = Some(snapshot_state.catalog_status());
break; break;
} }
Ok(Ok(Some(Event::Alert(_)))) => continue, Ok(Ok(Some(Event::Alert(_)))) => continue,
@@ -1507,7 +1507,7 @@ mod tests {
context_window: 0, context_window: 0,
context_tokens: 0, context_tokens: 0,
}, },
status: WorkerStatus::Idle, state: WorkerStatus::Idle.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}) })
@@ -1543,7 +1543,7 @@ mod tests {
context_window: 0, context_window: 0,
context_tokens: 0, context_tokens: 0,
}, },
status: WorkerStatus::Idle, state: WorkerStatus::Idle.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}) })
@@ -1638,7 +1638,7 @@ mod tests {
context_window: 0, context_window: 0,
context_tokens: 0, context_tokens: 0,
}, },
status: WorkerStatus::Idle, state: WorkerStatus::Idle.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}) })
@@ -1665,7 +1665,7 @@ mod tests {
context_window: 0, context_window: 0,
context_tokens: 0, context_tokens: 0,
}, },
status: WorkerStatus::Idle, state: WorkerStatus::Idle.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}) })
@@ -1773,7 +1773,7 @@ mod tests {
context_window: 0, context_window: 0,
context_tokens: 0, context_tokens: 0,
}, },
status: WorkerStatus::Paused, state: WorkerStatus::Paused.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}) })
@@ -1827,7 +1827,7 @@ mod tests {
context_window: 0, context_window: 0,
context_tokens: 0, context_tokens: 0,
}, },
status: WorkerStatus::Idle, state: WorkerStatus::Idle.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), internal_workers: Vec::new(),
}) })
+51 -19
View File
@@ -295,6 +295,38 @@ impl InternalWorkerSessionStatus {
} }
} }
fn send_internal_worker_state(
event_tx: &broadcast::Sender<Event>,
state_revision: &std::sync::atomic::AtomicU64,
status: InternalWorkerSessionStatus,
) {
let state = match status {
InternalWorkerSessionStatus::Idle
| InternalWorkerSessionStatus::Stopped
| InternalWorkerSessionStatus::Failed => protocol::WorkerState::Idle,
InternalWorkerSessionStatus::Paused => protocol::WorkerState::Busy(
protocol::WorkerBusyState::Run(protocol::WorkerRunState::Paused),
),
InternalWorkerSessionStatus::Running => protocol::WorkerState::Busy(
protocol::WorkerBusyState::Run(protocol::WorkerRunState::Running),
),
InternalWorkerSessionStatus::Stopping => protocol::WorkerState::Busy(
protocol::WorkerBusyState::Run(protocol::WorkerRunState::Cancelling),
),
};
let revision = state_revision
.fetch_add(1, std::sync::atomic::Ordering::AcqRel)
.saturating_add(1);
let _ = event_tx.send(Event::WorkerState {
snapshot: protocol::WorkerStateSnapshot {
execution_generation: 1,
revision,
last_command_id: 0,
state,
},
});
}
fn classify_internal_turn_result( fn classify_internal_turn_result(
result: Result<WorkerRunResult, WorkerError>, result: Result<WorkerRunResult, WorkerError>,
) -> (InternalWorkerSessionStatus, Option<String>) { ) -> (InternalWorkerSessionStatus, Option<String>) {
@@ -351,6 +383,7 @@ pub(crate) struct InternalWorkerSessionSnapshot {
pub(crate) struct InternalWorkerSessionHandle { pub(crate) struct InternalWorkerSessionHandle {
command_tx: tokio::sync::mpsc::Sender<InternalWorkerSessionCommand>, command_tx: tokio::sync::mpsc::Sender<InternalWorkerSessionCommand>,
status: Arc<std::sync::atomic::AtomicU8>, status: Arc<std::sync::atomic::AtomicU8>,
state_revision: Arc<std::sync::atomic::AtomicU64>,
store: EphemeralSessionStore, store: EphemeralSessionStore,
session_id: SessionId, session_id: SessionId,
segment_id: SegmentId, segment_id: SegmentId,
@@ -400,6 +433,10 @@ impl InternalWorkerSessionHandle {
self.in_flight.text_delta(block_id, text.to_owned()); 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 { pub(crate) fn protocol_snapshot(&self) -> InternalWorkerSessionSnapshot {
let (entries, in_flight) = { let (entries, in_flight) = {
let guard = self.in_flight.snapshot_guard(); let guard = self.in_flight.snapshot_guard();
@@ -473,9 +510,7 @@ impl InternalWorkerSessionHandle {
}); });
return Err(InternalWorkerSessionError::Unavailable); return Err(InternalWorkerSessionError::Unavailable);
} }
let _ = self.event_tx.send(Event::Status { self.emit_worker_state(InternalWorkerSessionStatus::Running);
status: WorkerStatus::Running,
});
Ok(()) Ok(())
} }
@@ -767,11 +802,13 @@ pub(crate) async fn prepare_internal_worker_session(
let status = Arc::new(std::sync::atomic::AtomicU8::new( let status = Arc::new(std::sync::atomic::AtomicU8::new(
InternalWorkerSessionStatus::Idle.encode(), InternalWorkerSessionStatus::Idle.encode(),
)); ));
let state_revision = Arc::new(std::sync::atomic::AtomicU64::new(0));
let state_changed = Arc::new(tokio::sync::Notify::new()); let state_changed = Arc::new(tokio::sync::Notify::new());
let last_error = Arc::new(Mutex::new(None)); let last_error = Arc::new(Mutex::new(None));
let handle = InternalWorkerSessionHandle { let handle = InternalWorkerSessionHandle {
command_tx, command_tx,
status: status.clone(), status: status.clone(),
state_revision: state_revision.clone(),
store, store,
session_id, session_id,
segment_id, segment_id,
@@ -807,19 +844,11 @@ pub(crate) async fn prepare_internal_worker_session(
message, message,
}); });
} }
let protocol_status = match turn_status { send_internal_worker_state(
InternalWorkerSessionStatus::Idle => WorkerStatus::Idle, &event_tx,
InternalWorkerSessionStatus::Paused => WorkerStatus::Paused, &state_revision,
InternalWorkerSessionStatus::Stopped turn_status,
| InternalWorkerSessionStatus::Failed => WorkerStatus::Stopped, );
InternalWorkerSessionStatus::Running
| InternalWorkerSessionStatus::Stopping => {
unreachable!("run completion cannot remain active")
}
};
let _ = event_tx.send(Event::Status {
status: protocol_status,
});
if let Some(callback) = &on_turn_end { if let Some(callback) = &on_turn_end {
callback(turn_status); callback(turn_status);
} }
@@ -861,9 +890,11 @@ pub(crate) async fn prepare_internal_worker_session(
InternalWorkerSessionStatus::Stopped.encode(), InternalWorkerSessionStatus::Stopped.encode(),
std::sync::atomic::Ordering::Release, std::sync::atomic::Ordering::Release,
); );
let _ = event_tx.send(Event::Status { send_internal_worker_state(
status: WorkerStatus::Stopped, &event_tx,
}); &state_revision,
InternalWorkerSessionStatus::Stopped,
);
let _ = event_tx.send(Event::Shutdown); let _ = event_tx.send(Event::Shutdown);
state_changed.notify_waiters(); state_changed.notify_waiters();
if let Some(done) = stop_done { 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( status: Arc::new(std::sync::atomic::AtomicU8::new(
InternalWorkerSessionStatus::Idle.encode(), InternalWorkerSessionStatus::Idle.encode(),
)), )),
state_revision: Arc::new(std::sync::atomic::AtomicU64::new(0)),
store, store,
session_id, session_id,
segment_id, segment_id,
+3 -2
View File
@@ -197,7 +197,6 @@ pub fn default_base() -> Result<PathBuf, io::Error> {
mod tests { mod tests {
use super::*; use super::*;
use crate::shared_state::WorkerSharedState; use crate::shared_state::WorkerSharedState;
use protocol::WorkerStatus;
fn test_state() -> WorkerSharedState { fn test_state() -> WorkerSharedState {
WorkerSharedState::new( WorkerSharedState::new(
@@ -247,7 +246,9 @@ mod tests {
let rt = RuntimeDir::create(tmp.path(), "my-worker").await.unwrap(); let rt = RuntimeDir::create(tmp.path(), "my-worker").await.unwrap();
let state = test_state(); 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(); rt.write_status(&state).await.unwrap();
let content = std::fs::read_to_string(rt.path().join("status.json")).unwrap(); let content = std::fs::read_to_string(rt.path().join("status.json")).unwrap();
+97 -44
View File
@@ -1,7 +1,12 @@
use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{
use std::sync::{OnceLock, RwLock}; OnceLock, RwLock,
atomic::{AtomicBool, AtomicU64, Ordering},
};
use protocol::WorkerStatus; use protocol::{
WorkerBusyState, WorkerMaintenanceState, WorkerRunState, WorkerState, WorkerStateSnapshot,
WorkerStatus,
};
use serde_json::json; use serde_json::json;
use session_store::SegmentId; use session_store::SegmentId;
@@ -9,20 +14,16 @@ use crate::fs_view::WorkerFsView;
/// Shared state between WorkerController and runtime directory. /// Shared state between WorkerController and runtime directory.
/// ///
/// Controller updates this in-memory; RuntimeDir writes the status /// `WorkerStateSnapshot` is the sole live execution-state authority. Runtime
/// snapshot to disk. Wrapped in `Arc` for sharing. /// catalog status remains a separate lifecycle projection because `Stopped`
/// /// describes the execution handle rather than a live controller state.
/// 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.
pub struct WorkerSharedState { pub struct WorkerSharedState {
pub worker_name: String, pub worker_name: String,
pub segment_id: SegmentId, pub segment_id: SegmentId,
pub manifest_toml: String, pub manifest_toml: String,
pub greeting: protocol::Greeting, pub greeting: protocol::Greeting,
pub status: RwLock<WorkerStatus>, state: RwLock<WorkerStateSnapshot>,
last_command_id: AtomicU64,
/// Worker-from-the-inside view of the filesystem. Set once in /// Worker-from-the-inside view of the filesystem. Set once in
/// `WorkerController::start` after the local WorkdirSession provider is /// `WorkerController::start` after the local WorkdirSession provider is
/// materialised, and read from the IPC server layer to answer /// materialised, and read from the IPC server layer to answer
@@ -38,13 +39,24 @@ impl WorkerSharedState {
segment_id: SegmentId, segment_id: SegmentId,
manifest_toml: String, manifest_toml: String,
greeting: protocol::Greeting, 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 {
Self { Self {
worker_name, worker_name,
segment_id, segment_id,
manifest_toml, manifest_toml,
greeting, greeting,
status: RwLock::new(WorkerStatus::Idle), state: RwLock::new(WorkerStateSnapshot::initial(execution_generation)),
last_command_id: AtomicU64::new(0),
fs_view: OnceLock::new(), fs_view: OnceLock::new(),
flow_transition_enabled: AtomicBool::new(false), flow_transition_enabled: AtomicBool::new(false),
} }
@@ -70,21 +82,57 @@ impl WorkerSharedState {
self.flow_transition_enabled.load(Ordering::Acquire) self.flow_transition_enabled.load(Ordering::Acquire)
} }
pub fn set_status(&self, status: WorkerStatus) { pub fn transition(&self, state: WorkerState) -> WorkerStateSnapshot {
if let Ok(mut s) = self.status.write() { let mut snapshot = self
*s = status; .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 { /// Serialize the runtime-directory lifecycle projection as JSON while
self.status.read().map(|s| *s).unwrap_or(WorkerStatus::Idle) /// retaining the full state snapshot for diagnostics and reconnects.
}
/// Serialize status as JSON.
pub fn status_json(&self) -> String { pub fn status_json(&self) -> String {
let status = self.get_status(); let snapshot = self.snapshot();
json!({ json!({
"state": status, "state": self.catalog_status(),
"worker_state": snapshot,
"segment_id": self.segment_id.to_string(), "segment_id": self.segment_id.to_string(),
"worker_name": self.worker_name, "worker_name": self.worker_name,
}) })
@@ -97,11 +145,12 @@ mod tests {
use super::*; use super::*;
fn test_state() -> WorkerSharedState { fn test_state() -> WorkerSharedState {
WorkerSharedState::new( WorkerSharedState::new_with_generation(
"test-worker".into(), "test-worker".into(),
session_store::new_segment_id(), session_store::new_segment_id(),
"[engine]\nname = \"test-worker\"".into(), "[engine]\nname = \"test-worker\"".into(),
test_greeting(), test_greeting(),
7,
) )
} }
@@ -119,36 +168,40 @@ mod tests {
} }
#[test] #[test]
fn initial_status_is_idle() { fn initial_snapshot_is_idle() {
let state = test_state(); 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] #[test]
fn set_and_get_status() { fn transitions_increment_revision_only_when_state_changes() {
let state = test_state(); let state = test_state();
state.set_status(WorkerStatus::Running); let running = WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running));
assert_eq!(state.get_status(), WorkerStatus::Running); let snapshot = state.transition(running.clone());
state.set_status(WorkerStatus::Paused); assert_eq!(snapshot.revision, 1);
assert_eq!(state.get_status(), WorkerStatus::Paused); 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] #[test]
fn status_json_contains_fields() { fn status_json_contains_full_snapshot_and_catalog_projection() {
let state = test_state(); let state = test_state();
let json = state.status_json(); state.transition(WorkerState::Busy(WorkerBusyState::Maintenance(
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap(); WorkerMaintenanceState::Compacting,
assert_eq!(parsed["state"], "idle"); )));
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_eq!(parsed["worker_name"], "test-worker");
assert!(parsed["segment_id"].is_string()); 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");
}
} }
+9 -3
View File
@@ -97,7 +97,7 @@ mod tests {
context_window: 200_000, context_window: 200_000,
context_tokens: 0, context_tokens: 0,
}, },
status: WorkerStatus::Idle, state: WorkerStatus::Idle.into(),
in_flight: Default::default(), in_flight: Default::default(),
internal_workers: Vec::new(), 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"); let method = received.await.unwrap().expect("expected method");
assert!(matches!(method, Method::Shutdown)); assert!(matches!(method, Method::Shutdown { .. }));
} }
#[tokio::test] #[tokio::test]
+186 -20
View File
@@ -4571,15 +4571,6 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
Ok(()) 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( fn persist_and_send_compact_failed(
&mut self, &mut self,
lifecycle: CompactionLifecycle, lifecycle: CompactionLifecycle,
@@ -4724,7 +4715,97 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
Ok(rewrite_guard) Ok(rewrite_guard)
} }
/// Terminalize and clean up any compaction that was left active by the
/// previous controller generation. This runs before the restored
/// controller publishes its first Idle state.
pub async fn recover_unfinished_compaction(&mut self) -> Result<(), WorkerError> {
let (entries, _) = self.sink.subscribe_with_snapshot();
let latest_payload = entries.iter().rev().find_map(|entry| match entry {
LogEntry::Extension {
domain, payload, ..
} if domain == COMPACTION_EXTENSION_DOMAIN => Some(payload.clone()),
_ => None,
});
let Some(payload) = latest_payload else {
return Ok(());
};
#[derive(serde::Deserialize)]
#[serde(deny_unknown_fields)]
struct CompactionLifecycleWire {
schema_version: u32,
compaction_id: String,
revision: u64,
#[serde(default)]
internal_worker: Option<protocol::InternalWorkerRef>,
state: CompactionLifecycleState,
started_at_ms: u64,
#[serde(default)]
ended_at_ms: Option<u64>,
#[serde(default)]
summary: Option<String>,
#[serde(default)]
error: Option<String>,
#[serde(default)]
new_segment_id: Option<String>,
}
let wire: CompactionLifecycleWire = serde_json::from_value(payload).map_err(|error| {
WorkerError::InvalidState(format!("decode compaction lifecycle: {error}"))
})?;
if !matches!(wire.schema_version, 2 | 3) {
return Err(WorkerError::InvalidState(format!(
"unsupported compaction lifecycle schema version {}",
wire.schema_version
)));
}
let mut lifecycle = CompactionLifecycle {
schema_version: wire.schema_version,
compaction_id: wire.compaction_id,
revision: wire.revision,
internal_worker: wire.internal_worker,
state: wire.state,
started_at_ms: wire.started_at_ms,
ended_at_ms: wire.ended_at_ms,
summary: wire.summary,
error: wire.error,
new_segment_id: wire.new_segment_id,
};
match lifecycle.state {
CompactionLifecycleState::Running => {
lifecycle.schema_version = 3;
lifecycle.revision = lifecycle.revision.saturating_add(1);
lifecycle.state = CompactionLifecycleState::Interrupted;
lifecycle.ended_at_ms = Some(segment_log::now_millis());
lifecycle.error =
Some("worker execution restarted before compaction completed".into());
self.persist_compaction_lifecycle(&lifecycle)?;
self.send_event(Event::CompactFailed {
lifecycle: lifecycle.clone(),
});
self.release_compaction_service(&lifecycle).await;
}
CompactionLifecycleState::Interrupted => {
self.release_compaction_service(&lifecycle).await;
}
CompactionLifecycleState::Done | CompactionLifecycleState::Failed => {}
}
Ok(())
}
pub async fn manual_compact(&mut self) -> Result<ManualCompactResult, WorkerError> { pub async fn manual_compact(&mut self) -> Result<ManualCompactResult, WorkerError> {
self.manual_compact_inner(None).await
}
pub async fn manual_compact_with_cancel(
&mut self,
cancel: tokio::sync::watch::Receiver<bool>,
) -> Result<ManualCompactResult, WorkerError> {
self.manual_compact_inner(Some(cancel)).await
}
async fn manual_compact_inner(
&mut self,
mut cancel: Option<tokio::sync::watch::Receiver<bool>>,
) -> Result<ManualCompactResult, WorkerError> {
if self.manifest.compaction.is_none() { if self.manifest.compaction.is_none() {
let message = let message =
"manual compact is unavailable because [compaction] is not configured".to_string(); "manual compact is unavailable because [compaction] is not configured".to_string();
@@ -4764,7 +4845,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
return Ok(ManualCompactResult::Skipped { message }); return Ok(ManualCompactResult::Skipped { message });
} }
match self.compact(retained).await { match self.compact_with_cancel(retained, cancel.take()).await {
Ok(new_segment_id) => { Ok(new_segment_id) => {
info!(new_segment_id = %new_segment_id, "Manual compaction succeeded"); info!(new_segment_id = %new_segment_id, "Manual compaction succeeded");
if let Some(ref state) = state { if let Some(ref state) = state {
@@ -4937,11 +5018,19 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// Runs one parent-owned observable compaction service and returns the new /// Runs one parent-owned observable compaction service and returns the new
/// Segment ID. Lifecycle revisions are committed before they are broadcast. /// Segment ID. Lifecycle revisions are committed before they are broadcast.
pub async fn compact(&mut self, retained_tokens: u64) -> Result<SegmentId, WorkerError> { pub async fn compact(&mut self, retained_tokens: u64) -> Result<SegmentId, WorkerError> {
self.compact_with_cancel(retained_tokens, None).await
}
async fn compact_with_cancel(
&mut self,
retained_tokens: u64,
mut cancel: Option<tokio::sync::watch::Receiver<bool>>,
) -> Result<SegmentId, WorkerError> {
let _rewrite_guard = self let _rewrite_guard = self
.prepare_session_rewrite(SessionRewriteKind::Compact) .prepare_session_rewrite(SessionRewriteKind::Compact)
.await?; .await?;
let mut lifecycle = CompactionLifecycle { let mut lifecycle = CompactionLifecycle {
schema_version: 2, schema_version: 3,
compaction_id: uuid::Uuid::now_v7().to_string(), compaction_id: uuid::Uuid::now_v7().to_string(),
revision: 1, revision: 1,
internal_worker: None, internal_worker: None,
@@ -4953,16 +5042,25 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
new_segment_id: None, new_segment_id: None,
}; };
self.persist_and_send_compact_start(lifecycle.clone())?; self.persist_and_send_compact_start(lifecycle.clone())?;
match self.compact_impl(retained_tokens, &mut lifecycle).await { let outcome = if let Some(cancel) = cancel.as_mut() {
Ok((new_segment_id, summary)) => { tokio::select! {
lifecycle.revision = lifecycle.revision.saturating_add(1); biased;
lifecycle.state = CompactionLifecycleState::Done; changed = cancel.changed() => {
lifecycle.ended_at_ms = Some(segment_log::now_millis()); let _ = changed;
lifecycle.summary = Some(summary); Err(WorkerError::CompactCancelled)
lifecycle.new_segment_id = Some(new_segment_id.to_string()); }
let terminal = self.persist_and_send_compact_done(lifecycle.clone()); 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; self.release_compaction_service(&lifecycle).await;
terminal?;
Ok(new_segment_id) Ok(new_segment_id)
} }
Err(error) => { Err(error) => {
@@ -5543,6 +5641,24 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
})?, })?,
}); });
} }
// Commit the terminal lifecycle in the same atomic replacement-segment
// creation as the rewritten history. Restore can therefore never see a
// replacement segment without the Done fact for the compaction that
// created it.
lifecycle.revision = lifecycle.revision.saturating_add(1);
lifecycle.state = CompactionLifecycleState::Done;
lifecycle.ended_at_ms = Some(segment_log::now_millis());
lifecycle.summary = Some(summary_text.clone());
lifecycle.new_segment_id = Some(new_segment_id.to_string());
initial_entries.push(LogEntry::Extension {
ts: segment_log::now_millis(),
domain: COMPACTION_EXTENSION_DOMAIN.to_string(),
payload: serde_json::to_value(&*lifecycle).map_err(|error| {
WorkerError::InvalidState(format!(
"serialize terminal compaction lifecycle: {error}"
))
})?,
});
self.store self.store
.create_segment(old_loc.session_id, new_segment_id, &initial_entries)?; .create_segment(old_loc.session_id, new_segment_id, &initial_entries)?;
self.segment_state.set_location(SegmentLocation { self.segment_state.set_location(SegmentLocation {
@@ -10165,6 +10281,56 @@ mod build_summary_prompt_tests {
assert_eq!(state.notification_receipts.len(), 1); assert_eq!(state.notification_receipts.len(), 1);
} }
#[tokio::test]
async fn restore_terminalizes_running_compaction_before_idle_publication() {
let (_dir, mut worker) = rewind_test_worker().await;
let lifecycle = CompactionLifecycle {
schema_version: 3,
compaction_id: "compact-before-restart".into(),
revision: 1,
internal_worker: None,
state: CompactionLifecycleState::Running,
started_at_ms: segment_log::now_millis(),
ended_at_ms: None,
summary: None,
error: None,
new_segment_id: None,
};
worker.persist_compaction_lifecycle(&lifecycle).unwrap();
worker.recover_unfinished_compaction().await.unwrap();
let (entries, _) = worker.sink.subscribe_with_snapshot();
let restored = entries.iter().rev().find_map(|entry| match entry {
LogEntry::Extension {
domain, payload, ..
} if domain == COMPACTION_EXTENSION_DOMAIN => {
serde_json::from_value::<CompactionLifecycle>(payload.clone()).ok()
}
_ => None,
});
let restored = restored.expect("terminal compaction lifecycle");
assert_eq!(restored.state, CompactionLifecycleState::Interrupted);
assert_eq!(restored.revision, 2);
assert!(
restored
.error
.as_deref()
.is_some_and(|error| error.contains("restarted"))
);
let mut future = lifecycle;
future.schema_version = 4;
future.compaction_id = "future-compaction".into();
worker.persist_compaction_lifecycle(&future).unwrap();
let error = worker.recover_unfinished_compaction().await.unwrap_err();
assert!(
error
.to_string()
.contains("unsupported compaction lifecycle schema version 4")
);
}
fn minimal_manifest() -> WorkerManifest { fn minimal_manifest() -> WorkerManifest {
let toml_str = r#" let toml_str = r#"
[worker] [worker]
+202 -6
View File
@@ -72,6 +72,41 @@ impl LlmClient for MockClient {
} }
} }
#[derive(Clone)]
struct BlockingCompactClient {
calls: Arc<AtomicUsize>,
}
impl BlockingCompactClient {
fn new() -> Self {
Self {
calls: Arc::new(AtomicUsize::new(0)),
}
}
}
#[async_trait]
impl LlmClient for BlockingCompactClient {
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
async fn stream(
&self,
_request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + Send>>, ClientError>
{
let call = self.calls.fetch_add(1, Ordering::SeqCst);
if call == 0 {
Ok(Box::pin(futures::stream::iter(
single_text_events("seed").into_iter().map(Ok),
)))
} else {
Ok(Box::pin(futures::stream::pending()))
}
}
}
fn single_text_events(text: &str) -> Vec<LlmEvent> { fn single_text_events(text: &str) -> Vec<LlmEvent> {
vec![ vec![
LlmEvent::text_block_start(0), LlmEvent::text_block_start(0),
@@ -156,10 +191,10 @@ target = "./"
permission = "write" permission = "write"
"#; "#;
async fn make_worker_with_manifest( async fn make_worker_with_manifest<C>(manifest_toml: &str, client: C) -> Worker<C, TestStore>
manifest_toml: &str, where
client: MockClient, C: LlmClient + Clone + Send + Sync + 'static,
) -> Worker<MockClient, TestStore> { {
let manifest = worker::WorkerManifest::from_toml(manifest_toml).unwrap(); let manifest = worker::WorkerManifest::from_toml(manifest_toml).unwrap();
let store_tmp = tempfile::tempdir().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] #[tokio::test]
async fn controller_compact_method_emits_start_and_done() { async fn controller_compact_method_emits_start_and_done() {
let client = MockClient::new(vec![ let client = MockClient::new(vec![
text_events_with_usage("hi", 1000), text_events_with_usage("hi", 1000),
write_summary_tool_use_events("manual-summary", "manual compact summary"), write_summary_tool_use_events("manual-summary", "manual compact summary"),
single_text_events("done"), single_text_events("done"),
single_text_events("follow-up"),
]); ]);
let worker = make_worker_with_manifest(POST_RUN_MANIFEST_TOML, client).await; let worker = make_worker_with_manifest(POST_RUN_MANIFEST_TOML, client).await;
let runtime_tmp = tempfile::tempdir().unwrap(); 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; let mut saw_start = false;
loop { loop {
match tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) 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"); 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;
} }
+152 -48
View File
@@ -1,5 +1,5 @@
use std::pin::Pin; use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use agen::Engine; use agen::Engine;
@@ -25,6 +25,15 @@ use worker::{
type TestStore = CombinedStore<FsStore, FsWorkerStore>; type TestStore = CombinedStore<FsStore, FsWorkerStore>;
static NEXT_COMMAND_ID: AtomicU64 = AtomicU64::new(1);
fn worker_command(handle: &WorkerHandle) -> protocol::WorkerCommandEnvelope {
protocol::WorkerCommandEnvelope::for_snapshot(
NEXT_COMMAND_ID.fetch_add(1, Ordering::Relaxed),
&handle.shared_state.snapshot(),
)
}
/// Reconstruct a worker-history-like `Vec<Item>` from the live session /// Reconstruct a worker-history-like `Vec<Item>` from the live session
/// log mirror held by the Worker's broadcast sink. Replaces the previous /// log mirror held by the Worker's broadcast sink. Replaces the previous
/// `WorkerSharedState.history()` test helper now that the mirror lives in /// `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()); 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(); shutdown_rx.await.unwrap();
} }
@@ -345,7 +359,12 @@ async fn shutdown_closes_bound_workdir_session() {
WorkerController::spawn(worker, runtime_base.path(), &bash_output_dir) WorkerController::spawn(worker, runtime_base.path(), &bash_output_dir)
.await .await
.unwrap(); .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) tokio::time::timeout(std::time::Duration::from_secs(5), shutdown_rx)
.await .await
.expect("controller should shut down") .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"), !durable_history.contains("ready") && !durable_history.contains("done"),
"operational command chunks must not be appended to Worker history: {durable_history}" "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] #[tokio::test]
@@ -530,7 +554,12 @@ async fn controller_refreshes_command_snapshot_after_high_output_provider_lag()
.await .await
.unwrap(); .unwrap();
assert_eq!(output.status, workdir::CommandStatus::Cancelled); 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] #[tokio::test]
@@ -571,13 +600,13 @@ async fn controller_startup_failure_closes_bound_workdir_session() {
async fn wait_for_status(handle: &WorkerHandle, status: WorkerStatus) { async fn wait_for_status(handle: &WorkerHandle, status: WorkerStatus) {
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2); let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop { loop {
if handle.shared_state.get_status() == status { if handle.shared_state.catalog_status() == status {
return; return;
} }
assert!( assert!(
tokio::time::Instant::now() < deadline, tokio::time::Instant::now() < deadline,
"timed out waiting for status {status:?}; current={:?}", "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; 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 }) => { Ok(Event::RunEnd { result: protocol::RunResult::Finished }) => {
saw_run_end = true; 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; saw_idle_status = true;
break; break;
} }
@@ -1046,7 +1076,7 @@ async fn run_end_returns_to_idle_without_busy_status() {
saw_idle_status, saw_idle_status,
"expected idle status immediately after RunEnd" "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] #[tokio::test]
@@ -1124,9 +1154,7 @@ async fn snapshot_includes_user_input_for_in_flight_turn() {
loop { loop {
if matches!( if matches!(
events.recv().await, events.recv().await,
Ok(Event::Status { Ok(Event::WorkerState { snapshot }) if snapshot.catalog_status() == WorkerStatus::Running
status: WorkerStatus::Running,
})
) { ) {
break; break;
} }
@@ -1201,8 +1229,8 @@ async fn attach_snapshot_includes_current_status() {
loop { loop {
let event = reader.next::<Event>().await.unwrap().unwrap(); let event = reader.next::<Event>().await.unwrap().unwrap();
match event { match event {
Event::Snapshot { status, .. } => { Event::Snapshot { state, .. } => {
assert_eq!(status, WorkerStatus::Running); assert_eq!(state.catalog_status(), WorkerStatus::Running);
return; return;
} }
Event::Alert(_) => continue, Event::Alert(_) => continue,
@@ -1217,7 +1245,7 @@ async fn shared_state_starts_idle() {
let worker = make_worker(client).await; let worker = make_worker(client).await;
let handle = spawn_controller(worker).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] #[tokio::test]
@@ -1237,7 +1265,7 @@ async fn run_updates_shared_state_to_idle_after_completion() {
// Wait for the run to complete // Wait for the run to complete
tokio::time::sleep(std::time::Duration::from_millis(100)).await; 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] #[tokio::test]
@@ -1360,7 +1388,12 @@ async fn submit_while_running_is_durably_queued() {
assert_eq!(accepted, Some(protocol::SubmissionDisposition::Queued)); assert_eq!(accepted, Some(protocol::SubmissionDisposition::Queued));
let pending_snapshot = pending_snapshot.expect("pending snapshot"); let pending_snapshot = pending_snapshot.expect("pending snapshot");
assert_eq!(pending_snapshot.submissions.len(), 1); 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; wait_for_status(&handle, WorkerStatus::Paused).await;
handle handle
.send(Method::ContinuePending { .send(Method::ContinuePending {
@@ -1382,17 +1415,22 @@ async fn submit_while_running_is_durably_queued() {
.await .await
.expect("paused ContinuePending rejection"); .expect("paused ContinuePending rejection");
assert!(rejection.contains("Resume or Cancel")); 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] #[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 client = MockClient::new(simple_text_events());
let worker = make_worker(client).await; let worker = make_worker(client).await;
let handle = spawn_controller(worker).await; let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe(); 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 mut saw_not_paused = false;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(1); 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! { tokio::select! {
event = rx.recv() => { event = rx.recv() => {
match event { 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; saw_not_paused = true;
break; 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] #[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 client = MockClient::new(simple_text_events());
let worker = make_worker(client).await; let worker = make_worker(client).await;
let handle = spawn_controller(worker).await; let handle = spawn_controller(worker).await;
let mut rx = handle.subscribe(); 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 mut saw_not_running = false;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(1); 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! { tokio::select! {
event = rx.recv() => { event = rx.recv() => {
match event { 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; saw_not_running = true;
break; 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] #[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; 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!( assert!(
client_for_assert.captured_requests().is_empty(), client_for_assert.captured_requests().is_empty(),
"weak Notify must not stage RunForNotification while idle" "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, saw_worker_event_in_mirror,
"Method::WorkerEvent should commit a SystemItem::WorkerEvent entry" "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(); let requests = client_for_assert.captured_requests();
assert_eq!( 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; tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert_eq!( assert_eq!(
handle.shared_state.get_status(), handle.shared_state.catalog_status(),
WorkerStatus::Idle, WorkerStatus::Idle,
"control-plane ScopeSubDelegated must not auto-start the parent LLM" "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 .await
.unwrap(); .unwrap();
} }
handle.send(Method::Cancel).await.unwrap(); handle
.send(Method::Cancel {
command: worker_command(&handle),
})
.await
.unwrap();
wait_for_status(&handle, WorkerStatus::Idle).await; wait_for_status(&handle, WorkerStatus::Idle).await;
let mut rx = handle.subscribe(); 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" "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 // The controller emits RunEnd { Paused } when the
// EngineError::Cancelled is translated under pause_requested. // 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; 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!( assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!( 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; 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 // History consistency: exactly [user "hello", assistant
// "resumed output"]. No artifacts from the aborted stream // "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" "tool_call_done should arrive before pause"
); );
handle.send(Method::Pause).await.unwrap(); handle
.send(Method::Pause {
command: worker_command(&handle),
})
.await
.unwrap();
assert!( assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!( drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
e, e,
@@ -2622,7 +2691,7 @@ async fn paused_then_run_closes_orphan_tool_use_for_next_request() {
"expected RunEnd::Paused" "expected RunEnd::Paused"
); );
tokio::time::sleep(std::time::Duration::from_millis(50)).await; 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 // New user input while Paused → `Worker::run` observes
// `last_run_interrupted` and runs its interrupt-prep step, which // `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" "tool_call_done should arrive before pause"
); );
handle.send(Method::Pause).await.unwrap(); handle
.send(Method::Pause {
command: worker_command(&handle),
})
.await
.unwrap();
assert!( assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!( drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
e, e,
@@ -2794,7 +2868,12 @@ async fn paused_cancel_abandons_resume_and_next_input_is_fresh_run() {
); );
wait_for_status(&handle, WorkerStatus::Paused).await; 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; wait_for_status(&handle, WorkerStatus::Idle).await;
let (entries_after_cancel, _rx_after_cancel) = handle.sink.subscribe_with_snapshot(); let (entries_after_cancel, _rx_after_cancel) = handle.sink.subscribe_with_snapshot();
assert!( 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" "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!( assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!( drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
e, e,
Event::Error { Event::CommandAcknowledged { acknowledgement }
code: worker::ErrorCode::NotPaused, if acknowledgement.command == protocol::WorkerCommandKind::Resume
.. && acknowledgement.disposition
} == protocol::WorkerCommandDisposition::InvalidState
)) ))
.await, .await,
"resume after paused cancel should be rejected as not paused" "resume after paused cancel should receive invalid-state acknowledgement"
); );
assert_eq!( assert_eq!(
client_for_assert.captured_requests().len(), client_for_assert.captured_requests().len(),
@@ -2939,7 +3023,12 @@ async fn empty_turn_cancel_rolls_back_submit_entries_and_emits_signal() {
.await .await
.unwrap(); .unwrap();
wait_for_status(&handle, WorkerStatus::Running).await; wait_for_status(&handle, WorkerStatus::Running).await;
handle.send(Method::Cancel).await.unwrap(); handle
.send(Method::Cancel {
command: worker_command(&handle),
})
.await
.unwrap();
assert!( assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!( 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 .await
.unwrap(); .unwrap();
wait_for_status(&handle, WorkerStatus::Running).await; wait_for_status(&handle, WorkerStatus::Running).await;
handle.send(Method::Pause).await.unwrap(); handle
.send(Method::Pause {
command: worker_command(&handle),
})
.await
.unwrap();
assert!( assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!( 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 .await
.unwrap(); .unwrap();
wait_for_status(&handle, WorkerStatus::Running).await; wait_for_status(&handle, WorkerStatus::Running).await;
handle.send(Method::Cancel).await.unwrap(); handle
.send(Method::Cancel {
command: worker_command(&handle),
})
.await
.unwrap();
assert!( assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!( drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
e, e,
@@ -3091,7 +3190,12 @@ async fn pause_after_assistant_token_does_not_rollback() {
.await, .await,
"assistant token should be visible before pause" "assistant token should be visible before pause"
); );
handle.send(Method::Pause).await.unwrap(); handle
.send(Method::Pause {
command: worker_command(&handle),
})
.await
.unwrap();
assert!( assert!(
drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!( drain_until(&mut rx, std::time::Duration::from_secs(2), |e| matches!(
+4 -9
View File
@@ -36,8 +36,6 @@ use worker_runtime::config_bundle::{
ConfigBundleMetadata, ConfigBundleProvenance, ConfigProfileDescriptor, ConfigBundleMetadata, ConfigBundleProvenance, ConfigProfileDescriptor,
}; };
use worker_runtime::error::RuntimeError as EmbeddedRuntimeError; use worker_runtime::error::RuntimeError as EmbeddedRuntimeError;
#[cfg(test)]
use worker_runtime::execution::WorkerExecutionRunState;
use worker_runtime::fs_store::FsRuntimeStoreOptions; use worker_runtime::fs_store::FsRuntimeStoreOptions;
use worker_runtime::http_server::{ use worker_runtime::http_server::{
RUNTIME_PING_PERMISSION, RUNTIME_WORKSPACE_SCOPE_HEADER, RUNTIME_PING_PERMISSION, RUNTIME_WORKSPACE_SCOPE_HEADER,
@@ -5170,7 +5168,6 @@ mod tests {
request.worker_ref, request.worker_ref,
self.backend_id(), self.backend_id(),
), ),
run_state: WorkerExecutionRunState::Idle,
working_directory: request working_directory: request
.working_directory .working_directory
.as_ref() .as_ref()
@@ -5199,8 +5196,8 @@ mod tests {
let content = input.content; let content = input.content;
std::thread::spawn(move || { std::thread::spawn(move || {
std::thread::sleep(std::time::Duration::from_millis(10)); std::thread::sleep(std::time::Duration::from_millis(10));
let _ = context.publish_protocol_event(protocol::Event::Status { let _ = context.publish_protocol_event(protocol::Event::WorkerState {
status: protocol::WorkerStatus::Running, snapshot: protocol::WorkerStatus::Running.into(),
}); });
let _ = context.publish_protocol_event(protocol::Event::TextDone { let _ = context.publish_protocol_event(protocol::Event::TextDone {
text: format!("echo: {content}"), text: format!("echo: {content}"),
@@ -5208,14 +5205,13 @@ mod tests {
let _ = context.publish_protocol_event(protocol::Event::RunEnd { let _ = context.publish_protocol_event(protocol::Event::RunEnd {
result: protocol::RunResult::Finished, result: protocol::RunResult::Finished,
}); });
let _ = context.publish_protocol_event(protocol::Event::Status { let _ = context.publish_protocol_event(protocol::Event::WorkerState {
status: protocol::WorkerStatus::Idle, snapshot: protocol::WorkerStatus::Idle.into(),
}); });
}); });
if let Some(submission_request_id) = submission_request_id { if let Some(submission_request_id) = submission_request_id {
worker_runtime::execution::WorkerExecutionResult::accepted_submission( worker_runtime::execution::WorkerExecutionResult::accepted_submission(
worker_runtime::execution::WorkerExecutionOperation::Input, worker_runtime::execution::WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
submission_request_id, submission_request_id,
uuid::Uuid::now_v7().to_string(), uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started, protocol::SubmissionDisposition::Started,
@@ -5223,7 +5219,6 @@ mod tests {
} else { } else {
worker_runtime::execution::WorkerExecutionResult::accepted( worker_runtime::execution::WorkerExecutionResult::accepted(
worker_runtime::execution::WorkerExecutionOperation::Input, worker_runtime::execution::WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
) )
} }
} }
@@ -6,7 +6,7 @@ use worker_runtime::catalog::{
}; };
use worker_runtime::execution::{ use worker_runtime::execution::{
WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, WorkerExecutionResult, WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, WorkerExecutionResult,
WorkerExecutionRunState, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult,
}; };
use worker_runtime::identity::WorkerId; use worker_runtime::identity::WorkerId;
use worker_runtime::profile_archive::{ProfileSourceArchiveRef, ProfileSourceGraphSummary}; use worker_runtime::profile_archive::{ProfileSourceArchiveRef, ProfileSourceGraphSummary};
@@ -22,7 +22,6 @@ impl WorkerExecutionBackend for TestExecutionBackend {
fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult {
WorkerExecutionSpawnResult::connected( WorkerExecutionSpawnResult::connected(
WorkerExecutionHandle::new(request.worker_ref, self.backend_id()), WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
WorkerExecutionRunState::Idle,
None, None,
) )
} }
@@ -35,24 +34,17 @@ impl WorkerExecutionBackend for TestExecutionBackend {
if let Some(submission_request_id) = input.submission_request_id { if let Some(submission_request_id) = input.submission_request_id {
WorkerExecutionResult::accepted_submission( WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input, WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
submission_request_id, submission_request_id,
uuid::Uuid::now_v7().to_string(), uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started, protocol::SubmissionDisposition::Started,
) )
} else { } else {
WorkerExecutionResult::accepted( WorkerExecutionResult::accepted(WorkerExecutionOperation::Input)
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
)
} }
} }
fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
WorkerExecutionResult::accepted( WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop)
WorkerExecutionOperation::Stop,
WorkerExecutionRunState::Stopped,
)
} }
} }
@@ -199,8 +191,8 @@ async fn equal_downstream_selectors_share_one_upstream_subscription() {
runtime runtime
.observe_worker_event( .observe_worker_event(
&worker.worker_ref, &worker.worker_ref,
protocol::Event::Status { protocol::Event::WorkerState {
status: protocol::WorkerStatus::Running, snapshot: protocol::WorkerStatus::Running.into(),
}, },
) )
.unwrap(); .unwrap();
@@ -339,8 +331,8 @@ async fn embedded_runtime_uses_in_process_subscription_source() {
runtime runtime
.observe_worker_event( .observe_worker_event(
&worker.worker_ref, &worker.worker_ref,
protocol::Event::Status { protocol::Event::WorkerState {
status: protocol::WorkerStatus::Running, snapshot: protocol::WorkerStatus::Running.into(),
}, },
) )
.unwrap(); .unwrap();
+11 -9
View File
@@ -18559,7 +18559,6 @@ mod tests {
request.worker_ref, request.worker_ref,
self.backend_id(), self.backend_id(),
), ),
run_state: worker_runtime::execution::WorkerExecutionRunState::Idle,
working_directory, working_directory,
} }
} }
@@ -18575,7 +18574,6 @@ mod tests {
.push((handle.worker_ref().clone(), method)); .push((handle.worker_ref().clone(), method));
worker_runtime::execution::WorkerExecutionResult::accepted( worker_runtime::execution::WorkerExecutionResult::accepted(
worker_runtime::execution::WorkerExecutionOperation::ProtocolMethod, worker_runtime::execution::WorkerExecutionOperation::ProtocolMethod,
worker_runtime::execution::WorkerExecutionRunState::Idle,
) )
} }
@@ -18585,7 +18583,6 @@ mod tests {
) -> worker_runtime::execution::WorkerExecutionResult { ) -> worker_runtime::execution::WorkerExecutionResult {
worker_runtime::execution::WorkerExecutionResult::accepted( worker_runtime::execution::WorkerExecutionResult::accepted(
worker_runtime::execution::WorkerExecutionOperation::Stop, worker_runtime::execution::WorkerExecutionOperation::Stop,
worker_runtime::execution::WorkerExecutionRunState::Stopped,
) )
} }
@@ -18595,7 +18592,6 @@ mod tests {
) -> worker_runtime::execution::WorkerExecutionResult { ) -> worker_runtime::execution::WorkerExecutionResult {
worker_runtime::execution::WorkerExecutionResult::accepted( worker_runtime::execution::WorkerExecutionResult::accepted(
worker_runtime::execution::WorkerExecutionOperation::Cancel, 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 { if let Some(submission_request_id) = submission_request_id {
worker_runtime::execution::WorkerExecutionResult::accepted_submission( worker_runtime::execution::WorkerExecutionResult::accepted_submission(
worker_runtime::execution::WorkerExecutionOperation::Input, worker_runtime::execution::WorkerExecutionOperation::Input,
worker_runtime::execution::WorkerExecutionRunState::Idle,
submission_request_id, submission_request_id,
uuid::Uuid::now_v7().to_string(), uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started, protocol::SubmissionDisposition::Started,
) )
.with_worker_state(protocol::WorkerStateSnapshot::initial(1))
} else { } else {
worker_runtime::execution::WorkerExecutionResult::accepted( worker_runtime::execution::WorkerExecutionResult::accepted(
worker_runtime::execution::WorkerExecutionOperation::Input, 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::SubscriptionFramePayload::WorkerProtocol(
protocol::subscription::SubscriptionWorkerProtocolMethod { protocol::subscription::SubscriptionWorkerProtocolMethod {
subscription_id: second_protocol_subscription_id, 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() .iter()
.any(|(worker_ref, method)| { .any(|(worker_ref, method)| {
worker_ref.worker_id.to_string() == worker_id worker_ref.worker_id.to_string() == worker_id
&& matches!(method, protocol::Method::Resume) && matches!(method, protocol::Method::Resume { .. })
}) })
{ {
break; break;
@@ -27527,7 +27529,7 @@ mod tests {
let protocol_methods = execution_backend.protocol_methods(); let protocol_methods = execution_backend.protocol_methods();
assert!(protocol_methods.iter().any(|(worker_ref, method)| { assert!(protocol_methods.iter().any(|(worker_ref, method)| {
worker_ref.worker_id.to_string() == worker_id worker_ref.worker_id.to_string() == worker_id
&& matches!(method, protocol::Method::Resume) && matches!(method, protocol::Method::Resume { .. })
})); }));
server.abort(); server.abort();
} }
+39 -3
View File
@@ -10,6 +10,37 @@ export type CompletionKind = "file";
export type WorkerStatus = "idle" | "running" | "paused" | "stopped"; 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 TurnResult = "finished" | "paused";
export type InvokeKind = "user_send" | "notify" | "worker_event" | "system_reminder" | "wakeup"; 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 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<Segment>, } } | { "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<Segment>, } } | { "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<Segment>, } } | { "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, 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<Segment>, } } | { "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 * Full tool output. Absent when the tool chose to return
* summary-only, or when the result was pruned. * 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 * Unfinished model output that has already streamed in the current
* run but is not yet represented by committed snapshot entries. * 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. * Parent-owned Internal Worker sessions visible to this client.
* Service-private Internal Workers are deliberately excluded. * Service-private Internal Workers are deliberately excluded.
*/ */
internal_workers?: Array<InternalWorkerSnapshot>, } } | { "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<CompletionEntry>, } } | { "event": "rewind_targets", "data": { head_entries: number, targets: Array<RewindTarget>, } } | { "event": "rewind_applied", "data": { session: SessionSnapshot, input: Array<Segment>, 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<InternalWorkerSnapshot>, } } | { "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<CompletionEntry>, } } | { "event": "rewind_targets", "data": { head_entries: number, targets: Array<RewindTarget>, } } | { "event": "rewind_applied", "data": { session: SessionSnapshot, input: Array<Segment>, 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" };
@@ -1,4 +1,4 @@
import type { Event } from "$lib/generated/protocol"; import type { Event, WorkerStateSnapshot, WorkerStatus } from "$lib/generated/protocol";
import { import {
type ConsoleEventInput, type ConsoleEventInput,
type ConsoleLine, type ConsoleLine,
@@ -19,6 +19,23 @@ declare const Deno: {
test(name: string, fn: () => void): void; 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 { function assert(condition: unknown, message: string): asserts condition {
if (!condition) { if (!condition) {
throw new Error(message); throw new Error(message);
@@ -131,7 +148,7 @@ function snapshotEvent(cwd: string, entries: unknown[] = []): Event {
context_window: 100, context_window: 100,
context_tokens: 20, context_tokens: 20,
}, },
status: "idle", state: workerState("idle"),
in_flight: { blocks: [] }, in_flight: { blocks: [] },
}, },
}; };
@@ -213,7 +230,7 @@ Deno.test("snapshot replaces a live error with one durable run_errored row", ()
}, },
{ {
eventId: "idle-after-error", 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", () => { Deno.test("snapshot restores bounded in-flight Bash command output", () => {
const snapshot = snapshotEvent("/repo"); const snapshot = snapshotEvent("/repo");
if (snapshot.event !== "snapshot") throw new Error("snapshot fixture expected"); if (snapshot.event !== "snapshot") throw new Error("snapshot fixture expected");
snapshot.data.status = "running"; snapshot.data.state = workerState("running");
snapshot.data.in_flight = { snapshot.data.in_flight = {
blocks: [{ blocks: [{
kind: "tool_call", kind: "tool_call",
@@ -1403,7 +1420,7 @@ Deno.test("projectConsole hides lifecycle events and renders system items", () =
const projection = projectConsole([ const projection = projectConsole([
{ {
eventId: "30", eventId: "30",
event: { event: "status", data: { status: "running" } } satisfies Event, event: { event: "worker_state", data: { snapshot: workerState("running") } } satisfies Event,
}, },
{ {
eventId: "31", eventId: "31",
@@ -1527,7 +1544,7 @@ Deno.test("projectConsole renders snapshot entries and in-flight output", () =>
context_window: 100, context_window: 100,
context_tokens: 20, context_tokens: 20,
}, },
status: "running", state: workerState("running"),
in_flight: { in_flight: {
blocks: [ blocks: [
{ kind: "text", text: "partial" }, { kind: "text", text: "partial" },
@@ -1578,7 +1595,7 @@ Deno.test("projectConsole restores system items from snapshot entries", () => {
context_window: 100, context_window: 100,
context_tokens: 20, context_tokens: 20,
}, },
status: "idle", state: workerState("idle"),
}, },
} satisfies Event, } satisfies Event,
}]); }]);
@@ -1922,7 +1939,7 @@ Deno.test("console Worker views expose only direct Internal Workers", () => {
kind: "sub_worker", kind: "sub_worker",
}, },
revision: 1, 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", kind: "sub_worker",
}, },
revision: 1, 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", kind: "sub_worker",
}, },
revision: 1, revision: 1,
event: { event: "status", data: { status: "running" } }, event: { event: "worker_state", data: { snapshot: workerState("running") } },
}, },
}, },
}]); }]);
@@ -10,6 +10,8 @@ import type {
InternalWorkerRef, InternalWorkerRef,
InternalWorkerSnapshot, InternalWorkerSnapshot,
Segment, Segment,
WorkerStateSnapshot,
WorkerStatus,
} from "$lib/generated/protocol"; } from "$lib/generated/protocol";
import { stringify as stringifyYaml } from "yaml"; import { stringify as stringifyYaml } from "yaml";
import { workspaceRoute } from "$lib/workspace/api/http"; import { workspaceRoute } from "$lib/workspace/api/http";
@@ -169,6 +171,7 @@ export type ConsoleProjection = {
tasks: ConsoleTask[]; tasks: ConsoleTask[];
taskNextId: number; taskNextId: number;
status: string | null; status: string | null;
workerState: WorkerStateSnapshot | null;
usage: string | null; usage: string | null;
runActivity: RunActivityStats; runActivity: RunActivityStats;
cwd: string | null; cwd: string | null;
@@ -251,12 +254,22 @@ export function isConsoleProjectionEvent(event: ProtocolEvent): boolean {
return event.event !== "completions"; 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 { export function emptyConsoleProjection(): ConsoleProjection {
return { return {
lines: [], lines: [],
tasks: [], tasks: [],
taskNextId: 1, taskNextId: 1,
status: null, status: null,
workerState: null,
usage: null, usage: null,
runActivity: emptyRunActivityStats(), runActivity: emptyRunActivityStats(),
cwd: null, cwd: null,
@@ -793,6 +806,7 @@ export function applyProtocolEvent(
tasks: [...projection.tasks], tasks: [...projection.tasks],
taskNextId: projection.taskNextId, taskNextId: projection.taskNextId,
status: projection.status, status: projection.status,
workerState: projection.workerState,
usage: projection.usage, usage: projection.usage,
runActivity: applyRunActivityEvent( runActivity: applyRunActivityEvent(
projection.runActivity, projection.runActivity,
@@ -903,7 +917,8 @@ export function applyProtocolEvent(
); );
break; break;
case "snapshot": { case "snapshot": {
next.status = event.data.status; next.workerState = event.data.state;
next.status = workerStatusFromState(event.data.state);
next.cwd = event.data.greeting.cwd; next.cwd = event.data.greeting.cwd;
const snapshot = snapshotProjectionFromSession( const snapshot = snapshotProjectionFromSession(
envelope.eventId, envelope.eventId,
@@ -1000,8 +1015,13 @@ export function applyProtocolEvent(
if (existingIndex >= 0) next.internalWorkers.splice(existingIndex, 1); if (existingIndex >= 0) next.internalWorkers.splice(existingIndex, 1);
break; break;
} }
case "status": case "worker_state":
next.status = event.data.status; 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; break;
case "command": case "command":
applyCommandEvent(next, envelope.eventId, event.data.event); applyCommandEvent(next, envelope.eventId, event.data.event);
@@ -1939,6 +1959,7 @@ function snapshotProjectionFromSession(
tasks: [], tasks: [],
taskNextId: 1, taskNextId: 1,
status: null, status: null,
workerState: null,
usage: null, usage: null,
runActivity: emptyRunActivityStats(), runActivity: emptyRunActivityStats(),
cwd, cwd,
@@ -75,7 +75,12 @@ Deno.test("new invoke and running snapshot reset run activity", () => {
data: { data: {
entries: [], entries: [],
greeting: { text: "", profile: "" }, greeting: { text: "", profile: "" },
status: "idle", state: {
execution_generation: 1,
revision: 0,
last_command_id: 0,
state: { kind: "idle" },
},
in_flight: {}, in_flight: {},
internal_workers: [], internal_workers: [],
}, },
@@ -25,7 +25,9 @@ export function applyRunActivityEvent(
case "invoke_start": case "invoke_start":
return { ...emptyRunActivityStats(), startedAtMs: observedAtMs }; return { ...emptyRunActivityStats(), startedAtMs: observedAtMs };
case "snapshot": 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(), startedAtMs: observedAtMs }
: emptyRunActivityStats(); : emptyRunActivityStats();
case "turn_start": case "turn_start":
@@ -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") { function sendWorkerControl(command: "pause" | "cancel" | "resume") {
const label = command[0].toUpperCase() + command.slice(1); 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 { function isEditableTarget(target: EventTarget | null): boolean {
@@ -627,8 +660,11 @@
auto_run: true, auto_run: true,
}, },
}; };
case "compact": case "compact": {
return { method: "compact" }; const method = lifecycleMethod("compact");
if (!method) throw new Error("Worker state snapshot is not available");
return method;
}
case "list_rewind_targets": case "list_rewind_targets":
return { method: "list_rewind_targets" }; return { method: "list_rewind_targets" };
case "register_peer": case "register_peer":
@@ -691,7 +727,7 @@
function handleComposerSubmit() { function handleComposerSubmit() {
if (workerRunning) { if (workerRunning) {
sendControl({ method: "cancel" }, "Stop"); sendWorkerControl("cancel");
return; return;
} }
void submitDraft(composerInputElement?.snapshot() ?? draft); void submitDraft(composerInputElement?.snapshot() ?? draft);
@@ -894,8 +930,26 @@
): string | null { ): string | null {
switch (event.event) { switch (event.event) {
case "snapshot": case "snapshot":
case "status": return event.data.state.state.kind === "idle"
return event.data.status; ? "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": case "shutdown":
return "shutdown"; return "shutdown";
default: default:
@@ -1620,7 +1674,10 @@
type="button" type="button"
class="secondary-button" class="secondary-button"
disabled={protocolState !== "open"} disabled={protocolState !== "open"}
onclick={() => sendControl({ method: "compact" }, "Compact")} onclick={() => {
const method = lifecycleMethod("compact");
if (method) sendControl(method, "Compact");
}}
> >
Compact Compact
</button> </button>