feat: add worker observation websocket proxy
This commit is contained in:
@@ -6,6 +6,7 @@
|
||||
|
||||
pub mod hosts;
|
||||
pub mod identity;
|
||||
pub mod observation;
|
||||
pub mod records;
|
||||
pub mod repositories;
|
||||
pub mod server;
|
||||
|
||||
@@ -0,0 +1,440 @@
|
||||
use std::collections::{BTreeMap, VecDeque};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio_tungstenite::connect_async;
|
||||
use tokio_tungstenite::tungstenite::Message as TungsteniteMessage;
|
||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
use worker_runtime::http_server::{RuntimeWorkerEventWsEnvelope, RuntimeWorkerEventWsFrame};
|
||||
|
||||
/// Backend-private source for a runtime worker observation stream.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct RuntimeObservationSourceConfig {
|
||||
pub runtime_id: String,
|
||||
pub worker_id: String,
|
||||
pub endpoint: String,
|
||||
pub bearer_token: Option<String>,
|
||||
}
|
||||
|
||||
/// Event consumed from a Runtime-owned worker observation WebSocket.
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct RuntimeObservationUpstreamEvent {
|
||||
pub runtime_id: String,
|
||||
pub worker_id: String,
|
||||
pub runtime_cursor: String,
|
||||
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 opaque event envelope. It intentionally omits Runtime endpoints,
|
||||
/// credentials, sockets and session paths.
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct ClientWorkerEventWsEnvelope {
|
||||
pub cursor: String,
|
||||
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, Default, Deserialize)]
|
||||
pub struct ClientWorkerEventsWsQuery {
|
||||
pub cursor: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum ObservationProxyError {
|
||||
RuntimeUnavailable(String),
|
||||
WorkerNotFound(String),
|
||||
CursorMalformed(String),
|
||||
CursorUnknownOrExpired(String),
|
||||
UpstreamDisconnect(String),
|
||||
MalformedFrame(String),
|
||||
ObservationOnly,
|
||||
}
|
||||
|
||||
impl ObservationProxyError {
|
||||
pub fn code(&self) -> &'static str {
|
||||
match self {
|
||||
ObservationProxyError::RuntimeUnavailable(_) => "backend.runtime_unavailable",
|
||||
ObservationProxyError::WorkerNotFound(_) => "backend.worker_not_found",
|
||||
ObservationProxyError::CursorMalformed(_) => "backend.cursor_malformed",
|
||||
ObservationProxyError::CursorUnknownOrExpired(_) => "backend.cursor_unknown_or_expired",
|
||||
ObservationProxyError::UpstreamDisconnect(_) => "backend.upstream_disconnect",
|
||||
ObservationProxyError::MalformedFrame(_) => "backend.malformed_frame",
|
||||
ObservationProxyError::ObservationOnly => "backend.observation_only",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn message(&self) -> &str {
|
||||
match self {
|
||||
ObservationProxyError::RuntimeUnavailable(message)
|
||||
| ObservationProxyError::WorkerNotFound(message)
|
||||
| ObservationProxyError::CursorMalformed(message)
|
||||
| ObservationProxyError::CursorUnknownOrExpired(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(),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
|
||||
pub struct BackendObservationCursor {
|
||||
pub sequence: u64,
|
||||
}
|
||||
|
||||
impl BackendObservationCursor {
|
||||
pub fn new(sequence: u64) -> Self {
|
||||
Self { sequence }
|
||||
}
|
||||
|
||||
pub fn zero() -> Self {
|
||||
Self { sequence: 0 }
|
||||
}
|
||||
|
||||
pub fn encode(self) -> String {
|
||||
format!("bo_{:016x}", self.sequence)
|
||||
}
|
||||
|
||||
pub fn decode(value: &str) -> Option<Self> {
|
||||
let encoded = value.strip_prefix("bo_")?;
|
||||
if encoded.len() != 16 {
|
||||
return None;
|
||||
}
|
||||
u64::from_str_radix(encoded, 16)
|
||||
.ok()
|
||||
.map(|sequence| Self { sequence })
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct BackendObservationState {
|
||||
next_sequence: u64,
|
||||
history: BTreeMap<ObservationKey, VecDeque<StoredBackendEvent>>,
|
||||
}
|
||||
|
||||
impl BackendObservationState {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
next_sequence: 1,
|
||||
history: BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
|
||||
struct ObservationKey {
|
||||
runtime_id: String,
|
||||
worker_id: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct StoredBackendEvent {
|
||||
sequence: u64,
|
||||
runtime_cursor: String,
|
||||
envelope: ClientWorkerEventWsEnvelope,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct BackendObservationOpen {
|
||||
pub replay: Vec<ClientWorkerEventWsEnvelope>,
|
||||
pub runtime_cursor: Option<String>,
|
||||
pub backend_cursor: BackendObservationCursor,
|
||||
}
|
||||
|
||||
/// Backend-owned in-memory v0 observation proxy state.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct BackendObservationProxy {
|
||||
sources: Arc<BTreeMap<ObservationKey, RuntimeObservationSourceConfig>>,
|
||||
state: Arc<Mutex<BackendObservationState>>,
|
||||
}
|
||||
|
||||
impl BackendObservationProxy {
|
||||
pub fn new(sources: Vec<RuntimeObservationSourceConfig>) -> Self {
|
||||
let sources = sources
|
||||
.into_iter()
|
||||
.map(|source| {
|
||||
(
|
||||
ObservationKey {
|
||||
runtime_id: source.runtime_id.clone(),
|
||||
worker_id: source.worker_id.clone(),
|
||||
},
|
||||
source,
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
Self {
|
||||
sources: Arc::new(sources),
|
||||
state: Arc::new(Mutex::new(BackendObservationState::new())),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn source(
|
||||
&self,
|
||||
runtime_id: &str,
|
||||
worker_id: &str,
|
||||
) -> Result<RuntimeObservationSourceConfig, ObservationProxyError> {
|
||||
self.sources
|
||||
.get(&ObservationKey {
|
||||
runtime_id: runtime_id.to_string(),
|
||||
worker_id: worker_id.to_string(),
|
||||
})
|
||||
.cloned()
|
||||
.ok_or_else(|| {
|
||||
ObservationProxyError::WorkerNotFound(format!(
|
||||
"worker {worker_id} is not registered for runtime {runtime_id}"
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn open(
|
||||
&self,
|
||||
runtime_id: &str,
|
||||
worker_id: &str,
|
||||
cursor: Option<&str>,
|
||||
) -> Result<BackendObservationOpen, ObservationProxyError> {
|
||||
let key = ObservationKey {
|
||||
runtime_id: runtime_id.to_string(),
|
||||
worker_id: worker_id.to_string(),
|
||||
};
|
||||
let cursor = match cursor {
|
||||
Some(raw) => BackendObservationCursor::decode(raw).ok_or_else(|| {
|
||||
ObservationProxyError::CursorMalformed(format!(
|
||||
"malformed backend observation cursor: {raw}"
|
||||
))
|
||||
})?,
|
||||
None => BackendObservationCursor::zero(),
|
||||
};
|
||||
let state = self.state.lock().map_err(|_| {
|
||||
ObservationProxyError::RuntimeUnavailable(
|
||||
"backend observation state lock poisoned".into(),
|
||||
)
|
||||
})?;
|
||||
let history = state.history.get(&key);
|
||||
let replay: Vec<_> = history
|
||||
.into_iter()
|
||||
.flat_map(|events| events.iter())
|
||||
.filter(|event| event.sequence > cursor.sequence)
|
||||
.cloned()
|
||||
.collect();
|
||||
if cursor.sequence != 0 {
|
||||
let found = history
|
||||
.into_iter()
|
||||
.flat_map(|events| events.iter())
|
||||
.any(|event| event.sequence == cursor.sequence);
|
||||
if !found {
|
||||
return Err(ObservationProxyError::CursorUnknownOrExpired(format!(
|
||||
"backend observation cursor {} is unknown or expired for runtime {runtime_id} worker {worker_id}",
|
||||
cursor.encode()
|
||||
)));
|
||||
}
|
||||
}
|
||||
let runtime_cursor = replay
|
||||
.last()
|
||||
.map(|event| event.runtime_cursor.clone())
|
||||
.or_else(|| {
|
||||
history.and_then(|events| {
|
||||
events
|
||||
.iter()
|
||||
.find(|event| event.sequence == cursor.sequence)
|
||||
.map(|event| event.runtime_cursor.clone())
|
||||
})
|
||||
});
|
||||
Ok(BackendObservationOpen {
|
||||
replay: replay.into_iter().map(|event| event.envelope).collect(),
|
||||
runtime_cursor,
|
||||
backend_cursor: cursor,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn store(
|
||||
&self,
|
||||
event: RuntimeObservationUpstreamEvent,
|
||||
) -> Result<ClientWorkerEventWsEnvelope, ObservationProxyError> {
|
||||
let mut state = self.state.lock().map_err(|_| {
|
||||
ObservationProxyError::RuntimeUnavailable(
|
||||
"backend observation state lock poisoned".into(),
|
||||
)
|
||||
})?;
|
||||
let sequence = state.next_sequence;
|
||||
state.next_sequence += 1;
|
||||
let cursor = BackendObservationCursor::new(sequence).encode();
|
||||
let envelope = ClientWorkerEventWsEnvelope {
|
||||
cursor: cursor.clone(),
|
||||
event_id: cursor,
|
||||
runtime_id: event.runtime_id.clone(),
|
||||
worker_id: event.worker_id.clone(),
|
||||
payload: event.payload,
|
||||
};
|
||||
let key = ObservationKey {
|
||||
runtime_id: event.runtime_id,
|
||||
worker_id: event.worker_id,
|
||||
};
|
||||
let history = state.history.entry(key).or_default();
|
||||
history.push_back(StoredBackendEvent {
|
||||
sequence,
|
||||
runtime_cursor: event.runtime_cursor,
|
||||
envelope: envelope.clone(),
|
||||
});
|
||||
while history.len() > 1024 {
|
||||
history.pop_front();
|
||||
}
|
||||
Ok(envelope)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct RuntimeWsObservationClient {
|
||||
runtime_id: String,
|
||||
worker_id: String,
|
||||
stream: tokio_tungstenite::WebSocketStream<
|
||||
tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
|
||||
>,
|
||||
}
|
||||
|
||||
impl RuntimeWsObservationClient {
|
||||
pub async fn connect(
|
||||
source: &RuntimeObservationSourceConfig,
|
||||
runtime_cursor: Option<&str>,
|
||||
) -> Result<Self, ObservationProxyError> {
|
||||
let mut endpoint = source.endpoint.clone();
|
||||
if let Some(cursor) = runtime_cursor {
|
||||
let separator = if endpoint.contains('?') { '&' } else { '?' };
|
||||
endpoint.push(separator);
|
||||
endpoint.push_str("cursor=");
|
||||
endpoint.push_str(cursor);
|
||||
}
|
||||
let mut request = endpoint.into_client_request().map_err(|error| {
|
||||
ObservationProxyError::RuntimeUnavailable(format!(
|
||||
"failed to build runtime WebSocket request: {error}"
|
||||
))
|
||||
})?;
|
||||
if let Some(token) = &source.bearer_token {
|
||||
request.headers_mut().insert(
|
||||
"authorization",
|
||||
format!("Bearer {token}").parse().map_err(|error| {
|
||||
ObservationProxyError::RuntimeUnavailable(format!(
|
||||
"failed to build runtime authorization header: {error}"
|
||||
))
|
||||
})?,
|
||||
);
|
||||
}
|
||||
let (stream, _) = connect_async(request).await.map_err(|error| {
|
||||
ObservationProxyError::RuntimeUnavailable(format!(
|
||||
"failed to connect runtime WebSocket: {error}"
|
||||
))
|
||||
})?;
|
||||
Ok(Self {
|
||||
runtime_id: source.runtime_id.clone(),
|
||||
worker_id: source.worker_id.clone(),
|
||||
stream,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn next_event(
|
||||
&mut self,
|
||||
) -> Result<RuntimeObservationUpstreamEvent, ObservationProxyError> {
|
||||
loop {
|
||||
let Some(message) = self.stream.next().await else {
|
||||
return Err(ObservationProxyError::UpstreamDisconnect(
|
||||
"runtime WebSocket closed".into(),
|
||||
));
|
||||
};
|
||||
let message = message.map_err(|error| {
|
||||
ObservationProxyError::UpstreamDisconnect(format!(
|
||||
"runtime WebSocket receive error: {error}"
|
||||
))
|
||||
})?;
|
||||
let text = match message {
|
||||
TungsteniteMessage::Text(text) => text,
|
||||
TungsteniteMessage::Close(_) => {
|
||||
return Err(ObservationProxyError::UpstreamDisconnect(
|
||||
"runtime WebSocket closed".into(),
|
||||
));
|
||||
}
|
||||
TungsteniteMessage::Ping(payload) => {
|
||||
self.stream
|
||||
.send(TungsteniteMessage::Pong(payload))
|
||||
.await
|
||||
.map_err(|error| {
|
||||
ObservationProxyError::UpstreamDisconnect(format!(
|
||||
"failed to reply to runtime ping: {error}"
|
||||
))
|
||||
})?;
|
||||
continue;
|
||||
}
|
||||
TungsteniteMessage::Pong(_) => continue,
|
||||
TungsteniteMessage::Binary(_) | TungsteniteMessage::Frame(_) => {
|
||||
return Err(ObservationProxyError::MalformedFrame(
|
||||
"runtime sent a non-text observation frame".into(),
|
||||
));
|
||||
}
|
||||
};
|
||||
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(ObservationProxyError::UpstreamDisconnect(format!(
|
||||
"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_cursor: envelope.cursor,
|
||||
payload: envelope.payload,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,12 +1,14 @@
|
||||
use std::path::{Component, Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade};
|
||||
use axum::extract::{Path as AxumPath, Query, State};
|
||||
use axum::http::header::CONTENT_TYPE;
|
||||
use axum::http::{StatusCode, Uri};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::get;
|
||||
use axum::{Json, Router};
|
||||
use futures::StreamExt;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
@@ -15,6 +17,10 @@ use crate::hosts::{
|
||||
RuntimeSummary, WorkerSummary,
|
||||
};
|
||||
use crate::identity::WorkspaceIdentity;
|
||||
use crate::observation::{
|
||||
BackendObservationProxy, ClientWorkerEventWsFrame, ClientWorkerEventsWsQuery,
|
||||
ObservationProxyError, RuntimeObservationSourceConfig, RuntimeWsObservationClient,
|
||||
};
|
||||
use crate::records::{
|
||||
LocalProjectRecordReader, ObjectiveDetail, ProjectRecordList, TicketDetail, TicketSummary,
|
||||
};
|
||||
@@ -39,6 +45,7 @@ pub struct ServerConfig {
|
||||
pub auth: AuthConfig,
|
||||
pub max_records: usize,
|
||||
pub local_runtime_data_dir: Option<PathBuf>,
|
||||
pub runtime_event_sources: Vec<RuntimeObservationSourceConfig>,
|
||||
}
|
||||
|
||||
impl ServerConfig {
|
||||
@@ -55,6 +62,7 @@ impl ServerConfig {
|
||||
},
|
||||
max_records: 200,
|
||||
local_runtime_data_dir: manifest::paths::data_dir(),
|
||||
runtime_event_sources: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -65,6 +73,7 @@ pub struct WorkspaceApi {
|
||||
store: Arc<dyn ControlPlaneStore>,
|
||||
records: LocalProjectRecordReader,
|
||||
runtime: Arc<RuntimeRegistry>,
|
||||
observation_proxy: BackendObservationProxy,
|
||||
}
|
||||
|
||||
impl WorkspaceApi {
|
||||
@@ -83,11 +92,13 @@ impl WorkspaceApi {
|
||||
config.workspace_root.clone(),
|
||||
config.local_runtime_data_dir.clone(),
|
||||
)));
|
||||
let observation_proxy = BackendObservationProxy::new(config.runtime_event_sources.clone());
|
||||
Ok(Self {
|
||||
records: LocalProjectRecordReader::new(config.workspace_root.clone()),
|
||||
config,
|
||||
store,
|
||||
runtime,
|
||||
observation_proxy,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -128,6 +139,10 @@ pub fn build_router(api: WorkspaceApi) -> Router {
|
||||
.route("/api/hosts", get(list_hosts))
|
||||
.route("/api/runtimes", get(list_runtimes))
|
||||
.route("/api/workers", get(list_workers))
|
||||
.route(
|
||||
"/api/runtimes/{runtime_id}/workers/{worker_id}/events/ws",
|
||||
get(worker_observation_ws),
|
||||
)
|
||||
.route("/api/hosts/{host_id}/workers", get(list_host_workers))
|
||||
.fallback(get(static_or_spa_fallback))
|
||||
.with_state(api)
|
||||
@@ -423,6 +438,144 @@ async fn list_workers(
|
||||
workers_response(api).map(Json)
|
||||
}
|
||||
|
||||
async fn worker_observation_ws(
|
||||
State(api): State<WorkspaceApi>,
|
||||
AxumPath((runtime_id, worker_id)): AxumPath<(String, String)>,
|
||||
Query(query): Query<ClientWorkerEventsWsQuery>,
|
||||
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, query, socket)
|
||||
}),
|
||||
Err(error) => {
|
||||
let status = match error {
|
||||
ObservationProxyError::WorkerNotFound(_) => StatusCode::NOT_FOUND,
|
||||
_ => 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: RuntimeObservationSourceConfig,
|
||||
query: ClientWorkerEventsWsQuery,
|
||||
mut socket: WebSocket,
|
||||
) {
|
||||
let open = match proxy.open(
|
||||
&source.runtime_id,
|
||||
&source.worker_id,
|
||||
query.cursor.as_deref(),
|
||||
) {
|
||||
Ok(open) => open,
|
||||
Err(error) => {
|
||||
let _ = send_client_ws_frame(&mut socket, ClientWorkerEventWsFrame::diagnostic(error))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let mut backend_cursor = open.backend_cursor;
|
||||
for envelope in open.replay {
|
||||
backend_cursor = crate::observation::BackendObservationCursor::decode(&envelope.cursor)
|
||||
.unwrap_or(backend_cursor);
|
||||
if !send_client_ws_frame(&mut socket, ClientWorkerEventWsFrame::event(envelope)).await {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
let mut upstream =
|
||||
match RuntimeWsObservationClient::connect(&source, open.runtime_cursor.as_deref()).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) => match proxy.store(event) {
|
||||
Ok(envelope) => {
|
||||
backend_cursor = crate::observation::BackendObservationCursor::decode(&envelope.cursor)
|
||||
.unwrap_or(backend_cursor);
|
||||
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;
|
||||
}
|
||||
},
|
||||
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>,
|
||||
@@ -636,7 +789,10 @@ mod tests {
|
||||
use super::*;
|
||||
use axum::body::{Body, to_bytes};
|
||||
use axum::http::Request;
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use serde_json::Value;
|
||||
use tokio_tungstenite::connect_async;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tower::ServiceExt;
|
||||
|
||||
use crate::store::SqliteWorkspaceStore;
|
||||
@@ -844,6 +1000,127 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxies_worker_observation_ws_with_backend_cursors_and_diagnostics() {
|
||||
let runtime = worker_runtime::Runtime::new_memory();
|
||||
let worker = runtime
|
||||
.create_worker(worker_runtime::catalog::CreateWorkerRequest::default())
|
||||
.unwrap();
|
||||
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 = ServerConfig::local_dev(dir.path(), test_identity());
|
||||
config.local_runtime_data_dir = Some(dir.path().join("data"));
|
||||
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.worker_ref.worker_id
|
||||
),
|
||||
bearer_token: None,
|
||||
});
|
||||
let api = WorkspaceApi::new(config, Arc::new(store)).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");
|
||||
};
|
||||
assert_eq!(snapshot.runtime_id, "runtime-a");
|
||||
assert_eq!(snapshot.worker_id, "worker-a");
|
||||
assert!(matches!(snapshot.payload, protocol::Event::Snapshot { .. }));
|
||||
|
||||
runtime
|
||||
.observe_worker_event(
|
||||
&worker.worker_ref,
|
||||
protocol::Event::TextDelta {
|
||||
text: "live".into(),
|
||||
},
|
||||
)
|
||||
.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 resumed, _) = connect_async(format!("{url}?cursor={}", live.cursor))
|
||||
.await
|
||||
.unwrap();
|
||||
let _snapshot = next_client_frame(&mut resumed).await;
|
||||
runtime
|
||||
.observe_worker_event(
|
||||
&worker.worker_ref,
|
||||
protocol::Event::TextDone {
|
||||
text: "done".into(),
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
let resumed_event = next_client_frame(&mut resumed).await;
|
||||
let ClientWorkerEventWsFrame::Event {
|
||||
envelope: resumed_event,
|
||||
} = resumed_event
|
||||
else {
|
||||
panic!("expected resumed live event");
|
||||
};
|
||||
assert_ne!(resumed_event.cursor, live.cursor);
|
||||
assert!(matches!(
|
||||
resumed_event.payload,
|
||||
protocol::Event::TextDone { .. }
|
||||
));
|
||||
|
||||
let (mut malformed, _) = connect_async(format!("{url}?cursor=bad")).await.unwrap();
|
||||
let diagnostic = next_client_frame(&mut malformed).await;
|
||||
let ClientWorkerEventWsFrame::Diagnostic { diagnostic } = diagnostic else {
|
||||
panic!("expected malformed cursor diagnostic");
|
||||
};
|
||||
assert_eq!(diagnostic.code, "backend.cursor_malformed");
|
||||
|
||||
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");
|
||||
}
|
||||
|
||||
async fn next_client_frame(
|
||||
stream: &mut tokio_tungstenite::WebSocketStream<
|
||||
tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
|
||||
>,
|
||||
) -> ClientWorkerEventWsFrame {
|
||||
let message = stream.next().await.unwrap().unwrap();
|
||||
let Message::Text(text) = message else {
|
||||
panic!("expected text frame");
|
||||
};
|
||||
serde_json::from_str(&text).unwrap()
|
||||
}
|
||||
|
||||
async fn get_json(app: Router, uri: &str) -> Value {
|
||||
let response = app
|
||||
.oneshot(Request::builder().uri(uri).body(Body::empty()).unwrap())
|
||||
|
||||
Reference in New Issue
Block a user