diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index 0aeb1d9c6..ada6f3e8e 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -5361,10 +5361,17 @@ fn embedded_profile_label(profile: &ProfileSelector) -> Option { } const MEMORY_CONSOLIDATION_PROFILE: &str = "memory-consolidation"; -const MEMORY_CONSOLIDATION_SINGLETON_KEY: &str = "workspace-memory-consolidation"; +pub(crate) const MEMORY_CONSOLIDATION_SINGLETON_KEY: &str = "workspace-memory-consolidation"; const WORKSPACE_ORCHESTRATOR_PROFILE: &str = "orchestrator"; pub(crate) const WORKSPACE_ORCHESTRATOR_SINGLETON_KEY: &str = "workspace-orchestrator"; +pub(crate) fn is_reserved_internal_worker_singleton_key(key: &str) -> bool { + matches!( + key, + MEMORY_CONSOLIDATION_SINGLETON_KEY | WORKSPACE_ORCHESTRATOR_SINGLETON_KEY + ) +} + struct WorkerDisplayMetadata { display_name: String, singleton_key: Option, diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index f1a2bcbe1..0d82d8560 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -16233,7 +16233,7 @@ fn find_online_workspace_orchestrator(api: &WorkspaceApi) -> Option Option bool { + worker.workspace.workspace_id.as_deref() == Some(workspace_id) + && worker.singleton_key.as_deref() + == Some(crate::hosts::WORKSPACE_ORCHESTRATOR_SINGLETON_KEY) + && worker.tags.iter().any(|tag| tag == "orchestrator") + && worker.tags.iter().any(|tag| tag == "singleton") } async fn scoped_get_memory_document( @@ -23056,6 +23064,14 @@ async fn create_workspace_worker_inner( workdir_attachments, control_operation_id: _, } = request; + if let Some(singleton_key) = singleton_key.as_deref() + && crate::hosts::is_reserved_internal_worker_singleton_key(singleton_key) + { + return Err(Error::InvalidInput(format!( + "Worker singleton key {singleton_key:?} is reserved for a Backend-managed Worker" + )) + .into()); + } let config_state = api .config_store .load_workspace_config(&api.config.workspace_id)? @@ -33372,7 +33388,7 @@ mod tests { } #[tokio::test] - async fn arbitrary_singleton_create_conflicts_while_idle_or_stopped_and_restore_keeps_key() { + async fn arbitrary_singleton_exact_retry_replays_while_competing_create_conflicts() { let workspace = tempfile::tempdir().unwrap(); init_clean_git_workspace(workspace.path()); let api = test_api(workspace.path()).await; @@ -33395,12 +33411,22 @@ mod tests { created.worker.singleton_key.as_deref(), Some("subjektiv:subject-42") ); - let duplicate = + let Json(duplicate) = create_workspace_worker(State(api.clone()), HeaderMap::new(), Json(request())) .await - .unwrap_err(); + .unwrap(); + assert_eq!(duplicate.worker_id, created.worker_id); + let mut competing_request = request(); + competing_request.display_name = "Competing Subject Worker".to_string(); + let competing = create_workspace_worker( + State(api.clone()), + HeaderMap::new(), + Json(competing_request), + ) + .await + .unwrap_err(); assert!( - matches!(duplicate.error, Error::RepositoryConflict(message) if message.contains("worker_singleton_owned")) + matches!(competing.error, Error::RepositoryConflict(message) if message.contains("worker_singleton_owned")) ); let worker = RuntimeWorkerRef::new(&created.runtime_id, &created.worker_id); @@ -33415,12 +33441,22 @@ mod tests { ) .unwrap(); assert_eq!(stopped.state, InternalWorkerOperationState::Accepted); - let stopped_duplicate = + let Json(stopped_retry) = create_workspace_worker(State(api.clone()), HeaderMap::new(), Json(request())) .await - .unwrap_err(); + .unwrap(); + assert_eq!(stopped_retry.worker_id, created.worker_id); + let mut stopped_competing_request = request(); + stopped_competing_request.display_name = "Stopped competing Worker".to_string(); + let stopped_competing = create_workspace_worker( + State(api.clone()), + HeaderMap::new(), + Json(stopped_competing_request), + ) + .await + .unwrap_err(); assert!(matches!( - stopped_duplicate.error, + stopped_competing.error, Error::RepositoryConflict(_) )); @@ -33478,6 +33514,37 @@ mod tests { )); assert!(runtime.spawn_requests().is_empty()); + // Simulate a Backend crash after the durable singleton reservation commits but before + // Runtime dispatch. The restarted request receives a fresh manual allocation key, yet its + // identical singleton intent replays the reserved Worker and dispatches that generation. + let mut restart_request = request("ignored-manual-operation"); + restart_request.singleton_key = Some("subjektiv:restart-retry".to_string()); + restart_request.resolved_control_operation = None; + let restart_fingerprint = worker_spawn_create_fingerprint(&restart_request).unwrap(); + let memory = api + .config_store + .get_workspace_memory_settings(&api.config.workspace_id) + .unwrap(); + let pre_crash_reservation = api + .config_store + .reserve_worker_create( + &api.config.workspace_id, + WorkdirlessFixtureRuntime::RUNTIME_ID, + "manual:lost-before-dispatch", + &restart_fingerprint, + restart_request.singleton_key.as_deref(), + &memory, + ) + .unwrap(); + let recovered = api + .spawn_workspace_worker(WorkdirlessFixtureRuntime::RUNTIME_ID, restart_request) + .unwrap(); + assert_eq!(recovered.state, InternalWorkerOperationState::Accepted); + assert_eq!( + recovered.worker.unwrap().worker.worker_id, + pre_crash_reservation.worker_id.to_string() + ); + assert!(matches!( api.spawn_workspace_worker("missing-runtime", request("singleton-create-undispatched")) .unwrap_err() @@ -33657,6 +33724,90 @@ mod tests { .unwrap(); assert_eq!(generic.worker.singleton_key, None); assert!(find_workspace_orchestrator(&api).is_none()); + + let forged_public = create_workspace_worker( + State(api.clone()), + HeaderMap::new(), + Json(CreateWorkspaceWorkerRequest { + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + display_name: "Generic singleton key claimant".to_string(), + singleton_key: Some(crate::hosts::WORKSPACE_ORCHESTRATOR_SINGLETON_KEY.to_string()), + profile: Some("builtin:companion".to_string()), + ticket_assignment: None, + initial_submit: Vec::new(), + workdir_attachments: Vec::new(), + control_operation_id: None, + }), + ) + .await + .unwrap_err(); + assert!(matches!(forged_public.error, Error::InvalidInput(_))); + assert!( + api.store + .current_worker_singleton_owner( + &workspace_id, + crate::hosts::WORKSPACE_ORCHESTRATOR_SINGLETON_KEY, + ) + .unwrap() + .is_none() + ); + + // Even a trusted internal caller that accidentally assigns the legacy key cannot turn + // generic Worker metadata into Orchestrator completion authority. + let forged_internal = api + .spawn_workspace_worker( + EMBEDDED_WORKER_RUNTIME_ID, + WorkerSpawnRequest { + requested_worker_name: Some("Forged Orchestrator claimant".to_string()), + singleton_key: Some( + crate::hosts::WORKSPACE_ORCHESTRATOR_SINGLETON_KEY.to_string(), + ), + intent: WorkerSpawnIntent::WorkspaceCompanion, + acceptance: WorkerSpawnAcceptanceRequirement::RunAccepted { + expected_segments: 0, + }, + profile: ProfileSelector::Builtin("builtin:companion".to_string()), + ticket_assignment: None, + initial_submit: Vec::new(), + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), + resolved_config_bundle: None, + resolved_worker_observation_enabled: false, + resolved_worker_observation_grants: Vec::new(), + resolved_workspace_api: None, + resolved_memory_settings: None, + resolved_control_operation: None, + }, + ) + .unwrap() + .worker + .unwrap(); + assert!(find_workspace_orchestrator(&api).is_none()); + assert!( + require_online_workspace_orchestrator_source(&api, &forged_internal.worker).is_err() + ); + let stopped = api + .runtime + .stop_worker( + &forged_internal.worker, + WorkerLifecycleRequest { + reason: Some("remove forged Orchestrator claimant".to_string()), + ticket_assignment: None, + }, + ) + .unwrap(); + assert_eq!(stopped.state, InternalWorkerOperationState::Accepted); + assert!( + api.runtime + .delete_worker(&forged_internal.worker) + .unwrap() + .deleted + ); + api.store + .delete_worker_registry(&workspace_id, &forged_internal.worker) + .unwrap(); + let reserved = create_workspace_worker( State(api.clone()), HeaderMap::new(), diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index c97acaecd..d8bf0bc46 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -2231,6 +2231,86 @@ impl SqliteWorkspaceStore { }); } + let singleton_owner = singleton_key + .map(|key| current_worker_singleton_owner(&tx, workspace_id, key)) + .transpose()? + .flatten(); + if let Some((owner, state)) = singleton_owner.as_ref() + && state.as_deref() != Some("removed") + { + let stored = tx + .query_row( + "SELECT allocation_key, request_fingerprint, create_fingerprint, \ + memory_settings_revision, memory_language, singleton_key, \ + singleton_generation \ + FROM worker_create_reservations \ + WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3", + params![ + workspace_id, + owner.worker.runtime_id, + owner.worker.worker_id + ], + |row| { + Ok(( + row.get::<_, String>(0)?, + row.get::<_, Option>(1)?, + row.get::<_, String>(2)?, + row.get::<_, Option>(3)?, + row.get::<_, Option>(4)?, + row.get::<_, Option>(5)?, + row.get::<_, Option>(6)?, + )) + }, + ) + .optional()?; + if let Some((stored_allocation_key, stored_request_fingerprint, create_fingerprint, revision, language, stored_singleton_key, stored_generation)) = stored + && allocation_key.starts_with("manual:") + && stored_allocation_key.starts_with("manual:") + && owner.worker.runtime_id == runtime_id + && stored_request_fingerprint.as_deref() == Some(request_fingerprint) + && stored_singleton_key.as_deref() == singleton_key + && stored_generation.and_then(|generation| u64::try_from(generation).ok()) + == Some(owner.generation) + { + let revision = revision.ok_or_else(|| { + Error::InvalidInput(format!( + "Worker singleton reservation {}:{} has no persisted Memory settings snapshot", + owner.worker.runtime_id, owner.worker.worker_id + )) + })?; + let language = language.ok_or_else(|| { + Error::InvalidInput(format!( + "Worker singleton reservation {}:{} has no persisted Memory language", + owner.worker.runtime_id, owner.worker.worker_id + )) + })?; + let worker_id = owner.worker.worker_id.parse::().map_err(|_| { + Error::Store(format!( + "Worker singleton reservation {}:{} has a non-UUIDv7 worker id", + owner.worker.runtime_id, owner.worker.worker_id + )) + })?; + let snapshot = manifest::WorkspaceMemorySettingsSnapshot { + workspace_id: workspace_id.to_string(), + settings_revision: revision.try_into().map_err(|_| { + rusqlite::Error::IntegralValueOutOfRange(2, revision) + })?, + language, + }; + validate_workspace_memory_settings_snapshot(&snapshot, workspace_id)?; + return Ok(WorkerCreateReservation { + worker_id, + create_fingerprint, + memory_settings: snapshot, + singleton: Some(owner.clone()), + }); + } + return Err(Error::RepositoryConflict(format!( + "worker_singleton_owned: singleton key is already owned by Worker {}:{}", + owner.worker.runtime_id, owner.worker.worker_id + ))); + } + let (authoritative_revision, authoritative_language) = tx .query_row( "SELECT settings_revision, language FROM workspace_memory_settings WHERE workspace_id = ?1", @@ -2262,10 +2342,6 @@ impl SqliteWorkspaceStore { let worker_id = WorkerId::now_v7(); let worker_id_text = worker_id.to_string(); let now = chrono::Utc::now().to_rfc3339(); - let singleton_owner = singleton_key - .map(|key| current_worker_singleton_owner(&tx, workspace_id, key)) - .transpose()? - .flatten(); let singleton_generation = match singleton_owner.as_ref() { None if singleton_key.is_some() => Some(1_u64), None => None, @@ -18220,6 +18296,75 @@ mod tests { assert!(store.get_workspace_memory_settings("workspace-a").is_err()); } + #[tokio::test] + async fn worker_singleton_reservation_retry_survives_restart_and_new_allocation_key() { + let dir = tempfile::tempdir().unwrap(); + let database = dir.path().join("server.db"); + let store = SqliteWorkspaceStore::open(&database).unwrap(); + store + .upsert_workspace(&WorkspaceRecord { + workspace_id: "workspace-a".to_string(), + owner_account_id: "owner-account".to_string(), + display_name: "Workspace A".to_string(), + state: "active".to_string(), + created_at: "2026-10-01T00:00:00Z".to_string(), + updated_at: "2026-10-01T00:00:00Z".to_string(), + }) + .await + .unwrap(); + let memory = store.get_workspace_memory_settings("workspace-a").unwrap(); + let reserved = store + .reserve_worker_create( + "workspace-a", + "runtime-a", + "manual:lost-before-dispatch", + "sha256:stable-request", + Some("subjektiv:restart-retry"), + &memory, + ) + .unwrap(); + drop(store); + + let reopened = SqliteWorkspaceStore::open(&database).unwrap(); + let retried = reopened + .reserve_worker_create( + "workspace-a", + "runtime-a", + "manual:new-after-restart", + "sha256:stable-request", + Some("subjektiv:restart-retry"), + &memory, + ) + .unwrap(); + assert_eq!(retried, reserved); + let reservation_count: i64 = reopened + .with_conn(|conn| { + conn.query_row( + "SELECT COUNT(*) FROM worker_create_reservations \ + WHERE workspace_id = 'workspace-a' AND singleton_key = 'subjektiv:restart-retry'", + [], + |row| row.get(0), + ) + .map_err(Error::from) + }) + .unwrap(); + assert_eq!(reservation_count, 1); + let conflict = reopened + .reserve_worker_create( + "workspace-a", + "runtime-a", + "manual:different-request", + "sha256:different-request", + Some("subjektiv:restart-retry"), + &memory, + ) + .unwrap_err(); + assert!(matches!( + conflict, + Error::RepositoryConflict(message) if message.contains("worker_singleton_owned") + )); + } + #[tokio::test] async fn worker_singleton_ownership_is_workspace_scoped_durable_and_fences_old_generations() { let dir = tempfile::tempdir().unwrap();