diff --git a/crates/protocol/src/subscription.rs b/crates/protocol/src/subscription.rs index 03459387..af2666c7 100644 --- a/crates/protocol/src/subscription.rs +++ b/crates/protocol/src/subscription.rs @@ -643,6 +643,8 @@ pub enum SubscriptionEventPayload { }, WorkerRemoved { worker_id: SubscriptionWorkerId, + #[serde(default, skip_serializing_if = "Option::is_none")] + runtime_id: Option, }, WorkerProtocol { worker_id: SubscriptionWorkerId, @@ -660,9 +662,17 @@ impl SubscriptionEventPayload { pub fn validate(&self) -> Result<(), SubscriptionValidationError> { match self { Self::WorkerUpserted { worker } => worker.validate(), - Self::WorkerRemoved { worker_id } | Self::WorkerProtocol { worker_id, .. } => { - worker_id.validate() + Self::WorkerRemoved { + 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::WorkerProtocol { worker_id, .. } => worker_id.validate(), Self::WorkdirUpserted { workdir } => workdir.validate(), Self::WorkdirRemoved { working_directory_id, @@ -688,7 +698,7 @@ impl SubscriptionEventPayload { ) if worker_ids.contains(&worker.worker_id) => Ok(()), ( EventSubscriptionSelector::WorkerLifecycle { worker_ids }, - Self::WorkerRemoved { worker_id }, + Self::WorkerRemoved { worker_id, .. }, ) if worker_ids.contains(worker_id) => Ok(()), ( EventSubscriptionSelector::WorkerProtocol { @@ -708,7 +718,7 @@ impl SubscriptionEventPayload { }), ( EventSubscriptionSelector::WorkerLifecycle { .. }, - Self::WorkerRemoved { worker_id }, + Self::WorkerRemoved { worker_id, .. }, ) => Err(SubscriptionValidationError::UnselectedWorker { worker_id: worker_id.to_string(), }), @@ -721,7 +731,7 @@ fn validate_workers(workers: &[SubscriptionWorker]) -> Result<(), SubscriptionVa let mut seen = HashSet::with_capacity(workers.len()); for worker in workers { worker.validate()?; - if !seen.insert(&worker.worker_id) { + if !seen.insert((worker.runtime_id.as_deref(), &worker.worker_id)) { return Err(SubscriptionValidationError::DuplicateWorkerId { worker_id: worker.worker_id.to_string(), }); @@ -868,6 +878,7 @@ mod tests { subject_revision: 8, payload: SubscriptionEventPayload::WorkerRemoved { worker_id: worker_id("worker-2"), + runtime_id: None, }, }; assert!(matches!( @@ -880,6 +891,19 @@ mod tests { )); } + #[test] + fn workspace_snapshot_allows_equal_local_worker_ids_from_distinct_runtimes() { + let mut first = worker("1"); + first.runtime_id = Some("runtime-a".to_string()); + let mut second = worker("1"); + second.runtime_id = Some("runtime-b".to_string()); + SubscriptionSnapshot::Workers { + workers: vec![first, second], + } + .validate_for_selector(&EventSubscriptionSelector::WorkspaceWorkers) + .unwrap(); + } + #[test] fn subscription_closed_is_a_typed_server_event() { let frame = SubscriptionFrame::new(SubscriptionFramePayload::Event( diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index ac91e8b4..eafbbeb5 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -2177,6 +2177,7 @@ impl RuntimeState { subject_revision, payload: SubscriptionEventPayload::WorkerRemoved { worker_id: worker_id.clone(), + runtime_id: None, }, }, ); diff --git a/crates/workspace-server/src/lib.rs b/crates/workspace-server/src/lib.rs index fad63c88..8a3f4d40 100644 --- a/crates/workspace-server/src/lib.rs +++ b/crates/workspace-server/src/lib.rs @@ -23,6 +23,7 @@ pub mod runtime_subscription; pub mod server; pub mod skills; pub mod store; +mod workspace_subscription; pub use authority::{ MemoryAuthority, MemoryDocument, MemoryStagingEntry, MemoryStagingResolution, diff --git a/crates/workspace-server/src/runtime_subscription.rs b/crates/workspace-server/src/runtime_subscription.rs index 69b178ee..1ab799d6 100644 --- a/crates/workspace-server/src/runtime_subscription.rs +++ b/crates/workspace-server/src/runtime_subscription.rs @@ -242,6 +242,7 @@ impl RuntimeSubscriptionBroker { })?; let downstream_id = self.next_downstream.fetch_add(1, Ordering::Relaxed); let (events, receiver) = mpsc::channel(DOWNSTREAM_QUEUE_CAPACITY); + let initial_events = events.clone(); registration .commands .send(Command::Subscribe { @@ -250,6 +251,17 @@ impl RuntimeSubscriptionBroker { events, }) .map_err(|_| RuntimeSubscriptionBrokerError::Closed)?; + let initial_status = registration + .status + .read() + .expect("broker status poisoned") + .clone(); + if !initial_status.connected { + let _ = initial_events.try_send(BrokerSubscriptionEvent::Disconnected { + connection_generation: initial_status.connection_generation, + message: "Runtime subscription connection is not currently available".to_string(), + }); + } Ok(BrokerSubscription { downstream_id, runtime_id: runtime_id.to_string(), @@ -291,6 +303,7 @@ impl SelectorState { } struct State { + runtime_id: String, generation: u64, next_request: u64, selectors: HashMap, @@ -299,8 +312,9 @@ struct State { upstream_index: HashMap, } impl State { - fn new(generation: u64) -> Self { + fn new(runtime_id: String, generation: u64) -> Self { Self { + runtime_id, generation, next_request: 1, selectors: HashMap::new(), @@ -433,8 +447,17 @@ 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()); + match &mut payload { + SubscriptionEventPayload::WorkerUpserted { worker } => { + worker.runtime_id = Some(runtime_id.to_string()); + } + SubscriptionEventPayload::WorkerRemoved { + runtime_id: projected_runtime_id, + .. + } => { + *projected_runtime_id = Some(runtime_id.to_string()); + } + _ => {} } payload } @@ -446,7 +469,7 @@ async fn run_connection( mut commands: mpsc::UnboundedReceiver, status: Arc>, ) { - let mut state = State::new(generation); + let mut state = State::new(config.runtime_id.clone(), generation); let mut disconnect_notified = false; loop { update_status(&status, &state, false); @@ -653,6 +676,7 @@ async fn handle_frame( if state.pending.remove(&request_id) != Some(selector.clone()) { return Err(()); } + let snapshot = project_snapshot_runtime(snapshot, &state.runtime_id); let entry = state.selectors.get_mut(&selector).ok_or(())?; entry.pending = false; entry.upstream_id = Some(subscription_id.clone()); @@ -704,6 +728,7 @@ async fn handle_frame( .cloned() .ok_or(())?; payload.validate_for_selector(&selector).map_err(|_| ())?; + let payload = project_payload_runtime(payload, &state.runtime_id); 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); @@ -793,7 +818,16 @@ fn snapshot_revisions(snapshot: &SubscriptionSnapshot) -> HashMap { match snapshot { SubscriptionSnapshot::Workers { workers } => workers .iter() - .map(|worker| (worker.worker_id.to_string(), worker.subject_revision)) + .map(|worker| { + ( + format!( + "{}:{}", + worker.runtime_id.as_deref().unwrap_or_default(), + worker.worker_id + ), + worker.subject_revision, + ) + }) .collect(), SubscriptionSnapshot::WorkerProtocol { worker_id, .. } => { HashMap::from([(worker_id.to_string(), 0)]) @@ -803,9 +837,20 @@ fn snapshot_revisions(snapshot: &SubscriptionSnapshot) -> HashMap { } fn event_subject(payload: &SubscriptionEventPayload) -> Option { Some(match payload { - SubscriptionEventPayload::WorkerUpserted { worker } => worker.worker_id.to_string(), - SubscriptionEventPayload::WorkerRemoved { worker_id } - | SubscriptionEventPayload::WorkerProtocol { worker_id, .. } => worker_id.to_string(), + SubscriptionEventPayload::WorkerUpserted { worker } => format!( + "{}:{}", + worker.runtime_id.as_deref().unwrap_or_default(), + worker.worker_id + ), + SubscriptionEventPayload::WorkerRemoved { + worker_id, + runtime_id, + } => format!( + "{}:{}", + runtime_id.as_deref().unwrap_or_default(), + worker_id + ), + SubscriptionEventPayload::WorkerProtocol { worker_id, .. } => worker_id.to_string(), SubscriptionEventPayload::WorkdirUpserted { workdir } => { workdir.working_directory_id.to_string() } diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 4941c47e..92b008ae 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -867,6 +867,10 @@ pub fn build_router(api: WorkspaceApi) -> Router { "/api/w/{workspace_id}/workers", get(scoped_list_workers).post(scoped_create_workspace_worker), ) + .route( + "/api/w/{workspace_id}/protocol/ws", + get(scoped_workspace_protocol_ws), + ) .route( "/api/workers/launch-options", get(get_worker_launch_options), @@ -3779,6 +3783,38 @@ async fn scoped_list_runtimes( list_runtimes(State(api)).await } +async fn scoped_workspace_protocol_ws( + State(api): State, + AxumPath(workspace_id): AxumPath, + headers: HeaderMap, + ws: axum::extract::ws::WebSocketUpgrade, +) -> std::result::Result { + validate_workspace_scope(&api, &workspace_id).map_err(|error| error.into_response())?; + let actor = resolve_actor(&api, &headers) + .await + .map_err(|error| error.into_response())? + .ok_or_else(|| StatusCode::UNAUTHORIZED.into_response())?; + let workspace = api + .store + .get_workspace(&workspace_id) + .await + .map_err(|error| ApiError::from(error).into_response())? + .ok_or_else(|| StatusCode::NOT_FOUND.into_response())?; + if workspace + .owner_account_id + .as_deref() + .is_some_and(|owner| owner != actor.account_id.as_str()) + { + 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) + }) + .into_response()) +} + async fn scoped_list_workers( State(api): State, AxumPath(path): AxumPath, @@ -8716,6 +8752,7 @@ mod tests { use std::{fs, sync::Arc}; use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::Message; + use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tower::ServiceExt; use worker_runtime::resource::BackendResourceClient; use worker_runtime::working_directory::WorkingDirectoryMaterializer; @@ -8725,8 +8762,9 @@ mod tests { WorkerSpawnIntent, }; use crate::store::{ - MemoryDocumentRecord, MemoryStagingRecord, ObjectiveRecord, ObjectiveResourceRecord, - ObjectiveTicketLinkRecord, SqliteWorkspaceStore, WorkspaceRecord, + AccountRecord, MemoryDocumentRecord, MemoryStagingRecord, ObjectiveRecord, + ObjectiveResourceRecord, ObjectiveTicketLinkRecord, SqliteWorkspaceStore, UserRecord, + WorkspaceRecord, }; const TEST_WORKSPACE_ID: &str = "0192f0e8-4d84-7d6e-a000-000000000001"; @@ -12623,6 +12661,105 @@ mod tests { ) } + #[tokio::test] + async fn workspace_subscription_requires_browser_session() { + let dir = tempfile::tempdir().unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = build_router(test_api(dir.path()).await); + let server = tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + let error = tokio_tungstenite::connect_async(format!( + "ws://{address}/api/w/{TEST_WORKSPACE_ID}/protocol/ws" + )) + .await + .unwrap_err(); + let tokio_tungstenite::tungstenite::Error::Http(response) = error else { + panic!("expected HTTP authentication rejection"); + }; + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + server.abort(); + } + + #[tokio::test] + async fn workspace_subscription_returns_authenticated_workspace_snapshot() { + let dir = tempfile::tempdir().unwrap(); + let api = test_api(dir.path()).await; + let account = AccountRecord { + account_id: "account-test".to_string(), + kind: "user".to_string(), + handle: "tester".to_string(), + display_name: "Tester".to_string(), + created_at: TEST_CREATED_AT.to_string(), + updated_at: TEST_CREATED_AT.to_string(), + }; + let user = UserRecord { + user_id: "user-test".to_string(), + account_id: account.account_id.clone(), + handle: account.handle.clone(), + display_name: account.display_name.clone(), + created_at: TEST_CREATED_AT.to_string(), + updated_at: TEST_CREATED_AT.to_string(), + }; + api.store.upsert_account(&account).unwrap(); + api.store.upsert_user(&user).unwrap(); + let session = issue_browser_session_response(&api, user).unwrap(); + let cookie = session + .headers() + .get(SET_COOKIE) + .unwrap() + .to_str() + .unwrap() + .split(';') + .next() + .unwrap() + .to_string(); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = build_router(api); + let server = tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + let mut request = format!("ws://{address}/api/w/{TEST_WORKSPACE_ID}/protocol/ws") + .into_client_request() + .unwrap(); + request + .headers_mut() + .insert(axum::http::header::COOKIE, cookie.parse().unwrap()); + let (mut socket, _) = connect_async(request).await.unwrap(); + let frame = protocol::subscription::SubscriptionFrame::new( + protocol::subscription::SubscriptionFramePayload::Request( + protocol::subscription::SubscriptionRequest::SubscribeEvents { + request_id: protocol::subscription::SubscriptionRequestId::new("request-1") + .unwrap(), + selector: protocol::subscription::EventSubscriptionSelector::WorkspaceWorkers, + }, + ), + ); + socket + .send(Message::Text(serde_json::to_string(&frame).unwrap().into())) + .await + .unwrap(); + let Message::Text(text) = socket.next().await.unwrap().unwrap() else { + panic!("expected subscription response"); + }; + let response: protocol::subscription::SubscriptionFrame = + serde_json::from_str(text.as_str()).unwrap(); + assert!(matches!( + response.payload, + protocol::subscription::SubscriptionFramePayload::Response( + protocol::subscription::SubscriptionResponse::Subscribed { + selector: protocol::subscription::EventSubscriptionSelector::WorkspaceWorkers, + snapshot: protocol::subscription::SubscriptionSnapshot::Workers { .. }, + .. + } + ) + )); + server.abort(); + } + #[tokio::test] async fn passkey_registration_rejects_unverified_credential_response() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/workspace-server/src/workspace_subscription.rs b/crates/workspace-server/src/workspace_subscription.rs new file mode 100644 index 00000000..10414816 --- /dev/null +++ b/crates/workspace-server/src/workspace_subscription.rs @@ -0,0 +1,389 @@ +use std::collections::{BTreeMap, HashMap, HashSet}; + +use axum::extract::ws::{Message as WsMessage, WebSocket}; +use futures::{SinkExt, StreamExt}; +use protocol::subscription::{ + EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame, + SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest, + SubscriptionResponse, SubscriptionSnapshot, SubscriptionTerminationCode, SubscriptionWorker, +}; +use tokio::sync::mpsc; + +use crate::runtime_subscription::{BrokerSubscriptionEvent, RuntimeSubscriptionBroker}; + +const OUTBOUND_CAPACITY: usize = 256; + +pub(crate) async fn serve_workspace_subscription( + broker: RuntimeSubscriptionBroker, + socket: WebSocket, +) { + let (mut socket_sender, mut socket_receiver) = socket.split(); + let (outbound, mut outbound_receiver) = mpsc::channel::(OUTBOUND_CAPACITY); + let writer = tokio::spawn(async move { + while let Some(message) = outbound_receiver.recv().await { + if socket_sender.send(message).await.is_err() { + break; + } + } + }); + let mut next_subscription_id = 1_u64; + let mut subscriptions = HashMap::>::new(); + + while let Some(message) = socket_receiver.next().await { + let Ok(message) = message else { break }; + match message { + WsMessage::Text(text) => { + let Ok(frame) = serde_json::from_str::(text.as_str()) else { + break; + }; + if frame.validate().is_err() { + break; + } + let SubscriptionFramePayload::Request(request) = frame.payload else { + break; + }; + 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( + broker.clone(), + request_id, + subscription_id.clone(), + outbound.clone(), + )); + subscriptions.insert(subscription_id, task); + } + SubscriptionRequest::UnsubscribeEvents { + request_id, + subscription_id, + } => { + if let Some(task) = subscriptions.remove(&subscription_id) { + task.abort(); + } + if send_frame( + &outbound, + SubscriptionFrame::new(SubscriptionFramePayload::Response( + SubscriptionResponse::Unsubscribed { + request_id, + subscription_id, + }, + )), + ) + .await + .is_err() + { + break; + } + } + } + } + WsMessage::Ping(value) => { + if outbound.send(WsMessage::Pong(value)).await.is_err() { + break; + } + } + WsMessage::Pong(_) => {} + WsMessage::Close(_) | WsMessage::Binary(_) => break, + } + } + + for (_, task) in subscriptions { + task.abort(); + } + drop(outbound); + let _ = writer.await; +} + +async fn run_workspace_workers( + broker: RuntimeSubscriptionBroker, + request_id: protocol::subscription::SubscriptionRequestId, + subscription_id: SubscriptionId, + outbound: mpsc::Sender, +) { + let runtime_ids = broker.runtime_ids(); + let mut pending = runtime_ids.iter().cloned().collect::>(); + let (events, mut event_receiver) = mpsc::channel(OUTBOUND_CAPACITY); + let mut upstreams = tokio::task::JoinSet::new(); + for runtime_id in runtime_ids { + let Ok(mut subscription) = + broker.subscribe(&runtime_id, EventSubscriptionSelector::RuntimeWorkers) + else { + pending.remove(&runtime_id); + continue; + }; + let sender = events.clone(); + upstreams.spawn(async move { + while let Some(event) = subscription.recv().await { + if sender.send((runtime_id.clone(), event)).await.is_err() { + break; + } + } + }); + } + drop(events); + + let mut workers = HashMap::>::new(); + while !pending.is_empty() { + let Some((runtime_id, event)) = event_receiver.recv().await else { + return; + }; + match event { + BrokerSubscriptionEvent::Snapshot { snapshot, .. } => { + install_snapshot(&mut workers, &runtime_id, snapshot); + pending.remove(&runtime_id); + } + BrokerSubscriptionEvent::Disconnected { .. } + | BrokerSubscriptionEvent::Rejected { .. } + | BrokerSubscriptionEvent::Closed { .. } => { + pending.remove(&runtime_id); + } + BrokerSubscriptionEvent::Event { .. } => {} + } + } + + let mut revisions = HashMap::::new(); + let mut initial_workers = workers + .values_mut() + .flat_map(|runtime| runtime.values_mut()) + .map(|worker| { + let key = worker_key(worker.runtime_id.as_deref(), worker.worker_id.as_str()); + worker.subject_revision = next_revision(&mut revisions, &key); + worker.clone() + }) + .collect::>(); + sort_workers(&mut initial_workers); + if send_frame( + &outbound, + SubscriptionFrame::new(SubscriptionFramePayload::Response( + SubscriptionResponse::Subscribed { + request_id, + subscription_id: subscription_id.clone(), + selector: EventSubscriptionSelector::WorkspaceWorkers, + snapshot_revision: 1, + snapshot: SubscriptionSnapshot::Workers { + workers: initial_workers, + }, + }, + )), + ) + .await + .is_err() + { + return; + } + + while let Some((runtime_id, event)) = event_receiver.recv().await { + match event { + BrokerSubscriptionEvent::Snapshot { snapshot, .. } => { + let removed = workers.remove(&runtime_id).unwrap_or_default(); + for worker in removed.values() { + let key = worker_key(Some(&runtime_id), worker.worker_id.as_str()); + let revision = next_revision(&mut revisions, &key); + if send_event( + &outbound, + &subscription_id, + revision, + SubscriptionEventPayload::WorkerRemoved { + worker_id: worker.worker_id.clone(), + runtime_id: Some(runtime_id.clone()), + }, + ) + .await + .is_err() + { + return; + } + } + install_snapshot(&mut workers, &runtime_id, snapshot); + if let Some(current) = workers.get_mut(&runtime_id) { + for worker in current.values_mut() { + let key = worker_key(Some(&runtime_id), worker.worker_id.as_str()); + let revision = next_revision(&mut revisions, &key); + worker.subject_revision = revision; + if send_event( + &outbound, + &subscription_id, + revision, + SubscriptionEventPayload::WorkerUpserted { + worker: worker.clone(), + }, + ) + .await + .is_err() + { + return; + } + } + } + } + BrokerSubscriptionEvent::Event { payload, .. } => match payload { + SubscriptionEventPayload::WorkerUpserted { mut worker } => { + worker.runtime_id = Some(runtime_id.clone()); + let key = worker_key(Some(&runtime_id), worker.worker_id.as_str()); + let revision = next_revision(&mut revisions, &key); + worker.subject_revision = revision; + workers + .entry(runtime_id) + .or_default() + .insert(worker.worker_id.to_string(), worker.clone()); + if send_event( + &outbound, + &subscription_id, + revision, + SubscriptionEventPayload::WorkerUpserted { worker }, + ) + .await + .is_err() + { + return; + } + } + SubscriptionEventPayload::WorkerRemoved { worker_id, .. } => { + workers + .entry(runtime_id.clone()) + .or_default() + .remove(worker_id.as_str()); + let key = worker_key(Some(&runtime_id), worker_id.as_str()); + let revision = next_revision(&mut revisions, &key); + if send_event( + &outbound, + &subscription_id, + revision, + SubscriptionEventPayload::WorkerRemoved { + worker_id, + runtime_id: Some(runtime_id), + }, + ) + .await + .is_err() + { + return; + } + } + _ => {} + }, + BrokerSubscriptionEvent::Disconnected { .. } => {} + BrokerSubscriptionEvent::Rejected { code, message, .. } => { + let _ = send_frame( + &outbound, + SubscriptionFrame::new(SubscriptionFramePayload::Event( + SubscriptionEvent::SubscriptionClosed { + subscription_id: subscription_id.clone(), + code: rejection_termination(code), + message, + }, + )), + ) + .await; + return; + } + BrokerSubscriptionEvent::Closed { code, message, .. } => { + let _ = send_frame( + &outbound, + SubscriptionFrame::new(SubscriptionFramePayload::Event( + SubscriptionEvent::SubscriptionClosed { + subscription_id: subscription_id.clone(), + code, + message, + }, + )), + ) + .await; + return; + } + } + } +} + +fn install_snapshot( + workers: &mut HashMap>, + runtime_id: &str, + snapshot: SubscriptionSnapshot, +) { + let SubscriptionSnapshot::Workers { + workers: snapshot_workers, + } = snapshot + else { + return; + }; + let mut projected = BTreeMap::new(); + for mut worker in snapshot_workers { + worker.runtime_id = Some(runtime_id.to_string()); + projected.insert(worker.worker_id.to_string(), worker); + } + workers.insert(runtime_id.to_string(), projected); +} + +async fn send_event( + outbound: &mpsc::Sender, + subscription_id: &SubscriptionId, + subject_revision: u64, + payload: SubscriptionEventPayload, +) -> Result<(), ()> { + send_frame( + outbound, + SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event { + subscription_id: subscription_id.clone(), + subject_revision, + payload, + })), + ) + .await +} + +async fn send_frame( + outbound: &mpsc::Sender, + frame: SubscriptionFrame, +) -> Result<(), ()> { + frame.validate().map_err(|_| ())?; + outbound + .send(WsMessage::Text( + serde_json::to_string(&frame).map_err(|_| ())?.into(), + )) + .await + .map_err(|_| ()) +} + +fn next_revision(revisions: &mut HashMap, key: &str) -> u64 { + let revision = revisions.entry(key.to_string()).or_insert(0); + *revision = revision.saturating_add(1); + *revision +} +fn worker_key(runtime_id: Option<&str>, worker_id: &str) -> String { + format!("{}:{worker_id}", runtime_id.unwrap_or_default()) +} +fn sort_workers(workers: &mut [SubscriptionWorker]) { + workers.sort_by(|left, right| { + left.runtime_id + .cmp(&right.runtime_id) + .then_with(|| left.worker_id.cmp(&right.worker_id)) + }); +} +fn rejection_termination(code: SubscriptionRejectionCode) -> SubscriptionTerminationCode { + match code { + SubscriptionRejectionCode::Unauthorized => SubscriptionTerminationCode::Unauthorized, + SubscriptionRejectionCode::ResourceNotFound => SubscriptionTerminationCode::ResourceGone, + _ => SubscriptionTerminationCode::ServerShutdown, + } +} diff --git a/web/workspace/src/lib/generated/protocol.ts b/web/workspace/src/lib/generated/protocol.ts index ea5637d5..3903aeee 100644 --- a/web/workspace/src/lib/generated/protocol.ts +++ b/web/workspace/src/lib/generated/protocol.ts @@ -128,7 +128,7 @@ export type SubscriptionWorkdir = { working_directory_id: SubscriptionWorkdirId, export type SubscriptionSnapshot = { "topic": "workers", "data": { workers: Array, } } | { "topic": "worker_protocol", "data": { worker_id: SubscriptionWorkerId, events: Array, } } | { "topic": "workspace_workdirs", "data": { workdirs: Array, } }; -export type SubscriptionEventPayload = { "event": "worker_upserted", "data": { worker: SubscriptionWorker, } } | { "event": "worker_removed", "data": { worker_id: SubscriptionWorkerId, } } | { "event": "worker_protocol", "data": { worker_id: SubscriptionWorkerId, event: Event, } } | { "event": "workdir_upserted", "data": { workdir: SubscriptionWorkdir, } } | { "event": "workdir_removed", "data": { working_directory_id: SubscriptionWorkdirId, } }; +export type SubscriptionEventPayload = { "event": "worker_upserted", "data": { worker: SubscriptionWorker, } } | { "event": "worker_removed", "data": { worker_id: SubscriptionWorkerId, runtime_id?: string | null, } } | { "event": "worker_protocol", "data": { worker_id: SubscriptionWorkerId, event: Event, } } | { "event": "workdir_upserted", "data": { workdir: SubscriptionWorkdir, } } | { "event": "workdir_removed", "data": { working_directory_id: SubscriptionWorkdirId, } }; export type SubscriptionRejectionCode = "invalid_request" | "unsupported_protocol_version" | "unsupported_selector" | "unauthorized" | "resource_not_found" | "capacity_exceeded" | "internal";