yoi/crates/pod/src/ipc/server.rs

281 lines
10 KiB
Rust

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::PodHandle;
use protocol::{Event, Method};
/// Unix socket server for Pod Protocol.
///
/// Listens on the Pod's runtime directory socket path.
/// Each client connection gets bidirectional JSONL:
/// - Client writes Method lines → forwarded to PodController
/// - Pod events → written as Event lines to all connected clients
pub struct SocketServer {
_accept_task: JoinHandle<()>,
path: PathBuf,
}
impl SocketServer {
/// Start listening on the PodHandle's socket path.
pub async fn start(handle: &PodHandle) -> 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: PodHandle) {
let (reader, writer) = stream.into_split();
let mut reader = JsonLineReader::new(reader);
let mut writer = JsonLineWriter::new(writer);
// Atomically subscribe to the session-log mirror first. The
// returned (snapshot, rx) pair partitions the entry timeline:
// entries committed before this call appear in `entries`, every
// entry after lands on `entry_rx`. Doing this before the alert
// snapshot keeps both ordering pairs internally consistent.
let (entries_snapshot, mut 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 — they appear in the snapshot, and any alert
// arriving afterwards reaches us through `rx`.
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;
}
}
// 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(),
};
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:?}"),
}
}
}