server: broker runtime event subscriptions

crates/workspace-server/Cargo.toml

@@ -31,7 +31,7 @@ sha2.workspace = true
 thiserror.workspace = true
 ticket.workspace = true
 memory.workspace = true
-tokio = { workspace = true, features = ["fs", "macros", "net", "rt-multi-thread", "sync"] }
+tokio = { workspace = true, features = ["fs", "macros", "net", "rt-multi-thread", "sync", "time"] }
 tokio-tungstenite.workspace = true
 worker.workspace = true
 worker-runtime.workspace = true

crates/workspace-server/src/lib.rs

@@ -19,6 +19,7 @@ pub mod records;
 pub use records::ticket_api_typescript;
 pub mod repositories;
 pub mod resource_broker;
+pub mod runtime_subscription;
 pub mod server;
 pub mod skills;
 pub mod store;

crates/workspace-server/src/runtime_subscription.rs

@@ -0,0 +1,737 @@
+use std::collections::{HashMap, HashSet};
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{Arc, RwLock};
+use std::time::Duration;
+
+use futures::{SinkExt, StreamExt};
+use protocol::subscription::{
+    EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
+    SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest,
+    SubscriptionRequestId, SubscriptionResponse, SubscriptionSnapshot, SubscriptionTerminationCode,
+};
+use tokio::sync::mpsc;
+use tokio_tungstenite::connect_async;
+use tokio_tungstenite::tungstenite::Message;
+use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+use worker_runtime::auth::{CapabilityTokenSigner, capability_claims};
+
+use crate::hosts::RemoteRuntimeConfig;
+
+const DOWNSTREAM_QUEUE_CAPACITY: usize = 256;
+const RECONNECT_DELAY: Duration = Duration::from_millis(100);
+
+type RuntimeSocket =
+    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
+
+#[derive(Debug, thiserror::Error)]
+pub enum RuntimeSubscriptionBrokerError {
+    #[error("unknown Runtime {0:?}")]
+    UnknownRuntime(String),
+    #[error("Runtime subscription broker command channel closed")]
+    Closed,
+}
+
+#[derive(Clone, Debug)]
+pub enum BrokerSubscriptionEvent {
+    Snapshot {
+        connection_generation: u64,
+        snapshot_revision: u64,
+        snapshot: SubscriptionSnapshot,
+    },
+    Event {
+        connection_generation: u64,
+        subject_revision: u64,
+        payload: SubscriptionEventPayload,
+    },
+    Disconnected {
+        connection_generation: u64,
+        message: String,
+    },
+    Rejected {
+        connection_generation: u64,
+        code: SubscriptionRejectionCode,
+        message: String,
+    },
+    Closed {
+        connection_generation: u64,
+        code: SubscriptionTerminationCode,
+        message: String,
+    },
+}
+
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub struct RuntimeSubscriptionBrokerStatus {
+    pub connection_generation: u64,
+    pub connected: bool,
+    pub desired_selectors: usize,
+    pub upstream_subscriptions: usize,
+}
+
+pub struct BrokerSubscription {
+    downstream_id: u64,
+    runtime_id: String,
+    selector: EventSubscriptionSelector,
+    receiver: mpsc::Receiver<BrokerSubscriptionEvent>,
+    commands: mpsc::UnboundedSender<Command>,
+}
+
+impl BrokerSubscription {
+    pub fn runtime_id(&self) -> &str {
+        &self.runtime_id
+    }
+    pub fn selector(&self) -> &EventSubscriptionSelector {
+        &self.selector
+    }
+    pub async fn recv(&mut self) -> Option<BrokerSubscriptionEvent> {
+        self.receiver.recv().await
+    }
+}
+
+impl Drop for BrokerSubscription {
+    fn drop(&mut self) {
+        let _ = self.commands.send(Command::Unsubscribe(self.downstream_id));
+    }
+}
+
+#[derive(Clone)]
+struct Registration {
+    generation: u64,
+    commands: mpsc::UnboundedSender<Command>,
+    status: Arc<RwLock<RuntimeSubscriptionBrokerStatus>>,
+}
+
+#[derive(Clone)]
+pub struct RuntimeSubscriptionBroker {
+    workspace_id: Arc<str>,
+    next_generation: Arc<AtomicU64>,
+    next_downstream: Arc<AtomicU64>,
+    registrations: Arc<RwLock<HashMap<String, Registration>>>,
+}
+
+impl RuntimeSubscriptionBroker {
+    pub fn new(workspace_id: impl Into<String>) -> Self {
+        Self {
+            workspace_id: Arc::from(workspace_id.into()),
+            next_generation: Arc::new(AtomicU64::new(1)),
+            next_downstream: Arc::new(AtomicU64::new(1)),
+            registrations: Arc::new(RwLock::new(HashMap::new())),
+        }
+    }
+
+    pub fn register_remote_runtime(&self, config: RemoteRuntimeConfig) -> u64 {
+        let generation = self.next_generation.fetch_add(1, Ordering::Relaxed);
+        let (commands, receiver) = mpsc::unbounded_channel();
+        let status = Arc::new(RwLock::new(RuntimeSubscriptionBrokerStatus {
+            connection_generation: generation,
+            ..Default::default()
+        }));
+        let registration = Registration {
+            generation,
+            commands: commands.clone(),
+            status: status.clone(),
+        };
+        let previous = self
+            .registrations
+            .write()
+            .expect("broker registry poisoned")
+            .insert(config.runtime_id.clone(), registration);
+        if let Some(previous) = previous {
+            let _ = previous.commands.send(Command::Shutdown(generation));
+        }
+        tokio::spawn(run_connection(
+            config,
+            self.workspace_id.to_string(),
+            generation,
+            receiver,
+            status,
+        ));
+        generation
+    }
+
+    pub fn unregister_runtime(&self, runtime_id: &str) {
+        if let Some(registration) = self
+            .registrations
+            .write()
+            .expect("broker registry poisoned")
+            .remove(runtime_id)
+        {
+            let _ = registration
+                .commands
+                .send(Command::Shutdown(registration.generation.saturating_add(1)));
+        }
+    }
+
+    pub fn status(&self, runtime_id: &str) -> Option<RuntimeSubscriptionBrokerStatus> {
+        let status = self
+            .registrations
+            .read()
+            .expect("broker registry poisoned")
+            .get(runtime_id)?
+            .status
+            .clone();
+        Some(status.read().expect("broker status poisoned").clone())
+    }
+
+    pub fn subscribe(
+        &self,
+        runtime_id: &str,
+        selector: EventSubscriptionSelector,
+    ) -> Result<BrokerSubscription, RuntimeSubscriptionBrokerError> {
+        selector
+            .validate()
+            .map_err(|_| RuntimeSubscriptionBrokerError::Closed)?;
+        let registration = self
+            .registrations
+            .read()
+            .expect("broker registry poisoned")
+            .get(runtime_id)
+            .cloned()
+            .ok_or_else(|| {
+                RuntimeSubscriptionBrokerError::UnknownRuntime(runtime_id.to_string())
+            })?;
+        let downstream_id = self.next_downstream.fetch_add(1, Ordering::Relaxed);
+        let (events, receiver) = mpsc::channel(DOWNSTREAM_QUEUE_CAPACITY);
+        registration
+            .commands
+            .send(Command::Subscribe {
+                downstream_id,
+                selector: selector.clone(),
+                events,
+            })
+            .map_err(|_| RuntimeSubscriptionBrokerError::Closed)?;
+        Ok(BrokerSubscription {
+            downstream_id,
+            runtime_id: runtime_id.to_string(),
+            selector,
+            receiver,
+            commands: registration.commands,
+        })
+    }
+}
+
+#[derive(Debug)]
+enum Command {
+    Subscribe {
+        downstream_id: u64,
+        selector: EventSubscriptionSelector,
+        events: mpsc::Sender<BrokerSubscriptionEvent>,
+    },
+    Unsubscribe(u64),
+    Shutdown(u64),
+}
+
+struct SelectorState {
+    downstreams: HashMap<u64, mpsc::Sender<BrokerSubscriptionEvent>>,
+    upstream_id: Option<SubscriptionId>,
+    pending: bool,
+    snapshot: Option<(u64, SubscriptionSnapshot)>,
+    revisions: HashMap<String, u64>,
+}
+impl SelectorState {
+    fn new() -> Self {
+        Self {
+            downstreams: HashMap::new(),
+            upstream_id: None,
+            pending: false,
+            snapshot: None,
+            revisions: HashMap::new(),
+        }
+    }
+}
+
+struct State {
+    generation: u64,
+    next_request: u64,
+    selectors: HashMap<EventSubscriptionSelector, SelectorState>,
+    downstream_index: HashMap<u64, EventSubscriptionSelector>,
+    pending: HashMap<SubscriptionRequestId, EventSubscriptionSelector>,
+    upstream_index: HashMap<SubscriptionId, EventSubscriptionSelector>,
+}
+impl State {
+    fn new(generation: u64) -> Self {
+        Self {
+            generation,
+            next_request: 1,
+            selectors: HashMap::new(),
+            downstream_index: HashMap::new(),
+            pending: HashMap::new(),
+            upstream_index: HashMap::new(),
+        }
+    }
+    fn request_id(&mut self) -> SubscriptionRequestId {
+        let id = self.next_request;
+        self.next_request = self.next_request.saturating_add(1);
+        SubscriptionRequestId::new(format!("server-{}-{id}", self.generation)).unwrap()
+    }
+    fn disconnected(&mut self, message: String) {
+        self.pending.clear();
+        self.upstream_index.clear();
+        for selector in self.selectors.values_mut() {
+            selector.upstream_id = None;
+            selector.pending = false;
+            selector.snapshot = None;
+            selector.revisions.clear();
+            broadcast(
+                &mut selector.downstreams,
+                BrokerSubscriptionEvent::Disconnected {
+                    connection_generation: self.generation,
+                    message: message.clone(),
+                },
+            );
+        }
+    }
+}
+
+async fn run_connection(
+    config: RemoteRuntimeConfig,
+    workspace_id: String,
+    generation: u64,
+    mut commands: mpsc::UnboundedReceiver<Command>,
+    status: Arc<RwLock<RuntimeSubscriptionBrokerStatus>>,
+) {
+    let mut state = State::new(generation);
+    let mut disconnect_notified = false;
+    loop {
+        update_status(&status, &state, false);
+        let connecting = connect_runtime(&config, &workspace_id);
+        tokio::pin!(connecting);
+        let connection = loop {
+            tokio::select! {
+                command = commands.recv() => match command {
+                    Some(Command::Shutdown(replacement)) => { close_all(&mut state, replacement); return; }
+                    Some(command) => { apply_offline(&mut state, command); update_status(&status, &state, false); }
+                    None => return,
+                },
+                connected = &mut connecting => break connected,
+            }
+        };
+        let mut socket = match connection {
+            Ok(socket) => socket,
+            Err(error) => {
+                if !disconnect_notified {
+                    state.disconnected(error);
+                    disconnect_notified = true;
+                }
+                tokio::time::sleep(RECONNECT_DELAY).await;
+                continue;
+            }
+        };
+        if resubscribe_all(&mut socket, &mut state).await.is_err() {
+            state.disconnected("failed to restore Runtime subscriptions".into());
+            disconnect_notified = true;
+            tokio::time::sleep(RECONNECT_DELAY).await;
+            continue;
+        }
+        update_status(&status, &state, true);
+        let reason = loop {
+            tokio::select! {
+                command = commands.recv() => match command {
+                    Some(Command::Shutdown(replacement)) => { let _ = socket.close(None).await; close_all(&mut state, replacement); return; }
+                    Some(command) => if apply_online(&mut socket, &mut state, command).await.is_err() { break "failed to apply Runtime subscription command".into(); },
+                    None => return,
+                },
+                message = socket.next() => match message {
+                    Some(Ok(Message::Text(text))) => match serde_json::from_str::<SubscriptionFrame>(text.as_str()) {
+                        Ok(frame) if frame.validate().is_ok() => if handle_frame(&mut socket, &mut state, frame).await.is_err() { break "invalid Runtime subscription transition".into(); },
+                        _ => break "Runtime returned an invalid subscription frame".into(),
+                    },
+                    Some(Ok(Message::Ping(value))) => if socket.send(Message::Pong(value)).await.is_err() { break "Runtime pong failed".into(); },
+                    Some(Ok(Message::Pong(_))) => {},
+                    Some(Ok(Message::Close(_))) | None => break "Runtime subscription connection closed".into(),
+                    Some(Ok(Message::Binary(_) | Message::Frame(_))) => break "Runtime returned a non-text subscription frame".into(),
+                    Some(Err(error)) => break format!("Runtime subscription connection failed: {error}"),
+                }
+            }
+            update_status(&status, &state, true);
+        };
+        state.disconnected(reason);
+        disconnect_notified = true;
+        update_status(&status, &state, false);
+        tokio::time::sleep(RECONNECT_DELAY).await;
+    }
+}
+
+fn apply_offline(state: &mut State, command: Command) {
+    match command {
+        Command::Subscribe {
+            downstream_id,
+            selector,
+            events,
+        } => {
+            state
+                .downstream_index
+                .insert(downstream_id, selector.clone());
+            state
+                .selectors
+                .entry(selector)
+                .or_insert_with(SelectorState::new)
+                .downstreams
+                .insert(downstream_id, events);
+        }
+        Command::Unsubscribe(id) => remove_downstream(state, id),
+        Command::Shutdown(_) => unreachable!(),
+    }
+}
+
+async fn apply_online(
+    socket: &mut RuntimeSocket,
+    state: &mut State,
+    command: Command,
+) -> Result<(), ()> {
+    match command {
+        Command::Subscribe {
+            downstream_id,
+            selector,
+            events,
+        } => {
+            state
+                .downstream_index
+                .insert(downstream_id, selector.clone());
+            let entry = state
+                .selectors
+                .entry(selector.clone())
+                .or_insert_with(SelectorState::new);
+            if let Some((revision, snapshot)) = &entry.snapshot {
+                let _ = events.try_send(BrokerSubscriptionEvent::Snapshot {
+                    connection_generation: state.generation,
+                    snapshot_revision: *revision,
+                    snapshot: snapshot.clone(),
+                });
+            }
+            entry.downstreams.insert(downstream_id, events);
+            if entry.upstream_id.is_none() && !entry.pending {
+                send_subscribe(socket, state, selector).await?;
+            }
+        }
+        Command::Unsubscribe(id) => {
+            let selector = state.downstream_index.get(&id).cloned();
+            remove_downstream(state, id);
+            if let Some(selector) = selector {
+                maybe_unsubscribe(socket, state, selector).await?;
+            }
+        }
+        Command::Shutdown(_) => unreachable!(),
+    }
+    Ok(())
+}
+
+async fn resubscribe_all(socket: &mut RuntimeSocket, state: &mut State) -> Result<(), ()> {
+    let selectors = state
+        .selectors
+        .iter()
+        .filter(|(_, value)| !value.downstreams.is_empty())
+        .map(|(key, _)| key.clone())
+        .collect::<Vec<_>>();
+    for selector in selectors {
+        send_subscribe(socket, state, selector).await?;
+    }
+    Ok(())
+}
+
+async fn send_subscribe(
+    socket: &mut RuntimeSocket,
+    state: &mut State,
+    selector: EventSubscriptionSelector,
+) -> Result<(), ()> {
+    let request_id = state.request_id();
+    send_frame(
+        socket,
+        SubscriptionFrame::new(SubscriptionFramePayload::Request(
+            SubscriptionRequest::SubscribeEvents {
+                request_id: request_id.clone(),
+                selector: selector.clone(),
+            },
+        )),
+    )
+    .await?;
+    state.pending.insert(request_id, selector.clone());
+    state.selectors.get_mut(&selector).unwrap().pending = true;
+    Ok(())
+}
+
+async fn maybe_unsubscribe(
+    socket: &mut RuntimeSocket,
+    state: &mut State,
+    selector: EventSubscriptionSelector,
+) -> Result<(), ()> {
+    let Some(entry) = state.selectors.get(&selector) else {
+        return Ok(());
+    };
+    if !entry.downstreams.is_empty() {
+        return Ok(());
+    }
+    if let Some(subscription_id) = entry.upstream_id.clone() {
+        let request_id = state.request_id();
+        send_frame(
+            socket,
+            SubscriptionFrame::new(SubscriptionFramePayload::Request(
+                SubscriptionRequest::UnsubscribeEvents {
+                    request_id,
+                    subscription_id: subscription_id.clone(),
+                },
+            )),
+        )
+        .await?;
+        state.upstream_index.remove(&subscription_id);
+        state.selectors.remove(&selector);
+    } else if !entry.pending {
+        state.selectors.remove(&selector);
+    }
+    Ok(())
+}
+
+async fn handle_frame(
+    socket: &mut RuntimeSocket,
+    state: &mut State,
+    frame: SubscriptionFrame,
+) -> Result<(), ()> {
+    match frame.payload {
+        SubscriptionFramePayload::Response(SubscriptionResponse::Subscribed {
+            request_id,
+            subscription_id,
+            selector,
+            snapshot_revision,
+            snapshot,
+        }) => {
+            if state.pending.remove(&request_id) != Some(selector.clone()) {
+                return Err(());
+            }
+            let entry = state.selectors.get_mut(&selector).ok_or(())?;
+            entry.pending = false;
+            entry.upstream_id = Some(subscription_id.clone());
+            entry.snapshot = Some((snapshot_revision, snapshot.clone()));
+            entry.revisions = snapshot_revisions(&snapshot);
+            state
+                .upstream_index
+                .insert(subscription_id, selector.clone());
+            broadcast(
+                &mut entry.downstreams,
+                BrokerSubscriptionEvent::Snapshot {
+                    connection_generation: state.generation,
+                    snapshot_revision,
+                    snapshot,
+                },
+            );
+            if entry.downstreams.is_empty() {
+                maybe_unsubscribe(socket, state, selector).await?;
+            }
+        }
+        SubscriptionFramePayload::Response(SubscriptionResponse::Unsubscribed { .. }) => {}
+        SubscriptionFramePayload::Response(SubscriptionResponse::SubscriptionRejected {
+            request_id,
+            code,
+            message,
+            ..
+        }) => {
+            if let Some(selector) = state.pending.remove(&request_id) {
+                if let Some(mut entry) = state.selectors.remove(&selector) {
+                    broadcast(
+                        &mut entry.downstreams,
+                        BrokerSubscriptionEvent::Rejected {
+                            connection_generation: state.generation,
+                            code,
+                            message,
+                        },
+                    );
+                }
+            }
+        }
+        SubscriptionFramePayload::Event(SubscriptionEvent::Event {
+            subscription_id,
+            subject_revision,
+            payload,
+        }) => {
+            let selector = state
+                .upstream_index
+                .get(&subscription_id)
+                .cloned()
+                .ok_or(())?;
+            payload.validate_for_selector(&selector).map_err(|_| ())?;
+            let entry = state.selectors.get_mut(&selector).ok_or(())?;
+            if let Some(subject) = event_subject(&payload) {
+                let revision = entry.revisions.entry(subject).or_insert(0);
+                if subject_revision <= *revision {
+                    return Ok(());
+                }
+                *revision = subject_revision;
+            }
+            broadcast(
+                &mut entry.downstreams,
+                BrokerSubscriptionEvent::Event {
+                    connection_generation: state.generation,
+                    subject_revision,
+                    payload,
+                },
+            );
+        }
+        SubscriptionFramePayload::Event(SubscriptionEvent::SubscriptionClosed {
+            subscription_id,
+            code,
+            message,
+        }) => {
+            let selector = state.upstream_index.remove(&subscription_id).ok_or(())?;
+            let should_resubscribe = if let Some(entry) = state.selectors.get_mut(&selector) {
+                entry.upstream_id = None;
+                entry.snapshot = None;
+                entry.revisions.clear();
+                broadcast(
+                    &mut entry.downstreams,
+                    BrokerSubscriptionEvent::Closed {
+                        connection_generation: state.generation,
+                        code,
+                        message,
+                    },
+                );
+                !entry.downstreams.is_empty()
+            } else {
+                false
+            };
+            if should_resubscribe {
+                send_subscribe(socket, state, selector).await?;
+            }
+        }
+        SubscriptionFramePayload::Request(_) => return Err(()),
+    }
+    Ok(())
+}
+
+fn remove_downstream(state: &mut State, id: u64) {
+    if let Some(selector) = state.downstream_index.remove(&id) {
+        if let Some(entry) = state.selectors.get_mut(&selector) {
+            entry.downstreams.remove(&id);
+        }
+    }
+}
+
+fn close_all(state: &mut State, replacement: u64) {
+    for entry in state.selectors.values_mut() {
+        broadcast(
+            &mut entry.downstreams,
+            BrokerSubscriptionEvent::Closed {
+                connection_generation: state.generation,
+                code: SubscriptionTerminationCode::ServerShutdown,
+                message: format!(
+                    "Runtime connection generation {} was fenced by generation {replacement}",
+                    state.generation
+                ),
+            },
+        );
+    }
+}
+
+fn broadcast(
+    downstreams: &mut HashMap<u64, mpsc::Sender<BrokerSubscriptionEvent>>,
+    event: BrokerSubscriptionEvent,
+) {
+    let mut closed = HashSet::new();
+    for (id, sender) in downstreams.iter() {
+        if sender.try_send(event.clone()).is_err() {
+            closed.insert(*id);
+        }
+    }
+    downstreams.retain(|id, _| !closed.contains(id));
+}
+
+fn snapshot_revisions(snapshot: &SubscriptionSnapshot) -> HashMap<String, u64> {
+    match snapshot {
+        SubscriptionSnapshot::Workers { workers } => workers
+            .iter()
+            .map(|worker| (worker.worker_id.to_string(), worker.subject_revision))
+            .collect(),
+        SubscriptionSnapshot::WorkerProtocol { worker_id, .. } => {
+            HashMap::from([(worker_id.to_string(), 0)])
+        }
+        SubscriptionSnapshot::WorkspaceWorkdirs { .. } => HashMap::new(),
+    }
+}
+fn event_subject(payload: &SubscriptionEventPayload) -> Option<String> {
+    Some(match payload {
+        SubscriptionEventPayload::WorkerUpserted { worker } => worker.worker_id.to_string(),
+        SubscriptionEventPayload::WorkerRemoved { worker_id }
+        | SubscriptionEventPayload::WorkerProtocol { worker_id, .. } => worker_id.to_string(),
+        SubscriptionEventPayload::WorkdirUpserted { workdir } => {
+            workdir.working_directory_id.to_string()
+        }
+        SubscriptionEventPayload::WorkdirRemoved {
+            working_directory_id,
+        } => working_directory_id.to_string(),
+    })
+}
+async fn send_frame(socket: &mut RuntimeSocket, frame: SubscriptionFrame) -> Result<(), ()> {
+    frame.validate().map_err(|_| ())?;
+    socket
+        .send(Message::Text(
+            serde_json::to_string(&frame).map_err(|_| ())?.into(),
+        ))
+        .await
+        .map_err(|_| ())
+}
+
+async fn connect_runtime(
+    config: &RemoteRuntimeConfig,
+    workspace_id: &str,
+) -> Result<RuntimeSocket, String> {
+    let endpoint = runtime_endpoint(&config.base_url);
+    let mut request = endpoint
+        .into_client_request()
+        .map_err(|error| format!("invalid Runtime subscription endpoint: {error}"))?;
+    if let Some(token) = runtime_token(config, workspace_id)? {
+        request.headers_mut().insert(
+            "authorization",
+            format!("Bearer {token}")
+                .parse()
+                .map_err(|error| format!("invalid Runtime authorization header: {error}"))?,
+        );
+    }
+    connect_async(request)
+        .await
+        .map(|(socket, _)| socket)
+        .map_err(|error| format!("failed to connect Runtime subscription endpoint: {error}"))
+}
+fn runtime_endpoint(base_url: &str) -> String {
+    let base = base_url.trim_end_matches('/');
+    if let Some(rest) = base.strip_prefix("https://") {
+        format!("wss://{rest}/v1/protocol/ws")
+    } else if let Some(rest) = base.strip_prefix("http://") {
+        format!("ws://{rest}/v1/protocol/ws")
+    } else {
+        format!("{base}/v1/protocol/ws")
+    }
+}
+fn runtime_token(
+    config: &RemoteRuntimeConfig,
+    workspace_id: &str,
+) -> Result<Option<String>, String> {
+    let Some(auth) = config.auth.as_ref() else {
+        return Ok(config.bearer_token.clone());
+    };
+    let signer = CapabilityTokenSigner::new(&auth.server_id, &auth.server_private_key);
+    let claims = capability_claims(
+        &auth.server_id,
+        &config.runtime_id,
+        workspace_id,
+        vec!["workers:list".into()],
+        300,
+    )
+    .map_err(|error| error.to_string())?;
+    signer
+        .sign(&claims)
+        .map(Some)
+        .map_err(|error| error.to_string())
+}
+fn update_status(status: &RwLock<RuntimeSubscriptionBrokerStatus>, state: &State, connected: bool) {
+    *status.write().expect("broker status poisoned") = RuntimeSubscriptionBrokerStatus {
+        connection_generation: state.generation,
+        connected,
+        desired_selectors: state
+            .selectors
+            .values()
+            .filter(|value| !value.downstreams.is_empty())
+            .count(),
+        upstream_subscriptions: state.upstream_index.len(),
+    };
+}
+
+#[cfg(test)]
+#[path = "runtime_subscription_tests.rs"]
+mod tests;

