refactor: rename pod crate to worker
This commit is contained in:
@@ -0,0 +1,168 @@
|
||||
//! User-facing alert channel for Worker → client.
|
||||
//!
|
||||
//! Separate from `tracing` (which is for developer logs). Alerts
|
||||
//! are short human-readable messages the Worker layer wants a client to
|
||||
//! see — for example "compaction failed", "tool output truncated".
|
||||
//!
|
||||
//! Each alert is broadcast on the shared `Event` channel and
|
||||
//! also appended to an in-memory buffer so that clients connecting
|
||||
//! after the fact still see everything emitted during the session.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
use protocol::{Alert, AlertLevel, AlertSource, Event};
|
||||
|
||||
/// Upper bound on buffered alerts. When exceeded, the oldest
|
||||
/// entries are discarded so a long-running session cannot leak
|
||||
/// memory through a pathological loop of recurring alerts
|
||||
/// (e.g. compaction failing every turn).
|
||||
const MAX_BUFFERED_ALERTS: usize = 512;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Alerter {
|
||||
inner: Arc<Inner>,
|
||||
}
|
||||
|
||||
struct Inner {
|
||||
event_tx: broadcast::Sender<Event>,
|
||||
buffer: Mutex<VecDeque<Alert>>,
|
||||
}
|
||||
|
||||
impl Alerter {
|
||||
pub fn new(event_tx: broadcast::Sender<Event>) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Inner {
|
||||
event_tx,
|
||||
buffer: Mutex::new(VecDeque::with_capacity(MAX_BUFFERED_ALERTS)),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Record and broadcast an alert.
|
||||
///
|
||||
/// The broadcast may have no subscribers (e.g. during Worker
|
||||
/// construction before any client has connected); the buffer
|
||||
/// guarantees the message is still delivered once a client
|
||||
/// attaches.
|
||||
///
|
||||
/// The buffer mutex is held across `broadcast::send` to make
|
||||
/// `subscribe_with_snapshot` race-free — a client that snapshots
|
||||
/// the buffer while holding the same lock sees every alert
|
||||
/// exactly once: older ones from the snapshot, newer ones from
|
||||
/// the freshly-subscribed receiver.
|
||||
pub fn alert(&self, level: AlertLevel, source: AlertSource, message: String) {
|
||||
let alert = Alert {
|
||||
level,
|
||||
source,
|
||||
message,
|
||||
timestamp_ms: now_ms(),
|
||||
};
|
||||
if let Ok(mut buf) = self.inner.buffer.lock() {
|
||||
if buf.len() >= MAX_BUFFERED_ALERTS {
|
||||
buf.pop_front();
|
||||
}
|
||||
buf.push_back(alert.clone());
|
||||
let _ = self.inner.event_tx.send(Event::Alert(alert));
|
||||
}
|
||||
}
|
||||
|
||||
/// Subscribe and atomically snapshot the current buffer.
|
||||
///
|
||||
/// The returned snapshot contains alerts emitted before
|
||||
/// this call; the receiver will deliver alerts emitted
|
||||
/// after. An alert cannot appear in both.
|
||||
pub fn subscribe_with_snapshot(&self) -> (Vec<Alert>, broadcast::Receiver<Event>) {
|
||||
let buf = self
|
||||
.inner
|
||||
.buffer
|
||||
.lock()
|
||||
.expect("alerter buffer mutex poisoned");
|
||||
let rx = self.inner.event_tx.subscribe();
|
||||
let snapshot: Vec<Alert> = buf.iter().cloned().collect();
|
||||
(snapshot, rx)
|
||||
}
|
||||
}
|
||||
|
||||
fn now_ms() -> i64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as i64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn alert_broadcasts_to_existing_subscriber() {
|
||||
let (tx, _keep) = broadcast::channel::<Event>(8);
|
||||
let alerter = Alerter::new(tx);
|
||||
let (_snapshot, mut rx) = alerter.subscribe_with_snapshot();
|
||||
|
||||
alerter.alert(
|
||||
AlertLevel::Warn,
|
||||
AlertSource::Compactor,
|
||||
"test message".into(),
|
||||
);
|
||||
|
||||
match rx.try_recv() {
|
||||
Ok(Event::Alert(a)) => assert_eq!(a.message, "test message"),
|
||||
other => panic!("unexpected event: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn late_subscriber_sees_earlier_alerts_via_snapshot() {
|
||||
let (tx, _keep) = broadcast::channel::<Event>(8);
|
||||
let alerter = Alerter::new(tx);
|
||||
|
||||
alerter.alert(AlertLevel::Error, AlertSource::Worker, "first".into());
|
||||
alerter.alert(AlertLevel::Warn, AlertSource::AgentsMd, "second".into());
|
||||
|
||||
let (snapshot, mut rx) = alerter.subscribe_with_snapshot();
|
||||
assert_eq!(snapshot.len(), 2);
|
||||
assert_eq!(snapshot[0].message, "first");
|
||||
assert_eq!(snapshot[1].message, "second");
|
||||
assert!(rx.try_recv().is_err()); // nothing pending on the receiver
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn buffer_discards_oldest_past_cap() {
|
||||
let (tx, _keep) = broadcast::channel::<Event>(1024);
|
||||
let alerter = Alerter::new(tx);
|
||||
|
||||
for i in 0..(MAX_BUFFERED_ALERTS + 50) {
|
||||
alerter.alert(AlertLevel::Warn, AlertSource::Engine, format!("msg-{i}"));
|
||||
}
|
||||
|
||||
let (snapshot, _rx) = alerter.subscribe_with_snapshot();
|
||||
assert_eq!(snapshot.len(), MAX_BUFFERED_ALERTS);
|
||||
// First 50 were evicted; the oldest remaining is msg-50.
|
||||
assert_eq!(snapshot.first().unwrap().message, "msg-50");
|
||||
let last = format!("msg-{}", MAX_BUFFERED_ALERTS + 49);
|
||||
assert_eq!(snapshot.last().unwrap().message, last);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn subscribe_snapshot_and_live_do_not_overlap() {
|
||||
let (tx, _keep) = broadcast::channel::<Event>(8);
|
||||
let alerter = Alerter::new(tx);
|
||||
|
||||
alerter.alert(AlertLevel::Warn, AlertSource::Engine, "historic".into());
|
||||
let (snapshot, mut rx) = alerter.subscribe_with_snapshot();
|
||||
alerter.alert(AlertLevel::Error, AlertSource::Engine, "live".into());
|
||||
|
||||
assert_eq!(snapshot.len(), 1);
|
||||
assert_eq!(snapshot[0].message, "historic");
|
||||
match rx.try_recv() {
|
||||
Ok(Event::Alert(a)) => assert_eq!(a.message, "live"),
|
||||
other => panic!("unexpected: {other:?}"),
|
||||
}
|
||||
assert!(rx.try_recv().is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,170 @@
|
||||
//! `WorkerEvent` send / receive helpers.
|
||||
//!
|
||||
//! This module owns the parent-facing lifecycle-event primitive
|
||||
//! (`WorkerEvent`) that children fire upward on turn-end / error /
|
||||
//! shutdown / scope-sub-delegation. Three responsibilities live here:
|
||||
//!
|
||||
//! - **Send** a `Method::WorkerEvent` to the parent socket, fire-and-forget,
|
||||
//! logging failures without blocking the child.
|
||||
//! - **Render** agent-visible variants into human-readable strings for the
|
||||
//! parent's notification buffer. Control-plane-only variants may still have
|
||||
//! a renderer for diagnostics, but receive-side classification keeps them
|
||||
//! out of LLM history/context.
|
||||
//! - **Apply side effects** on the parent (registry / pod-registry
|
||||
//! updates) so that the receive path is idempotent and tolerant of
|
||||
//! out-of-order delivery.
|
||||
//!
|
||||
//! Transport is fire-and-forget — the ticket's decision is that
|
||||
//! callbacks are an optimisation and `ListWorkers` + `reclaim_stale` are
|
||||
//! the real fallback. This module is allowed to drop events on the
|
||||
//! floor (with a warn log) rather than retry.
|
||||
//!
|
||||
//! `apply_event_side_effects` takes its dependencies (registry, scope
|
||||
//! lock path, self identity) by reference so the caller owns lifetime
|
||||
//! and locking concerns.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use protocol::{Method, ScopeRule, WorkerEvent};
|
||||
|
||||
use crate::runtime::dir::SpawnedWorkerRecord;
|
||||
use crate::spawn::comm_tools::connect_and_send;
|
||||
use crate::spawn::registry::SpawnedWorkerRegistry;
|
||||
|
||||
/// Connect to `socket`, send a single `Method::WorkerEvent(event)`, and
|
||||
/// return. Used by children to report up to their parent.
|
||||
///
|
||||
/// This is a synchronous helper — callers that want fire-and-forget
|
||||
/// semantics should wrap the call in `tokio::spawn` themselves.
|
||||
pub async fn send_worker_event(socket: &Path, event: WorkerEvent) -> std::io::Result<()> {
|
||||
connect_and_send(socket, &Method::WorkerEvent(event)).await
|
||||
}
|
||||
|
||||
/// Spawn a fire-and-forget task that sends `event` to `socket`. If
|
||||
/// `socket` is `None`, no send happens (top-level Workers have no parent).
|
||||
/// Any send failure is logged at warn level but otherwise ignored —
|
||||
/// the parent is treated as best-effort.
|
||||
pub fn fire_and_forget(socket: Option<PathBuf>, event: WorkerEvent) {
|
||||
let Some(socket) = socket else { return };
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = send_worker_event(&socket, event).await {
|
||||
tracing::warn!(error = %e, socket = %socket.display(), "WorkerEvent send failed");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Render a variant into a one-line human-readable string.
|
||||
///
|
||||
/// Only events classified by `WorkerEvent::should_notify_agent` are injected
|
||||
/// into the parent's LLM context as system messages; control-plane-only events
|
||||
/// keep this renderer for diagnostics/tests. Agent-visible summaries are kept
|
||||
/// deliberately short — the LLM can always call `ReadWorkerOutput` to fetch more
|
||||
/// detail if the event summary is not enough.
|
||||
pub fn render_event(event: &WorkerEvent) -> String {
|
||||
match event {
|
||||
WorkerEvent::TurnEnded { worker_name } => {
|
||||
format!("Worker `{worker_name}` finished a turn.")
|
||||
}
|
||||
WorkerEvent::Errored {
|
||||
worker_name,
|
||||
message,
|
||||
} => {
|
||||
format!("Worker `{worker_name}` reported an error: {message}")
|
||||
}
|
||||
WorkerEvent::ShutDown { worker_name } => {
|
||||
format!("Worker `{worker_name}` has stopped.")
|
||||
}
|
||||
WorkerEvent::ScopeSubDelegated {
|
||||
parent_worker,
|
||||
sub_worker,
|
||||
..
|
||||
} => {
|
||||
format!("Worker `{parent_worker}` spawned `{sub_worker}` and delegated scope to it.")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply the variant-specific side effect on the parent side.
|
||||
///
|
||||
/// All operations are idempotent so that out-of-order delivery (e.g.
|
||||
/// `TurnEnded` arriving after `ShutDown`) does not produce errors:
|
||||
///
|
||||
/// - `TurnEnded` / `Errored`: no system work; the LLM handles the
|
||||
/// semantic response.
|
||||
/// - `ShutDown`: remove the child from `spawned_workers.json`, Worker state,
|
||||
/// and reclaim its delegated scope/allocation. Missing entries are swallowed.
|
||||
/// - `ScopeSubDelegated`: register the grandchild locally and re-emit
|
||||
/// upward to our own parent if we have one. Duplicate grandchild
|
||||
/// entries (re-delivery) are swallowed.
|
||||
pub async fn apply_event_side_effects(
|
||||
event: &WorkerEvent,
|
||||
registry: &Arc<SpawnedWorkerRegistry>,
|
||||
self_name: &str,
|
||||
self_parent_socket: &Option<PathBuf>,
|
||||
) {
|
||||
match event {
|
||||
WorkerEvent::TurnEnded { .. } | WorkerEvent::Errored { .. } => {}
|
||||
|
||||
WorkerEvent::ShutDown { worker_name } => {
|
||||
if let Err(e) = registry.remove(worker_name).await {
|
||||
tracing::warn!(error = %e, worker = %worker_name, "registry remove on ShutDown failed");
|
||||
}
|
||||
}
|
||||
|
||||
WorkerEvent::ScopeSubDelegated {
|
||||
parent_worker,
|
||||
sub_worker,
|
||||
sub_socket,
|
||||
scope,
|
||||
} => {
|
||||
if registry.get(sub_worker).await.is_some() {
|
||||
return;
|
||||
}
|
||||
let callback_address = registry
|
||||
.get(parent_worker)
|
||||
.await
|
||||
.map(|r| r.socket_path)
|
||||
.unwrap_or_else(PathBuf::new);
|
||||
let record = SpawnedWorkerRecord {
|
||||
worker_name: sub_worker.clone(),
|
||||
socket_path: sub_socket.clone(),
|
||||
scope_delegated: scope.clone(),
|
||||
callback_address,
|
||||
};
|
||||
if let Err(e) = registry.add(record).await {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
sub_worker = %sub_worker,
|
||||
"registry add on ScopeSubDelegated failed"
|
||||
);
|
||||
}
|
||||
reemit_scope_sub_delegated(
|
||||
self_parent_socket,
|
||||
self_name,
|
||||
sub_worker.clone(),
|
||||
sub_socket.clone(),
|
||||
scope.clone(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn reemit_scope_sub_delegated(
|
||||
self_parent_socket: &Option<PathBuf>,
|
||||
self_name: &str,
|
||||
sub_worker: String,
|
||||
sub_socket: PathBuf,
|
||||
scope: Vec<ScopeRule>,
|
||||
) {
|
||||
let Some(parent_socket) = self_parent_socket.clone() else {
|
||||
return;
|
||||
};
|
||||
let event = WorkerEvent::ScopeSubDelegated {
|
||||
parent_worker: self_name.to_string(),
|
||||
sub_worker,
|
||||
sub_socket,
|
||||
scope,
|
||||
};
|
||||
fire_and_forget(Some(parent_socket), event);
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,6 @@
|
||||
pub mod alerter;
|
||||
pub mod event;
|
||||
pub mod server;
|
||||
|
||||
pub(crate) mod interceptor;
|
||||
pub(crate) mod notify_buffer;
|
||||
@@ -0,0 +1,198 @@
|
||||
//! Pending-notify buffer for `Method::Notify` and `Method::WorkerEvent`.
|
||||
//!
|
||||
//! Entries are queued here by the Controller (on receipt of the
|
||||
//! corresponding IPC method) and drained by
|
||||
//! `WorkerInterceptor::pending_history_appends`, which the Engine calls
|
||||
//! at the head of each turn loop iteration. The drain renders each
|
||||
//! pending entry into a typed `SystemItem` (with the `notify_wrapper`
|
||||
//! prompt applied), commits a `LogEntry::SystemItem` per entry through
|
||||
//! the session-log sink, and returns the corresponding
|
||||
//! `Item::system_message`s for the worker to append to its
|
||||
//! persistent history.
|
||||
//!
|
||||
//! This is the **single lane** for "system messages produced by Worker
|
||||
//! state that should land in the next LLM request": Notify,
|
||||
//! agent-visible WorkerEvent variants, and any future `<system-reminder>`
|
||||
//! injection all ride this queue.
|
||||
//! Per `tickets/notify-history-persist.md` and `AGENTS.md` (LLM
|
||||
//! context の加工原則), there is **no** "transient, history-skipping"
|
||||
//! lane — everything injected into a request is also committed to
|
||||
//! history so any LLM reaction has a visible trigger across turns,
|
||||
//! resume, and compaction, and so the Anthropic prompt cache prefix
|
||||
//! stays stable across requests.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use protocol::WorkerEvent;
|
||||
use session_store::SystemItem;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::prompt::catalog::{CatalogError, PromptCatalog};
|
||||
|
||||
/// Maximum queued pending entries. Oldest entries are dropped beyond this.
|
||||
const CAPACITY: usize = 128;
|
||||
|
||||
/// One pending entry awaiting drain into the next LLM request.
|
||||
///
|
||||
/// The buffer keeps the raw input shape so the drain step can decide
|
||||
/// the right `SystemItem` kind (and apply `notify_wrapper` to the
|
||||
/// rendered body) at the moment of commit, when the prompt catalog
|
||||
/// is available.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PendingNotify {
|
||||
Notify { message: String },
|
||||
WorkerEvent { event: WorkerEvent },
|
||||
}
|
||||
|
||||
/// Shared, mutex-guarded buffer of pending entries.
|
||||
///
|
||||
/// Cloned between the Worker (producer) and WorkerInterceptor (consumer).
|
||||
#[derive(Clone, Default)]
|
||||
pub struct NotifyBuffer {
|
||||
inner: Arc<Mutex<VecDeque<PendingNotify>>>,
|
||||
}
|
||||
|
||||
impl NotifyBuffer {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Push a notify entry onto the queue. If the queue is full, the
|
||||
/// oldest entry is dropped and a `tracing::warn` is emitted — the
|
||||
/// caller should never hit this in normal operation.
|
||||
pub fn push_notify(&self, message: String) {
|
||||
self.push_entry(PendingNotify::Notify { message });
|
||||
}
|
||||
|
||||
/// Push a typed worker-event entry onto the queue.
|
||||
pub fn push_worker_event(&self, event: WorkerEvent) {
|
||||
self.push_entry(PendingNotify::WorkerEvent { event });
|
||||
}
|
||||
|
||||
fn push_entry(&self, entry: PendingNotify) {
|
||||
let mut q = self.inner.lock().expect("notify buffer poisoned");
|
||||
if q.len() >= CAPACITY {
|
||||
let dropped = q.pop_front();
|
||||
warn!(
|
||||
capacity = CAPACITY,
|
||||
dropped = ?dropped,
|
||||
"notify buffer overflow; dropped oldest"
|
||||
);
|
||||
}
|
||||
q.push_back(entry);
|
||||
}
|
||||
|
||||
/// Remove and return all pending entries in FIFO order.
|
||||
pub fn drain(&self) -> Vec<PendingNotify> {
|
||||
let mut q = self.inner.lock().expect("notify buffer poisoned");
|
||||
q.drain(..).collect()
|
||||
}
|
||||
|
||||
/// Number of pending entries. Primarily for tests.
|
||||
pub fn len(&self) -> usize {
|
||||
self.inner.lock().expect("notify buffer poisoned").len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
}
|
||||
|
||||
/// Render one pending entry into a typed `SystemItem`. The
|
||||
/// `notify_wrapper` prompt produces the LLM-context body for both
|
||||
/// `Notify` (raw message) and `WorkerEvent` (rendered event line).
|
||||
pub(crate) fn build_system_item(
|
||||
entry: &PendingNotify,
|
||||
prompts: &PromptCatalog,
|
||||
) -> Result<SystemItem, CatalogError> {
|
||||
match entry {
|
||||
PendingNotify::Notify { message } => {
|
||||
let body = prompts.notify_wrapper(message)?;
|
||||
Ok(SystemItem::Notification {
|
||||
message: message.clone(),
|
||||
body,
|
||||
})
|
||||
}
|
||||
PendingNotify::WorkerEvent { event } => {
|
||||
let rendered = session_store::render_worker_event(event);
|
||||
let body = prompts.notify_wrapper(&rendered)?;
|
||||
Ok(SystemItem::WorkerEvent {
|
||||
event: event.clone(),
|
||||
body,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn push_then_drain_preserves_order() {
|
||||
let buf = NotifyBuffer::new();
|
||||
buf.push_notify("one".into());
|
||||
buf.push_notify("two".into());
|
||||
let drained = buf.drain();
|
||||
assert_eq!(drained.len(), 2);
|
||||
match &drained[0] {
|
||||
PendingNotify::Notify { message } => assert_eq!(message, "one"),
|
||||
other => panic!("unexpected: {other:?}"),
|
||||
}
|
||||
assert!(buf.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capacity_drops_oldest() {
|
||||
let buf = NotifyBuffer::new();
|
||||
for i in 0..(CAPACITY + 5) {
|
||||
buf.push_notify(format!("msg{i}"));
|
||||
}
|
||||
let drained = buf.drain();
|
||||
assert_eq!(drained.len(), CAPACITY);
|
||||
match &drained[0] {
|
||||
PendingNotify::Notify { message } => assert_eq!(message, "msg5"),
|
||||
other => panic!("unexpected: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_system_item_for_notify_carries_wrapper_body() {
|
||||
let entry = PendingNotify::Notify {
|
||||
message: "hello".into(),
|
||||
};
|
||||
let catalog = PromptCatalog::builtins_only().unwrap();
|
||||
let item = build_system_item(&entry, &catalog).unwrap();
|
||||
match item {
|
||||
SystemItem::Notification { message, body } => {
|
||||
assert_eq!(message, "hello");
|
||||
assert!(body.contains("[Notification]"));
|
||||
assert!(body.contains("hello"));
|
||||
assert!(body.contains("not a blocking request"));
|
||||
}
|
||||
other => panic!("unexpected: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_system_item_for_worker_event_wraps_rendered_event_text() {
|
||||
let entry = PendingNotify::WorkerEvent {
|
||||
event: WorkerEvent::TurnEnded {
|
||||
worker_name: "child".into(),
|
||||
},
|
||||
};
|
||||
let catalog = PromptCatalog::builtins_only().unwrap();
|
||||
let item = build_system_item(&entry, &catalog).unwrap();
|
||||
match item {
|
||||
SystemItem::WorkerEvent { event, body } => {
|
||||
assert!(
|
||||
matches!(event, WorkerEvent::TurnEnded { ref worker_name } if worker_name == "child")
|
||||
);
|
||||
assert!(body.contains("[Notification]"));
|
||||
assert!(body.contains("`child`"));
|
||||
}
|
||||
other => panic!("unexpected: {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,286 @@
|
||||
use std::io;
|
||||
use std::io::ErrorKind;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use protocol::stream::{JsonLineReader, JsonLineWriter};
|
||||
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.
|
||||
///
|
||||
/// Listens on the Worker's runtime directory socket path.
|
||||
/// Each client connection gets bidirectional JSONL:
|
||||
/// - Client writes Method lines → forwarded to WorkerController
|
||||
/// - Worker events → written as Event lines to all connected clients
|
||||
pub struct SocketServer {
|
||||
_accept_task: JoinHandle<()>,
|
||||
path: PathBuf,
|
||||
}
|
||||
|
||||
impl SocketServer {
|
||||
/// Start listening on the WorkerHandle's socket path.
|
||||
pub async fn start(handle: &WorkerHandle) -> Result<Self, io::Error> {
|
||||
let path = handle.runtime_dir.socket_path();
|
||||
|
||||
// Remove stale socket file if it exists
|
||||
let _ = tokio::fs::remove_file(&path).await;
|
||||
|
||||
let listener = UnixListener::bind(&path)?;
|
||||
let handle = handle.clone();
|
||||
|
||||
let _accept_task = tokio::spawn(async move {
|
||||
loop {
|
||||
match listener.accept().await {
|
||||
Ok((stream, _)) => {
|
||||
let handle = handle.clone();
|
||||
tokio::spawn(handle_connection(stream, handle));
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(Self { _accept_task, path })
|
||||
}
|
||||
|
||||
/// The socket file path.
|
||||
pub fn path(&self) -> &std::path::Path {
|
||||
&self.path
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for SocketServer {
|
||||
fn drop(&mut self) {
|
||||
let _ = std::fs::remove_file(&self.path);
|
||||
}
|
||||
}
|
||||
|
||||
fn is_peer_disconnect_read_error(error: &io::Error) -> bool {
|
||||
matches!(
|
||||
error.kind(),
|
||||
ErrorKind::ConnectionReset
|
||||
| ErrorKind::ConnectionAborted
|
||||
| ErrorKind::BrokenPipe
|
||||
| ErrorKind::UnexpectedEof
|
||||
)
|
||||
}
|
||||
|
||||
fn live_entry_event(entry: session_store::LogEntry) -> Option<Event> {
|
||||
match entry {
|
||||
session_store::LogEntry::SegmentStart { .. } => {
|
||||
let value = serde_json::to_value(&entry).expect("LogEntry is Serialize");
|
||||
Some(Event::SegmentRotated { entry: value })
|
||||
}
|
||||
session_store::LogEntry::UserInput { segments, .. } => {
|
||||
Some(Event::UserMessage { segments })
|
||||
}
|
||||
session_store::LogEntry::SystemItem { item, .. } => {
|
||||
let value = serde_json::to_value(&item).expect("SystemItem is Serialize");
|
||||
Some(Event::SystemItem { item: value })
|
||||
}
|
||||
session_store::LogEntry::Invoke { trigger, .. } => {
|
||||
Some(Event::InvokeStart { kind: trigger })
|
||||
}
|
||||
other => {
|
||||
// `SegmentLogSink::is_live_relevant` keeps non-live-relevant
|
||||
// variants off the broadcast lane; reaching here means the two
|
||||
// are out of sync and we silently dropped a wire event. Log so a
|
||||
// future regression surfaces instead of vanishing.
|
||||
tracing::error!(
|
||||
entry_kind = ?std::mem::discriminant(&other),
|
||||
"session-log broadcast emitted a non-live-relevant entry; \
|
||||
sink filter and IPC dispatch are out of sync"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_connection(stream: tokio::net::UnixStream, handle: WorkerHandle) {
|
||||
let (reader, writer) = stream.into_split();
|
||||
let mut reader = JsonLineReader::new(reader);
|
||||
let mut writer = JsonLineWriter::new(writer);
|
||||
|
||||
// Hold the in-flight stream lock while taking the session-log mirror
|
||||
// snapshot. `LogEntry::AssistantItem` is mirror-only for live clients,
|
||||
// so a finalized assistant block must be observed either as an already
|
||||
// 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)
|
||||
};
|
||||
for alert in alert_snapshot {
|
||||
if writer.write(&Event::Alert(alert)).await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// 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;
|
||||
}
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
// Live session-log entries → dispatched as the role-specific
|
||||
// wire events. `SegmentLogSink` only broadcasts committed log
|
||||
// entries with live UI meaning; `UserInput` travels this lane so
|
||||
// the visible user line is ordered with `SegmentStart` rotation.
|
||||
entry = entry_rx.recv() => {
|
||||
match entry {
|
||||
Ok(entry) => {
|
||||
if let Some(event) = live_entry_event(entry) {
|
||||
if writer.write(&event).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
|
||||
// Slow client fell behind the broadcast buffer.
|
||||
// Drop the connection so the next reconnect
|
||||
// re-seeds the prefix via subscribe_with_snapshot.
|
||||
break;
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
|
||||
}
|
||||
}
|
||||
// Broadcast events → this client
|
||||
event = rx.recv() => {
|
||||
match event {
|
||||
Ok(event) => {
|
||||
if writer.write(&event).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
// Client methods → handle or forward to controller
|
||||
method = reader.next::<Method>() => {
|
||||
match method {
|
||||
Ok(Some(Method::ListCompletions { kind, prefix })) => {
|
||||
let entries = match kind {
|
||||
protocol::CompletionKind::File => handle
|
||||
.shared_state
|
||||
.fs_view()
|
||||
.map(|view| view.list_file_completions(&prefix))
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.map(|c| protocol::CompletionEntry {
|
||||
value: c.path,
|
||||
is_dir: c.is_dir,
|
||||
})
|
||||
.collect(),
|
||||
protocol::CompletionKind::Knowledge => handle
|
||||
.shared_state
|
||||
.list_knowledge_completions(&prefix)
|
||||
.into_iter()
|
||||
.map(|c| protocol::CompletionEntry {
|
||||
value: c.slug,
|
||||
is_dir: false,
|
||||
})
|
||||
.collect(),
|
||||
protocol::CompletionKind::Workflow => handle
|
||||
.shared_state
|
||||
.list_workflow_completions(&prefix)
|
||||
.into_iter()
|
||||
.map(|c| protocol::CompletionEntry {
|
||||
value: c.slug,
|
||||
is_dir: false,
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
if writer
|
||||
.write(&Event::Completions { kind, entries })
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(Some(method)) => {
|
||||
let _ = handle.send(method).await;
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(e) if is_peer_disconnect_read_error(&e) => break,
|
||||
Err(e) => {
|
||||
if writer
|
||||
.write(&Event::Error {
|
||||
code: protocol::ErrorCode::InvalidRequest,
|
||||
message: format!("invalid method: {e}"),
|
||||
})
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn peer_disconnect_read_errors_are_connection_close() {
|
||||
for kind in [
|
||||
ErrorKind::ConnectionReset,
|
||||
ErrorKind::ConnectionAborted,
|
||||
ErrorKind::BrokenPipe,
|
||||
ErrorKind::UnexpectedEof,
|
||||
] {
|
||||
let error = io::Error::new(kind, "peer disconnected");
|
||||
assert!(
|
||||
is_peer_disconnect_read_error(&error),
|
||||
"{kind:?} should be treated as a normal peer disconnect"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_data_is_not_peer_disconnect() {
|
||||
let error = io::Error::new(ErrorKind::InvalidData, "malformed method");
|
||||
assert!(!is_peer_disconnect_read_error(&error));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_input_log_entry_maps_to_user_message_event() {
|
||||
let segments = vec![protocol::Segment::text("hello from log")];
|
||||
let event = live_entry_event(session_store::LogEntry::UserInput {
|
||||
ts: session_store::segment_log::now_millis(),
|
||||
segments: segments.clone(),
|
||||
})
|
||||
.expect("UserInput must be live-relevant");
|
||||
|
||||
match event {
|
||||
Event::UserMessage { segments: echoed } => assert_eq!(echoed, segments),
|
||||
other => panic!("expected UserMessage, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user