use futures::{SinkExt, StreamExt}; use protocol::stream::{decode_event, encode_method}; 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; #[derive(Debug, Clone, PartialEq, Eq)] pub struct BackendRuntimeTarget { /// Workspace Backend API root URL, for example `http://127.0.0.1:8787`. /// This is intentionally the Backend endpoint, not a Runtime endpoint. pub base_url: String, /// Backend-owned Runtime identity used as path authority. pub runtime_id: String, /// Backend-owned Worker identity used as path authority. pub worker_id: String, } impl BackendRuntimeTarget { pub fn new( base_url: impl Into, runtime_id: impl Into, worker_id: impl Into, ) -> Self { Self { base_url: base_url.into(), runtime_id: runtime_id.into(), worker_id: worker_id.into(), } } pub fn display_label(&self) -> String { format!("{}:{}", self.runtime_id, self.worker_id) } } #[derive(Debug, Clone, PartialEq, Eq)] pub struct BackendRuntimeListTarget { pub base_url: String, pub workspace_id: Option, pub runtime_id: Option, } impl BackendRuntimeListTarget { pub fn new( base_url: impl Into, workspace_id: Option, runtime_id: Option, ) -> Self { Self { base_url: base_url.into(), workspace_id, runtime_id, } } } #[derive(Debug, Clone, Deserialize)] pub struct BackendRuntimeListResponse { pub workspace_id: String, pub limit: usize, pub items: Vec, pub source: String, #[serde(default)] pub diagnostics: Vec, } #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] pub struct BackendRuntimeSummary { pub runtime_id: String, pub label: String, pub kind: String, pub status: String, #[serde(default)] pub host_ids: Vec, #[serde(default)] pub diagnostics: Vec, } #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] pub struct BackendWorkerWorkspaceSummary { pub visibility: String, pub identity: String, } #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] pub struct BackendWorkerImplementationSummary { pub kind: String, pub display_hint: String, } #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] pub struct BackendWorkerCapabilitySummary { pub can_stop: bool, pub can_spawn_followup: bool, } #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] pub struct BackendWorkingDirectoryCleanupTarget { pub kind: String, pub working_directory_id: String, pub repository_id: String, } #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] pub struct BackendWorkingDirectoryOccupancy { pub runtime_id: String, pub runtime_worker_id: u64, pub worker_id: String, pub display_name: String, pub linked_at: String, } #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] pub struct BackendWorkingDirectorySummary { pub working_directory_id: String, pub repository_id: String, #[serde(default)] pub creation_selector: Option, #[serde(default)] pub creation_ref: Option, #[serde(default)] pub current_selector: Option, #[serde(default)] pub current_ref: Option, pub materializer_kind: String, #[serde(default)] pub cleanup_target: Option, pub status: String, #[serde(default)] pub cleanliness: Option, #[serde(default)] pub primary_worker_id: Option, #[serde(default)] pub occupied_by: Option, } #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] pub struct BackendWorkerSummary { pub runtime_id: String, pub worker_id: String, pub host_id: String, #[serde(default)] pub display_name: String, pub label: String, #[serde(default)] pub profile: Option, #[serde(default)] pub singleton_key: Option, #[serde(default)] pub tags: Vec, pub workspace: BackendWorkerWorkspaceSummary, pub state: String, #[serde(default)] pub last_seen_at: Option, #[serde(default)] pub pinned: bool, #[serde(default)] pub retention_state: String, pub implementation: BackendWorkerImplementationSummary, pub capabilities: BackendWorkerCapabilitySummary, #[serde(default)] pub working_directory: Option, #[serde(default)] pub diagnostics: Vec, } #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] pub struct BackendWorkerRestoreResult { pub state: String, #[serde(default)] pub worker: Option, #[serde(default)] pub diagnostics: Vec, } #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] pub struct BackendWorkerRestoreResponse { pub workspace_id: String, pub runtime_id: String, pub worker_id: String, pub result: BackendWorkerRestoreResult, } #[derive(Debug)] pub struct BackendRuntimeClient { target: BackendRuntimeTarget, command_tx: mpsc::UnboundedSender, events: mpsc::UnboundedReceiver, diagnostics: VecDeque, _protocol_task: tokio::task::JoinHandle<()>, } #[derive(Debug)] pub enum BackendRuntimeClientError { InvalidTarget(String), Http(reqwest::Error), } impl fmt::Display for BackendRuntimeClientError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::InvalidTarget(message) => f.write_str(message), Self::Http(error) => write!(f, "{error}"), } } } impl std::error::Error for BackendRuntimeClientError {} impl From for BackendRuntimeClientError { fn from(error: reqwest::Error) -> Self { Self::Http(error) } } pub async fn list_backend_workers( target: &BackendRuntimeListTarget, ) -> Result, BackendRuntimeClientError> { validate_list_target(target)?; let http = reqwest::Client::new(); if let Some(runtime_id) = target.runtime_id.as_deref() { let path = backend_runtime_workers_path(target.workspace_id.as_deref(), runtime_id); let url = join_base_and_path(&target.base_url, &path); return Ok(http .get(url) .send() .await? .error_for_status()? .json::>() .await?); } let runtime_path = backend_runtimes_path(target.workspace_id.as_deref()); let runtime_url = join_base_and_path(&target.base_url, &runtime_path); let runtimes = http .get(runtime_url) .send() .await? .error_for_status()? .json::>() .await?; let mut items = Vec::new(); let mut diagnostics = runtimes.diagnostics; for runtime in runtimes.items { let path = backend_runtime_workers_path(target.workspace_id.as_deref(), &runtime.runtime_id); let url = join_base_and_path(&target.base_url, &path); match http .get(url) .send() .await .and_then(|response| response.error_for_status()) { Ok(response) => { let response = response .json::>() .await?; diagnostics.extend(response.diagnostics); items.extend(response.items); } Err(error) => diagnostics.push(BackendDiagnostic { code: "runtime_worker_list_failed".to_string(), severity: Some("error".to_string()), message: format!( "failed to list workers for runtime {}: {error}", runtime.runtime_id ), }), } } Ok(BackendRuntimeListResponse { workspace_id: runtimes.workspace_id, limit: runtimes.limit, items, source: "backend_runtime_worker_summary".to_string(), diagnostics, }) } pub async fn list_backend_stopped_workers( target: &BackendRuntimeListTarget, ) -> Result, BackendRuntimeClientError> { validate_list_target(target)?; let Some(runtime_id) = target.runtime_id.as_deref() else { return Err(BackendRuntimeClientError::InvalidTarget( "stopped worker listing requires a runtime id".to_string(), )); }; let http = reqwest::Client::new(); let path = backend_runtime_workers_path(target.workspace_id.as_deref(), runtime_id); let url = join_base_and_path(&target.base_url, &format!("{path}?status=stopped")); Ok(http .get(url) .send() .await? .error_for_status()? .json::>() .await?) } pub async fn restore_backend_worker( target: &BackendRuntimeTarget, ) -> Result { validate_target(target)?; let http = reqwest::Client::new(); let path = backend_runtime_worker_restore_path(None, &target.runtime_id, &target.worker_id); let url = join_base_and_path(&target.base_url, &path); Ok(http .post(url) .json(&serde_json::json!({})) .send() .await? .error_for_status()? .json::() .await?) } impl BackendRuntimeClient { pub async fn connect(target: BackendRuntimeTarget) -> Result { validate_target(&target)?; let (event_tx, rx) = mpsc::unbounded_channel(); let (command_tx, command_rx) = mpsc::unbounded_channel(); 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 { target, command_tx, events: rx, diagnostics: VecDeque::new(), _protocol_task: protocol_task, }) } pub fn try_next_event(&mut self) -> Option { if let Some(event) = self.diagnostics.pop_front() { return Some(event); } self.events.try_recv().ok() } pub async fn next_event(&mut self) -> Option { if let Some(event) = self.diagnostics.pop_front() { return Some(event); } self.events.recv().await } pub async fn send(&mut self, method: &Method) -> Result<(), BackendRuntimeClientError> { self.command_tx.send(method.clone()).map_err(|_| { BackendRuntimeClientError::InvalidTarget(format!( "Backend protocol command stream is closed for {}", self.target.display_label() )) })?; Ok(()) } } impl Drop for BackendRuntimeClient { fn drop(&mut self) { self._protocol_task.abort(); } } async fn run_worker_protocol_transport( target: BackendRuntimeTarget, mut commands: mpsc::UnboundedReceiver, tx: mpsc::UnboundedSender, ) { let url = protocol_ws_url(&target); match connect_async(&url).await { Ok((ws, _)) => { let (mut sink, mut stream) = ws.split(); loop { tokio::select! { maybe_method = commands.recv() => { let Some(method) = maybe_method else { break; }; match encode_method(&method) { Ok(text) => { if let Err(error) = sink.send(TungsteniteMessage::Text(text.into())).await { let _ = tx.send(diagnostic_event(format!( "Backend protocol command send failed for {}: {error}", target.display_label() ))); break; } } Err(error) => { let _ = tx.send(diagnostic_event(format!( "Backend protocol command could not serialize method for {}: {error}", target.display_label() ))); } } } frame = stream.next() => { match frame { Some(Ok(TungsteniteMessage::Text(text))) => { match decode_event(&text) { Ok(event) => { let _ = tx.send(event); } Err(error) => { let _ = tx.send(diagnostic_event(format!( "Backend protocol response was not valid Event JSON for {}: {error}", target.display_label() ))); } } } Some(Ok(TungsteniteMessage::Close(_))) | None => { let _ = tx.send(diagnostic_event(format!( "Backend protocol command stream closed for {}", target.display_label() ))); break; } Some(Ok(TungsteniteMessage::Ping(_))) | Some(Ok(TungsteniteMessage::Pong(_))) | Some(Ok(TungsteniteMessage::Binary(_))) | Some(Ok(TungsteniteMessage::Frame(_))) => {} Some(Err(error)) => { let _ = tx.send(diagnostic_event(format!( "Backend protocol WebSocket error for {}: {error}", target.display_label() ))); break; } } } } } } Err(error) => { let _ = tx.send(diagnostic_event(format!( "Backend protocol WebSocket connect failed for {}: {error}", target.display_label() ))); while commands.recv().await.is_some() { let _ = tx.send(diagnostic_event(format!( "Backend protocol command was not sent because command stream is unavailable for {}", target.display_label() ))); } } } } fn diagnostic_event(message: impl Into) -> Event { Event::Error { code: ErrorCode::Internal, message: message.into(), } } fn validate_target(target: &BackendRuntimeTarget) -> Result<(), BackendRuntimeClientError> { if target.base_url.trim().is_empty() { return Err(BackendRuntimeClientError::InvalidTarget( "Backend API base URL is required".to_string(), )); } if !(target.base_url.starts_with("http://") || target.base_url.starts_with("https://")) { return Err(BackendRuntimeClientError::InvalidTarget( "Backend API base URL must start with http:// or https://".to_string(), )); } if target.runtime_id.is_empty() { return Err(BackendRuntimeClientError::InvalidTarget( "runtime_id is required".to_string(), )); } if target.worker_id.is_empty() { return Err(BackendRuntimeClientError::InvalidTarget( "worker_id is required".to_string(), )); } Ok(()) } fn validate_list_target( target: &BackendRuntimeListTarget, ) -> Result<(), BackendRuntimeClientError> { if target.base_url.trim().is_empty() { return Err(BackendRuntimeClientError::InvalidTarget( "Backend API base URL is required".to_string(), )); } if !(target.base_url.starts_with("http://")) && !(target.base_url.starts_with("https://")) { return Err(BackendRuntimeClientError::InvalidTarget( "Backend API base URL must start with http:// or https://".to_string(), )); } if target.workspace_id.as_deref().is_some_and(str::is_empty) { return Err(BackendRuntimeClientError::InvalidTarget( "workspace_id must not be empty when provided".to_string(), )); } if target.runtime_id.as_deref().is_some_and(str::is_empty) { return Err(BackendRuntimeClientError::InvalidTarget( "runtime_id must not be empty when provided".to_string(), )); } Ok(()) } fn backend_runtimes_path(workspace_id: Option<&str>) -> String { match workspace_id { Some(workspace_id) => format!("/api/w/{}/runtimes", path_segment_encode(workspace_id)), None => "/api/runtimes".to_string(), } } fn backend_runtime_workers_path(workspace_id: Option<&str>, runtime_id: &str) -> String { match workspace_id { Some(workspace_id) => format!( "/api/w/{}/runtimes/{}/workers", path_segment_encode(workspace_id), path_segment_encode(runtime_id) ), None => format!("/api/runtimes/{}/workers", path_segment_encode(runtime_id)), } } fn backend_runtime_worker_restore_path( workspace_id: Option<&str>, runtime_id: &str, worker_id: &str, ) -> String { match workspace_id { Some(workspace_id) => format!( "/api/w/{}/runtimes/{}/workers/{}/restore", path_segment_encode(workspace_id), path_segment_encode(runtime_id), path_segment_encode(worker_id) ), None => format!( "/api/runtimes/{}/workers/{}/restore", path_segment_encode(runtime_id), path_segment_encode(worker_id) ), } } fn protocol_ws_url(target: &BackendRuntimeTarget) -> String { let path = format!( "/api/runtimes/{}/workers/{}/protocol/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 http_base_to_ws(base: &str) -> String { if let Some(rest) = base.strip_prefix("https://") { format!("wss://{rest}") } else if let Some(rest) = base.strip_prefix("http://") { format!("ws://{rest}") } else { base.to_string() } } fn join_base_and_path(base: &str, path: &str) -> String { format!("{}{}", base.trim_end_matches('/'), path) } fn path_segment_encode(input: &str) -> String { percent_encode(input, |byte| { byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_' | b'~') }) } fn percent_encode(input: &str, keep: impl Fn(u8) -> bool) -> String { let mut encoded = String::with_capacity(input.len()); for byte in input.bytes() { if keep(byte) { encoded.push(byte as char); } else { encoded.push('%'); encoded.push_str(&format!("{byte:02X}")); } } encoded } #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] pub struct BackendDiagnostic { pub code: String, #[serde(default)] pub severity: Option, pub message: String, } #[cfg(test)] mod tests { use super::*; #[test] fn protocol_url_uses_backend_runtime_worker_identity() { let target = BackendRuntimeTarget::new("http://127.0.0.1:8787/", "runtime/one", "worker one"); assert_eq!( protocol_ws_url(&target), "ws://127.0.0.1:8787/api/runtimes/runtime%2Fone/workers/worker%20one/protocol/ws" ); } #[test] fn workers_path_can_be_workspace_scoped_for_status_queries() { let path = backend_runtime_workers_path(Some("team main"), "runtime/one"); assert_eq!( format!("{path}?status=stopped"), "/api/w/team%20main/runtimes/runtime%2Fone/workers?status=stopped" ); } #[test] fn restore_worker_path_uses_backend_runtime_worker_identity() { assert_eq!( backend_runtime_worker_restore_path(None, "runtime/one", "worker one"), "/api/runtimes/runtime%2Fone/workers/worker%20one/restore" ); } }