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