fix: use live worker observation snapshots

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<protocol::Event> {
+        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<protocol::Event> {
+        self.backend.worker_snapshot(handle)
+    }
 }
 
 impl fmt::Debug for WorkerExecutionBackendRef {

crates/worker-runtime/src/runtime.rs

@@ -783,8 +783,19 @@ impl Runtime {
         &self,
         worker_ref: &WorkerRef,
     ) -> Result<protocol::Event, RuntimeError> {
-        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<Option<WorkerExecutionSpawnResult>>,
         restore_count: Mutex<u64>,
         contexts: Mutex<BTreeMap<WorkerId, WorkerExecutionContext>>,
+        #[cfg(feature = "ws-server")]
+        snapshots: Mutex<BTreeMap<WorkerId, protocol::Event>>,
     }
 
     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<protocol::Event> {
+            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 {

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<protocol::Event> {
+        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)]