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/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())