fix: use live worker observation snapshots
This commit is contained in:
parent
ec951ee5c0
commit
6ca0d48327
|
|
@ -393,6 +393,11 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static {
|
||||||
"execution backend does not support cancelling workers",
|
"execution backend does not support cancelling workers",
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(feature = "ws-server")]
|
||||||
|
fn worker_snapshot(&self, _handle: &WorkerExecutionHandle) -> Option<protocol::Event> {
|
||||||
|
None
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
|
|
@ -472,6 +477,14 @@ impl WorkerExecutionBackendRef {
|
||||||
pub(crate) fn cancel_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
|
pub(crate) fn cancel_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
|
||||||
self.backend.cancel_worker(handle)
|
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 {
|
impl fmt::Debug for WorkerExecutionBackendRef {
|
||||||
|
|
|
||||||
|
|
@ -783,8 +783,19 @@ impl Runtime {
|
||||||
&self,
|
&self,
|
||||||
worker_ref: &WorkerRef,
|
worker_ref: &WorkerRef,
|
||||||
) -> Result<protocol::Event, RuntimeError> {
|
) -> Result<protocol::Event, RuntimeError> {
|
||||||
let state = self.lock()?;
|
let (backend, handle) = {
|
||||||
let _worker = state.worker(worker_ref)?;
|
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 {
|
Ok(protocol::Event::Snapshot {
|
||||||
entries: Vec::new(),
|
entries: Vec::new(),
|
||||||
greeting: protocol::Greeting {
|
greeting: protocol::Greeting {
|
||||||
|
|
@ -1826,6 +1837,8 @@ mod tests {
|
||||||
restore_result: Mutex<Option<WorkerExecutionSpawnResult>>,
|
restore_result: Mutex<Option<WorkerExecutionSpawnResult>>,
|
||||||
restore_count: Mutex<u64>,
|
restore_count: Mutex<u64>,
|
||||||
contexts: Mutex<BTreeMap<WorkerId, WorkerExecutionContext>>,
|
contexts: Mutex<BTreeMap<WorkerId, WorkerExecutionContext>>,
|
||||||
|
#[cfg(feature = "ws-server")]
|
||||||
|
snapshots: Mutex<BTreeMap<WorkerId, protocol::Event>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl TestExecutionBackend {
|
impl TestExecutionBackend {
|
||||||
|
|
@ -1833,6 +1846,14 @@ mod tests {
|
||||||
*self.dispatch_result.lock().unwrap() = Some(result);
|
*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")]
|
#[cfg(feature = "ws-server")]
|
||||||
fn publish_text_delta(
|
fn publish_text_delta(
|
||||||
&self,
|
&self,
|
||||||
|
|
@ -1917,6 +1938,15 @@ mod tests {
|
||||||
WorkerExecutionRunState::Stopped,
|
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 {
|
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;
|
struct InputOnlyBackend;
|
||||||
|
|
||||||
impl WorkerExecutionBackend for InputOnlyBackend {
|
impl WorkerExecutionBackend for InputOnlyBackend {
|
||||||
|
|
|
||||||
|
|
@ -1082,6 +1082,17 @@ where
|
||||||
WorkerExecutionRunState::Idle,
|
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)]
|
#[cfg(test)]
|
||||||
|
|
|
||||||
|
|
@ -5,8 +5,8 @@ use std::sync::atomic::Ordering;
|
||||||
use llm_engine::EngineError;
|
use llm_engine::EngineError;
|
||||||
use llm_engine::llm_client::client::LlmClient;
|
use llm_engine::llm_client::client::LlmClient;
|
||||||
use manifest::TicketFeatureAccessConfig;
|
use manifest::TicketFeatureAccessConfig;
|
||||||
use session_store::Store;
|
|
||||||
use session_store::WorkerMetadataStore;
|
use session_store::WorkerMetadataStore;
|
||||||
|
use session_store::{LogEntry, Store};
|
||||||
use ticket::LocalTicketBackend;
|
use ticket::LocalTicketBackend;
|
||||||
use ticket::config::TicketConfig;
|
use ticket::config::TicketConfig;
|
||||||
use tokio::sync::{broadcast, mpsc, oneshot};
|
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,
|
WorkerDiscovery, list_workers_tool, restore_worker_tool, send_to_peer_worker_tool,
|
||||||
};
|
};
|
||||||
use crate::feature::FeatureRegistryBuilder;
|
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::alerter::Alerter;
|
||||||
use crate::ipc::notify_buffer::NotifyBuffer;
|
use crate::ipc::notify_buffer::NotifyBuffer;
|
||||||
use crate::ipc::server::SocketServer;
|
use crate::ipc::server::SocketServer;
|
||||||
|
|
@ -66,6 +66,31 @@ impl WorkerHandle {
|
||||||
self.event_tx.subscribe()
|
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<LogEntry>) {
|
||||||
|
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).
|
/// Broadcast an event to all listeners (including socket clients).
|
||||||
pub fn send_event(&self, event: Event) -> Result<usize, broadcast::error::SendError<Event>> {
|
pub fn send_event(&self, event: Event) -> Result<usize, broadcast::error::SendError<Event>> {
|
||||||
self.event_tx.send(event)
|
self.event_tx.send(event)
|
||||||
|
|
|
||||||
|
|
@ -7,7 +7,6 @@ use tokio::net::UnixListener;
|
||||||
use tokio::task::JoinHandle;
|
use tokio::task::JoinHandle;
|
||||||
|
|
||||||
use crate::controller::WorkerHandle;
|
use crate::controller::WorkerHandle;
|
||||||
use crate::in_flight::snapshot_from_guard;
|
|
||||||
use protocol::{Event, Method};
|
use protocol::{Event, Method};
|
||||||
|
|
||||||
/// Unix socket server for Worker Protocol.
|
/// 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
|
// committed entry or as the still-present in-flight block. This lock
|
||||||
// order matches `append_entry` (in-flight clear before sink publish) and
|
// order matches `append_entry` (in-flight clear before sink publish) and
|
||||||
// keeps the snapshot/live boundary gap-free.
|
// keeps the snapshot/live boundary gap-free.
|
||||||
let (entries_snapshot, mut entry_rx, alert_snapshot, mut rx, in_flight) = {
|
let (snapshot_event, mut entry_rx) = handle.snapshot_event_with_entry_subscription();
|
||||||
let in_flight_guard = handle.in_flight.snapshot_guard();
|
// Atomically subscribe and snapshot buffered alerts so that warnings
|
||||||
let (entries_snapshot, entry_rx) = handle.sink.subscribe_with_snapshot();
|
// emitted before this client connected are replayed exactly once.
|
||||||
|
let (alert_snapshot, mut rx) = handle.alerter.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)
|
|
||||||
};
|
|
||||||
for alert in alert_snapshot {
|
for alert in alert_snapshot {
|
||||||
if writer.write(&Event::Alert(alert)).await.is_err() {
|
if writer.write(&Event::Alert(alert)).await.is_err() {
|
||||||
return;
|
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
|
// Send the typed snapshot up front so late attachers can
|
||||||
// reconstruct view state without an extra round trip.
|
// 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() {
|
if writer.write(&snapshot_event).await.is_err() {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue
Block a user