From 8aafd388fa382e03955f74df126ca88d9f4b7155 Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 21 Jul 2026 20:42:49 +0900 Subject: [PATCH] fix: proxy backend protocol websocket frames --- .yoi/tickets/00001KY1ZKCPC/item.md | 2 +- .yoi/tickets/00001KY1ZKCPC/thread.md | 31 +++++ crates/workspace-server/src/server.rs | 169 +++++++++++++++++++++++--- 3 files changed, 186 insertions(+), 16 deletions(-) diff --git a/.yoi/tickets/00001KY1ZKCPC/item.md b/.yoi/tickets/00001KY1ZKCPC/item.md index ed5c8037..b42bf5ab 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-21T11:29:08Z' +updated_at: '2026-07-21T11:42:36Z' assignee: null readiness: 'draft' --- diff --git a/.yoi/tickets/00001KY1ZKCPC/thread.md b/.yoi/tickets/00001KY1ZKCPC/thread.md index 8c7e31ad..fdae1157 100644 --- a/.yoi/tickets/00001KY1ZKCPC/thread.md +++ b/.yoi/tickets/00001KY1ZKCPC/thread.md @@ -135,3 +135,34 @@ Verification: - `git diff --check` --- + + + +## Implementation report + +Bug fix after manual TUI send failure: + +Observed symptom: +- TUI send over Backend protocol reported a `does not support Worker protocol command transport` style error. + +Root cause: +- `/events/ws` had been removed and `/protocol/ws` unified, but workspace-server still handled inbound client `protocol::Method` by calling `RuntimeRegistry::send_protocol_method(...)`. +- For remote runtime-backed workers, Backend must be only a protocol WS proxy. Calling the registry default method hits the unsupported path instead of tunneling the frame to the runtime's `/protocol/ws`. + +Fix: +- Workspace-server now branches by observation/protocol source: + - `RemoteWs`: proxy raw WebSocket frames bidirectionally between client `/api/.../protocol/ws` and runtime `/v1/.../protocol/ws`. Backend does not deserialize or interpret method frames in this path. + - `Embedded`: use in-process runtime adapter for method dispatch while streaming embedded observation events. +- Added/updated test so sending `Method::ListCompletions` through the workspace-server protocol proxy returns `Event::Completions` from the runtime, proving method frames reach the runtime instead of the Backend unsupported fallback. + +Verification: +- `nix develop -c cargo fmt -- --check` +- `nix develop -c cargo check -p yoi-workspace-server -p client -p worker-runtime` +- `nix develop -c cargo test -p yoi-workspace-server proxies_worker_protocol_ws_as_raw_events --lib` +- `nix develop -c cargo test -p client backend_runtime --lib` +- `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` +- `deno test --allow-read web/workspace/src/lib/workspace/console/worker-console.ui.test.ts` +- `git diff --check` + +--- diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 1c7c0e49..3ffe726a 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -11,7 +11,7 @@ use axum::response::{IntoResponse, Response}; use axum::routing::{delete, get, post, put}; use axum::{Json, Router}; use chrono::{SecondsFormat, Utc}; -use futures::StreamExt; +use futures::{SinkExt, StreamExt}; use memory::backend::{ MemoryBackendHttpResponse, MemoryBackendOperation, execute_memory_backend_operation, }; @@ -22,6 +22,9 @@ use ticket::{ execute_ticket_backend_operation, }; use tokio::net::TcpListener; +use tokio_tungstenite::connect_async; +use tokio_tungstenite::tungstenite::Message as TungsteniteMessage; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; use worker_runtime::resource::{BackendResourceError, BackendResourceFetchRequest}; use worker_runtime::worker_backend::{ProfileRuntimeWorkerFactory, WorkerRuntimeExecutionBackend}; @@ -42,7 +45,7 @@ use crate::hosts::{ use crate::identity::WorkspaceIdentity; use crate::observation::{ BackendObservationProxy, ObservationProxyError, RuntimeObservationClient, - RuntimeObservationSourceConfig, + RuntimeObservationSource, RuntimeObservationSourceConfig, }; use crate::profile_settings::{ CreateWorkspaceProfileSourceRequest, DeleteWorkspaceProfileSourceRequest, @@ -3412,19 +3415,140 @@ async fn worker_protocol_ws( .into_response(); } }; - ws.on_upgrade(move |socket| { - worker_protocol_ws_session(api.runtime, source, runtime_id, worker_id, socket) - }) + ws.on_upgrade(move |socket| worker_protocol_ws_session(source, socket)) } -async fn worker_protocol_ws_session( - runtime: Arc, - source: crate::observation::RuntimeObservationSource, - runtime_id: String, - worker_id: String, +async fn worker_protocol_ws_session(source: RuntimeObservationSource, socket: WebSocket) { + match source { + RuntimeObservationSource::RemoteWs(config) => { + remote_worker_protocol_ws_session(config, socket).await; + } + RuntimeObservationSource::Embedded(source) => { + embedded_worker_protocol_ws_session(source, socket).await; + } + } +} + +async fn remote_worker_protocol_ws_session( + config: RuntimeObservationSourceConfig, + socket: WebSocket, +) { + let mut request = match config.endpoint.clone().into_client_request() { + Ok(request) => request, + Err(error) => { + let mut socket = socket; + let event = protocol_error_event(format!( + "failed to build runtime protocol WebSocket request: {error}" + )); + let _ = send_protocol_event(&mut socket, &event).await; + return; + } + }; + if let Some(token) = &config.bearer_token { + match format!("Bearer {token}").parse() { + Ok(value) => { + request.headers_mut().insert("authorization", value); + } + Err(error) => { + let mut socket = socket; + let event = protocol_error_event(format!( + "failed to build runtime authorization header: {error}" + )); + let _ = send_protocol_event(&mut socket, &event).await; + return; + } + } + } + + let (upstream, _) = match connect_async(request).await { + Ok(connection) => connection, + Err(error) => { + let mut socket = socket; + let event = protocol_error_event(format!( + "failed to connect runtime protocol WebSocket: {error}" + )); + let _ = send_protocol_event(&mut socket, &event).await; + return; + } + }; + + let (mut client_sink, mut client_stream) = socket.split(); + let (mut upstream_sink, mut upstream_stream) = upstream.split(); + + loop { + tokio::select! { + inbound = client_stream.next() => { + match inbound { + Some(Ok(WsMessage::Text(text))) => { + if upstream_sink.send(TungsteniteMessage::Text(text.to_string().into())).await.is_err() { + break; + } + } + Some(Ok(WsMessage::Binary(binary))) => { + if upstream_sink.send(TungsteniteMessage::Binary(binary.to_vec().into())).await.is_err() { + break; + } + } + Some(Ok(WsMessage::Close(_))) | None => { + let _ = upstream_sink.send(TungsteniteMessage::Close(None)).await; + break; + } + Some(Ok(WsMessage::Ping(payload))) => { + if upstream_sink.send(TungsteniteMessage::Ping(payload.to_vec().into())).await.is_err() { + break; + } + } + Some(Ok(WsMessage::Pong(payload))) => { + if upstream_sink.send(TungsteniteMessage::Pong(payload.to_vec().into())).await.is_err() { + break; + } + } + Some(Err(_)) => break, + } + } + outbound = upstream_stream.next() => { + match outbound { + Some(Ok(TungsteniteMessage::Text(text))) => { + if client_sink.send(WsMessage::Text(text.to_string().into())).await.is_err() { + break; + } + } + Some(Ok(TungsteniteMessage::Binary(binary))) => { + if client_sink.send(WsMessage::Binary(binary.to_vec().into())).await.is_err() { + break; + } + } + Some(Ok(TungsteniteMessage::Close(_))) | None => { + let _ = client_sink.send(WsMessage::Close(None)).await; + break; + } + Some(Ok(TungsteniteMessage::Ping(payload))) => { + if client_sink.send(WsMessage::Ping(payload.to_vec().into())).await.is_err() { + break; + } + } + Some(Ok(TungsteniteMessage::Pong(payload))) => { + if client_sink.send(WsMessage::Pong(payload.to_vec().into())).await.is_err() { + break; + } + } + Some(Ok(TungsteniteMessage::Frame(_))) => {} + Some(Err(_)) => break, + } + } + } + } +} + +async fn embedded_worker_protocol_ws_session( + source: crate::observation::EmbeddedRuntimeObservationSource, mut socket: WebSocket, ) { - let mut upstream = match RuntimeObservationClient::connect(&source).await { + let mut upstream = match RuntimeObservationClient::connect(&RuntimeObservationSource::Embedded( + source.clone(), + )) + .await + { Ok(client) => client, Err(error) => { let event = protocol_error_event(error.message()); @@ -3438,7 +3562,7 @@ async fn worker_protocol_ws_session( inbound = socket.next() => { match inbound { Some(Ok(WsMessage::Text(text))) => match serde_json::from_str::(&text) { - Ok(method) => match runtime.send_protocol_method(&runtime_id, &worker_id, method) { + Ok(method) => match source.runtime.send_protocol_method(&source.worker_ref, method) { Ok(events) => { for event in events { if !send_protocol_event(&mut socket, &event).await { @@ -3447,7 +3571,7 @@ async fn worker_protocol_ws_session( } } Err(error) => { - let event = protocol_error_event(error.message()); + let event = protocol_error_event(error.to_string()); if !send_protocol_event(&mut socket, &event).await { return; } @@ -3492,7 +3616,6 @@ async fn worker_protocol_ws_session( } } } - async fn send_protocol_event(socket: &mut WebSocket, event: &protocol::Event) -> bool { match serde_json::to_string(event) { Ok(text) => socket.send(WsMessage::Text(text.into())).await.is_ok(), @@ -5349,7 +5472,7 @@ mod tests { use super::*; use axum::body::{Body, to_bytes}; use axum::http::Request; - use futures::StreamExt; + use futures::{SinkExt, StreamExt}; use serde_json::{Value, json}; use std::{fs, sync::Arc}; use tokio_tungstenite::connect_async; @@ -7788,6 +7911,22 @@ mod tests { protocol::Event::Snapshot { .. } )); + stream + .send(Message::Text( + serde_json::to_string(&protocol::Method::ListCompletions { + kind: protocol::CompletionKind::File, + prefix: String::new(), + }) + .unwrap() + .into(), + )) + .await + .unwrap(); + assert!(matches!( + next_client_frame(&mut stream).await, + protocol::Event::Completions { .. } + )); + runtime .observe_worker_event( &worker_ref,