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

feat: route backend worker controls over protocol

Browse Source
This commit is contained in:
Hare
2026-07-21 19:33:24 +09:00
parent bfa9346de2
commit 70a26a3042
10 changed files with 957 additions and 292 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/client/src/backend_runtime.rs
+119 -222
View File
@@ -2,9 +2,9 @@ use std::collections::VecDeque;
use std::fmt;
use std::time::Duration;
use futures::StreamExt;
use protocol::{ErrorCode, Event, Method, Segment};
use serde::{Deserialize, Serialize};
use futures::{SinkExt, StreamExt};
use protocol::{ErrorCode, Event, Method};
use serde::Deserialize;
use tokio::sync::mpsc;
use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite::Message as TungsteniteMessage;
@@ -160,10 +160,11 @@ pub struct BackendWorkerSummary {
#[derive(Debug)]
pub struct BackendRuntimeClient {
target: BackendRuntimeTarget,
http: reqwest::Client,
command_tx: mpsc::UnboundedSender<Method>,
events: mpsc::UnboundedReceiver<Event>,
diagnostics: VecDeque<Event>,
_observation_task: tokio::task::JoinHandle<()>,
_command_task: tokio::task::JoinHandle<()>,
}
#[derive(Debug)]
@@ -258,21 +259,28 @@ pub async fn list_backend_workers(
impl BackendRuntimeClient {
pub async fn connect(target: BackendRuntimeTarget) -> Result<Self, BackendRuntimeClientError> {
validate_target(&target)?;
let http = reqwest::Client::new();
let (tx, rx) = mpsc::unbounded_channel();
let (event_tx, rx) = mpsc::unbounded_channel();
let (command_tx, command_rx) = mpsc::unbounded_channel();
let observation_target = target.clone();
let observation_tx = tx.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;
});
Ok(Self {
target,
http,
command_tx,
events: rx,
diagnostics: VecDeque::new(),
_observation_task: observation_task,
_command_task: command_task,
})
}
@@ -291,163 +299,106 @@ impl BackendRuntimeClient {
}
pub async fn send(&mut self, method: &Method) -> Result<(), BackendRuntimeClientError> {
match backend_command_from_method(method) {
BackendCommand::Input { kind, content } => {
let url = self.worker_api_url("input");
match self
.http
.post(url)
.json(&WorkerInputRequest { kind, content })
.send()
.await
.and_then(|response| response.error_for_status())
{
Ok(response) => match response.json::<WorkerInputResult>().await {
Ok(result) => self.enqueue_operation_diagnostics(
"input",
result.state,
result.diagnostics,
),
Err(error) => self.enqueue_diagnostic(format!(
"Backend runtime input response could not be decoded for {}: {error}",
self.target.display_label()
)),
},
Err(error) => self.enqueue_diagnostic(format!(
"Backend runtime input failed for {}: {error}",
self.target.display_label()
)),
}
}
BackendCommand::Lifecycle { action, reason } => {
let url = self.worker_api_url(action);
match self
.http
.post(url)
.json(&WorkerLifecycleRequest { reason })
.send()
.await
.and_then(|response| response.error_for_status())
{
Ok(response) => match response.json::<WorkerLifecycleResult>().await {
Ok(result) => self.enqueue_operation_diagnostics(
action,
result.state,
result.diagnostics,
),
Err(error) => self.enqueue_diagnostic(format!(
"Backend runtime {action} response could not be decoded for {}: {error}",
self.target.display_label()
)),
},
Err(error) => self.enqueue_diagnostic(format!(
"Backend runtime {action} failed for {}: {error}",
self.target.display_label()
)),
}
}
BackendCommand::Unsupported(message) => {
self.enqueue_diagnostic(message);
}
}
Ok(())
}
fn worker_api_url(&self, suffix: &str) -> String {
let path = format!(
"/api/runtimes/{}/workers/{}/{}",
path_segment_encode(&self.target.runtime_id),
path_segment_encode(&self.target.worker_id),
suffix
);
join_base_and_path(&self.target.base_url, &path)
}
fn enqueue_operation_diagnostics(
&mut self,
operation: &str,
state: String,
diagnostics: Vec<BackendDiagnostic>,
) {
if state != "accepted" {
self.enqueue_diagnostic(format!(
"Backend runtime {operation} was {state} for {}",
self.command_tx.send(method.clone()).map_err(|_| {
BackendRuntimeClientError::InvalidTarget(format!(
"Backend protocol command stream is closed for {}",
self.target.display_label()
));
}
for diagnostic in diagnostics {
self.enqueue_diagnostic(format!(
"Backend runtime {operation} diagnostic [{}]: {}",
diagnostic.code, diagnostic.message
));
}
}
fn enqueue_diagnostic(&mut self, message: impl Into<String>) {
self.diagnostics.push_back(diagnostic_event(message));
))
})?;
Ok(())
}
}
impl Drop for BackendRuntimeClient {
fn drop(&mut self) {
self._observation_task.abort();
self._command_task.abort();
}
}
#[derive(Debug, PartialEq, Eq)]
enum BackendCommand {
Input {
kind: WorkerInputKind,
content: String,
},
Lifecycle {
action: &'static str,
reason: Option<String>,
},
Unsupported(String),
}
fn backend_command_from_method(method: &Method) -> BackendCommand {
match method {
Method::Run { input } => BackendCommand::Input {
kind: WorkerInputKind::User,
content: Segment::flatten_to_text(input),
},
Method::Notify { message, .. } => BackendCommand::Input {
kind: WorkerInputKind::System,
content: message.clone(),
},
Method::Cancel => BackendCommand::Lifecycle {
action: "cancel",
reason: Some("requested from TUI Backend Runtime API client".to_string()),
},
Method::Shutdown => BackendCommand::Lifecycle {
action: "stop",
reason: Some("requested from TUI Backend Runtime API client".to_string()),
},
Method::Pause => BackendCommand::Unsupported(
"Backend Runtime API does not expose pause/resume for the TUI client yet; command was not sent".to_string(),
),
Method::Resume => BackendCommand::Unsupported(
"Backend Runtime API does not expose resume for the TUI client yet; command was not sent".to_string(),
),
Method::Compact => BackendCommand::Unsupported(
"Backend Runtime API does not expose compaction for the TUI client yet; command was not sent".to_string(),
),
Method::ListCompletions { .. } => BackendCommand::Unsupported(
"Backend Runtime API does not expose completion lookup for the TUI client yet".to_string(),
),
Method::ListRewindTargets | Method::RewindTo { .. } => BackendCommand::Unsupported(
"Backend Runtime API does not expose rewind controls for the TUI client yet; command was not sent".to_string(),
),
Method::ListWorkers | Method::RestoreWorker { .. } | Method::RegisterPeer { .. } => {
BackendCommand::Unsupported(
"Backend Runtime API worker-management controls are not available from this Console connection".to_string(),
)
async fn run_worker_protocol_commands(
target: BackendRuntimeTarget,
mut commands: mpsc::UnboundedReceiver<Method>,
tx: mpsc::UnboundedSender<Event>,
) {
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 serde_json::to_string(&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 serde_json::from_str::<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()
)));
}
}
Method::WorkerEvent(_) => BackendCommand::Unsupported(
"Backend Runtime API does not accept child Worker lifecycle events from this Console connection".to_string(),
),
}
}
@@ -610,6 +561,15 @@ fn observation_ws_url(target: &BackendRuntimeTarget) -> String {
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",
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}")
@@ -643,38 +603,6 @@ fn percent_encode(input: &str, keep: impl Fn(u8) -> bool) -> String {
encoded
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
enum WorkerInputKind {
User,
System,
}
#[derive(Debug, Serialize)]
struct WorkerInputRequest {
kind: WorkerInputKind,
content: String,
}
#[derive(Debug, Serialize)]
struct WorkerLifecycleRequest {
reason: Option<String>,
}
#[derive(Debug, Deserialize)]
struct WorkerInputResult {
state: String,
#[serde(default)]
diagnostics: Vec<BackendDiagnostic>,
}
#[derive(Debug, Deserialize)]
struct WorkerLifecycleResult {
state: String,
#[serde(default)]
diagnostics: Vec<BackendDiagnostic>,
}
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
pub struct BackendDiagnostic {
pub code: String,
@@ -712,47 +640,16 @@ mod tests {
use super::*;
#[test]
fn backend_command_maps_run_to_user_input_without_runtime_endpoint() {
let method = Method::Run {
input: vec![
Segment::text("hello"),
Segment::FileRef {
path: "src/lib.rs".into(),
},
],
};
assert_eq!(
backend_command_from_method(&method),
BackendCommand::Input {
kind: WorkerInputKind::User,
content: "hello@src/lib.rs".to_string(),
}
);
}
#[test]
fn backend_worker_list_paths_use_scoped_workspace_when_available() {
assert_eq!(
backend_runtimes_path(Some("workspace/one")),
"/api/w/workspace%2Fone/runtimes"
);
assert_eq!(
backend_runtime_workers_path(Some("workspace/one"), "runtime one"),
"/api/w/workspace%2Fone/runtimes/runtime%20one/workers"
);
assert_eq!(
backend_runtime_workers_path(None, "runtime one"),
"/api/runtimes/runtime%20one/workers"
);
}
#[test]
fn observation_url_uses_backend_runtime_worker_identity() {
fn command_and_observation_urls_use_backend_protocol_paths() {
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/execution.rs
+21
View File
@@ -6,6 +6,7 @@ use crate::interaction::WorkerInput;
#[cfg(feature = "ws-server")]
use crate::observation::WorkerObservationEvent;
use crate::working_directory::{WorkingDirectoryBinding, WorkingDirectoryDiagnostic};
use protocol::Method;
use serde::{Deserialize, Serialize};
use std::fmt;
use std::sync::Arc;
@@ -63,6 +64,7 @@ pub enum WorkerExecutionOperation {
Spawn,
Restore,
Input,
ProtocolMethod,
Stop,
Cancel,
}
@@ -380,6 +382,17 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static {
input: WorkerInput,
) -> WorkerExecutionResult;
fn dispatch_method(
&self,
_handle: &WorkerExecutionHandle,
_method: Method,
) -> WorkerExecutionResult {
WorkerExecutionResult::unsupported(
WorkerExecutionOperation::ProtocolMethod,
"execution backend does not support direct Worker protocol methods",
)
}
fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
WorkerExecutionResult::unsupported(
WorkerExecutionOperation::Stop,
@@ -479,6 +492,14 @@ impl WorkerExecutionBackendRef {
self.backend.dispatch_input(handle, input)
}
pub(crate) fn dispatch_method(
&self,
handle: &WorkerExecutionHandle,
method: Method,
) -> WorkerExecutionResult {
self.backend.dispatch_method(handle, method)
}
pub(crate) fn stop_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
self.backend.stop_worker(handle)
}
crates/worker-runtime/src/http_server.rs
+94 -1
View File
@@ -154,7 +154,12 @@ 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));
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),
);
router
.with_state(state.clone())
@@ -506,6 +511,94 @@ async fn create_worker(
Ok(Json(RuntimeHttpWorkerResponse { worker }))
}
#[cfg(feature = "ws-server")]
async fn worker_protocol_ws(
State(state): State<RuntimeHttpState>,
Path(worker_id): Path<String>,
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_protocol_ws_session(state.runtime, worker_ref, socket))
.into_response())
}
#[cfg(feature = "ws-server")]
async fn worker_protocol_ws_session(
runtime: Runtime,
worker_ref: WorkerRef,
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 {
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;
}
}
},
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;
}
}
}
}
#[cfg(feature = "ws-server")]
async fn send_protocol_event(socket: &mut WebSocket, event: &protocol::Event) -> bool {
match serde_json::to_string(event) {
Ok(text) => socket.send(WsMessage::Text(text.into())).await.is_ok(),
Err(error) => {
let fallback = protocol_error_event(format!(
"failed to serialize protocol response event: {error}"
));
let Ok(text) = serde_json::to_string(&fallback) else {
return false;
};
socket.send(WsMessage::Text(text.into())).await.is_ok()
}
}
}
#[cfg(feature = "ws-server")]
fn protocol_error_event(message: impl Into<String>) -> protocol::Event {
protocol::Event::Error {
code: protocol::ErrorCode::Internal,
message: message.into(),
}
}
#[cfg(feature = "ws-server")]
async fn worker_events_ws(
State(state): State<RuntimeHttpState>,
crates/worker-runtime/src/runtime.rs
+68 -1
View File
@@ -35,6 +35,7 @@ use crate::observation::{
};
#[cfg(feature = "ws-server")]
use crate::observation::{WorkerObservationCursor, WorkerObservationEvent};
use protocol::{Event, Method};
use std::collections::BTreeMap;
#[cfg(feature = "ws-server")]
use std::collections::VecDeque;
@@ -559,6 +560,71 @@ impl Runtime {
Ok(backend.worker_completions(&handle, kind, prefix))
}
/// Accept a protocol method for a Worker through a Backend/runtime transport.
///
/// Most methods are delivered to the execution backend unchanged. Methods with
/// direct same-connection replies in the local socket protocol return those
/// events from this function so WebSocket transports can write them back to the
/// requesting client without rebroadcasting them.
pub fn send_protocol_method(
&self,
worker_ref: &WorkerRef,
method: Method,
) -> Result<Vec<Event>, RuntimeError> {
if let Method::ListCompletions { kind, prefix } = method {
let entries = self.worker_completions(worker_ref, kind, &prefix)?;
return Ok(vec![Event::Completions { kind, entries }]);
}
let (backend, handle) = {
let mut state = self.lock()?;
state.ensure_running()?;
state.ensure_worker_ref(worker_ref)?;
let worker = state.worker(worker_ref)?;
if !worker.status.is_active() {
return Err(RuntimeError::InvalidRequest(format!(
"worker {} is not running",
worker_ref.worker_id
)));
}
let backend = state.execution_backend.clone();
let handle = worker.execution_handle.clone();
match (backend, handle) {
(Some(backend), Some(handle)) => (backend, handle),
_ => {
let result = WorkerExecutionResult::rejected(
WorkerExecutionOperation::ProtocolMethod,
"worker has no execution backend",
);
let worker = state.worker_mut(worker_ref)?;
let mut execution = WorkerExecutionStatus::unconnected().with_result(result);
execution.binding = worker.execution.binding.clone();
worker.execution = execution;
state.persist_worker(&worker_ref.worker_id)?;
return Err(RuntimeError::WorkerExecutionUnavailable {
worker_id: worker_ref.worker_id.clone(),
message: "worker has no execution backend".to_string(),
});
}
}
};
let dispatch_result = backend.dispatch_method(&handle, method);
if !dispatch_result.is_accepted() {
self.record_execution_result(worker_ref, dispatch_result.clone())?;
return Err(RuntimeError::WorkerExecutionRejected {
worker_id: worker_ref.worker_id.clone(),
operation: dispatch_result.operation,
outcome: dispatch_result.outcome,
message: dispatch_result.message_or_default(),
result: dispatch_result,
});
}
self.record_execution_result(worker_ref, dispatch_result)?;
Ok(Vec::new())
}
fn commit_created_worker(
&self,
worker_ref: &WorkerRef,
@@ -641,7 +707,8 @@ impl Runtime {
WorkerExecutionOperation::Cancel => backend.cancel_worker(&handle),
WorkerExecutionOperation::Spawn
| WorkerExecutionOperation::Restore
| WorkerExecutionOperation::Input => return Ok(()),
| WorkerExecutionOperation::Input
| WorkerExecutionOperation::ProtocolMethod => return Ok(()),
};
if result.is_accepted() {
self.record_execution_result(worker_ref, result)?;
crates/worker-runtime/src/worker_backend.rs
+63
View File
@@ -722,6 +722,27 @@ impl<F> Drop for WorkerRuntimeExecutionBackend<F> {
}
}
fn method_starts_turn(method: &Method) -> bool {
matches!(
method,
Method::Run { .. }
| Method::Notify { auto_run: true, .. }
| Method::Resume
| Method::Compact
)
}
fn accepted_run_state_for_method(method: &Method) -> WorkerExecutionRunState {
match method {
Method::Run { .. }
| Method::Notify { auto_run: true, .. }
| Method::Resume
| Method::Compact => WorkerExecutionRunState::Busy,
Method::Shutdown => WorkerExecutionRunState::Stopped,
_ => WorkerExecutionRunState::Idle,
}
}
impl<F> WorkerExecutionBackend for WorkerRuntimeExecutionBackend<F>
where
F: RuntimeWorkerFactory,
@@ -1039,6 +1060,48 @@ where
result
}
fn dispatch_method(
&self,
handle: &WorkerExecutionHandle,
method: Method,
) -> WorkerExecutionResult {
let (worker, busy) = match self.get_execution(handle) {
Ok(execution) => execution,
Err(mut result) => {
result.operation = WorkerExecutionOperation::ProtocolMethod;
return result;
}
};
let starts_turn = method_starts_turn(&method);
if starts_turn
&& (worker.shared_state.get_status() != WorkerStatus::Idle
|| busy
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_err())
{
return WorkerExecutionResult::busy(
WorkerExecutionOperation::ProtocolMethod,
"Worker is already running; runtime adapter v0 does not queue protocol methods",
);
}
let accepted_run_state = accepted_run_state_for_method(&method);
let accepted_is_idle = accepted_run_state == WorkerExecutionRunState::Idle;
let result = self.send_method(
WorkerExecutionOperation::ProtocolMethod,
worker,
method,
accepted_run_state,
);
if (starts_turn && accepted_is_idle)
|| (starts_turn && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted)
{
busy.store(false, Ordering::SeqCst);
}
result
}
fn stop_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
if handle.backend_id() != self.backend_id() {
return WorkerExecutionResult::rejected(
crates/workspace-server/src/hosts.rs
+76
View File
@@ -514,6 +514,21 @@ pub enum RuntimeRegistryError {
}
impl RuntimeRegistryError {
pub fn message(&self) -> String {
match self {
Self::InvalidIdentifier { kind, value } => {
format!("invalid {kind} identifier `{value}`")
}
Self::UnknownRuntime(runtime_id) => format!("unknown runtime `{runtime_id}`"),
Self::UnknownHost(host_id) => format!("unknown host `{host_id}`"),
Self::UnknownWorker {
runtime_id,
worker_id,
} => format!("unknown worker `{worker_id}` in runtime `{runtime_id}`"),
Self::RuntimeOperationFailed { message, .. } => message.clone(),
}
}
pub fn into_error(self) -> Error {
match self {
Self::InvalidIdentifier { kind, value } => Error::InvalidRuntimeIdentifier {
@@ -658,6 +673,18 @@ pub trait WorkspaceWorkerRuntime: Send + Sync {
}
}
fn send_protocol_method(
&self,
_worker_id: &str,
_method: protocol::Method,
) -> Result<Vec<protocol::Event>, RuntimeRegistryError> {
Err(RuntimeRegistryError::RuntimeOperationFailed {
runtime_id: self.runtime_id().to_string(),
code: "worker_protocol_method_unsupported".to_string(),
message: "runtime does not support Worker protocol command transport".to_string(),
})
}
fn stop_worker(
&self,
worker_id: &str,
@@ -1062,6 +1089,26 @@ impl RuntimeRegistry {
Ok(runtime.list_config_bundles())
}
pub fn send_protocol_method(
&self,
runtime_id: &str,
worker_id: &str,
method: protocol::Method,
) -> Result<Vec<protocol::Event>, RuntimeRegistryError> {
validate_backend_identifier("runtime_id", runtime_id)?;
validate_backend_identifier("worker_id", worker_id)?;
let runtime = self.runtime(runtime_id)?;
let lookup = runtime.worker(worker_id);
if lookup.worker.is_none() {
return Err(operation_failed_or_unknown_worker(
runtime_id,
worker_id,
lookup.diagnostics,
));
}
runtime.send_protocol_method(worker_id, method)
}
pub fn send_input(
&self,
runtime_id: &str,
@@ -1807,6 +1854,35 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
))
}
fn send_protocol_method(
&self,
worker_id: &str,
method: protocol::Method,
) -> Result<Vec<protocol::Event>, RuntimeRegistryError> {
if !self.execution_enabled {
return Err(RuntimeRegistryError::RuntimeOperationFailed {
runtime_id: self.runtime_id.clone(),
code: "embedded_worker_execution_unavailable".to_string(),
message: format!(
"worker protocol command for '{worker_id}' requires an embedded execution backend"
),
});
}
let Some(worker_ref) = self.worker_ref(worker_id) else {
return Err(RuntimeRegistryError::UnknownWorker {
runtime_id: self.runtime_id.clone(),
worker_id: worker_id.to_string(),
});
};
self.runtime
.send_protocol_method(&worker_ref, method)
.map_err(|error| RuntimeRegistryError::RuntimeOperationFailed {
runtime_id: self.runtime_id.clone(),
code: "embedded_worker_protocol_command_failed".to_string(),
message: error.to_string(),
})
}
fn send_input(&self, worker_id: &str, request: WorkerInputRequest) -> WorkerInputResult {
if !self.execution_enabled {
return embedded_input_rejected(
crates/workspace-server/src/server.rs
+104
View File
@@ -555,6 +555,14 @@ pub fn build_router(api: WorkspaceApi) -> Router {
"/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),
)
.route(
"/api/w/{workspace_id}/runtimes/{runtime_id}/workers/{worker_id}/protocol/ws",
get(scoped_worker_protocol_ws),
)
.route("/api/hosts/{host_id}/workers", get(list_host_workers))
.route(
"/api/w/{workspace_id}/hosts/{host_id}/workers",
@@ -2457,6 +2465,19 @@ async fn scoped_worker_observation_ws(
.into_response()
}
async fn scoped_worker_protocol_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_protocol_ws(State(api), AxumPath((path.runtime_id, path.worker_id)), ws)
.await
.into_response()
}
async fn scoped_list_host_workers(
State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedHostPath>,
@@ -3388,6 +3409,89 @@ async fn cancel_runtime_worker(
Ok(Json(result))
}
async fn worker_protocol_ws(
State(api): State<WorkspaceApi>,
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(),
}
}
async fn worker_protocol_ws_session(
runtime: Arc<RuntimeRegistry>,
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 {
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;
}
}
},
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;
}
}
}
}
async fn send_protocol_event(socket: &mut WebSocket, event: &protocol::Event) -> bool {
match serde_json::to_string(event) {
Ok(text) => socket.send(WsMessage::Text(text.into())).await.is_ok(),
Err(error) => {
let fallback = protocol_error_event(format!(
"failed to serialize protocol response event: {error}"
));
let Ok(text) = serde_json::to_string(&fallback) else {
return false;
};
socket.send(WsMessage::Text(text.into())).await.is_ok()
}
}
}
fn protocol_error_event(message: impl Into<String>) -> protocol::Event {
protocol::Event::Error {
code: protocol::ErrorCode::Internal,
message: message.into(),
}
}
async fn worker_observation_ws(
State(api): State<WorkspaceApi>,
AxumPath((runtime_id, worker_id)): AxumPath<(String, String)>,