server: multiplex worker protocol subscriptions

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<String>,
     },
     /// 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");

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::<SubscriptionRejectionCode>(&cfg, &mut output);
     push_decl::<SubscriptionTerminationCode>(&cfg, &mut output);
     push_decl::<SubscriptionRequest>(&cfg, &mut output);
+    push_decl::<SubscriptionWorkerProtocolMethod>(&cfg, &mut output);
     push_decl::<SubscriptionResponse>(&cfg, &mut output);
     push_decl::<SubscriptionEvent>(&cfg, &mut output);
     push_decl::<SubscriptionFramePayload>(&cfg, &mut output);

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

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<protocol::Method>,
+    pub(crate) events: tokio::sync::mpsc::Receiver<protocol::Event>,
+}
+
+pub(crate) async fn connect_workspace_worker_protocol(
+    api: &WorkspaceApi,
+    runtime_id: &str,
+    worker_id: &str,
+) -> Result<WorkspaceWorkerProtocolConnection> {
+    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<WorkspaceWorkerProtocolConnection> {
+    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<WorkspaceWorkerProtocolConnection> {
+    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();
     }
 

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(
+struct ActiveSubscription {
-    broker: RuntimeSubscriptionBroker,
+    task: tokio::task::JoinHandle<()>,
-    socket: WebSocket,
+    methods: Option<mpsc::Sender<protocol::Method>>,
-) {
+}
+
+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::<WsMessage>(OUTBOUND_CAPACITY);
+    let (control_outbound, mut control_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY);
+    let (protocol_outbound, mut protocol_receiver) = mpsc::channel::<WsMessage>(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::<SubscriptionId, tokio::task::JoinHandle<()>>::new();
+    let mut subscriptions = HashMap::<SubscriptionId, ActiveSubscription>::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 {
+                subscriptions.retain(|_, subscription| !subscription.task.is_finished());
-                    break;
+                match frame.payload {
-                };
+                    SubscriptionFramePayload::Request(SubscriptionRequest::SubscribeEvents {
-                subscriptions.retain(|_, task| !task.is_finished());
-                match 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(
+                        match selector {
-                            broker.clone(),
+                            EventSubscriptionSelector::WorkspaceWorkers => {
-                            request_id,
+                                let task = tokio::spawn(run_workspace_workers(
-                            subscription_id.clone(),
+                                    broker.clone(),
-                            outbound.clone(),
+                                    request_id,
-                        ));
+                                    subscription_id.clone(),
-                        subscriptions.insert(subscription_id, task);
+                                    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) {
+                        if let Some(subscription) = subscriptions.remove(&subscription_id) {
-                            task.abort();
+                            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 {
+    for (_, subscription) in subscriptions {
-        task.abort();
+        subscription.task.abort();
     }
-    drop(outbound);
+    drop(control_outbound);
+    drop(protocol_outbound);
     let _ = writer.await;
 }
 
+async fn send_rejected(
+    outbound: &mpsc::Sender<WsMessage>,
+    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<protocol::Event>,
+    control_outbound: mpsc::Sender<WsMessage>,
+    protocol_outbound: mpsc::Sender<WsMessage>,
+) {
+    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<WsMessage>, 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<WsMessage>,
     frame: SubscriptionFrame,

web/workspace/src/lib/generated/protocol.ts

@@ -111,7 +111,7 @@ export type SubscriptionWorkerIds = Array<SubscriptionWorkerId>;
 
 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<Segment>, } } | { "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, } };