From 945716d0100031deec84c319b1b88cac5f33e28a Mon Sep 17 00:00:00 2001 From: Hare Date: Wed, 22 Jul 2026 12:51:19 +0900 Subject: [PATCH] fix: share worker protocol session logic --- .yoi/tickets/00001KY1ZKCPC/item.md | 2 +- .yoi/tickets/00001KY1ZKCPC/thread.md | 66 ++++++++++++++ crates/worker-runtime/src/worker_backend.rs | 42 ++------- crates/worker/src/ipc/mod.rs | 1 + crates/worker/src/ipc/protocol_session.rs | 96 +++++++++++++++++++++ crates/worker/src/ipc/server.rs | 80 +++-------------- 6 files changed, 185 insertions(+), 102 deletions(-) create mode 100644 crates/worker/src/ipc/protocol_session.rs diff --git a/.yoi/tickets/00001KY1ZKCPC/item.md b/.yoi/tickets/00001KY1ZKCPC/item.md index 82c9096e..4ecbd35e 100644 --- a/.yoi/tickets/00001KY1ZKCPC/item.md +++ b/.yoi/tickets/00001KY1ZKCPC/item.md @@ -2,7 +2,7 @@ title: 'Backend runtime経由の操作をprotocol transportへ統一しTUI同等にする' state: 'planning' created_at: '2026-07-21T09:20:07Z' -updated_at: '2026-07-22T02:54:28Z' +updated_at: '2026-07-22T03:51:04Z' assignee: null readiness: 'draft' --- diff --git a/.yoi/tickets/00001KY1ZKCPC/thread.md b/.yoi/tickets/00001KY1ZKCPC/thread.md index 5b134e97..c2bd896e 100644 --- a/.yoi/tickets/00001KY1ZKCPC/thread.md +++ b/.yoi/tickets/00001KY1ZKCPC/thread.md @@ -235,3 +235,69 @@ Verification: - `git diff --check` --- + + + +## Decision + +Design correction / implementation target: + +Runtime protocol WS should be thin. `/v1/workers/{worker_id}/protocol/ws` is already worker-id scoped, so after worker lookup it should not reimplement Worker protocol semantics. It should connect the WebSocket framing to the same Worker protocol session behavior used by local IPC. + +Current issue: +- Local IPC protocol session in `worker/src/ipc/server.rs` already knows how to combine: + - `WorkerHandle::event_tx` protocol events + - `SegmentLogSink` live log entries (`LogEntry::UserInput` -> `Event::UserMessage`, etc.) + - `Method` frame reading and event writing +- Runtime WS has been recreating similar logic in `worker-runtime/src/http_server.rs` / `worker_backend.rs`. +- This duplicated logic caused drift: runtime protocol streaming initially missed `SegmentLogSink` live entries, and a bad synthetic echo workaround was added before being corrected. + +Decision: +- Do not make runtime/backend interpret normal Worker protocol operations. +- Backend remains a raw protocol WS proxy for remote runtime workers. +- Runtime remains a worker-id lookup + lifecycle boundary; for protocol WS it should delegate to a shared Worker protocol session implementation. +- Local IPC and runtime WS should use the same session logic so user echo, system items, invoke events, completions, rewind, pause/resume, etc. cannot diverge. + +Implementation target: +- Extract the local IPC protocol session logic into a reusable Worker protocol session helper that is not Unix-socket-specific. +- Local IPC server uses the helper. +- Runtime protocol WS uses the same helper or a thin adapter around it. +- Remove duplicated live-entry/event bridging from runtime-specific code where possible. +- Keep lifecycle/runtime management outside the protocol session helper. + +--- + + + +## Implementation report + +Implementation update for thin runtime protocol session: + +- Added `worker::ipc::protocol_session` as the shared Worker protocol session helper. + - `subscribe_worker_protocol_session(&WorkerHandle)` centralizes snapshot + alert + live log entry + event subscriptions. + - `live_log_entry_event(LogEntry)` centralizes the local IPC mapping: + - `SegmentStart` -> `Event::SegmentRotated` + - `UserInput` -> `Event::UserMessage` + - `SystemItem` -> `Event::SystemItem` + - `Invoke` -> `Event::InvokeStart` + - `dispatch_worker_protocol_method(...)` centralizes same-connection direct response handling for `ListCompletions`, forwarding other methods to `WorkerHandle::send`. +- Updated local IPC server to use the shared helper instead of owning its own mapping/subscription logic. +- Updated `worker-runtime` embedded backend bridge to use the same shared helper for live Worker/log event bridging. +- Removed the duplicated `live_log_entry_event` mapping from runtime-specific code. + +Scope note: +- This keeps Backend as raw WS proxy for remote runtime workers. +- Runtime still performs worker-id lookup and execution/lifecycle boundary work, but protocol stream composition now uses the same Worker session helper as local IPC instead of maintaining a separate mapping implementation. +- Full WS socket pump and Unix JSONL pump still have transport-specific loops, but the protocol subscription/mapping/direct-response semantics are no longer duplicated. + +Verification: +- `nix develop -c cargo fmt -- --check` +- `nix develop -c cargo check -p worker-runtime -p worker -p yoi-workspace-server -p client -p tui` +- `nix develop -c cargo test -p worker --lib ipc::protocol_session` +- `nix develop -c cargo test -p worker --lib ipc::server` +- `nix develop -c cargo test -p worker-runtime --features ws-server --lib protocol_ws` +- `nix develop -c cargo test -p yoi-workspace-server protocol_ws --lib` +- `nix develop -c cargo test -p client backend_runtime --lib` +- `git diff --check` + +--- diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 89fa049d..360d8089 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -30,13 +30,15 @@ use crate::working_directory::{ }; use async_trait::async_trait; use manifest::paths; -use protocol::{Event, Method, Segment, WorkerStatus}; +use protocol::{Method, Segment, WorkerStatus}; use session_store::FsStore; -use session_store::{CombinedStore, FsWorkerStore, LogEntry}; +use session_store::{CombinedStore, FsWorkerStore}; use tokio::runtime::Runtime; #[cfg(feature = "ws-server")] use tokio::sync::broadcast; +#[cfg(feature = "ws-server")] +use worker::ipc::protocol_session::{live_log_entry_event, subscribe_worker_protocol_session}; use worker::{ PromptLoader, Worker, WorkerController, WorkerError, WorkerFilesystemAuthority, WorkerHandle, WorkerWorkspaceContext, WorkspaceClient, WorkspaceId, @@ -666,8 +668,9 @@ where let busy = Arc::new(AtomicBool::new(false)); #[cfg(feature = "ws-server")] { - let mut events = handle.subscribe(); - let (_entries, mut entry_events) = handle.sink.subscribe_with_snapshot(); + let streams = subscribe_worker_protocol_session(&handle); + let mut events = streams.events; + let mut entry_events = streams.log_entries; let bridge_handle = handle.clone(); let bridge_busy = busy.clone(); if let Err(message) = self.spawn_on_adapter_runtime(async move { @@ -738,22 +741,6 @@ impl Drop for WorkerRuntimeExecutionBackend { } } -fn live_log_entry_event(entry: LogEntry) -> Option { - 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 }), - _ => None, - } -} - fn method_starts_turn(method: &Method) -> bool { matches!( method, @@ -1240,21 +1227,6 @@ mod tests { use llm_engine::llm_client::{ClientError, LlmClient, Request}; use manifest::{Scope, WorkerManifest}; - #[test] - fn runtime_bridge_maps_live_user_input_log_entry_to_user_message() { - let segments = vec![Segment::text("hello through normal bridge")]; - 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:?}"), - } - } - #[derive(Clone)] struct MockClient { responses: Arc>>, diff --git a/crates/worker/src/ipc/mod.rs b/crates/worker/src/ipc/mod.rs index 7072fb0a..053bbe35 100644 --- a/crates/worker/src/ipc/mod.rs +++ b/crates/worker/src/ipc/mod.rs @@ -1,5 +1,6 @@ pub mod alerter; pub mod event; +pub mod protocol_session; pub mod server; pub(crate) mod interceptor; diff --git a/crates/worker/src/ipc/protocol_session.rs b/crates/worker/src/ipc/protocol_session.rs new file mode 100644 index 00000000..99345225 --- /dev/null +++ b/crates/worker/src/ipc/protocol_session.rs @@ -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, + pub log_entries: broadcast::Receiver, + pub events: broadcast::Receiver, +} + +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 { + 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 { + 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:?}"), + } + } +} diff --git a/crates/worker/src/ipc/server.rs b/crates/worker/src/ipc/server.rs index 885ddfc7..27f0676d 100644 --- a/crates/worker/src/ipc/server.rs +++ b/crates/worker/src/ipc/server.rs @@ -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 { - 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::() => { 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:?}"), - } - } }