server: broker embedded runtime subscriptions

crates/workspace-server/src/hosts.rs

@@ -1427,6 +1427,10 @@ impl EmbeddedWorkerRuntime {
         self
     }
 
+    pub(crate) fn subscription_runtime(&self) -> worker_runtime::Runtime {
+        self.runtime.clone()
+    }
+
     pub fn from_runtime(workspace_id: impl AsRef<str>, runtime: worker_runtime::Runtime) -> Self {
         let workspace_id = workspace_id.as_ref().to_string();
         Self {

crates/workspace-server/src/runtime_subscription.rs

@@ -148,6 +148,45 @@ impl RuntimeSubscriptionBroker {
         generation
     }
 
+    pub fn register_embedded_runtime(
+        &self,
+        runtime_id: impl Into<String>,
+        runtime: worker_runtime::Runtime,
+    ) -> u64 {
+        let runtime_id = runtime_id.into();
+        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,
+            connected: true,
+            ..Default::default()
+        }));
+        let previous = self
+            .registrations
+            .write()
+            .expect("broker registry poisoned")
+            .insert(
+                runtime_id.clone(),
+                Registration {
+                    generation,
+                    commands: commands.clone(),
+                    status: status.clone(),
+                },
+            );
+        if let Some(previous) = previous {
+            let _ = previous.commands.send(Command::Shutdown(generation));
+        }
+        tokio::spawn(run_embedded_connection(
+            runtime_id,
+            runtime,
+            self.workspace_id.to_string(),
+            generation,
+            receiver,
+            status,
+        ));
+        generation
+    }
+
     pub fn unregister_runtime(&self, runtime_id: &str) {
         if let Some(registration) = self
             .registrations
@@ -282,6 +321,112 @@ impl State {
     }
 }
 