crates/workspace-server/src/runtime_subscription_tests.rs

@@ -0,0 +1,278 @@
+use super::*;
+use protocol::subscription::{SubscriptionWorkerIds, SubscriptionWorkerState};
+use worker_runtime::Runtime;
+use worker_runtime::catalog::{
+    CreateWorkerRequest, ProfileSelector, ProfileSourceArchiveHttpRef, ProfileSourceArchiveSource,
+};
+use worker_runtime::execution::{
+    WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, WorkerExecutionResult,
+    WorkerExecutionRunState, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult,
+};
+use worker_runtime::profile_archive::{ProfileSourceArchiveRef, ProfileSourceGraphSummary};
+
+#[derive(Debug)]
+struct TestExecutionBackend;
+
+impl WorkerExecutionBackend for TestExecutionBackend {
+    fn backend_id(&self) -> &str {
+        "runtime-subscription-test"
+    }
+
+    fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult {
+        WorkerExecutionSpawnResult::connected(
+            WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
+            WorkerExecutionRunState::Idle,
+            None,
+        )
+    }
+
+    fn dispatch_input(
+        &self,
+        _handle: &WorkerExecutionHandle,
+        _input: worker_runtime::interaction::WorkerInput,
+    ) -> WorkerExecutionResult {
+        WorkerExecutionResult::accepted(
+            WorkerExecutionOperation::Input,
+            WorkerExecutionRunState::Busy,
+        )
+    }
+}
+
+const TOKEN: &str = "runtime-subscription-test-token";
+
+fn create_request(name: &str) -> CreateWorkerRequest {
+    CreateWorkerRequest {
+        idempotency_key: None,
+        idempotency_fingerprint: None,
+        profile: ProfileSelector::Builtin("builtin:companion".to_string()),
+        display_name: Some(name.to_string()),
+        config_bundle: None,
+        profile_source: ProfileSourceArchiveSource::Http {
+            location: ProfileSourceArchiveHttpRef {
+                url: "http://127.0.0.1/profiles/test".to_string(),
+                etag: None,
+                archive: ProfileSourceArchiveRef {
+                    id: "test-profile-source".to_string(),
+                    digest: "test-digest".to_string(),
+                    size_bytes: 0,
+                    source_graph: ProfileSourceGraphSummary {
+                        source_count: 0,
+                        total_source_bytes: 0,
+                        entrypoints: std::collections::BTreeMap::new(),
+                        import_count: 0,
+                    },
+                },
+            },
+        },
+        initial_input: None,
+        working_directory_request: None,
+        working_directory: None,
+        workspace_api: None,
+    }
+}
+
+async fn start_runtime_server(
+    listener: tokio::net::TcpListener,
+    runtime: Runtime,
+) -> tokio::task::JoinHandle<()> {
+    let router = worker_runtime::http_server::runtime_http_router(runtime, TOKEN.to_string());
+    tokio::spawn(async move {
+        let _ = axum::serve(listener, router).await;
+    })
+}
+
+async fn fixture() -> (Runtime, RemoteRuntimeConfig, tokio::task::JoinHandle<()>) {
+    let runtime = Runtime::with_execution_backend(
+        worker_runtime::RuntimeOptions::default(),
+        std::sync::Arc::new(TestExecutionBackend),
+    )
+    .unwrap();
+    runtime
+        .create_worker_scoped(
+            &worker_runtime::RuntimeWorkspaceScope::new("local", "local-token"),
+            create_request("fixture"),
+        )
+        .unwrap();
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let address = listener.local_addr().unwrap();
+    let task = start_runtime_server(listener, runtime.clone()).await;
+    let config = RemoteRuntimeConfig::new(
+        "runtime-test",
+        "Runtime test",
+        format!("http://{address}"),
+        Some(TOKEN.to_string()),
+    );
+    (runtime, config, task)
+}
+
+async fn next_event(subscription: &mut BrokerSubscription) -> BrokerSubscriptionEvent {
+    tokio::time::timeout(Duration::from_secs(5), subscription.recv())
+        .await
+        .expect("subscription event timed out")
+        .expect("subscription closed")
+}
+
+async fn next_snapshot(subscription: &mut BrokerSubscription) -> BrokerSubscriptionEvent {
+    loop {
+        let event = next_event(subscription).await;
+        if matches!(event, BrokerSubscriptionEvent::Snapshot { .. }) {
+            return event;
+        }
+    }
+}
+
+async fn wait_for_status(
+    broker: &RuntimeSubscriptionBroker,
+    runtime_id: &str,
+    predicate: impl Fn(&RuntimeSubscriptionBrokerStatus) -> bool,
+) -> RuntimeSubscriptionBrokerStatus {
+    tokio::time::timeout(Duration::from_secs(5), async {
+        loop {
+            if let Some(status) = broker.status(runtime_id) {
+                if predicate(&status) {
+                    return status;
+                }
+            }
+            tokio::time::sleep(Duration::from_millis(10)).await;
+        }
+    })
+    .await
+    .expect("broker status timed out")
+}
+
+#[tokio::test]
+async fn equal_downstream_selectors_share_one_upstream_subscription() {
+    let (runtime, config, server) = fixture().await;
+    let worker = runtime.list_workers().unwrap().remove(0);
+    let broker = RuntimeSubscriptionBroker::new("local");
+    broker.register_remote_runtime(config);
+    let selector = EventSubscriptionSelector::WorkerLifecycle {
+        worker_ids: SubscriptionWorkerIds::new([
+            protocol::subscription::SubscriptionWorkerId::new(
+                worker.worker_ref.worker_id.to_string(),
+            )
+            .unwrap(),
+        ])
+        .unwrap(),
+    };
+    let mut first = broker.subscribe("runtime-test", selector.clone()).unwrap();
+    let mut second = broker.subscribe("runtime-test", selector).unwrap();
+    assert!(matches!(
+        next_snapshot(&mut first).await,
+        BrokerSubscriptionEvent::Snapshot { .. }
+    ));
+    assert!(matches!(
+        next_snapshot(&mut second).await,
+        BrokerSubscriptionEvent::Snapshot { .. }
+    ));
+    let status = wait_for_status(&broker, "runtime-test", |status| {
+        status.upstream_subscriptions == 1
+    })
+    .await;
+    assert_eq!(status.desired_selectors, 1);
+
+    runtime
+        .observe_worker_event(
+            &worker.worker_ref,
+            protocol::Event::Status {
+                status: protocol::WorkerStatus::Running,
+            },
+        )
+        .unwrap();
+    for subscription in [&mut first, &mut second] {
+        assert!(matches!(
+            next_event(subscription).await,
+            BrokerSubscriptionEvent::Event {
+                payload: SubscriptionEventPayload::WorkerUpserted { ref worker },
+                ..
+            } if worker.state == SubscriptionWorkerState::Running
+        ));
+    }
+
+    drop(first);
+    tokio::task::yield_now().await;
+    assert_eq!(
+        broker
+            .status("runtime-test")
+            .unwrap()
+            .upstream_subscriptions,
+        1
+    );
+    drop(second);
+    wait_for_status(&broker, "runtime-test", |status| {
+        status.upstream_subscriptions == 0 && status.desired_selectors == 0
+    })
+    .await;
+    server.abort();
+}
+
+#[tokio::test]
+async fn replacing_runtime_registration_fences_the_old_generation() {
+    let (_runtime, config, server) = fixture().await;
+    let broker = RuntimeSubscriptionBroker::new("local");
+    let first_generation = broker.register_remote_runtime(config.clone());
+    let mut old = broker
+        .subscribe("runtime-test", EventSubscriptionSelector::RuntimeWorkers)
+        .unwrap();
+    assert!(matches!(
+        next_snapshot(&mut old).await,
+        BrokerSubscriptionEvent::Snapshot { .. }
+    ));
+
+    let second_generation = broker.register_remote_runtime(config);
+    assert!(second_generation > first_generation);
+    assert!(matches!(
+        next_event(&mut old).await,
+        BrokerSubscriptionEvent::Closed {
+            connection_generation,
+            ..
+        } if connection_generation == first_generation
+    ));
+    let status = wait_for_status(&broker, "runtime-test", |status| {
+        status.connection_generation == second_generation && status.connected
+    })
+    .await;
+    assert_eq!(status.connection_generation, second_generation);
+    server.abort();
+}
+
+#[tokio::test]
+async fn reconnect_resubscribes_and_replaces_state_from_fresh_snapshot() {
+    let (runtime, config, server) = fixture().await;
+    let address = config
+        .base_url
+        .strip_prefix("http://")
+        .unwrap()
+        .parse::<std::net::SocketAddr>()
+        .unwrap();
+    server.abort();
+    tokio::task::yield_now().await;
+
+    let broker = RuntimeSubscriptionBroker::new("local");
+    broker.register_remote_runtime(config);
+    let mut subscription = broker
+        .subscribe("runtime-test", EventSubscriptionSelector::RuntimeWorkers)
+        .unwrap();
+    assert!(matches!(
+        next_event(&mut subscription).await,
+        BrokerSubscriptionEvent::Disconnected { .. }
+    ));
+
+    runtime
+        .create_worker_scoped(
+            &worker_runtime::RuntimeWorkspaceScope::new("local", "local-token"),
+            create_request("after-reconnect"),
+        )
+        .unwrap();
+    let listener = tokio::net::TcpListener::bind(address).await.unwrap();
+    let restarted = start_runtime_server(listener, runtime).await;
+    let event = next_snapshot(&mut subscription).await;
+    let BrokerSubscriptionEvent::Snapshot { snapshot, .. } = event else {
+        panic!("expected fresh snapshot after reconnect");
+    };
+    let SubscriptionSnapshot::Workers { workers } = snapshot else {
+        panic!("expected Worker snapshot");
+    };
+    assert_eq!(workers.len(), 2);
+    restarted.abort();
+}

