fix: accept reserved workers during initial flow resolution

This commit is contained in:
2026-08-26 07:52:04 +09:00
parent f7852e8034
commit 9d572d18bc
4 changed files with 121 additions and 2 deletions
+4 -1
View File
@@ -2460,7 +2460,10 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
.map(ToOwned::to_owned)
})
.unwrap_or_else(|| "Workspace rejected Flow source resolution".to_string());
return Err(WorkerError::FlowInput(message));
return Err(WorkerError::FlowInput(format!(
"{message} (HTTP {})",
response.status
)));
}
let source: flow::ResolvedFlowSource = serde_json::from_str(&response.body)
.map_err(|error| WorkerError::FlowInput(format!("decode Flow source: {error}")))?;
+75
View File
@@ -20387,6 +20387,81 @@ mod tests {
));
}
#[tokio::test]
async fn runtime_request_proof_accepts_active_worker_create_reservation() {
let workspace = tempfile::tempdir().unwrap();
let mut api = test_api(workspace.path()).await;
let identity =
worker_runtime::auth::RuntimeIdentityMaterial::generate("runtime-test").unwrap();
configure_runtime_request_auth(&mut api, &identity, "runtime-test");
let store = SqliteWorkspaceStore::open(&api.config.database_path).unwrap();
let memory_settings = store
.get_workspace_memory_settings(TEST_WORKSPACE_ID)
.unwrap();
let reserved = store
.reserve_worker_create(
TEST_WORKSPACE_ID,
"runtime-test",
"spawn-flow-race",
&"f".repeat(64),
&memory_settings,
)
.unwrap();
let worker_id = reserved.worker_id.to_string();
let path = format!("/api/w/{TEST_WORKSPACE_ID}/flows/resolve");
let signer = worker_runtime::auth::RuntimeRequestSourceSigner::from_identity(&identity);
let issue = || {
signer
.issue(
"server-test",
TEST_WORKSPACE_ID,
Some(worker_id.as_str()),
worker_runtime::auth::WORKSPACE_REQUEST_PERMISSION,
"POST",
&path,
b"{}",
i64::try_from(worker_runtime::auth::unix_now_seconds()).unwrap_or(i64::MAX),
30,
)
.unwrap()
};
let proof = issue();
let source = crate::worker_source::verify_runtime_request_source_proof_with_store(
api.store.as_ref(),
&api.config,
&proof,
TEST_WORKSPACE_ID,
worker_runtime::auth::WORKSPACE_REQUEST_PERMISSION,
"POST",
&path,
&worker_runtime::auth::request_body_digest(b"{}"),
)
.await
.expect("active create reservation should establish provisional Worker membership");
assert_eq!(source.worker_id, Some(worker_id.clone()));
store
.complete_worker_create_reservation(TEST_WORKSPACE_ID, reserved.worker_id)
.unwrap();
let proof = issue();
let result = crate::worker_source::verify_runtime_request_source_proof_with_store(
api.store.as_ref(),
&api.config,
&proof,
TEST_WORKSPACE_ID,
worker_runtime::auth::WORKSPACE_REQUEST_PERMISSION,
"POST",
&path,
&worker_runtime::auth::request_body_digest(b"{}"),
)
.await;
assert!(matches!(
result,
Err(crate::worker_source::WorkerMutationSourceProofError::WorkerCatalogMembership)
));
}
#[tokio::test]
async fn runtime_request_proof_verifies_path_and_query_for_ticket_search() {
let workspace = tempfile::tempdir().unwrap();
+38
View File
@@ -1002,6 +1002,11 @@ pub trait ControlPlaneStore: Send + Sync {
workspace_id: &str,
worker: &RuntimeWorkerRef,
) -> Result<Option<WorkerRegistryRecord>>;
fn has_active_worker_create_reservation(
&self,
workspace_id: &str,
worker: &RuntimeWorkerRef,
) -> Result<bool>;
fn list_worker_registry(
&self,
workspace_id: &str,
@@ -3268,6 +3273,28 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
})
}
fn has_active_worker_create_reservation(
&self,
workspace_id: &str,
worker: &RuntimeWorkerRef,
) -> Result<bool> {
self.with_conn(|conn| {
conn.query_row(
r#"SELECT EXISTS(
SELECT 1
FROM worker_create_reservations
WHERE workspace_id = ?1
AND runtime_id = ?2
AND worker_id = ?3
AND state = 'reserved'
)"#,
params![workspace_id, worker.runtime_id, worker.worker_id],
|row| row.get::<_, bool>(0),
)
.map_err(Error::from)
})
}
fn list_worker_registry(
&self,
workspace_id: &str,
@@ -10476,6 +10503,12 @@ INSERT INTO workdir_registry (
);
assert_eq!(reserved.memory_settings.settings_revision, 1);
assert_eq!(reserved.memory_settings.language, "English");
let reserved_worker = RuntimeWorkerRef::new("arcadia", reserved.worker_id.to_string());
assert!(
store
.has_active_worker_create_reservation("workspace-a", &reserved_worker)
.unwrap()
);
let unchanged_memory_settings = store
.update_workspace_memory_settings("workspace-a", 1, " English ")
.unwrap();
@@ -10549,6 +10582,11 @@ INSERT INTO workdir_registry (
store
.complete_worker_create_reservation("workspace-a", reserved.worker_id)
.unwrap();
assert!(
!store
.has_active_worker_create_reservation("workspace-a", &reserved_worker)
.unwrap()
);
let state: String = store
.with_conn(|conn| {
conn.query_row(
+4 -1
View File
@@ -106,7 +106,10 @@ pub async fn verify_runtime_request_source_proof_with_store(
let member = store
.get_worker_registry(workspace_id, &worker)
.map_err(|error| WorkerMutationSourceProofError::Authority(error.to_string()))?;
if member.is_none() {
let reserved = store
.has_active_worker_create_reservation(workspace_id, &worker)
.map_err(|error| WorkerMutationSourceProofError::Authority(error.to_string()))?;
if member.is_none() && !reserved {
return Err(WorkerMutationSourceProofError::WorkerCatalogMembership);
}
}