Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: share worker protocol session logic

Browse Source
This commit is contained in:
Hare
2026-07-22 12:51:19 +09:00
parent a164017432
commit 945716d010
6 changed files with 185 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/ipc/mod.rs
+1
View File
@@ -1,5 +1,6 @@
pub mod alerter;
pub mod event;
pub mod protocol_session;
pub mod server;
pub(crate) mod interceptor;
crates/worker/src/ipc/protocol_session.rs
+96
View File
@@ -0,0 +1,96 @@
use protocol::{Alert, Event, Method};
use session_store::LogEntry;
use tokio::sync::broadcast;
use crate::controller::WorkerHandle;
/// Live channels and initial replay data for a Worker protocol session.
///
/// This is intentionally transport-agnostic: Unix JSONL sockets and Runtime
/// WebSocket transports should subscribe through this helper so they cannot
/// drift on which Worker/log events make up the protocol stream.
pub struct WorkerProtocolSessionStreams {
pub snapshot_event: Event,
pub alert_snapshot: Vec<Alert>,
pub log_entries: broadcast::Receiver<LogEntry>,
pub events: broadcast::Receiver<Event>,
}
pub fn subscribe_worker_protocol_session(handle: &WorkerHandle) -> WorkerProtocolSessionStreams {
let (snapshot_event, log_entries) = handle.snapshot_event_with_entry_subscription();
let (alert_snapshot, events) = handle.alerter.subscribe_with_snapshot();
WorkerProtocolSessionStreams {
snapshot_event,
alert_snapshot,
log_entries,
events,
}
}
pub fn live_log_entry_event(entry: LogEntry) -> Option<Event> {
match entry {
LogEntry::SegmentStart { .. } => {
let value = serde_json::to_value(&entry).expect("LogEntry is Serialize");
Some(Event::SegmentRotated { entry: value })
}
LogEntry::UserInput { segments, .. } => Some(Event::UserMessage { segments }),
LogEntry::SystemItem { item, .. } => {
let value = serde_json::to_value(&item).expect("SystemItem is Serialize");
Some(Event::SystemItem { item: value })
}
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 protocol dispatch are out of sync"
);
None
}
}
}
/// Dispatch a client Method that has same-connection response semantics.
///
/// Methods returning `Some(Event)` are handled by the protocol session and must
/// be written back only to the requesting transport. Other methods are sent to
/// the Worker controller and their results appear through the normal protocol
/// event/log streams.
pub async fn dispatch_worker_protocol_method(
handle: &WorkerHandle,
method: Method,
) -> Option<Event> {
match method {
Method::ListCompletions { kind, prefix } => {
let entries = handle.completion_entries(kind, &prefix);
Some(Event::Completions { kind, entries })
}
method => {
let _ = handle.send(method).await;
None
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn user_input_log_entry_maps_to_user_message_event() {
let segments = vec![protocol::Segment::text("hello from log")];
let event = live_log_entry_event(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:?}"),
}
}
}
crates/worker/src/ipc/server.rs
+14 -66
View File
@@ -7,7 +7,10 @@ use tokio::net::UnixListener;
use tokio::task::JoinHandle;
use crate::controller::WorkerHandle;
use protocol::{Event, Method};
use crate::ipc::protocol_session::{
dispatch_worker_protocol_method, live_log_entry_event, subscribe_worker_protocol_session,
};
use protocol::{ErrorCode, Event, Method};
/// Unix socket server for Worker Protocol.
///
@@ -68,37 +71,6 @@ fn is_peer_disconnect_read_error(error: &io::Error) -> bool {
)
}
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);
@@ -110,11 +82,8 @@ 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 (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 {
let mut streams = subscribe_worker_protocol_session(&handle);
for alert in streams.alert_snapshot {
if writer.write(&Event::Alert(alert)).await.is_err() {
return;
}
@@ -122,7 +91,7 @@ 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.
if writer.write(&snapshot_event).await.is_err() {
if writer.write(&streams.snapshot_event).await.is_err() {
return;
}
@@ -132,10 +101,10 @@ async fn handle_connection(stream: tokio::net::UnixStream, handle: WorkerHandle)
// 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() => {
entry = streams.log_entries.recv() => {
match entry {
Ok(entry) => {
if let Some(event) = live_entry_event(entry) {
if let Some(event) = live_log_entry_event(entry) {
if writer.write(&event).await.is_err() {
break;
}
@@ -151,7 +120,7 @@ async fn handle_connection(stream: tokio::net::UnixStream, handle: WorkerHandle)
}
}
// Broadcast events → this client
event = rx.recv() => {
event = streams.events.recv() => {
match event {
Ok(event) => {
if writer.write(&event).await.is_err() {
@@ -164,25 +133,19 @@ async fn handle_connection(stream: tokio::net::UnixStream, handle: WorkerHandle)
// Client methods → handle or forward to controller
method = reader.next::<Method>() => {
match method {
Ok(Some(Method::ListCompletions { kind, prefix })) => {
let entries = handle.completion_entries(kind, &prefix);
if writer
.write(&Event::Completions { kind, entries })
.await
.is_err()
Ok(Some(method)) => {
if let Some(response) = dispatch_worker_protocol_method(&handle, method).await
&& writer.write(&response).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,
code: ErrorCode::InvalidRequest,
message: format!("invalid method: {e}"),
})
.await
@@ -222,19 +185,4 @@ mod tests {
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:?}"),
}
}
}