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

fix: unify worker protocol websocket

Browse Source
This commit is contained in:
Hare
2026-07-21 20:29:30 +09:00
parent 70a26a3042
commit 8841d063be
10 changed files with 377 additions and 1046 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/client/src/backend_runtime.rs
+11 -137
View File
@@ -1,17 +1,12 @@
use std::collections::VecDeque;
use std::fmt;
use std::time::Duration;
use futures::{SinkExt, StreamExt};
use protocol::{ErrorCode, Event, Method};
use serde::Deserialize;
use std::collections::VecDeque;
use std::fmt;
use tokio::sync::mpsc;
use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite::Message as TungsteniteMessage;
const RECONNECT_DELAY: Duration = Duration::from_millis(500);
const MAX_RECONNECT_ATTEMPTS: usize = 3;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BackendRuntimeTarget {
/// Workspace Backend API root URL, for example `http://127.0.0.1:8787`.
@@ -163,8 +158,7 @@ pub struct BackendRuntimeClient {
command_tx: mpsc::UnboundedSender<Method>,
events: mpsc::UnboundedReceiver<Event>,
diagnostics: VecDeque<Event>,
_observation_task: tokio::task::JoinHandle<()>,
_command_task: tokio::task::JoinHandle<()>,
_protocol_task: tokio::task::JoinHandle<()>,
}
#[derive(Debug)]
@@ -262,16 +256,10 @@ impl BackendRuntimeClient {
let (event_tx, rx) = mpsc::unbounded_channel();
let (command_tx, command_rx) = mpsc::unbounded_channel();
let observation_target = target.clone();
let observation_tx = event_tx.clone();
let observation_task = tokio::spawn(async move {
observe_worker_events(observation_target, observation_tx).await;
});
let command_target = target.clone();
let command_event_tx = event_tx.clone();
let command_task = tokio::spawn(async move {
run_worker_protocol_commands(command_target, command_rx, command_event_tx).await;
let protocol_target = target.clone();
let protocol_event_tx = event_tx.clone();
let protocol_task = tokio::spawn(async move {
run_worker_protocol_transport(protocol_target, command_rx, protocol_event_tx).await;
});
Ok(Self {
@@ -279,8 +267,7 @@ impl BackendRuntimeClient {
command_tx,
events: rx,
diagnostics: VecDeque::new(),
_observation_task: observation_task,
_command_task: command_task,
_protocol_task: protocol_task,
})
}
@@ -311,12 +298,11 @@ impl BackendRuntimeClient {
impl Drop for BackendRuntimeClient {
fn drop(&mut self) {
self._observation_task.abort();
self._command_task.abort();
self._protocol_task.abort();
}
}
async fn run_worker_protocol_commands(
async fn run_worker_protocol_transport(
target: BackendRuntimeTarget,
mut commands: mpsc::UnboundedReceiver<Method>,
tx: mpsc::UnboundedSender<Event>,
@@ -402,81 +388,6 @@ async fn run_worker_protocol_commands(
}
}
async fn observe_worker_events(target: BackendRuntimeTarget, tx: mpsc::UnboundedSender<Event>) {
let mut attempts = 0_usize;
loop {
let url = observation_ws_url(&target);
match connect_async(&url).await {
Ok((mut ws, _)) => {
attempts = 0;
while let Some(frame) = ws.next().await {
match frame {
Ok(TungsteniteMessage::Text(text)) => {
match serde_json::from_str::<ClientWorkerEventWsFrame>(&text) {
Ok(ClientWorkerEventWsFrame::Event { envelope }) => {
if envelope.runtime_id != target.runtime_id
|| envelope.worker_id != target.worker_id
{
let _ = tx.send(diagnostic_event(format!(
"Backend observation frame target mismatch: got {}:{}, expected {}",
envelope.runtime_id,
envelope.worker_id,
target.display_label()
)));
continue;
}
let _ = tx.send(envelope.payload);
}
Ok(ClientWorkerEventWsFrame::Diagnostic { diagnostic }) => {
let message = format!(
"Backend observation diagnostic [{}]: {}",
diagnostic.code, diagnostic.message
);
let _ = tx.send(diagnostic_event(message));
}
Err(error) => {
let _ = tx.send(diagnostic_event(format!(
"Backend observation frame was not valid JSON: {error}"
)));
}
}
}
Ok(TungsteniteMessage::Close(_)) => break,
Ok(TungsteniteMessage::Ping(_))
| Ok(TungsteniteMessage::Pong(_))
| Ok(TungsteniteMessage::Binary(_))
| Ok(TungsteniteMessage::Frame(_)) => {}
Err(error) => {
let _ = tx.send(diagnostic_event(format!(
"Backend observation WebSocket error for {}: {error}",
target.display_label()
)));
break;
}
}
}
}
Err(error) => {
let _ = tx.send(diagnostic_event(format!(
"Backend observation WebSocket connect failed for {}: {error}",
target.display_label()
)));
}
}
attempts += 1;
if attempts > MAX_RECONNECT_ATTEMPTS {
let _ = tx.send(diagnostic_event(format!(
"Backend observation stream for {} stopped after {MAX_RECONNECT_ATTEMPTS} reconnect attempts",
target.display_label()
)));
break;
}
tokio::time::sleep(RECONNECT_DELAY).await;
}
}
fn diagnostic_event(message: impl Into<String>) -> Event {
Event::Error {
code: ErrorCode::Internal,
@@ -552,15 +463,6 @@ fn backend_runtime_workers_path(workspace_id: Option<&str>, runtime_id: &str) ->
}
}
fn observation_ws_url(target: &BackendRuntimeTarget) -> String {
let path = format!(
"/api/runtimes/{}/workers/{}/events/ws",
path_segment_encode(&target.runtime_id),
path_segment_encode(&target.worker_id)
);
join_base_and_path(&http_base_to_ws(&target.base_url), &path)
}
fn protocol_ws_url(target: &BackendRuntimeTarget) -> String {
let path = format!(
"/api/runtimes/{}/workers/{}/protocol/ws",
@@ -611,42 +513,14 @@ pub struct BackendDiagnostic {
pub message: String,
}
#[derive(Debug, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
enum ClientWorkerEventWsFrame {
Event {
envelope: ClientWorkerEventWsEnvelope,
},
Diagnostic {
diagnostic: ClientWorkerEventWsDiagnostic,
},
}
#[derive(Debug, Deserialize)]
struct ClientWorkerEventWsEnvelope {
runtime_id: String,
worker_id: String,
payload: Event,
}
#[derive(Debug, Deserialize)]
struct ClientWorkerEventWsDiagnostic {
code: String,
message: String,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn command_and_observation_urls_use_backend_protocol_paths() {
fn protocol_url_uses_backend_runtime_worker_identity() {
let target =
BackendRuntimeTarget::new("http://127.0.0.1:8787/", "runtime/one", "worker one");
assert_eq!(
observation_ws_url(&target),
"ws://127.0.0.1:8787/api/runtimes/runtime%2Fone/workers/worker%20one/events/ws"
);
assert_eq!(
protocol_ws_url(&target),
"ws://127.0.0.1:8787/api/runtimes/runtime%2Fone/workers/worker%20one/protocol/ws"
crates/worker-runtime/src/http_server.rs
+148 -342
View File
@@ -154,12 +154,10 @@ pub fn runtime_http_router(runtime: Runtime, local_token: Option<String>) -> Rou
.route("/v1/workers/{worker_id}/cancel", post(cancel_worker));
#[cfg(feature = "ws-server")]
let router = router
.route("/v1/workers/{worker_id}/events/ws", get(worker_events_ws))
.route(
"/v1/workers/{worker_id}/protocol/ws",
get(worker_protocol_ws),
);
let router = router.route(
"/v1/workers/{worker_id}/protocol/ws",
get(worker_protocol_ws),
);
router
.with_state(state.clone())
@@ -277,88 +275,12 @@ pub struct RuntimeHttpErrorDetail {
pub message: String,
}
/// Runtime-owned WebSocket frame for worker-scoped observation.
#[cfg(feature = "ws-server")]
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum RuntimeWorkerEventWsFrame {
Event {
envelope: RuntimeWorkerEventWsEnvelope,
},
Diagnostic {
diagnostic: RuntimeWorkerEventWsDiagnostic,
},
}
/// Runtime-local protocol event envelope.
#[cfg(feature = "ws-server")]
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RuntimeWorkerEventWsEnvelope {
pub cursor: String,
pub event_id: String,
pub worker_id: WorkerId,
pub payload: protocol::Event,
}
/// Runtime-local observation diagnostic.
#[cfg(feature = "ws-server")]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct RuntimeWorkerEventWsDiagnostic {
pub code: String,
pub message: String,
}
#[cfg(feature = "ws-server")]
#[derive(Clone, Debug, Default, Deserialize)]
struct RuntimeWorkerEventsWsQuery {
cursor: Option<String>,
}
#[cfg(feature = "ws-server")]
impl RuntimeWorkerEventWsFrame {
fn event(
cursor: String,
event_id: String,
worker_id: WorkerId,
payload: protocol::Event,
) -> Self {
Self::Event {
envelope: RuntimeWorkerEventWsEnvelope {
cursor,
event_id,
worker_id,
payload,
},
}
}
fn diagnostic(code: impl Into<String>, message: impl Into<String>) -> Self {
Self::Diagnostic {
diagnostic: RuntimeWorkerEventWsDiagnostic {
code: code.into(),
message: message.into(),
},
}
}
}
#[cfg(feature = "ws-server")]
async fn send_ws_frame(socket: &mut WebSocket, frame: &RuntimeWorkerEventWsFrame) -> bool {
match serde_json::to_string(frame) {
Ok(text) => socket.send(WsMessage::Text(text.into())).await.is_ok(),
Err(error) => {
let fallback = RuntimeWorkerEventWsFrame::diagnostic(
"runtime.serialize_failed",
format!("failed to serialize observation frame: {error}"),
);
let Ok(text) = serde_json::to_string(&fallback) else {
return false;
};
socket.send(WsMessage::Text(text.into())).await.is_ok()
}
}
}
type RestResult<T> = Result<Json<T>, RuntimeHttpRestError>;
async fn get_runtime(
@@ -515,6 +437,7 @@ async fn create_worker(
async fn worker_protocol_ws(
State(state): State<RuntimeHttpState>,
Path(worker_id): Path<String>,
Query(query): Query<RuntimeWorkerEventsWsQuery>,
ws: WebSocketUpgrade,
) -> Result<Response, RuntimeHttpRestError> {
let worker_ref = worker_ref_for(&state.runtime, worker_id)?;
@@ -523,7 +446,9 @@ async fn worker_protocol_ws(
.worker_detail(&worker_ref)
.map_err(RuntimeHttpRestError::runtime)?;
Ok(ws
.on_upgrade(move |socket| worker_protocol_ws_session(state.runtime, worker_ref, socket))
.on_upgrade(move |socket| {
worker_protocol_ws_session(state.runtime, worker_ref, query, socket)
})
.into_response())
}
@@ -531,45 +456,130 @@ async fn worker_protocol_ws(
async fn worker_protocol_ws_session(
runtime: Runtime,
worker_ref: WorkerRef,
query: RuntimeWorkerEventsWsQuery,
mut socket: WebSocket,
) {
while let Some(frame) = socket.next().await {
match frame {
Ok(WsMessage::Text(text)) => match serde_json::from_str::<protocol::Method>(&text) {
Ok(method) => match runtime.send_protocol_method(&worker_ref, method) {
Ok(events) => {
for event in events {
let mut cursor = match query.cursor.as_deref() {
Some(raw) => match WorkerObservationCursor::decode(raw) {
Some(cursor) => cursor,
None => {
let event =
protocol_error_event(format!("malformed worker observation cursor: {raw}"));
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
},
None => match runtime.worker_observation_cursor_now(&worker_ref) {
Ok(cursor) => cursor,
Err(error) => {
let event = protocol_error_event(error.to_string());
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
},
};
let mut receiver = match runtime.subscribe_worker_observation() {
Ok(receiver) => receiver,
Err(error) => {
let event =
protocol_error_event(format!("runtime observation bus unavailable: {error}"));
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
};
let snapshot = match runtime.worker_observation_snapshot(&worker_ref) {
Ok(snapshot) => snapshot,
Err(error) => {
let event = protocol_error_event(error.to_string());
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
};
if !send_protocol_event(&mut socket, &snapshot).await {
return;
}
match runtime.read_worker_observation_events(&worker_ref, cursor) {
Ok(backlog) => {
for event in backlog {
cursor = WorkerObservationCursor::new(event.sequence);
if !send_protocol_event(&mut socket, &event.payload).await {
return;
}
}
}
Err(error) => {
let event = protocol_error_event(error.to_string());
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
}
loop {
tokio::select! {
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(&worker_ref, method) {
Ok(events) => {
for event in events {
if !send_protocol_event(&mut socket, &event).await {
return;
}
}
}
Err(error) => {
let event = protocol_error_event(error.to_string());
if !send_protocol_event(&mut socket, &event).await {
return;
}
}
},
Err(error) => {
let event = protocol_error_event(format!(
"malformed protocol method frame: {error}"
));
if !send_protocol_event(&mut socket, &event).await {
return;
}
}
}
Err(error) => {
let event = protocol_error_event(error.to_string());
if !send_protocol_event(&mut socket, &event).await {
},
Some(Ok(WsMessage::Close(_))) | None => return,
Some(Ok(WsMessage::Ping(payload))) => {
if socket.send(WsMessage::Pong(payload)).await.is_err() {
return;
}
}
},
Err(error) => {
let event =
protocol_error_event(format!("malformed protocol method frame: {error}"));
if !send_protocol_event(&mut socket, &event).await {
Some(Ok(WsMessage::Pong(_))) | Some(Ok(WsMessage::Binary(_))) => {}
Some(Err(error)) => {
let event = protocol_error_event(format!("protocol WebSocket error: {error}"));
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
}
},
Ok(WsMessage::Close(_)) => return,
Ok(WsMessage::Ping(payload)) => {
if socket.send(WsMessage::Pong(payload)).await.is_err() {
return;
}
}
Ok(WsMessage::Pong(_)) | Ok(WsMessage::Binary(_)) => {}
Err(error) => {
let event = protocol_error_event(format!("protocol WebSocket error: {error}"));
let _ = send_protocol_event(&mut socket, &event).await;
return;
event = receiver.recv() => {
match event {
Ok(event) if event.worker_ref == worker_ref && event.sequence > cursor.sequence => {
cursor = WorkerObservationCursor::new(event.sequence);
if !send_protocol_event(&mut socket, &event.payload).await {
return;
}
}
Ok(_) => {}
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
let event = protocol_error_event("runtime observation backlog was overrun");
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
let event = protocol_error_event("runtime observation bus closed");
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
}
}
}
}
@@ -599,182 +609,6 @@ fn protocol_error_event(message: impl Into<String>) -> protocol::Event {
}
}
#[cfg(feature = "ws-server")]
async fn worker_events_ws(
State(state): State<RuntimeHttpState>,
Path(worker_id): Path<String>,
Query(query): Query<RuntimeWorkerEventsWsQuery>,
ws: WebSocketUpgrade,
) -> Result<Response, RuntimeHttpRestError> {
let worker_ref = worker_ref_for(&state.runtime, worker_id)?;
state
.runtime
.worker_detail(&worker_ref)
.map_err(RuntimeHttpRestError::runtime)?;
Ok(ws
.on_upgrade(move |socket| {
worker_events_ws_session(state.runtime, worker_ref, query, socket)
})
.into_response())
}
#[cfg(feature = "ws-server")]
async fn worker_events_ws_session(
runtime: Runtime,
worker_ref: WorkerRef,
query: RuntimeWorkerEventsWsQuery,
mut socket: WebSocket,
) {
let mut cursor = match query.cursor.as_deref() {
Some(raw) => match WorkerObservationCursor::decode(raw) {
Some(cursor) => cursor,
None => {
let frame = RuntimeWorkerEventWsFrame::diagnostic(
"runtime.cursor_malformed",
format!("malformed worker observation cursor: {raw}"),
);
let _ = send_ws_frame(&mut socket, &frame).await;
return;
}
},
None => match runtime.worker_observation_cursor_now(&worker_ref) {
Ok(cursor) => cursor,
Err(error) => {
let frame = RuntimeWorkerEventWsFrame::diagnostic(
"runtime.worker_not_found",
error.to_string(),
);
let _ = send_ws_frame(&mut socket, &frame).await;
return;
}
},
};
let mut receiver = match runtime.subscribe_worker_observation() {
Ok(receiver) => receiver,
Err(error) => {
let frame = RuntimeWorkerEventWsFrame::diagnostic(
"runtime.unavailable",
format!("runtime observation bus unavailable: {error}"),
);
let _ = send_ws_frame(&mut socket, &frame).await;
return;
}
};
let snapshot = match runtime.worker_observation_snapshot(&worker_ref) {
Ok(snapshot) => snapshot,
Err(error) => {
let frame = RuntimeWorkerEventWsFrame::diagnostic(
"runtime.worker_not_found",
error.to_string(),
);
let _ = send_ws_frame(&mut socket, &frame).await;
return;
}
};
let snapshot_cursor = cursor.encode();
let snapshot_frame = RuntimeWorkerEventWsFrame::event(
snapshot_cursor.clone(),
format!("snapshot:{snapshot_cursor}"),
worker_ref.worker_id.clone(),
snapshot,
);
if !send_ws_frame(&mut socket, &snapshot_frame).await {
return;
}
match runtime.read_worker_observation_events(&worker_ref, cursor) {
Ok(backlog) => {
for event in backlog {
cursor = WorkerObservationCursor::new(event.sequence);
let frame = RuntimeWorkerEventWsFrame::event(
event.cursor,
event.event_id,
event.worker_ref.worker_id,
event.payload,
);
if !send_ws_frame(&mut socket, &frame).await {
return;
}
}
}
Err(error) => {
let frame = RuntimeWorkerEventWsFrame::diagnostic(
"runtime.cursor_unknown_or_expired",
error.to_string(),
);
let _ = send_ws_frame(&mut socket, &frame).await;
return;
}
}
loop {
tokio::select! {
inbound = socket.next() => {
match inbound {
Some(Ok(WsMessage::Close(_))) | None => return,
Some(Ok(WsMessage::Ping(payload))) => {
if socket.send(WsMessage::Pong(payload)).await.is_err() {
return;
}
}
Some(Ok(WsMessage::Pong(_))) => {}
Some(Ok(_)) => {
let frame = RuntimeWorkerEventWsFrame::diagnostic(
"runtime.observation_only",
"runtime worker event WebSocket is observation-only",
);
let _ = send_ws_frame(&mut socket, &frame).await;
return;
}
Some(Err(error)) => {
let frame = RuntimeWorkerEventWsFrame::diagnostic(
"runtime.websocket_error",
format!("runtime WebSocket receive error: {error}"),
);
let _ = send_ws_frame(&mut socket, &frame).await;
return;
}
}
}
event = receiver.recv() => {
match event {
Ok(event) if event.worker_ref == worker_ref && event.sequence > cursor.sequence => {
cursor = WorkerObservationCursor::new(event.sequence);
let frame = RuntimeWorkerEventWsFrame::event(
event.cursor,
event.event_id,
event.worker_ref.worker_id,
event.payload,
);
if !send_ws_frame(&mut socket, &frame).await {
return;
}
}
Ok(_) => {}
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
let frame = RuntimeWorkerEventWsFrame::diagnostic(
"runtime.cursor_expired",
"runtime observation backlog was overrun",
);
let _ = send_ws_frame(&mut socket, &frame).await;
return;
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
let frame = RuntimeWorkerEventWsFrame::diagnostic(
"runtime.upstream_closed",
"runtime observation bus closed",
);
let _ = send_ws_frame(&mut socket, &frame).await;
return;
}
}
}
}
}
}
async fn send_worker_input(
State(state): State<RuntimeHttpState>,
Path(worker_id): Path<String>,
@@ -1409,7 +1243,7 @@ mod ws_tests {
runtime,
worker.worker_ref.clone(),
format!(
"ws://{addr}/v1/workers/{}/events/ws",
"ws://{addr}/v1/workers/{}/protocol/ws",
worker.worker_ref.worker_id
),
)
@@ -1419,7 +1253,7 @@ mod ws_tests {
stream: &mut tokio_tungstenite::WebSocketStream<
tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
>,
) -> RuntimeWorkerEventWsFrame {
) -> protocol::Event {
let message = stream.next().await.unwrap().unwrap();
let Message::Text(text) = message else {
panic!("expected text frame");
@@ -1428,21 +1262,16 @@ mod ws_tests {
}
#[tokio::test]
async fn runtime_ws_connect_sends_snapshot_and_live_worker_events() {
async fn protocol_ws_connect_sends_snapshot_and_live_worker_events() {
let (runtime, worker_ref, url) = spawn_runtime_server().await;
let (mut stream, _) = connect_async(&url).await.unwrap();
match next_frame(&mut stream).await {
RuntimeWorkerEventWsFrame::Event { envelope } => {
assert_eq!(envelope.worker_id, worker_ref.worker_id);
assert!(matches!(envelope.payload, protocol::Event::Snapshot { .. }));
}
RuntimeWorkerEventWsFrame::Diagnostic { diagnostic } => {
panic!("unexpected diagnostic: {diagnostic:?}");
}
}
assert!(matches!(
next_frame(&mut stream).await,
protocol::Event::Snapshot { .. }
));
let stored = runtime
runtime
.observe_worker_event(
&worker_ref,
protocol::Event::TextDelta {
@@ -1450,23 +1279,14 @@ mod ws_tests {
},
)
.unwrap();
match next_frame(&mut stream).await {
RuntimeWorkerEventWsFrame::Event { envelope } => {
assert_eq!(envelope.worker_id, worker_ref.worker_id);
assert_eq!(envelope.cursor, stored.cursor);
assert!(matches!(
envelope.payload,
protocol::Event::TextDelta { .. }
));
}
RuntimeWorkerEventWsFrame::Diagnostic { diagnostic } => {
panic!("unexpected diagnostic: {diagnostic:?}");
}
}
assert!(matches!(
next_frame(&mut stream).await,
protocol::Event::TextDelta { .. }
));
}
#[tokio::test]
async fn runtime_ws_cursor_resume_is_duplicate_safe_and_filters_workers() {
async fn protocol_ws_cursor_resume_is_duplicate_safe_and_filters_workers() {
let (runtime, worker_ref, url) = spawn_runtime_server().await;
let other = runtime.create_worker(ws_create_request()).unwrap();
let first = runtime
@@ -1481,7 +1301,7 @@ mod ws_tests {
.observe_worker_event(
&other.worker_ref,
protocol::Event::TextDelta {
text: "started".into(),
text: "other".into(),
},
)
.unwrap();
@@ -1491,10 +1311,10 @@ mod ws_tests {
.unwrap();
assert!(matches!(
next_frame(&mut stream).await,
RuntimeWorkerEventWsFrame::Event { envelope } if matches!(envelope.payload, protocol::Event::Snapshot { .. })
protocol::Event::Snapshot { .. }
));
let second = runtime
runtime
.observe_worker_event(
&worker_ref,
protocol::Event::TextDone {
@@ -1502,41 +1322,27 @@ mod ws_tests {
},
)
.unwrap();
match next_frame(&mut stream).await {
RuntimeWorkerEventWsFrame::Event { envelope } => {
assert_eq!(envelope.cursor, second.cursor);
assert_ne!(envelope.cursor, first.cursor);
assert!(matches!(envelope.payload, protocol::Event::TextDone { .. }));
}
RuntimeWorkerEventWsFrame::Diagnostic { diagnostic } => {
panic!("unexpected diagnostic: {diagnostic:?}");
}
}
assert!(matches!(
next_frame(&mut stream).await,
protocol::Event::TextDone { .. }
));
}
#[tokio::test]
async fn runtime_ws_reports_malformed_cursor_and_observation_only_input() {
async fn protocol_ws_reports_malformed_cursor_and_method_frame() {
let (_runtime, _worker_ref, url) = spawn_runtime_server().await;
let (mut malformed, _) = connect_async(format!("{url}?cursor=bad")).await.unwrap();
match next_frame(&mut malformed).await {
RuntimeWorkerEventWsFrame::Diagnostic { diagnostic } => {
assert_eq!(diagnostic.code, "runtime.cursor_malformed");
}
RuntimeWorkerEventWsFrame::Event { envelope } => {
panic!("unexpected event: {envelope:?}");
}
}
assert!(matches!(
next_frame(&mut malformed).await,
protocol::Event::Error { .. }
));
let (mut stream, _) = connect_async(&url).await.unwrap();
let _ = next_frame(&mut stream).await;
stream.send(Message::Text("{}".into())).await.unwrap();
match next_frame(&mut stream).await {
RuntimeWorkerEventWsFrame::Diagnostic { diagnostic } => {
assert_eq!(diagnostic.code, "runtime.observation_only");
}
RuntimeWorkerEventWsFrame::Event { envelope } => {
panic!("unexpected event: {envelope:?}");
}
}
assert!(matches!(
next_frame(&mut stream).await,
protocol::Event::Error { .. }
));
}
}
crates/workspace-server/src/hosts.rs
+2 -2
View File
@@ -2128,7 +2128,7 @@ impl RemoteWorkerRuntime {
} else if let Some(rest) = base.strip_prefix("http://") {
base = format!("ws://{rest}");
}
format!("{base}/v1/workers/{worker_id}/events/ws")
format!("{base}/v1/workers/{worker_id}/protocol/ws")
}
fn get_json<T>(&self, path: &str) -> Result<T, RuntimeDiagnostic>
@@ -4285,7 +4285,7 @@ mod tests {
panic!("remote runtime should expose a remote WS observation source");
};
assert!(observation.endpoint.starts_with("ws://127.0.0.1:"));
assert!(observation.endpoint.ends_with("/v1/workers/1/events/ws"));
assert!(observation.endpoint.ends_with("/v1/workers/1/protocol/ws"));
assert_eq!(observation.bearer_token.as_deref(), Some(secret.as_str()));
let workers = registry.list_workers(10);
crates/workspace-server/src/observation.rs
+15 -106
View File
@@ -10,7 +10,6 @@ use serde::{Deserialize, Serialize};
use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::{Error as TungsteniteError, Message as TungsteniteMessage};
use worker_runtime::http_server::{RuntimeWorkerEventWsEnvelope, RuntimeWorkerEventWsFrame};
/// Backend-private source for a runtime worker observation stream.
#[derive(Clone, PartialEq, Eq)]
@@ -104,42 +103,12 @@ pub struct RuntimeObservationUpstreamEvent {
pub payload: protocol::Event,
}
/// Backend-local frame exposed to browser/future-TUI clients.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum ClientWorkerEventWsFrame {
Event {
envelope: ClientWorkerEventWsEnvelope,
},
Diagnostic {
diagnostic: ClientWorkerEventWsDiagnostic,
},
}
/// Backend-owned event envelope. It intentionally omits Runtime endpoints,
/// credentials, sockets and session paths.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ClientWorkerEventWsEnvelope {
pub event_id: String,
pub runtime_id: String,
pub worker_id: String,
pub payload: protocol::Event,
}
/// Client-facing typed observation diagnostic.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ClientWorkerEventWsDiagnostic {
pub code: String,
pub message: String,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ObservationProxyError {
RuntimeUnavailable(String),
WorkerNotFound(String),
UpstreamDisconnect(String),
MalformedFrame(String),
ObservationOnly,
}
impl ObservationProxyError {
@@ -149,7 +118,6 @@ impl ObservationProxyError {
ObservationProxyError::WorkerNotFound(_) => "backend.worker_not_found",
ObservationProxyError::UpstreamDisconnect(_) => "backend.upstream_disconnect",
ObservationProxyError::MalformedFrame(_) => "backend.malformed_frame",
ObservationProxyError::ObservationOnly => "backend.observation_only",
}
}
@@ -159,24 +127,6 @@ impl ObservationProxyError {
| ObservationProxyError::WorkerNotFound(message)
| ObservationProxyError::UpstreamDisconnect(message)
| ObservationProxyError::MalformedFrame(message) => message,
ObservationProxyError::ObservationOnly => {
"backend worker event WebSocket is observation-only"
}
}
}
}
impl ClientWorkerEventWsFrame {
pub fn event(envelope: ClientWorkerEventWsEnvelope) -> Self {
Self::Event { envelope }
}
pub fn diagnostic(error: ObservationProxyError) -> Self {
Self::Diagnostic {
diagnostic: ClientWorkerEventWsDiagnostic {
code: error.code().to_string(),
message: error.message().to_string(),
},
}
}
}
@@ -238,15 +188,6 @@ impl BackendObservationProxy {
))
})
}
pub fn map_event(&self, event: RuntimeObservationUpstreamEvent) -> ClientWorkerEventWsEnvelope {
ClientWorkerEventWsEnvelope {
event_id: event.runtime_event_id,
runtime_id: event.runtime_id,
worker_id: event.worker_id,
payload: event.payload,
}
}
}
fn map_runtime_connect_error(error: TungsteniteError) -> ObservationProxyError {
@@ -266,24 +207,6 @@ fn map_runtime_connect_error(error: TungsteniteError) -> ObservationProxyError {
}
}
fn map_runtime_diagnostic(code: String, message: String) -> ObservationProxyError {
match code.as_str() {
"runtime.worker_not_found" => ObservationProxyError::WorkerNotFound(message),
"runtime.cursor_malformed"
| "runtime.cursor_unknown_or_expired"
| "runtime.cursor_expired" => ObservationProxyError::RuntimeUnavailable(message),
"runtime.unavailable" => ObservationProxyError::RuntimeUnavailable(message),
"runtime.upstream_closed" | "runtime.websocket_error" => {
ObservationProxyError::UpstreamDisconnect(message)
}
"runtime.serialize_failed" => ObservationProxyError::MalformedFrame(message),
"runtime.observation_only" => ObservationProxyError::ObservationOnly,
_ => ObservationProxyError::RuntimeUnavailable(format!(
"runtime diagnostic {code}: {message}"
)),
}
}
pub struct RuntimeWsObservationClient {
runtime_id: String,
worker_id: String,
@@ -361,32 +284,17 @@ impl RuntimeWsObservationClient {
));
}
};
let frame: RuntimeWorkerEventWsFrame =
serde_json::from_str(&text).map_err(|error| {
ObservationProxyError::MalformedFrame(format!(
"failed to decode runtime observation frame: {error}"
))
})?;
match frame {
RuntimeWorkerEventWsFrame::Event { envelope } => {
return Ok(self.map_envelope(envelope));
}
RuntimeWorkerEventWsFrame::Diagnostic { diagnostic } => {
return Err(map_runtime_diagnostic(diagnostic.code, diagnostic.message));
}
}
}
}
fn map_envelope(
&self,
envelope: RuntimeWorkerEventWsEnvelope,
) -> RuntimeObservationUpstreamEvent {
RuntimeObservationUpstreamEvent {
runtime_id: self.runtime_id.clone(),
worker_id: self.worker_id.clone(),
runtime_event_id: envelope.event_id,
payload: envelope.payload,
let payload: protocol::Event = serde_json::from_str(&text).map_err(|error| {
ObservationProxyError::MalformedFrame(format!(
"failed to decode runtime protocol event frame: {error}"
))
})?;
return Ok(RuntimeObservationUpstreamEvent {
runtime_id: self.runtime_id.clone(),
worker_id: self.worker_id.clone(),
runtime_event_id: "protocol".to_string(),
payload,
});
}
}
}
@@ -548,7 +456,8 @@ mod tests {
RuntimeObservationSourceConfig {
runtime_id: "remote-runtime".to_string(),
worker_id: "worker-1".to_string(),
endpoint: "wss://remote.example.invalid/private/workers/worker-1/events/ws".to_string(),
endpoint: "wss://remote.example.invalid/private/workers/worker-1/protocol/ws"
.to_string(),
bearer_token: Some("top-secret-bearer-token".to_string()),
}
}
@@ -563,7 +472,7 @@ mod tests {
assert!(debug.contains("<redacted>"));
for forbidden in [
"remote.example.invalid",
"/private/workers/worker-1/events/ws",
"/private/workers/worker-1/protocol/ws",
"top-secret-bearer-token",
] {
assert!(
@@ -582,7 +491,7 @@ mod tests {
assert!(debug.contains("source_count"));
for forbidden in [
"remote.example.invalid",
"/private/workers/worker-1/events/ws",
"/private/workers/worker-1/protocol/ws",
"top-secret-bearer-token",
] {
assert!(
crates/workspace-server/src/server.rs
+103 -305
View File
@@ -41,8 +41,8 @@ use crate::hosts::{
};
use crate::identity::WorkspaceIdentity;
use crate::observation::{
BackendObservationProxy, ClientWorkerEventWsFrame, ObservationProxyError,
RuntimeObservationClient, RuntimeObservationSourceConfig,
BackendObservationProxy, ObservationProxyError, RuntimeObservationClient,
RuntimeObservationSourceConfig,
};
use crate::profile_settings::{
CreateWorkspaceProfileSourceRequest, DeleteWorkspaceProfileSourceRequest,
@@ -547,14 +547,6 @@ pub fn build_router(api: WorkspaceApi) -> Router {
"/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/cancel",
post(scoped_cancel_runtime_worker),
)
.route(
"/api/runtimes/{runtime_id}/workers/{worker_id}/events/ws",
get(worker_observation_ws),
)
.route(
"/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/events/ws",
get(scoped_worker_observation_ws),
)
.route(
"/api/runtimes/{runtime_id}/workers/{worker_id}/protocol/ws",
get(worker_protocol_ws),
@@ -2452,19 +2444,6 @@ async fn scoped_cancel_runtime_worker(
.await
}
async fn scoped_worker_observation_ws(
ws: WebSocketUpgrade,
State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedRuntimeWorkerPath>,
) -> Response {
if let Err(err) = validate_workspace_scope(&api, &path.workspace_id) {
return err.into_response();
}
worker_observation_ws(State(api), AxumPath((path.runtime_id, path.worker_id)), ws)
.await
.into_response()
}
async fn scoped_worker_protocol_ws(
ws: WebSocketUpgrade,
State(api): State<WorkspaceApi>,
@@ -3414,57 +3393,101 @@ async fn worker_protocol_ws(
AxumPath((runtime_id, worker_id)): AxumPath<(String, String)>,
ws: WebSocketUpgrade,
) -> impl IntoResponse {
match api.runtime.worker(&runtime_id, &worker_id) {
Ok(_) => ws.on_upgrade(move |socket| {
worker_protocol_ws_session(api.runtime, runtime_id, worker_id, socket)
}),
Err(error) => ApiError::from(error.into_error()).into_response(),
}
let source = match api.observation_proxy.source(&runtime_id, &worker_id) {
Ok(source) => source,
Err(ObservationProxyError::WorkerNotFound(_)) => {
match api.runtime.observation_source(&runtime_id, &worker_id) {
Ok(source) => source,
Err(error) => return ApiError::from(error.into_error()).into_response(),
}
}
Err(error) => {
return (
StatusCode::BAD_REQUEST,
Json(serde_json::json!({
"error": error.code(),
"message": error.message(),
})),
)
.into_response();
}
};
ws.on_upgrade(move |socket| {
worker_protocol_ws_session(api.runtime, source, runtime_id, worker_id, socket)
})
}
async fn worker_protocol_ws_session(
runtime: Arc<RuntimeRegistry>,
source: crate::observation::RuntimeObservationSource,
runtime_id: String,
worker_id: String,
mut socket: WebSocket,
) {
while let Some(frame) = socket.next().await {
match frame {
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(events) => {
for event in events {
let mut upstream = match RuntimeObservationClient::connect(&source).await {
Ok(client) => client,
Err(error) => {
let event = protocol_error_event(error.message());
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
};
loop {
tokio::select! {
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(events) => {
for event in events {
if !send_protocol_event(&mut socket, &event).await {
return;
}
}
}
Err(error) => {
let event = protocol_error_event(error.message());
if !send_protocol_event(&mut socket, &event).await {
return;
}
}
},
Err(error) => {
let event =
protocol_error_event(format!("malformed protocol method frame: {error}"));
if !send_protocol_event(&mut socket, &event).await {
return;
}
}
}
Err(error) => {
let event = protocol_error_event(error.message());
if !send_protocol_event(&mut socket, &event).await {
},
Some(Ok(WsMessage::Close(_))) | None => return,
Some(Ok(WsMessage::Ping(payload))) => {
if socket.send(WsMessage::Pong(payload)).await.is_err() {
return;
}
}
},
Err(error) => {
let event =
protocol_error_event(format!("malformed protocol method frame: {error}"));
if !send_protocol_event(&mut socket, &event).await {
Some(Ok(WsMessage::Pong(_))) | Some(Ok(WsMessage::Binary(_))) => {}
Some(Err(error)) => {
let event = protocol_error_event(format!("protocol WebSocket error: {error}"));
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
}
},
Ok(WsMessage::Close(_)) => return,
Ok(WsMessage::Ping(payload)) => {
if socket.send(WsMessage::Pong(payload)).await.is_err() {
return;
}
}
Ok(WsMessage::Pong(_)) | Ok(WsMessage::Binary(_)) => {}
Err(error) => {
let event = protocol_error_event(format!("protocol WebSocket error: {error}"));
let _ = send_protocol_event(&mut socket, &event).await;
return;
upstream_event = upstream.next_event() => {
match upstream_event {
Ok(event) => {
if !send_protocol_event(&mut socket, &event.payload).await {
return;
}
}
Err(error) => {
let event = protocol_error_event(error.message());
let _ = send_protocol_event(&mut socket, &event).await;
return;
}
}
}
}
}
@@ -3492,116 +3515,6 @@ fn protocol_error_event(message: impl Into<String>) -> protocol::Event {
}
}
async fn worker_observation_ws(
State(api): State<WorkspaceApi>,
AxumPath((runtime_id, worker_id)): AxumPath<(String, String)>,
ws: WebSocketUpgrade,
) -> impl IntoResponse {
match api.observation_proxy.source(&runtime_id, &worker_id) {
Ok(source) => ws.on_upgrade(move |socket| {
worker_observation_ws_session(api.observation_proxy, source, socket)
}),
Err(ObservationProxyError::WorkerNotFound(_)) => {
match api.runtime.observation_source(&runtime_id, &worker_id) {
Ok(source) => ws.on_upgrade(move |socket| {
worker_observation_ws_session(api.observation_proxy, source, socket)
}),
Err(error) => ApiError::from(error.into_error()).into_response(),
}
}
Err(error) => {
let status = StatusCode::BAD_REQUEST;
(
status,
Json(serde_json::json!({
"error": error.code(),
"message": error.message(),
})),
)
.into_response()
}
}
}
async fn worker_observation_ws_session(
proxy: BackendObservationProxy,
source: crate::observation::RuntimeObservationSource,
mut socket: WebSocket,
) {
let mut upstream = match RuntimeObservationClient::connect(&source).await {
Ok(client) => client,
Err(error) => {
let _ = send_client_ws_frame(&mut socket, ClientWorkerEventWsFrame::diagnostic(error))
.await;
return;
}
};
loop {
tokio::select! {
inbound = socket.next() => {
match inbound {
Some(Ok(WsMessage::Close(_))) | None => return,
Some(Ok(WsMessage::Ping(payload))) => {
if socket.send(WsMessage::Pong(payload)).await.is_err() {
return;
}
}
Some(Ok(WsMessage::Pong(_))) => {}
Some(Ok(_)) => {
let _ = send_client_ws_frame(
&mut socket,
ClientWorkerEventWsFrame::diagnostic(ObservationProxyError::ObservationOnly),
).await;
return;
}
Some(Err(error)) => {
let _ = send_client_ws_frame(
&mut socket,
ClientWorkerEventWsFrame::diagnostic(
ObservationProxyError::MalformedFrame(format!(
"client WebSocket receive error: {error}"
)),
),
).await;
return;
}
}
}
upstream_event = upstream.next_event() => {
match upstream_event {
Ok(event) => {
let envelope = proxy.map_event(event);
if !send_client_ws_frame(&mut socket, ClientWorkerEventWsFrame::event(envelope)).await {
return;
}
},
Err(error) => {
let _ = send_client_ws_frame(&mut socket, ClientWorkerEventWsFrame::diagnostic(error)).await;
return;
}
}
}
}
}
}
async fn send_client_ws_frame(socket: &mut WebSocket, frame: ClientWorkerEventWsFrame) -> bool {
match serde_json::to_string(&frame) {
Ok(text) => socket.send(WsMessage::Text(text.into())).await.is_ok(),
Err(error) => {
let fallback =
ClientWorkerEventWsFrame::diagnostic(ObservationProxyError::MalformedFrame(
format!("failed to serialize backend observation frame: {error}"),
));
let Ok(text) = serde_json::to_string(&fallback) else {
return false;
};
socket.send(WsMessage::Text(text.into())).await.is_ok()
}
}
}
async fn list_host_workers(
State(api): State<WorkspaceApi>,
AxumPath(host_id): AxumPath<String>,
@@ -5436,7 +5349,7 @@ mod tests {
use super::*;
use axum::body::{Body, to_bytes};
use axum::http::Request;
use futures::{SinkExt, StreamExt};
use futures::StreamExt;
use serde_json::{Value, json};
use std::{fs, sync::Arc};
use tokio_tungstenite::connect_async;
@@ -5449,7 +5362,6 @@ mod tests {
TicketWorkerRole, WorkerInputKind, WorkerOperationState, WorkerSpawnAcceptanceRequirement,
WorkerSpawnIntent,
};
use crate::observation::ClientWorkerEventWsDiagnostic;
use crate::store::SqliteWorkspaceStore;
const TEST_WORKSPACE_ID: &str = "0192f0e8-4d84-7d6e-a000-000000000001";
@@ -7861,53 +7773,20 @@ mod tests {
}
#[tokio::test]
async fn proxies_worker_observation_ws_with_backend_cursors_and_diagnostics() {
let (runtime, worker_ref) = runtime_with_worker();
let runtime_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let runtime_addr = runtime_listener.local_addr().unwrap();
tokio::spawn({
let runtime = runtime.clone();
async move {
worker_runtime::http_server::serve_runtime_http(runtime, runtime_listener, None)
.await
.unwrap()
}
});
let dir = tempfile::tempdir().unwrap();
let store = SqliteWorkspaceStore::in_memory().unwrap();
let mut config = test_server_config(dir.path());
config
.runtime_event_sources
.push(RuntimeObservationSourceConfig {
runtime_id: "runtime-a".into(),
worker_id: "worker-a".into(),
endpoint: format!(
"ws://{runtime_addr}/v1/workers/{}/events/ws",
worker_ref.worker_id
),
bearer_token: None,
});
let api = WorkspaceApi::new_with_execution_backend(
config,
Arc::new(store),
Arc::new(DeterministicExecutionBackend::default()),
)
.await
.unwrap();
let app_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let app_addr = app_listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(app_listener, build_router(api)).await.unwrap() });
let url = format!("ws://{app_addr}/api/runtimes/runtime-a/workers/worker-a/events/ws");
let (mut stream, _) = connect_async(&url).await.unwrap();
let snapshot = next_client_frame(&mut stream).await;
let ClientWorkerEventWsFrame::Event { envelope: snapshot } = snapshot else {
panic!("expected snapshot event");
async fn proxies_worker_protocol_ws_as_raw_events() {
let (runtime, worker_ref, endpoint) = spawn_runtime_worker().await;
let source = RuntimeObservationSourceConfig {
runtime_id: "runtime-a".into(),
worker_id: "worker-a".into(),
endpoint,
bearer_token: None,
};
assert_eq!(snapshot.runtime_id, "runtime-a");
assert_eq!(snapshot.worker_id, "worker-a");
assert!(matches!(snapshot.payload, protocol::Event::Snapshot { .. }));
let (url, _dir) = spawn_workspace_proxy(source).await;
let (mut stream, _) = connect_async(&url).await.unwrap();
assert!(matches!(
next_client_frame(&mut stream).await,
protocol::Event::Snapshot { .. }
));
runtime
.observe_worker_event(
@@ -7917,59 +7796,16 @@ mod tests {
},
)
.unwrap();
let live = next_client_frame(&mut stream).await;
let ClientWorkerEventWsFrame::Event { envelope: live } = live else {
panic!("expected live event");
};
assert_eq!(live.runtime_id, "runtime-a");
assert_eq!(live.worker_id, "worker-a");
assert!(matches!(live.payload, protocol::Event::TextDelta { .. }));
let (mut fresh, _) = connect_async(&url).await.unwrap();
let fresh_snapshot = next_client_frame(&mut fresh).await;
assert!(matches!(
fresh_snapshot,
ClientWorkerEventWsFrame::Event { envelope } if matches!(envelope.payload, protocol::Event::Snapshot { .. })
next_client_frame(&mut stream).await,
protocol::Event::TextDelta { .. }
));
runtime
.observe_worker_event(
&worker_ref,
protocol::Event::TextDone {
text: "fresh".into(),
},
)
.unwrap();
let fresh_event = next_client_frame(&mut fresh).await;
assert!(matches!(
fresh_event,
ClientWorkerEventWsFrame::Event { envelope } if matches!(envelope.payload, protocol::Event::TextDone { .. })
));
let (mut query_stream, _) = connect_async(format!("{url}?cursor=bad")).await.unwrap();
let query_snapshot = next_client_frame(&mut query_stream).await;
assert!(matches!(
query_snapshot,
ClientWorkerEventWsFrame::Event { envelope } if matches!(envelope.payload, protocol::Event::Snapshot { .. })
));
stream.send(Message::Text("{}".into())).await.unwrap();
let mut saw_observation_only = false;
for _ in 0..3 {
if let ClientWorkerEventWsFrame::Diagnostic { diagnostic } =
next_client_frame(&mut stream).await
{
assert_eq!(diagnostic.code, "backend.observation_only");
saw_observation_only = true;
break;
}
}
assert!(saw_observation_only, "expected observation-only diagnostic");
}
#[tokio::test]
async fn proxy_maps_runtime_worker_not_found_http_404_to_typed_backend_diagnostic() {
async fn proxy_maps_runtime_worker_not_found_http_404_to_protocol_error_event() {
let (_runtime, _worker_ref, endpoint) = spawn_runtime_worker().await;
let endpoint = endpoint.replace("/events/ws", "/missing-worker/events/ws");
let endpoint = endpoint.replace("/protocol/ws", "/missing-worker/protocol/ws");
let source = RuntimeObservationSourceConfig {
runtime_id: "runtime-a".into(),
worker_id: "worker-a".into(),
@@ -7978,30 +7814,17 @@ mod tests {
};
let (url, _dir) = spawn_workspace_proxy(source).await;
let (mut stream, _) = connect_async(&url).await.unwrap();
let diagnostic = next_client_diagnostic(&mut stream).await;
assert_eq!(diagnostic.code, "backend.worker_not_found");
}
#[tokio::test]
async fn proxy_reports_actual_upstream_disconnect_separately() {
let endpoint = spawn_closing_runtime_ws().await;
let source = RuntimeObservationSourceConfig {
runtime_id: "runtime-a".into(),
worker_id: "worker-a".into(),
endpoint,
bearer_token: None,
};
let (url, _dir) = spawn_workspace_proxy(source).await;
let (mut stream, _) = connect_async(&url).await.unwrap();
let diagnostic = next_client_diagnostic(&mut stream).await;
assert_eq!(diagnostic.code, "backend.upstream_disconnect");
assert!(matches!(
next_client_frame(&mut stream).await,
protocol::Event::Error { .. }
));
}
async fn next_client_frame(
stream: &mut tokio_tungstenite::WebSocketStream<
tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
>,
) -> ClientWorkerEventWsFrame {
) -> protocol::Event {
let message = stream.next().await.unwrap().unwrap();
let Message::Text(text) = message else {
panic!("expected text frame");
@@ -8009,19 +7832,6 @@ mod tests {
serde_json::from_str(&text).unwrap()
}
async fn next_client_diagnostic(
stream: &mut tokio_tungstenite::WebSocketStream<
tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
>,
) -> ClientWorkerEventWsDiagnostic {
match next_client_frame(stream).await {
ClientWorkerEventWsFrame::Diagnostic { diagnostic } => diagnostic,
ClientWorkerEventWsFrame::Event { envelope } => {
panic!("expected diagnostic, got event: {envelope:?}")
}
}
}
async fn spawn_runtime_worker() -> (
worker_runtime::Runtime,
worker_runtime::identity::WorkerRef,
@@ -8039,7 +7849,7 @@ mod tests {
}
});
let endpoint = format!(
"ws://{runtime_addr}/v1/workers/{}/events/ws",
"ws://{runtime_addr}/v1/workers/{}/protocol/ws",
worker_ref.worker_id
);
(runtime, worker_ref, endpoint)
@@ -8065,22 +7875,10 @@ mod tests {
let app_addr = app_listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(app_listener, build_router(api)).await.unwrap() });
(
format!("ws://{app_addr}/api/runtimes/{runtime_id}/workers/{worker_id}/events/ws"),
format!("ws://{app_addr}/api/runtimes/{runtime_id}/workers/{worker_id}/protocol/ws"),
dir,
)
}
async fn spawn_closing_runtime_ws() -> String {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = tokio_tungstenite::accept_async(stream).await.unwrap();
let _ = websocket.close(None).await;
});
format!("ws://{addr}/events/ws")
}
async fn get_json(app: Router, uri: &str) -> Value {
let response = app
.oneshot(Request::builder().uri(uri).body(Body::empty()).unwrap())