fix: proxy backend protocol websocket frames

This commit is contained in:
Keisuke Hirata 2026-07-21 20:42:49 +09:00
parent 8841d063be
commit 8aafd388fa
No known key found for this signature in database
3 changed files with 186 additions and 16 deletions

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-21T11:29:08Z'
updated_at: '2026-07-21T11:42:36Z'
assignee: null
readiness: 'draft'
---

View File

@ -135,3 +135,34 @@ Verification:
- `git diff --check`
---
<!-- event: implementation_report author: assistant at: 2026-07-21T11:42:36Z -->
## 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`
---

View File

@ -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<RuntimeRegistry>,
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::<protocol::Method>(&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,