Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: harden singleton retries and authority

Browse Source
This commit is contained in:
Hare
2026-10-01 15:46:18 +09:00
parent 1afddc148d
commit f351996d2c
3 changed files with 317 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/workspace-server/src/hosts.rs
+8 -1
View File
@@ -5361,10 +5361,17 @@ fn embedded_profile_label(profile: &ProfileSelector) -> Option<String> {
}
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<String>,
crates/workspace-server/src/server.rs
+160 -9
View File
@@ -16233,7 +16233,7 @@ fn find_online_workspace_orchestrator(api: &WorkspaceApi) -> Option<InternalWork
.flatten()?;
let mut worker = api.runtime.worker(&owner.worker).ok()?;
worker.singleton_key = Some(owner.key);
(worker.workspace.workspace_id.as_deref() == Some(api.config.workspace_id.as_str())
(is_dedicated_workspace_orchestrator(&worker, &api.config.workspace_id)
&& matches!(worker.state.as_str(), "idle" | "running" | "paused"))
.then_some(worker)
}
@@ -16249,7 +16249,15 @@ fn find_workspace_orchestrator(api: &WorkspaceApi) -> Option<InternalWorkerSumma
.flatten()?;
let mut worker = api.runtime.worker(&owner.worker).ok()?;
worker.singleton_key = Some(owner.key);
Some(worker)
is_dedicated_workspace_orchestrator(&worker, &api.config.workspace_id).then_some(worker)
}
fn is_dedicated_workspace_orchestrator(worker: &InternalWorkerSummary, workspace_id: &str) -> 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(),
crates/workspace-server/src/store.rs
+149 -4
View File
@@ -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<String>>(1)?,
row.get::<_, String>(2)?,
row.get::<_, Option<i64>>(3)?,
row.get::<_, Option<String>>(4)?,
row.get::<_, Option<String>>(5)?,
row.get::<_, Option<i64>>(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::<WorkerId>().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();