+struct EmbeddedEntry {
+    downstreams: HashMap<u64, mpsc::Sender<BrokerSubscriptionEvent>>,
+    snapshot_revision: u64,
+    snapshot: SubscriptionSnapshot,
+    task: tokio::task::JoinHandle<()>,
+}
+
+async fn run_embedded_connection(
+    runtime_id: String,
+    runtime: worker_runtime::Runtime,
+    _workspace_id: String,
+    generation: u64,
+    mut commands: mpsc::UnboundedReceiver<Command>,
+    status: Arc<RwLock<RuntimeSubscriptionBrokerStatus>>,
+) {
+    let (updates, mut update_receiver) = mpsc::unbounded_channel();
+    let mut entries = HashMap::<EventSubscriptionSelector, EmbeddedEntry>::new();
+    let mut downstream_index = HashMap::<u64, EventSubscriptionSelector>::new();
+    loop {
+        tokio::select! {
+            command = commands.recv() => match command {
+                Some(Command::Subscribe { downstream_id, selector, events }) => {
+                    downstream_index.insert(downstream_id, selector.clone());
+                    if let Some(entry) = entries.get_mut(&selector) {
+                        let _ = events.try_send(BrokerSubscriptionEvent::Snapshot {
+                            connection_generation: generation,
+                            snapshot_revision: entry.snapshot_revision,
+                            snapshot: entry.snapshot.clone(),
+                        });
+                        entry.downstreams.insert(downstream_id, events);
+                    } else {
+                        match runtime.subscribe_event_selector(selector.clone()) {
+                            Ok(mut subscription) => {
+                                let snapshot_revision = subscription.snapshot_revision();
+                                let snapshot = project_snapshot_runtime(subscription.snapshot().clone(), &runtime_id);
+                                let _ = events.try_send(BrokerSubscriptionEvent::Snapshot { connection_generation: generation, snapshot_revision, snapshot: snapshot.clone() });
+                                let sender = updates.clone();
+                                let task_selector = selector.clone();
+                                let task_runtime_id = runtime_id.clone();
+                                let task = tokio::spawn(async move {
+                                    while let Ok(update) = subscription.recv().await {
+                                        let payload = project_payload_runtime(update.payload, &task_runtime_id);
+                                        if sender.send((task_selector.clone(), update.subject_revision, payload)).is_err() { break; }
+                                    }
+                                });
+                                entries.insert(selector, EmbeddedEntry { downstreams: HashMap::from([(downstream_id, events)]), snapshot_revision, snapshot, task });
+                            }
+                            Err(error) => {
+                                let _ = events.try_send(BrokerSubscriptionEvent::Rejected { connection_generation: generation, code: SubscriptionRejectionCode::UnsupportedSelector, message: error.to_string() });
+                            }
+                        }
+                    }
+                }
+                Some(Command::Unsubscribe(id)) => {
+                    if let Some(selector) = downstream_index.remove(&id) {
+                        let empty = entries.get_mut(&selector).is_some_and(|entry| { entry.downstreams.remove(&id); entry.downstreams.is_empty() });
+                        if empty { if let Some(entry) = entries.remove(&selector) { entry.task.abort(); } }
+                    }
+                }
+                Some(Command::Shutdown(replacement)) => {
+                    for entry in entries.values_mut() {
+                        broadcast(&mut entry.downstreams, BrokerSubscriptionEvent::Closed { connection_generation: generation, code: SubscriptionTerminationCode::ServerShutdown, message: format!("embedded Runtime generation {generation} was fenced by {replacement}") });
+                        entry.task.abort();
+                    }
+                    return;
+                }
+                None => return,
+            },
+            update = update_receiver.recv() => {
+                let Some((selector, subject_revision, payload)) = update else { return; };
+                if let Some(entry) = entries.get_mut(&selector) {
+                    broadcast(&mut entry.downstreams, BrokerSubscriptionEvent::Event { connection_generation: generation, subject_revision, payload });
+                }
+            }
+        }
+        *status.write().expect("broker status poisoned") = RuntimeSubscriptionBrokerStatus {
+            connection_generation: generation,
+            connected: true,
+            desired_selectors: entries.len(),
+            upstream_subscriptions: entries.len(),
+        };
+    }
+}
+
+fn project_snapshot_runtime(
+    mut snapshot: SubscriptionSnapshot,
+    runtime_id: &str,
+) -> SubscriptionSnapshot {
+    if let SubscriptionSnapshot::Workers { workers } = &mut snapshot {
+        for worker in workers {
+            worker.runtime_id = Some(runtime_id.to_string());
+        }
+    }
+    snapshot
+}
+
+fn project_payload_runtime(
+    mut payload: SubscriptionEventPayload,
+    runtime_id: &str,
+) -> SubscriptionEventPayload {
+    if let SubscriptionEventPayload::WorkerUpserted { worker } = &mut payload {
+        worker.runtime_id = Some(runtime_id.to_string());
+    }
+    payload
+}
+
 async fn run_connection(
     config: RemoteRuntimeConfig,
     workspace_id: String,

crates/workspace-server/src/runtime_subscription_tests.rs

@@ -276,3 +276,40 @@ async fn reconnect_resubscribes_and_replaces_state_from_fresh_snapshot() {
     assert_eq!(workers.len(), 2);
     restarted.abort();
 }
+
+#[tokio::test]
+async fn embedded_runtime_uses_in_process_subscription_source() {
+    let (runtime, _config, server) = fixture().await;
+    let worker = runtime.list_workers().unwrap().remove(0);
+    let broker = RuntimeSubscriptionBroker::new("local");
+    broker.register_embedded_runtime("embedded-worker-runtime", runtime.clone());
+    let mut subscription = broker
+        .subscribe(
+            "embedded-worker-runtime",
+            EventSubscriptionSelector::RuntimeWorkers,
+        )
+        .unwrap();
+    let BrokerSubscriptionEvent::Snapshot { snapshot, .. } = next_snapshot(&mut subscription).await
+    else {
+        panic!("expected embedded snapshot");
+    };
+    let SubscriptionSnapshot::Workers { workers } = snapshot else {
+        panic!("expected Worker snapshot");
+    };
+    assert_eq!(
+        workers[0].runtime_id.as_deref(),
+        Some("embedded-worker-runtime")
+    );
+    runtime
+        .observe_worker_event(
+            &worker.worker_ref,
+            protocol::Event::Status {
+                status: protocol::WorkerStatus::Running,
+            },
+        )
+        .unwrap();
+    assert!(matches!(next_event(&mut subscription).await,
+        BrokerSubscriptionEvent::Event { payload: SubscriptionEventPayload::WorkerUpserted { worker }, .. }
+        if worker.runtime_id.as_deref() == Some("embedded-worker-runtime") && worker.state == SubscriptionWorkerState::Running));
+    server.abort();
+}

crates/workspace-server/src/server.rs

@@ -308,19 +308,20 @@ impl WorkspaceApi {
             .await?;
         import_configured_repositories(store.as_ref(), &config)?;
         config.repositories = load_configured_repositories_from_store(store.as_ref(), &config)?;
-        let runtime = RuntimeRegistry::for_workspace(
+        let embedded_runtime = EmbeddedWorkerRuntime::new_fs_store_with_execution_backend(
-            EmbeddedWorkerRuntime::new_fs_store_with_execution_backend(
+            config.workspace_id.clone(),
-                config.workspace_id.clone(),
+            config.embedded_runtime_store_root.clone(),
-                config.embedded_runtime_store_root.clone(),
+            execution_backend,
-                execution_backend,
+        )
-            )
+        .map(|runtime| runtime.with_resource_broker(resource_broker.clone()))
-            .map(|runtime| runtime.with_resource_broker(resource_broker.clone()))
+        .map_err(|err| crate::Error::Store(format!("invalid embedded Worker backend: {err}")))?;
-            .map_err(|err| {
+        let embedded_subscription_runtime = embedded_runtime.subscription_runtime();
-                crate::Error::Store(format!("invalid embedded Worker backend: {err}"))
+        let embedded_runtime_id = EMBEDDED_WORKER_RUNTIME_ID.to_string();
-            })?,
+        let runtime = RuntimeRegistry::for_workspace(embedded_runtime);
-        );
         let runtime_subscription_broker =
             RuntimeSubscriptionBroker::new(config.workspace_id.clone());
+        runtime_subscription_broker
+            .register_embedded_runtime(embedded_runtime_id, embedded_subscription_runtime);
         for remote_config in config.remote_runtime_sources.iter().cloned() {
             let remote_runtime = RemoteWorkerRuntime::new(
                 remote_config.clone(),