diff --git a/crates/workspace-server/src/runtime_subscription.rs b/crates/workspace-server/src/runtime_subscription.rs index b28755a0..88a3a47e 100644 --- a/crates/workspace-server/src/runtime_subscription.rs +++ b/crates/workspace-server/src/runtime_subscription.rs @@ -418,6 +418,8 @@ async fn run_embedded_connection( update = update_receiver.recv() => { let Some((selector, subject_revision, payload)) = update else { return; }; if let Some(entry) = entries.get_mut(&selector) { + entry.snapshot_revision = entry.snapshot_revision.saturating_add(1); + apply_event_to_cached_snapshot(&mut entry.snapshot, &payload); broadcast(&mut entry.downstreams, BrokerSubscriptionEvent::Event { connection_generation: generation, subject_revision, payload }); } } @@ -431,6 +433,31 @@ async fn run_embedded_connection( } } +fn apply_event_to_cached_snapshot( + snapshot: &mut SubscriptionSnapshot, + payload: &SubscriptionEventPayload, +) { + let SubscriptionSnapshot::Workers { workers } = snapshot else { + return; + }; + match payload { + SubscriptionEventPayload::WorkerUpserted { worker } => { + if let Some(existing) = workers + .iter_mut() + .find(|existing| existing.worker_id == worker.worker_id) + { + *existing = worker.clone(); + } else { + workers.push(worker.clone()); + } + } + SubscriptionEventPayload::WorkerRemoved { worker_id, .. } => { + workers.retain(|worker| worker.worker_id != *worker_id); + } + _ => {} + } +} + fn project_snapshot_runtime( mut snapshot: SubscriptionSnapshot, runtime_id: &str, @@ -737,6 +764,10 @@ async fn handle_frame( } *revision = subject_revision; } + if let Some((snapshot_revision, snapshot)) = entry.snapshot.as_mut() { + *snapshot_revision = snapshot_revision.saturating_add(1); + apply_event_to_cached_snapshot(snapshot, &payload); + } broadcast( &mut entry.downstreams, BrokerSubscriptionEvent::Event { diff --git a/crates/workspace-server/src/runtime_subscription_tests.rs b/crates/workspace-server/src/runtime_subscription_tests.rs index 48a2287e..921af63f 100644 --- a/crates/workspace-server/src/runtime_subscription_tests.rs +++ b/crates/workspace-server/src/runtime_subscription_tests.rs @@ -36,6 +36,13 @@ impl WorkerExecutionBackend for TestExecutionBackend { WorkerExecutionRunState::Busy, ) } + + fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { + WorkerExecutionResult::accepted( + WorkerExecutionOperation::Stop, + WorkerExecutionRunState::Stopped, + ) + } } const TOKEN: &str = "runtime-subscription-test-token"; @@ -156,7 +163,7 @@ async fn equal_downstream_selectors_share_one_upstream_subscription() { .unwrap(), }; let mut first = broker.subscribe("runtime-test", selector.clone()).unwrap(); - let mut second = broker.subscribe("runtime-test", selector).unwrap(); + let mut second = broker.subscribe("runtime-test", selector.clone()).unwrap(); assert!(matches!( next_snapshot(&mut first).await, BrokerSubscriptionEvent::Snapshot { .. } @@ -188,6 +195,16 @@ async fn equal_downstream_selectors_share_one_upstream_subscription() { } if worker.state == SubscriptionWorkerState::Running )); } + let mut late = broker.subscribe("runtime-test", selector.clone()).unwrap(); + let BrokerSubscriptionEvent::Snapshot { snapshot, .. } = next_snapshot(&mut late).await else { + panic!("expected cached snapshot for late subscriber"); + }; + assert!(matches!( + snapshot, + SubscriptionSnapshot::Workers { workers } + if workers.iter().any(|worker| worker.state == SubscriptionWorkerState::Running) + )); + drop(late); drop(first); tokio::task::yield_now().await; @@ -312,5 +329,53 @@ async fn embedded_runtime_uses_in_process_subscription_source() { 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)); + let mut late = broker + .subscribe( + "embedded-worker-runtime", + EventSubscriptionSelector::RuntimeWorkers, + ) + .unwrap(); + let BrokerSubscriptionEvent::Snapshot { snapshot, .. } = next_snapshot(&mut late).await else { + panic!("expected cached embedded snapshot for late subscriber"); + }; + assert!(matches!( + snapshot, + SubscriptionSnapshot::Workers { workers } + if workers.iter().any(|worker| worker.state == SubscriptionWorkerState::Running) + )); + + runtime + .stop_worker(&worker.worker_ref, Some("done".to_string())) + .unwrap(); + assert!(matches!( + next_event(&mut subscription).await, + BrokerSubscriptionEvent::Event { + payload: SubscriptionEventPayload::WorkerUpserted { .. }, + .. + } + )); + runtime.delete_worker(&worker.worker_ref).unwrap(); + assert!(matches!( + next_event(&mut subscription).await, + BrokerSubscriptionEvent::Event { + payload: SubscriptionEventPayload::WorkerRemoved { .. }, + .. + } + )); + let mut after_remove = broker + .subscribe( + "embedded-worker-runtime", + EventSubscriptionSelector::RuntimeWorkers, + ) + .unwrap(); + let BrokerSubscriptionEvent::Snapshot { snapshot, .. } = next_snapshot(&mut after_remove).await + else { + panic!("expected cached embedded snapshot after remove"); + }; + assert!(matches!( + snapshot, + SubscriptionSnapshot::Workers { workers } + if workers.iter().all(|candidate| candidate.worker_id.as_str() != worker.worker_ref.worker_id.to_string()) + )); server.abort(); } diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index e1e2c0a4..bfcc517c 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -129,9 +129,14 @@ const MIGRATIONS: &[Migration] = &[ }, Migration { version: 23, - name: "enforce exclusive Worker Workdir attachments and spawn reservations", + name: "enforce exclusive active Worker Workdir attachments", apply: enforce_exclusive_worker_workdir_attachments, }, + Migration { + version: 24, + name: "create Worker Workdir attachment reservations", + apply: create_worker_workdir_attachment_reservations, + }, ]; struct Migration { @@ -3415,6 +3420,20 @@ DROP TABLE IF EXISTS ticket_notification_outbox; fn enforce_exclusive_worker_workdir_attachments(conn: &Connection) -> Result<()> { conn.execute_batch( r#" +CREATE UNIQUE INDEX IF NOT EXISTS ux_worker_workdir_links_active_worker + ON worker_workdir_links(workspace_id, runtime_id, runtime_worker_id) + WHERE unlinked_at IS NULL; +CREATE UNIQUE INDEX IF NOT EXISTS ux_worker_workdir_links_active_workdir + ON worker_workdir_links(workspace_id, workdir_id) + WHERE unlinked_at IS NULL; +"#, + )?; + Ok(()) +} + +fn create_worker_workdir_attachment_reservations(conn: &Connection) -> Result<()> { + conn.execute_batch( + r#" CREATE TABLE IF NOT EXISTS worker_workdir_attachment_reservations ( workspace_id TEXT NOT NULL, workdir_id TEXT NOT NULL, @@ -3424,12 +3443,6 @@ CREATE TABLE IF NOT EXISTS worker_workdir_attachment_reservations ( ); CREATE UNIQUE INDEX IF NOT EXISTS ux_worker_workdir_attachment_reservation_id ON worker_workdir_attachment_reservations(workspace_id, reservation_id); -CREATE UNIQUE INDEX IF NOT EXISTS ux_worker_workdir_links_active_worker - ON worker_workdir_links(workspace_id, runtime_id, runtime_worker_id) - WHERE unlinked_at IS NULL; -CREATE UNIQUE INDEX IF NOT EXISTS ux_worker_workdir_links_active_workdir - ON worker_workdir_links(workspace_id, workdir_id) - WHERE unlinked_at IS NULL; "#, )?; Ok(()) @@ -4114,13 +4127,39 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT); } } + #[test] + fn schema_v24_adds_attachment_reservations_to_already_applied_v23() { + let conn = Connection::open_in_memory().unwrap(); + configure_sqlite(&conn).unwrap(); + for migration in MIGRATIONS + .iter() + .filter(|migration| migration.version <= 23) + { + let tx = conn.unchecked_transaction().unwrap(); + (migration.apply)(&tx).unwrap(); + tx.execute( + "INSERT INTO __yoi_schema_migrations (version, name) VALUES (?1, ?2)", + params![migration.version, migration.name], + ) + .unwrap(); + tx.commit().unwrap(); + } + assert_eq!(current_schema_version(&conn).unwrap(), 23); + assert!(!table_exists(&conn, "worker_workdir_attachment_reservations").unwrap()); + + apply_migrations(&conn).unwrap(); + + assert_eq!(current_schema_version(&conn).unwrap(), 24); + assert!(table_exists(&conn, "worker_workdir_attachment_reservations").unwrap()); + } + #[tokio::test] async fn migrates_sqlite_and_preserves_workspace_record() { let dir = tempfile::tempdir().unwrap(); let db = dir.path().join("control-plane.sqlite"); let store = SqliteWorkspaceStore::open(&db).unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 23); + assert_eq!(store.schema_version().await.unwrap(), 24); assert!( !store .with_conn(|conn| table_exists(conn, "worker_workspace_credentials")) @@ -4137,7 +4176,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT); store.upsert_workspace(&record).await.unwrap(); let reopened = SqliteWorkspaceStore::open(&db).unwrap(); - assert_eq!(reopened.schema_version().await.unwrap(), 23); + assert_eq!(reopened.schema_version().await.unwrap(), 24); assert_eq!( reopened.get_workspace("local-dev").await.unwrap(), Some(record) @@ -4588,7 +4627,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT); .unwrap(); let store = SqliteWorkspaceStore::from_connection(conn).unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 23); + assert_eq!(store.schema_version().await.unwrap(), 24); store .with_conn(|conn| { @@ -4777,7 +4816,7 @@ CREATE TABLE ticket_assignment_operations ( #[tokio::test] async fn repository_records_round_trip() { let store = SqliteWorkspaceStore::in_memory().unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 23); + assert_eq!(store.schema_version().await.unwrap(), 24); let workspace = WorkspaceRecord { workspace_id: "local-dev".to_string(), owner_account_id: None, @@ -4815,7 +4854,7 @@ CREATE TABLE ticket_assignment_operations ( #[tokio::test] async fn memory_authority_records_round_trip_and_close_staging() { let store = SqliteWorkspaceStore::in_memory().unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 23); + assert_eq!(store.schema_version().await.unwrap(), 24); let workspace = WorkspaceRecord { workspace_id: "local-dev".to_string(), owner_account_id: None, @@ -5065,7 +5104,7 @@ CREATE TABLE ticket_assignment_operations ( #[tokio::test] async fn account_and_login_records_round_trip() { let store = SqliteWorkspaceStore::in_memory().unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 23); + assert_eq!(store.schema_version().await.unwrap(), 24); let now = "2026-07-22T00:00:00Z".to_string(); let account = AccountRecord { account_id: "acct-user-alice".to_string(),