From 11f26a109023e8ee6186732c51d3c6b8dc2721aa Mon Sep 17 00:00:00 2001 From: Hare Date: Sat, 1 Aug 2026 21:07:45 +0900 Subject: [PATCH] server: multiplex worker protocol subscriptions --- crates/protocol/src/subscription.rs | 62 +++- crates/protocol/src/typescript.rs | 4 +- .../src/runtime_subscription.rs | 4 +- crates/workspace-server/src/server.rs | 286 +++++++++++++++++- .../src/workspace_subscription.rs | 255 +++++++++++++--- web/workspace/src/lib/generated/protocol.ts | 8 +- 6 files changed, 568 insertions(+), 51 deletions(-) diff --git a/crates/protocol/src/subscription.rs b/crates/protocol/src/subscription.rs index af2666c7..9fc4181a 100644 --- a/crates/protocol/src/subscription.rs +++ b/crates/protocol/src/subscription.rs @@ -248,6 +248,8 @@ pub enum EventSubscriptionSelector { }, WorkerProtocol { worker_id: SubscriptionWorkerId, + #[serde(default, skip_serializing_if = "Option::is_none")] + runtime_id: Option, }, /// Server-derived Workspace projection. Workspace identity comes from the /// authenticated connection and is deliberately absent from this selector. @@ -260,7 +262,16 @@ impl EventSubscriptionSelector { pub fn validate(&self) -> Result<(), SubscriptionValidationError> { match self { Self::WorkerLifecycle { worker_ids } => worker_ids.validate(), - Self::WorkerProtocol { worker_id } => worker_id.validate(), + Self::WorkerProtocol { + worker_id, + runtime_id, + } => { + worker_id.validate()?; + if let Some(runtime_id) = runtime_id { + validate_identifier("runtime_id", runtime_id, MAX_RESOURCE_ID_BYTES)?; + } + Ok(()) + } Self::RuntimeWorkers | Self::WorkspaceWorkers | Self::WorkspaceWorkdirs => Ok(()), } } @@ -271,6 +282,7 @@ impl EventSubscriptionSelector { Self::WorkerLifecycle { worker_ids } => worker_ids.contains(worker_id), Self::WorkerProtocol { worker_id: selected, + .. } => selected == worker_id, Self::WorkspaceWorkdirs => false, } @@ -312,6 +324,7 @@ pub enum SubscriptionFramePayload { Request(SubscriptionRequest), Response(SubscriptionResponse), Event(SubscriptionEvent), + WorkerProtocol(SubscriptionWorkerProtocolMethod), } impl SubscriptionFramePayload { @@ -320,10 +333,24 @@ impl SubscriptionFramePayload { Self::Request(request) => request.validate(), Self::Response(response) => response.validate(), Self::Event(event) => event.validate(), + Self::WorkerProtocol(message) => message.validate(), } } } +#[derive(Debug, Clone, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +pub struct SubscriptionWorkerProtocolMethod { + pub subscription_id: SubscriptionId, + pub method: crate::Method, +} + +impl SubscriptionWorkerProtocolMethod { + pub fn validate(&self) -> Result<(), SubscriptionValidationError> { + self.subscription_id.validate() + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "typescript", derive(ts_rs::TS))] #[serde(tag = "method", content = "params", rename_all = "snake_case")] @@ -617,6 +644,7 @@ impl SubscriptionSnapshot { ( EventSubscriptionSelector::WorkerProtocol { worker_id: selected, + .. }, Self::WorkerProtocol { worker_id, .. }, ) if selected == worker_id => worker_id.validate(), @@ -703,6 +731,7 @@ impl SubscriptionEventPayload { ( EventSubscriptionSelector::WorkerProtocol { worker_id: selected, + .. }, Self::WorkerProtocol { worker_id, .. }, ) if selected == worker_id => Ok(()), @@ -891,6 +920,37 @@ mod tests { )); } + #[test] + fn worker_protocol_method_uses_subscription_lane() { + let frame = SubscriptionFrame::new(SubscriptionFramePayload::WorkerProtocol( + SubscriptionWorkerProtocolMethod { + subscription_id: subscription_id(), + method: crate::Method::ListCompletions { + kind: crate::CompletionKind::File, + prefix: "src/".to_string(), + }, + }, + )); + frame.validate().unwrap(); + assert_eq!( + serde_json::to_value(frame).unwrap(), + serde_json::json!({ + "protocol_version": 1, + "frame": "worker_protocol", + "message": { + "subscription_id": "subscription-1", + "method": { + "method": "list_completions", + "params": { + "kind": "file", + "prefix": "src/" + } + } + } + }) + ); + } + #[test] fn workspace_snapshot_allows_equal_local_worker_ids_from_distinct_runtimes() { let mut first = worker("1"); diff --git a/crates/protocol/src/typescript.rs b/crates/protocol/src/typescript.rs index 607fd864..3bb2c778 100644 --- a/crates/protocol/src/typescript.rs +++ b/crates/protocol/src/typescript.rs @@ -12,7 +12,8 @@ use crate::{ SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest, SubscriptionRequestId, SubscriptionResponse, SubscriptionSnapshot, SubscriptionTerminationCode, SubscriptionWorkdir, SubscriptionWorkdirId, - SubscriptionWorker, SubscriptionWorkerId, SubscriptionWorkerIds, SubscriptionWorkerState, + SubscriptionWorker, SubscriptionWorkerId, SubscriptionWorkerIds, + SubscriptionWorkerProtocolMethod, SubscriptionWorkerState, }, }; @@ -71,6 +72,7 @@ pub fn generated_protocol_types() -> String { push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); diff --git a/crates/workspace-server/src/runtime_subscription.rs b/crates/workspace-server/src/runtime_subscription.rs index 1ab799d6..b28755a0 100644 --- a/crates/workspace-server/src/runtime_subscription.rs +++ b/crates/workspace-server/src/runtime_subscription.rs @@ -772,7 +772,9 @@ async fn handle_frame( send_subscribe(socket, state, selector).await?; } } - SubscriptionFramePayload::Request(_) => return Err(()), + SubscriptionFramePayload::Request(_) | SubscriptionFramePayload::WorkerProtocol(_) => { + return Err(()); + } } Ok(()) } diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 92b008ae..3d74536d 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -3807,10 +3807,9 @@ async fn scoped_workspace_protocol_ws( { return Err(StatusCode::FORBIDDEN.into_response()); } - let broker = api.runtime_subscription_broker().clone(); Ok(ws .on_upgrade(move |socket| { - crate::workspace_subscription::serve_workspace_subscription(broker, socket) + crate::workspace_subscription::serve_workspace_subscription(api, socket) }) .into_response()) } @@ -6685,6 +6684,131 @@ async fn worker_protocol_ws( ws.on_upgrade(move |socket| worker_protocol_ws_session(source, socket)) } +pub(crate) struct WorkspaceWorkerProtocolConnection { + pub(crate) methods: tokio::sync::mpsc::Sender, + pub(crate) events: tokio::sync::mpsc::Receiver, +} + +pub(crate) async fn connect_workspace_worker_protocol( + api: &WorkspaceApi, + runtime_id: &str, + worker_id: &str, +) -> Result { + let source = match api.observation_proxy.source(runtime_id, worker_id) { + Ok(source) => source, + Err(ObservationProxyError::WorkerNotFound(_)) => api + .runtime + .observation_source(runtime_id, worker_id) + .map_err(|error| error.into_error())?, + Err(error) => { + return Err(Error::RuntimeOperationFailed { + runtime_id: runtime_id.to_string(), + code: error.code().to_string(), + message: error.message().to_string(), + }); + } + }; + match source { + RuntimeObservationSource::RemoteWs(config) => connect_remote_worker_protocol(config).await, + RuntimeObservationSource::Embedded(source) => { + connect_embedded_worker_protocol(source).await + } + } +} + +async fn connect_remote_worker_protocol( + config: RuntimeObservationSourceConfig, +) -> Result { + let mut request = config + .endpoint + .clone() + .into_client_request() + .map_err(|error| Error::Config(format!("invalid Runtime protocol endpoint: {error}")))?; + if let Some(token) = &config.bearer_token { + request.headers_mut().insert( + "authorization", + format!("Bearer {token}").parse().map_err(|error| { + Error::Config(format!("invalid Runtime authorization: {error}")) + })?, + ); + } + let (socket, _) = + connect_async(request) + .await + .map_err(|error| Error::RuntimeOperationFailed { + runtime_id: config.runtime_id.clone(), + code: "worker_protocol_connect_failed".to_string(), + message: error.to_string(), + })?; + let (mut sink, mut stream) = socket.split(); + let (methods, mut method_receiver) = tokio::sync::mpsc::channel(256); + let (event_sender, events) = tokio::sync::mpsc::channel(512); + tokio::spawn(async move { + loop { + tokio::select! { + method = method_receiver.recv() => { + let Some(method) = method else { break }; + let Ok(text) = protocol::stream::encode_method(&method) else { break }; + if sink.send(TungsteniteMessage::Text(text.into())).await.is_err() { break; } + } + message = stream.next() => match message { + Some(Ok(TungsteniteMessage::Text(text))) => { + let Ok(event) = protocol::stream::decode_event(text.as_ref()) else { break; }; + if event_sender.send(event).await.is_err() { break; } + } + Some(Ok(TungsteniteMessage::Ping(value))) => { + if sink.send(TungsteniteMessage::Pong(value)).await.is_err() { break; } + } + Some(Ok(TungsteniteMessage::Pong(_))) => {} + _ => break, + } + } + } + }); + Ok(WorkspaceWorkerProtocolConnection { methods, events }) +} + +async fn connect_embedded_worker_protocol( + source: crate::observation::EmbeddedRuntimeObservationSource, +) -> Result { + let mut upstream = + RuntimeObservationClient::connect(&RuntimeObservationSource::Embedded(source.clone())) + .await + .map_err(|error| Error::RuntimeOperationFailed { + runtime_id: source.runtime_id.clone(), + code: error.code().to_string(), + message: error.message().to_string(), + })?; + let (methods, mut method_receiver) = tokio::sync::mpsc::channel(256); + let (event_sender, events) = tokio::sync::mpsc::channel(512); + tokio::spawn(async move { + loop { + tokio::select! { + method = method_receiver.recv() => { + let Some(method) = method else { break }; + match source.runtime.send_protocol_method(&source.worker_ref, method) { + Ok(direct_events) => { + for event in direct_events { + if event_sender.send(event).await.is_err() { return; } + } + } + Err(error) => { + if event_sender.send(protocol_error_event(error.to_string())).await.is_err() { return; } + } + } + } + event = upstream.next_event() => match event { + Ok(event) => { + if event_sender.send(event.payload).await.is_err() { break; } + } + Err(_) => break, + } + } + } + }); + Ok(WorkspaceWorkerProtocolConnection { methods, events }) +} + async fn worker_protocol_ws_session(source: RuntimeObservationSource, socket: WebSocket) { match source { RuntimeObservationSource::RemoteWs(config) => { @@ -12716,6 +12840,29 @@ mod tests { .unwrap() .to_string(); + let spawn_request = WorkerSpawnRequest { + intent: WorkerSpawnIntent::WorkspaceCompanion, + requested_worker_name: Some("multiplexed-console".to_string()), + acceptance: WorkerSpawnAcceptanceRequirement::RunAccepted { + expected_segments: 0, + }, + profile: worker_runtime::catalog::ProfileSelector::Builtin( + "builtin:companion".to_string(), + ), + ticket_assignment: None, + initial_input: None, + working_directory_request: None, + resolved_working_directory_request: None, + resolved_working_directory: None, + resolved_config_bundle: Some(runtime_test_bundle()), + resolved_workspace_api: None, + }; + let spawned = api + .spawn_workspace_worker(EMBEDDED_WORKER_RUNTIME_ID, spawn_request) + .unwrap(); + assert_eq!(spawned.state, WorkerOperationState::Accepted); + let worker_id = spawned.worker.unwrap().worker_id; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); let app = build_router(api); @@ -12757,6 +12904,141 @@ mod tests { } ) )); + + let subscribe_protocol = protocol::subscription::SubscriptionFrame::new( + protocol::subscription::SubscriptionFramePayload::Request( + protocol::subscription::SubscriptionRequest::SubscribeEvents { + request_id: protocol::subscription::SubscriptionRequestId::new("request-2") + .unwrap(), + selector: protocol::subscription::EventSubscriptionSelector::WorkerProtocol { + worker_id: protocol::subscription::SubscriptionWorkerId::new( + worker_id.clone(), + ) + .unwrap(), + runtime_id: Some(EMBEDDED_WORKER_RUNTIME_ID.to_string()), + }, + }, + ), + ); + socket + .send(Message::Text( + serde_json::to_string(&subscribe_protocol).unwrap().into(), + )) + .await + .unwrap(); + let protocol_subscription_id = loop { + let Message::Text(text) = socket.next().await.unwrap().unwrap() else { + continue; + }; + let frame: protocol::subscription::SubscriptionFrame = + serde_json::from_str(text.as_str()).unwrap(); + if let protocol::subscription::SubscriptionFramePayload::Response( + protocol::subscription::SubscriptionResponse::Subscribed { + subscription_id, + selector: + protocol::subscription::EventSubscriptionSelector::WorkerProtocol { .. }, + .. + }, + ) = frame.payload + { + break subscription_id; + } + }; + let second_subscribe = protocol::subscription::SubscriptionFrame::new( + protocol::subscription::SubscriptionFramePayload::Request( + protocol::subscription::SubscriptionRequest::SubscribeEvents { + request_id: protocol::subscription::SubscriptionRequestId::new("request-3") + .unwrap(), + selector: protocol::subscription::EventSubscriptionSelector::WorkerProtocol { + worker_id: protocol::subscription::SubscriptionWorkerId::new( + worker_id.clone(), + ) + .unwrap(), + runtime_id: Some(EMBEDDED_WORKER_RUNTIME_ID.to_string()), + }, + }, + ), + ); + socket + .send(Message::Text( + serde_json::to_string(&second_subscribe).unwrap().into(), + )) + .await + .unwrap(); + let second_protocol_subscription_id = loop { + let Message::Text(text) = socket.next().await.unwrap().unwrap() else { + continue; + }; + let frame: protocol::subscription::SubscriptionFrame = + serde_json::from_str(text.as_str()).unwrap(); + if let protocol::subscription::SubscriptionFramePayload::Response( + protocol::subscription::SubscriptionResponse::Subscribed { + subscription_id, + selector: + protocol::subscription::EventSubscriptionSelector::WorkerProtocol { .. }, + .. + }, + ) = frame.payload + { + break subscription_id; + } + }; + assert_ne!(protocol_subscription_id, second_protocol_subscription_id); + let unsubscribe = protocol::subscription::SubscriptionFrame::new( + protocol::subscription::SubscriptionFramePayload::Request( + protocol::subscription::SubscriptionRequest::UnsubscribeEvents { + request_id: protocol::subscription::SubscriptionRequestId::new("request-4") + .unwrap(), + subscription_id: protocol_subscription_id, + }, + ), + ); + socket + .send(Message::Text( + serde_json::to_string(&unsubscribe).unwrap().into(), + )) + .await + .unwrap(); + + let method = protocol::subscription::SubscriptionFrame::new( + protocol::subscription::SubscriptionFramePayload::WorkerProtocol( + protocol::subscription::SubscriptionWorkerProtocolMethod { + subscription_id: second_protocol_subscription_id.clone(), + method: protocol::Method::ListCompletions { + kind: protocol::CompletionKind::File, + prefix: String::new(), + }, + }, + ), + ); + socket + .send(Message::Text( + serde_json::to_string(&method).unwrap().into(), + )) + .await + .unwrap(); + loop { + let Message::Text(text) = socket.next().await.unwrap().unwrap() else { + continue; + }; + let frame: protocol::subscription::SubscriptionFrame = + serde_json::from_str(text.as_str()).unwrap(); + if matches!( + frame.payload, + protocol::subscription::SubscriptionFramePayload::Event( + protocol::subscription::SubscriptionEvent::Event { + subscription_id, + payload: protocol::subscription::SubscriptionEventPayload::WorkerProtocol { + event: protocol::Event::Completions { .. }, + .. + }, + .. + } + ) if subscription_id == second_protocol_subscription_id + ) { + break; + } + } server.abort(); } diff --git a/crates/workspace-server/src/workspace_subscription.rs b/crates/workspace-server/src/workspace_subscription.rs index 10414816..9584578a 100644 --- a/crates/workspace-server/src/workspace_subscription.rs +++ b/crates/workspace-server/src/workspace_subscription.rs @@ -10,24 +10,35 @@ use protocol::subscription::{ use tokio::sync::mpsc; use crate::runtime_subscription::{BrokerSubscriptionEvent, RuntimeSubscriptionBroker}; +use crate::server::{WorkspaceApi, connect_workspace_worker_protocol}; const OUTBOUND_CAPACITY: usize = 256; -pub(crate) async fn serve_workspace_subscription( - broker: RuntimeSubscriptionBroker, - socket: WebSocket, -) { +struct ActiveSubscription { + task: tokio::task::JoinHandle<()>, + methods: Option>, +} + +pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebSocket) { + let broker = api.runtime_subscription_broker().clone(); let (mut socket_sender, mut socket_receiver) = socket.split(); - let (outbound, mut outbound_receiver) = mpsc::channel::(OUTBOUND_CAPACITY); + let (control_outbound, mut control_receiver) = mpsc::channel::(OUTBOUND_CAPACITY); + let (protocol_outbound, mut protocol_receiver) = mpsc::channel::(OUTBOUND_CAPACITY); let writer = tokio::spawn(async move { - while let Some(message) = outbound_receiver.recv().await { + loop { + let message = tokio::select! { + biased; + message = control_receiver.recv() => message, + message = protocol_receiver.recv() => message, + }; + let Some(message) = message else { break }; if socket_sender.send(message).await.is_err() { break; } } }); let mut next_subscription_id = 1_u64; - let mut subscriptions = HashMap::>::new(); + let mut subscriptions = HashMap::::new(); while let Some(message) = socket_receiver.next().await { let Ok(message) = message else { break }; @@ -39,50 +50,91 @@ pub(crate) async fn serve_workspace_subscription( if frame.validate().is_err() { break; } - let SubscriptionFramePayload::Request(request) = frame.payload else { - break; - }; - subscriptions.retain(|_, task| !task.is_finished()); - match request { - SubscriptionRequest::SubscribeEvents { + subscriptions.retain(|_, subscription| !subscription.task.is_finished()); + match frame.payload { + SubscriptionFramePayload::Request(SubscriptionRequest::SubscribeEvents { request_id, selector, - } => { - if selector != EventSubscriptionSelector::WorkspaceWorkers { - let _ = send_frame(&outbound, SubscriptionFrame::new( - SubscriptionFramePayload::Response( - SubscriptionResponse::SubscriptionRejected { - request_id, - subscription_id: None, - code: SubscriptionRejectionCode::UnsupportedSelector, - message: "Workspace clients may subscribe only to workspace_workers on this endpoint".to_string(), - }, - ), - )).await; - continue; - } + }) => { let subscription_id = SubscriptionId::new(format!( "workspace-subscription-{next_subscription_id}" )) .expect("generated Workspace subscription id is valid"); next_subscription_id = next_subscription_id.saturating_add(1); - let task = tokio::spawn(run_workspace_workers( - broker.clone(), - request_id, - subscription_id.clone(), - outbound.clone(), - )); - subscriptions.insert(subscription_id, task); + match selector { + EventSubscriptionSelector::WorkspaceWorkers => { + let task = tokio::spawn(run_workspace_workers( + broker.clone(), + request_id, + subscription_id.clone(), + control_outbound.clone(), + )); + subscriptions.insert( + subscription_id, + ActiveSubscription { + task, + methods: None, + }, + ); + } + EventSubscriptionSelector::WorkerProtocol { + worker_id, + runtime_id: Some(runtime_id), + } => { + match connect_workspace_worker_protocol( + &api, + &runtime_id, + worker_id.as_str(), + ) + .await + { + Ok(connection) => { + let methods = connection.methods.clone(); + let task = tokio::spawn(run_worker_protocol( + request_id, + subscription_id.clone(), + runtime_id, + worker_id, + connection.events, + control_outbound.clone(), + protocol_outbound.clone(), + )); + subscriptions.insert( + subscription_id, + ActiveSubscription { + task, + methods: Some(methods), + }, + ); + } + Err(error) => { + let _ = send_rejected( + &control_outbound, + request_id, + SubscriptionRejectionCode::ResourceNotFound, + error.to_string(), + ) + .await; + } + } + } + _ => { + let _ = send_rejected( + &control_outbound, request_id, SubscriptionRejectionCode::UnsupportedSelector, + "Workspace clients may subscribe only to workspace_workers or a runtime-scoped worker_protocol selector".to_string(), + ).await; + } + } } - SubscriptionRequest::UnsubscribeEvents { + SubscriptionFramePayload::Request(SubscriptionRequest::UnsubscribeEvents { request_id, subscription_id, - } => { - if let Some(task) = subscriptions.remove(&subscription_id) { - task.abort(); + }) => { + if let Some(subscription) = subscriptions.remove(&subscription_id) { + subscription.task.abort(); } if send_frame( - &outbound, + &control_outbound, SubscriptionFrame::new(SubscriptionFramePayload::Response( SubscriptionResponse::Unsubscribed { request_id, @@ -96,10 +148,24 @@ pub(crate) async fn serve_workspace_subscription( break; } } + SubscriptionFramePayload::WorkerProtocol(message) => { + let Some(methods) = subscriptions + .get(&message.subscription_id) + .and_then(|value| value.methods.clone()) + else { + break; + }; + if methods.send(message.method).await.is_err() { + break; + } + } + SubscriptionFramePayload::Response(_) | SubscriptionFramePayload::Event(_) => { + break; + } } } WsMessage::Ping(value) => { - if outbound.send(WsMessage::Pong(value)).await.is_err() { + if control_outbound.send(WsMessage::Pong(value)).await.is_err() { break; } } @@ -108,13 +174,108 @@ pub(crate) async fn serve_workspace_subscription( } } - for (_, task) in subscriptions { - task.abort(); + for (_, subscription) in subscriptions { + subscription.task.abort(); } - drop(outbound); + drop(control_outbound); + drop(protocol_outbound); let _ = writer.await; } +async fn send_rejected( + outbound: &mpsc::Sender, + request_id: protocol::subscription::SubscriptionRequestId, + code: SubscriptionRejectionCode, + message: String, +) -> Result<(), ()> { + send_frame( + outbound, + SubscriptionFrame::new(SubscriptionFramePayload::Response( + SubscriptionResponse::SubscriptionRejected { + request_id, + subscription_id: None, + code, + message, + }, + )), + ) + .await +} + +async fn run_worker_protocol( + request_id: protocol::subscription::SubscriptionRequestId, + subscription_id: SubscriptionId, + runtime_id: String, + worker_id: protocol::subscription::SubscriptionWorkerId, + mut events: mpsc::Receiver, + control_outbound: mpsc::Sender, + protocol_outbound: mpsc::Sender, +) { + if send_frame( + &control_outbound, + SubscriptionFrame::new(SubscriptionFramePayload::Response( + SubscriptionResponse::Subscribed { + request_id, + subscription_id: subscription_id.clone(), + selector: EventSubscriptionSelector::WorkerProtocol { + worker_id: worker_id.clone(), + runtime_id: Some(runtime_id), + }, + snapshot_revision: 0, + snapshot: SubscriptionSnapshot::WorkerProtocol { + worker_id: worker_id.clone(), + events: Vec::new(), + }, + }, + )), + ) + .await + .is_err() + { + return; + } + let mut subject_revision = 0_u64; + while let Some(event) = events.recv().await { + subject_revision = subject_revision.saturating_add(1); + let frame = + SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event { + subscription_id: subscription_id.clone(), + subject_revision, + payload: SubscriptionEventPayload::WorkerProtocol { + worker_id: worker_id.clone(), + event, + }, + })); + if try_send_frame(&protocol_outbound, frame).is_err() { + let _ = send_frame( + &control_outbound, + SubscriptionFrame::new(SubscriptionFramePayload::Event( + SubscriptionEvent::SubscriptionClosed { + subscription_id: subscription_id.clone(), + code: SubscriptionTerminationCode::Lagged, + message: + "Worker protocol subscriber lagged; resubscribe for a fresh snapshot" + .to_string(), + }, + )), + ) + .await; + return; + } + } + let _ = send_frame( + &control_outbound, + SubscriptionFrame::new(SubscriptionFramePayload::Event( + SubscriptionEvent::SubscriptionClosed { + subscription_id, + code: SubscriptionTerminationCode::ResourceGone, + message: "Worker protocol stream closed".to_string(), + }, + )), + ) + .await; +} + async fn run_workspace_workers( broker: RuntimeSubscriptionBroker, request_id: protocol::subscription::SubscriptionRequestId, @@ -352,6 +513,14 @@ async fn send_event( .await } +fn try_send_frame(outbound: &mpsc::Sender, frame: SubscriptionFrame) -> Result<(), ()> { + frame.validate().map_err(|_| ())?; + let text = serde_json::to_string(&frame).map_err(|_| ())?; + outbound + .try_send(WsMessage::Text(text.into())) + .map_err(|_| ()) +} + async fn send_frame( outbound: &mpsc::Sender, frame: SubscriptionFrame, diff --git a/web/workspace/src/lib/generated/protocol.ts b/web/workspace/src/lib/generated/protocol.ts index 3903aeee..7a0ff0af 100644 --- a/web/workspace/src/lib/generated/protocol.ts +++ b/web/workspace/src/lib/generated/protocol.ts @@ -111,7 +111,7 @@ export type SubscriptionWorkerIds = Array; export type SubscriptionWorkerState = "idle" | "running" | "paused" | "stopped" | "cancelled"; -export type EventSubscriptionSelector = { "topic": "runtime_workers" } | { "topic": "worker_lifecycle", worker_ids: SubscriptionWorkerIds, } | { "topic": "worker_protocol", worker_id: SubscriptionWorkerId, } | { "topic": "workspace_workers" } | { "topic": "workspace_workdirs" }; +export type EventSubscriptionSelector = { "topic": "runtime_workers" } | { "topic": "worker_lifecycle", worker_ids: SubscriptionWorkerIds, } | { "topic": "worker_protocol", worker_id: SubscriptionWorkerId, runtime_id?: string | null, } | { "topic": "workspace_workers" } | { "topic": "workspace_workdirs" }; export type SubscriptionWorker = { worker_id: SubscriptionWorkerId, /** @@ -136,13 +136,15 @@ export type SubscriptionTerminationCode = "lagged" | "resource_gone" | "unauthor export type SubscriptionRequest = { "method": "subscribe_events", "params": { request_id: SubscriptionRequestId, selector: EventSubscriptionSelector, } } | { "method": "unsubscribe_events", "params": { request_id: SubscriptionRequestId, subscription_id: SubscriptionId, } }; +export type SubscriptionWorkerProtocolMethod = { subscription_id: SubscriptionId, method: Method, }; + export type SubscriptionResponse = { "result": "subscribed", "payload": { request_id: SubscriptionRequestId, subscription_id: SubscriptionId, selector: EventSubscriptionSelector, snapshot_revision: number, snapshot: SubscriptionSnapshot, } } | { "result": "unsubscribed", "payload": { request_id: SubscriptionRequestId, subscription_id: SubscriptionId, } } | { "result": "subscription_rejected", "payload": { request_id: SubscriptionRequestId, subscription_id?: SubscriptionId | null, code: SubscriptionRejectionCode, message: string, } }; export type SubscriptionEvent = { "event": "event", "data": { subscription_id: SubscriptionId, subject_revision: number, payload: SubscriptionEventPayload, } } | { "event": "subscription_closed", "data": { subscription_id: SubscriptionId, code: SubscriptionTerminationCode, message: string, } }; -export type SubscriptionFramePayload = { "frame": "request", "message": SubscriptionRequest } | { "frame": "response", "message": SubscriptionResponse } | { "frame": "event", "message": SubscriptionEvent }; +export type SubscriptionFramePayload = { "frame": "request", "message": SubscriptionRequest } | { "frame": "response", "message": SubscriptionResponse } | { "frame": "event", "message": SubscriptionEvent } | { "frame": "worker_protocol", "message": SubscriptionWorkerProtocolMethod }; -export type SubscriptionFrame = { protocol_version: number, } & ({ "frame": "request", "message": SubscriptionRequest } | { "frame": "response", "message": SubscriptionResponse } | { "frame": "event", "message": SubscriptionEvent }); +export type SubscriptionFrame = { protocol_version: number, } & ({ "frame": "request", "message": SubscriptionRequest } | { "frame": "response", "message": SubscriptionResponse } | { "frame": "event", "message": SubscriptionEvent } | { "frame": "worker_protocol", "message": SubscriptionWorkerProtocolMethod }); export type Method = { "method": "run", "params": { input: Array, } } | { "method": "notify", "params": { message: string, auto_run?: boolean, } } | { "method": "worker_event", "params": WorkerEvent } | { "method": "resume" } | { "method": "cancel" } | { "method": "pause" } | { "method": "compact" } | { "method": "list_rewind_targets" } | { "method": "rewind_to", "params": { target: RewindTargetId, expected_head_entries: number, } } | { "method": "shutdown" } | { "method": "list_completions", "params": { kind: CompletionKind, prefix: string, } } | { "method": "list_workers" } | { "method": "restore_worker", "params": { name: string, } } | { "method": "register_peer", "params": { name: string, } };