From 6ca0d48327316fab999fdc38b6cacd017980f1d1 Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 14 Jul 2026 02:59:52 +0900 Subject: [PATCH] fix: use live worker observation snapshots --- crates/worker-runtime/src/execution.rs | 13 ++++ crates/worker-runtime/src/runtime.rs | 79 ++++++++++++++++++++- crates/worker-runtime/src/worker_backend.rs | 11 +++ crates/worker/src/controller.rs | 29 +++++++- crates/worker/src/ipc/server.rs | 24 ++----- 5 files changed, 132 insertions(+), 24 deletions(-) diff --git a/crates/worker-runtime/src/execution.rs b/crates/worker-runtime/src/execution.rs index 872d2d0e..faa452aa 100644 --- a/crates/worker-runtime/src/execution.rs +++ b/crates/worker-runtime/src/execution.rs @@ -393,6 +393,11 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static { "execution backend does not support cancelling workers", ) } + + #[cfg(feature = "ws-server")] + fn worker_snapshot(&self, _handle: &WorkerExecutionHandle) -> Option { + None + } } #[derive(Clone)] @@ -472,6 +477,14 @@ impl WorkerExecutionBackendRef { pub(crate) fn cancel_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult { self.backend.cancel_worker(handle) } + + #[cfg(feature = "ws-server")] + pub(crate) fn worker_snapshot( + &self, + handle: &WorkerExecutionHandle, + ) -> Option { + self.backend.worker_snapshot(handle) + } } impl fmt::Debug for WorkerExecutionBackendRef { diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index f8e31dcd..f3866fd3 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -783,8 +783,19 @@ impl Runtime { &self, worker_ref: &WorkerRef, ) -> Result { - let state = self.lock()?; - let _worker = state.worker(worker_ref)?; + let (backend, handle) = { + let state = self.lock()?; + let worker = state.worker(worker_ref)?; + ( + state.execution_backend.clone(), + worker.execution_handle.clone(), + ) + }; + if let (Some(backend), Some(handle)) = (backend, handle) { + if let Some(snapshot) = backend.worker_snapshot(&handle) { + return Ok(snapshot); + } + } Ok(protocol::Event::Snapshot { entries: Vec::new(), greeting: protocol::Greeting { @@ -1826,6 +1837,8 @@ mod tests { restore_result: Mutex>, restore_count: Mutex, contexts: Mutex>, + #[cfg(feature = "ws-server")] + snapshots: Mutex>, } impl TestExecutionBackend { @@ -1833,6 +1846,14 @@ mod tests { *self.dispatch_result.lock().unwrap() = Some(result); } + #[cfg(feature = "ws-server")] + fn set_worker_snapshot(&self, worker_ref: &WorkerRef, snapshot: protocol::Event) { + self.snapshots + .lock() + .unwrap() + .insert(worker_ref.worker_id.clone(), snapshot); + } + #[cfg(feature = "ws-server")] fn publish_text_delta( &self, @@ -1917,6 +1938,15 @@ mod tests { WorkerExecutionRunState::Stopped, ) } + + #[cfg(feature = "ws-server")] + fn worker_snapshot(&self, handle: &WorkerExecutionHandle) -> Option { + self.snapshots + .lock() + .unwrap() + .get(&handle.worker_ref().worker_id) + .cloned() + } } fn runtime_with_backend() -> Runtime { @@ -2123,6 +2153,51 @@ mod tests { )); } + #[cfg(feature = "ws-server")] + #[test] + fn observation_snapshot_prefers_live_backend_snapshot() { + let (runtime, backend) = runtime_and_backend(); + let detail = runtime + .create_worker(task_request("observe snapshot")) + .unwrap(); + let expected_entry = serde_json::json!({"kind": "restored-log-entry"}); + backend.set_worker_snapshot( + &detail.worker_ref, + protocol::Event::Snapshot { + entries: vec![expected_entry.clone()], + greeting: protocol::Greeting { + worker_name: "live-worker".to_string(), + cwd: "/tmp/live".to_string(), + provider: "test-provider".to_string(), + model: "test-model".to_string(), + scope_summary: "live snapshot".to_string(), + tools: Vec::new(), + context_window: 128, + context_tokens: 64, + }, + status: protocol::WorkerStatus::Running, + in_flight: protocol::InFlightSnapshot { blocks: Vec::new() }, + }, + ); + + let snapshot = runtime + .worker_observation_snapshot(&detail.worker_ref) + .unwrap(); + match snapshot { + protocol::Event::Snapshot { + entries, + greeting, + status, + .. + } => { + assert_eq!(entries, vec![expected_entry]); + assert_eq!(greeting.worker_name, "live-worker"); + assert_eq!(status, protocol::WorkerStatus::Running); + } + other => panic!("expected snapshot, got {other:?}"), + } + } + struct InputOnlyBackend; impl WorkerExecutionBackend for InputOnlyBackend { diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 9c33652c..2a5d08c4 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -1082,6 +1082,17 @@ where WorkerExecutionRunState::Idle, ) } + + #[cfg(feature = "ws-server")] + fn worker_snapshot(&self, handle: &WorkerExecutionHandle) -> Option { + if handle.backend_id() != self.backend_id() { + return None; + } + let workers = self.workers.lock().ok()?; + workers + .get(handle.worker_ref()) + .map(|execution| execution.handle.snapshot_event()) + } } #[cfg(test)] diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index b25ddb46..ad5cdae6 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -5,8 +5,8 @@ use std::sync::atomic::Ordering; use llm_engine::EngineError; use llm_engine::llm_client::client::LlmClient; use manifest::TicketFeatureAccessConfig; -use session_store::Store; use session_store::WorkerMetadataStore; +use session_store::{LogEntry, Store}; use ticket::LocalTicketBackend; use ticket::config::TicketConfig; use tokio::sync::{broadcast, mpsc, oneshot}; @@ -16,7 +16,7 @@ use crate::discovery::{ WorkerDiscovery, list_workers_tool, restore_worker_tool, send_to_peer_worker_tool, }; use crate::feature::FeatureRegistryBuilder; -use crate::in_flight::InFlightEvents; +use crate::in_flight::{InFlightEvents, snapshot_from_guard}; use crate::ipc::alerter::Alerter; use crate::ipc::notify_buffer::NotifyBuffer; use crate::ipc::server::SocketServer; @@ -66,6 +66,31 @@ impl WorkerHandle { self.event_tx.subscribe() } + pub fn snapshot_event(&self) -> Event { + self.snapshot_event_with_entry_subscription().0 + } + + pub(crate) fn snapshot_event_with_entry_subscription( + &self, + ) -> (Event, broadcast::Receiver) { + let (entries, entry_rx, in_flight) = { + let in_flight_guard = self.in_flight.snapshot_guard(); + let (entries, entry_rx) = self.sink.subscribe_with_snapshot(); + let in_flight = snapshot_from_guard(&in_flight_guard); + (entries, entry_rx, in_flight) + }; + let event = Event::Snapshot { + entries: entries + .into_iter() + .map(|entry| serde_json::to_value(entry).expect("log entry serializes")) + .collect(), + greeting: self.shared_state.greeting.clone(), + status: self.shared_state.get_status(), + in_flight, + }; + (event, entry_rx) + } + /// Broadcast an event to all listeners (including socket clients). pub fn send_event(&self, event: Event) -> Result> { self.event_tx.send(event) diff --git a/crates/worker/src/ipc/server.rs b/crates/worker/src/ipc/server.rs index dbcea450..d366f029 100644 --- a/crates/worker/src/ipc/server.rs +++ b/crates/worker/src/ipc/server.rs @@ -7,7 +7,6 @@ use tokio::net::UnixListener; use tokio::task::JoinHandle; use crate::controller::WorkerHandle; -use crate::in_flight::snapshot_from_guard; use protocol::{Event, Method}; /// Unix socket server for Worker Protocol. @@ -111,16 +110,10 @@ async fn handle_connection(stream: tokio::net::UnixStream, handle: WorkerHandle) // committed entry or as the still-present in-flight block. This lock // order matches `append_entry` (in-flight clear before sink publish) and // keeps the snapshot/live boundary gap-free. - let (entries_snapshot, mut entry_rx, alert_snapshot, mut rx, in_flight) = { - let in_flight_guard = handle.in_flight.snapshot_guard(); - let (entries_snapshot, entry_rx) = handle.sink.subscribe_with_snapshot(); - - // Atomically subscribe and snapshot buffered alerts so that warnings - // emitted before this client connected are replayed exactly once. - let (alert_snapshot, rx) = handle.alerter.subscribe_with_snapshot(); - let in_flight = snapshot_from_guard(&in_flight_guard); - (entries_snapshot, entry_rx, alert_snapshot, rx, in_flight) - }; + let (snapshot_event, mut entry_rx) = handle.snapshot_event_with_entry_subscription(); + // Atomically subscribe and snapshot buffered alerts so that warnings + // emitted before this client connected are replayed exactly once. + let (alert_snapshot, mut rx) = handle.alerter.subscribe_with_snapshot(); for alert in alert_snapshot { if writer.write(&Event::Alert(alert)).await.is_err() { return; @@ -129,15 +122,6 @@ async fn handle_connection(stream: tokio::net::UnixStream, handle: WorkerHandle) // Send the typed snapshot up front so late attachers can // reconstruct view state without an extra round trip. - let snapshot_event = Event::Snapshot { - entries: entries_snapshot - .into_iter() - .map(|e| serde_json::to_value(&e).expect("LogEntry is Serialize")) - .collect(), - greeting: handle.shared_state.greeting.clone(), - status: handle.shared_state.get_status(), - in_flight, - }; if writer.write(&snapshot_event).await.is_err() { return; }