crates/workspace-server/src/server.rs

@@ -90,6 +90,7 @@ use crate::repositories::{
     RepositoryRegistryReader, RepositorySummary,
 };
 use crate::resource_broker::BackendResourceBroker;
+use crate::runtime_subscription::RuntimeSubscriptionBroker;
 use crate::skills;
 use crate::store::{
     AccountRecord, ApiTokenRecord, AuthChallengeRecord, BrowserSessionRecord, ControlPlaneStore,
@@ -248,6 +249,7 @@ pub struct WorkspaceApi {
     runtime: Arc<RuntimeRegistry>,
     companion: Arc<CompanionConsole>,
     observation_proxy: BackendObservationProxy,
+    runtime_subscription_broker: RuntimeSubscriptionBroker,
     resource_broker: BackendResourceBroker,
     credential_operation_lock: Arc<std::sync::Mutex<()>>,
 }
@@ -317,10 +319,11 @@ impl WorkspaceApi {
                 crate::Error::Store(format!("invalid embedded Worker backend: {err}"))
             })?,
         );
+        let runtime_subscription_broker =
+            RuntimeSubscriptionBroker::new(config.workspace_id.clone());
         for remote_config in config.remote_runtime_sources.iter().cloned() {
-            runtime.register(
+            let remote_runtime = RemoteWorkerRuntime::new(
-                RemoteWorkerRuntime::new(
+                remote_config.clone(),
-                    remote_config,
                 config.workspace_id.clone(),
                 config
                     .backend_base_url
@@ -328,8 +331,9 @@ impl WorkspaceApi {
                     .unwrap_or_else(|| "http://127.0.0.1:8787".to_string()),
             )
             .map(|host| host.with_resource_broker(resource_broker.clone()))
-                .map_err(|err| err.into_error())?,
+            .map_err(|err| err.into_error())?;
-            );
+            runtime.register(remote_runtime);
+            runtime_subscription_broker.register_remote_runtime(remote_config);
         }
         let runtime = Arc::new(runtime);
         let companion = Arc::new(CompanionConsole::disabled());
@@ -344,6 +348,7 @@ impl WorkspaceApi {
             runtime,
             companion,
             observation_proxy,
+            runtime_subscription_broker,
             resource_broker,
             credential_operation_lock: Arc::new(std::sync::Mutex::new(())),
         })
@@ -353,6 +358,10 @@ impl WorkspaceApi {
         self.config.workspace_id.as_str()
     }
 
+    pub fn runtime_subscription_broker(&self) -> &RuntimeSubscriptionBroker {
+        &self.runtime_subscription_broker
+    }
+
     fn mint_worker_workspace_credential(
         &self,
         runtime_id: &str,