Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

feat: add revisioned worker execution state

Browse Source
This commit is contained in:
Hare
2026-09-06 06:54:29 +09:00
parent 3ed1545c3c
commit e8b9adcde4
32 changed files with 2168 additions and 750 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/protocol/src/lib.rs
+222 -52
View File
@@ -85,6 +85,142 @@ impl AuthenticatedInputSource {
}
}
/// Immutable identity and revision fence for one state-changing Worker command.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
pub struct WorkerCommandEnvelope {
/// Caller-owned sequence. A controller accepts command ids in strictly
/// increasing order for one execution generation.
pub command_id: u64,
pub expected_execution_generation: u64,
pub expected_worker_state_revision: u64,
}
impl WorkerCommandEnvelope {
pub fn for_snapshot(command_id: u64, snapshot: &WorkerStateSnapshot) -> Self {
Self {
command_id,
expected_execution_generation: snapshot.execution_generation,
expected_worker_state_revision: snapshot.revision,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(rename_all = "snake_case")]
pub enum WorkerCommandKind {
Resume,
Cancel,
Pause,
Compact,
Shutdown,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(rename_all = "snake_case")]
pub enum WorkerCommandDisposition {
Accepted,
StaleExecutionGeneration,
StaleWorkerStateRevision,
StaleCommandId,
InvalidState,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
pub struct WorkerCommandAcknowledgement {
pub command_id: u64,
pub command: WorkerCommandKind,
pub disposition: WorkerCommandDisposition,
/// The complete authoritative state observed after command admission.
pub state: WorkerStateSnapshot,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(tag = "kind", content = "state", rename_all = "snake_case")]
pub enum WorkerState {
Idle,
Busy(WorkerBusyState),
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(tag = "kind", content = "state", rename_all = "snake_case")]
pub enum WorkerBusyState {
Run(WorkerRunState),
Maintenance(WorkerMaintenanceState),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(rename_all = "snake_case")]
pub enum WorkerRunState {
Running,
Pausing,
Paused,
Cancelling,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(rename_all = "snake_case")]
pub enum WorkerMaintenanceState {
Compacting,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
pub struct WorkerStateSnapshot {
pub execution_generation: u64,
pub revision: u64,
/// Highest lifecycle command id observed by this controller generation.
pub last_command_id: u64,
pub state: WorkerState,
}
impl WorkerStateSnapshot {
pub fn initial(execution_generation: u64) -> Self {
Self {
execution_generation,
revision: 0,
last_command_id: 0,
state: WorkerState::Idle,
}
}
/// Compatibility projection for Runtime catalog lifecycle. This value is
/// never command-admission authority and cannot produce `Stopped`.
pub fn catalog_status(&self) -> WorkerStatus {
match self.state {
WorkerState::Idle => WorkerStatus::Idle,
WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Paused)) => WorkerStatus::Paused,
WorkerState::Busy(WorkerBusyState::Run(_))
| WorkerState::Busy(WorkerBusyState::Maintenance(_)) => WorkerStatus::Running,
}
}
}
impl From<WorkerStatus> for WorkerStateSnapshot {
fn from(status: WorkerStatus) -> Self {
let state = match status {
WorkerStatus::Idle | WorkerStatus::Stopped => WorkerState::Idle,
WorkerStatus::Running => {
WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running))
}
WorkerStatus::Paused => WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Paused)),
};
Self {
execution_generation: 1,
revision: 0,
last_command_id: 0,
state,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[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);
}
#[test]
fn method_compact_roundtrip() {
let json = r#"{"method":"compact"}"#;
let method: Method = serde_json::from_str(json).unwrap();
assert!(matches!(method, Method::Compact));
let serialized = serde_json::to_string(&method).unwrap();
assert_eq!(serialized, json);
fn lifecycle_methods_roundtrip_with_fences() {
for method in [
Method::Pause {
command: WorkerCommandEnvelope {
command_id: 11,
expected_execution_generation: 4,
expected_worker_state_revision: 8,
},
},
Method::Compact {
command: WorkerCommandEnvelope {
command_id: 12,
expected_execution_generation: 4,
expected_worker_state_revision: 9,
},
},
] {
let json = serde_json::to_string(&method).unwrap();
let decoded: Method = serde_json::from_str(&json).unwrap();
match decoded {
Method::Pause { command } | Method::Compact { command } => {
assert_eq!(command.expected_execution_generation, 4);
assert!(command.command_id >= 11);
}
other => panic!("unexpected lifecycle method: {other:?}"),
}
}
}
#[test]
@@ -1902,7 +2063,7 @@ mod tests {
context_window: 200_000,
context_tokens: 42_000,
},
status: WorkerStatus::Paused,
state: WorkerStatus::Paused.into(),
in_flight: InFlightSnapshot::default(),
internal_workers: Vec::new(),
};
@@ -1919,12 +2080,13 @@ mod tests {
assert_eq!(parsed["data"]["greeting"]["tools"][0], "Read");
assert_eq!(parsed["data"]["greeting"]["context_window"], 200_000);
assert_eq!(parsed["data"]["greeting"]["context_tokens"], 42_000);
assert_eq!(parsed["data"]["status"], "paused");
assert_eq!(parsed["data"]["state"]["state"]["kind"], "busy");
assert_eq!(parsed["data"]["state"]["state"]["state"]["state"], "paused");
}
#[test]
fn event_snapshot_in_flight_roundtrip_and_default() {
let inbound = r#"{"event":"snapshot","data":{"session":{"entries":[]},"greeting":{"worker_name":"test","cwd":"/tmp","provider":"p","model":"m","scope_summary":"s","tools":[]},"status":"running"}}"#;
let inbound = r#"{"event":"snapshot","data":{"session":{"entries":[]},"greeting":{"worker_name":"test","cwd":"/tmp","provider":"p","model":"m","scope_summary":"s","tools":[]},"state":{"execution_generation":1,"revision":1,"last_command_id":0,"state":{"kind":"busy","state":{"kind":"run","state":"running"}}}}}"#;
let decoded: Event = serde_json::from_str(inbound).unwrap();
match decoded {
Event::Snapshot { in_flight, .. } => assert!(in_flight.is_empty()),
@@ -1946,7 +2108,7 @@ mod tests {
context_window: 0,
context_tokens: 0,
},
status: WorkerStatus::Running,
state: WorkerStatus::Running.into(),
in_flight: InFlightSnapshot {
blocks: vec![
InFlightBlock::Text {
@@ -2034,20 +2196,32 @@ mod tests {
}
#[test]
fn event_status_format() {
let event = Event::Status {
status: WorkerStatus::Running,
fn event_worker_state_format() {
let event = Event::WorkerState {
snapshot: WorkerStateSnapshot {
execution_generation: 7,
revision: 3,
last_command_id: 9,
state: WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running)),
},
};
let json = serde_json::to_string(&event).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(parsed["event"], "status");
assert_eq!(parsed["data"]["status"], "running");
assert_eq!(parsed["event"], "worker_state");
assert_eq!(parsed["data"]["snapshot"]["execution_generation"], 7);
assert_eq!(parsed["data"]["snapshot"]["revision"], 3);
assert_eq!(parsed["data"]["snapshot"]["state"]["kind"], "busy");
let decoded: Event = serde_json::from_str(&json).unwrap();
assert!(matches!(
decoded,
Event::Status {
status: WorkerStatus::Running
Event::WorkerState {
snapshot: WorkerStateSnapshot {
execution_generation: 7,
revision: 3,
state: WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running)),
..
}
}
));
}
@@ -2088,19 +2262,10 @@ mod tests {
}
#[test]
fn event_snapshot_without_status_defaults_to_idle() {
fn event_snapshot_without_worker_state_fails_closed() {
let json = r#"{"event":"snapshot","data":{"session":{"entries":[]},"greeting":{"worker_name":"test","cwd":"/tmp","provider":"anthropic","model":"claude","scope_summary":"","tools":[]}}}"#;
let decoded: Event = serde_json::from_str(json).unwrap();
match decoded {
Event::Snapshot {
status, greeting, ..
} => {
assert_eq!(status, WorkerStatus::Idle);
assert_eq!(greeting.context_window, 0);
assert_eq!(greeting.context_tokens, 0);
}
other => panic!("expected Snapshot, got {other:?}"),
}
let error = serde_json::from_str::<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();
crates/protocol/src/typescript.rs
+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);