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

fix: share worker protocol session logic

Browse Source
This commit is contained in:
Keisuke Hirata 2026-07-22 12:51:19 +09:00
parent a164017432
commit 945716d010
No known key found for this signature in database
6 changed files with 185 additions and 102 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.yoi/tickets/00001KY1ZKCPC/item.md
View File

@ -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'
---

66
.yoi/tickets/00001KY1ZKCPC/thread.md
View File

@ -235,3 +235,69 @@ Verification:
- `git diff --check`
---
<!-- event: decision author: assistant at: 2026-07-22T03:41:01Z -->
## 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.
---
<!-- event: implementation_report author: assistant at: 2026-07-22T03:51:04Z -->
## 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`
---

42
crates/worker-runtime/src/worker_backend.rs
View File

@ -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<F> Drop for WorkerRuntimeExecutionBackend<F> {
}
}
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 }),
_ => 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<Vec<Vec<LlmEvent>>>,

1
crates/worker/src/ipc/mod.rs
View File

@ -1,5 +1,6 @@
pub mod alerter;
pub mod event;
pub mod protocol_session;
pub mod server;
pub(crate) mod interceptor;

96
crates/worker/src/ipc/protocol_session.rs Normal file
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:?}"),
}
}
}

80
crates/worker/src/ipc/server.rs
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:?}"),
}
}
}