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

worker: persist known-worker control grants

Browse Source
This commit is contained in:
Hare
2026-08-16 23:08:32 +09:00
parent 46a44b232b
commit f41ab0e277
7 changed files with 1082 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/workspace-server/src/server.rs
+379 -73
View File
@@ -103,8 +103,8 @@ use crate::skills;
use crate::store::{
AccountRecord, ApiTokenRecord, AuthChallengeRecord, BrowserSessionRecord, ControlPlaneStore,
DeviceLoginFlowRecord, FlowSourceRecord, PasskeyCredentialRecord, RepositoryRecord,
TicketWorkerAssignmentRecord, UserRecord, WorkdirRegistryRecord, WorkerRegistryRecord,
WorkerWorkdirLinkRecord, WorkspaceRecord,
TicketWorkerAssignmentRecord, UserRecord, WorkdirRegistryRecord, WorkerControlGrantRecord,
WorkerRegistryRecord, WorkerWorkdirLinkRecord, WorkspaceRecord,
};
use crate::{Error, Result};
use worker_runtime::catalog::{
@@ -335,20 +335,10 @@ impl WorkspaceWorkerRemoveExecutor {
let runtime = self.runtime.upgrade().ok_or_else(|| {
"Workspace Runtime registry is unavailable during WorkerRemove".to_string()
})?;
let source_is_current_orchestrator =
runtime.list_workers(1_000).items.into_iter().any(|worker| {
worker.worker.runtime_id == source.runtime_id
&& worker.worker.worker_id == source.worker_id
&& worker.singleton_key.as_deref()
== Some(crate::hosts::WORKSPACE_ORCHESTRATOR_SINGLETON_KEY)
});
if !source_is_current_orchestrator {
return Ok(worker_remove_error_response(
StatusCode::FORBIDDEN,
"orchestrator_required",
"WorkerRemove is restricted to the current Workspace Orchestrator",
));
}
let target = RuntimeWorkerRef {
runtime_id: target_runtime_id.to_string(),
worker_id: target_worker_id.to_string(),
};
if source.runtime_id == target_runtime_id && source.worker_id == target_worker_id {
return Ok(worker_remove_error_response(
StatusCode::CONFLICT,
@@ -356,11 +346,25 @@ impl WorkspaceWorkerRemoveExecutor {
"The current Orchestrator cannot remove itself",
));
}
let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id);
let granted = self
.store
.get_active_worker_control_grant(&self.workspace_id, &controller, &target)
.map_err(|_| "Worker control grant authority is unavailable".to_string())?
.is_some_and(|grant| {
grant
.permissions
.iter()
.any(|permission| permission == "remove")
});
if !granted {
return Ok(worker_remove_error_response(
StatusCode::NOT_FOUND,
"unknown_worker",
"The target Worker is not known to the current Worker",
));
}
let target = RuntimeWorkerRef {
runtime_id: target_runtime_id.to_string(),
worker_id: target_worker_id.to_string(),
};
let remove_lock = {
let mut locks = self
.worker_remove_locks
@@ -1435,6 +1439,26 @@ pub fn build_router(api: WorkspaceApi) -> Router {
get(scoped_workspace_orchestrator_status)
.post(scoped_start_workspace_orchestrator),
)
.route(
"/api/w/{workspace_id}/worker-control/workers",
get(list_known_workers).post(spawn_known_worker),
)
.route(
"/api/w/{workspace_id}/worker-control/workers/{runtime_id}/{worker_id}/input",
post(send_known_worker_input),
)
.route(
"/api/w/{workspace_id}/worker-control/workers/{runtime_id}/{worker_id}/cancel",
post(cancel_known_worker),
)
.route(
"/api/w/{workspace_id}/worker-control/workers/{runtime_id}/{worker_id}/stop",
post(stop_known_worker),
)
.route(
"/api/w/{workspace_id}/worker-control/workers/{runtime_id}/{worker_id}/restore",
post(restore_known_worker),
)
.route(
"/api/w/{workspace_id}/worker-observation/sessions",
get(scoped_list_worker_observation_sessions),
@@ -2066,6 +2090,9 @@ pub struct CreateWorkspaceWorkerRequest {
pub initial_submit: Vec<Segment>,
#[serde(default)]
pub working_directory: Option<BrowserWorkerWorkingDirectorySelection>,
/// Backend idempotency key used only for authenticated Worker-owned spawn/control.
#[serde(default)]
pub control_operation_id: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
@@ -5832,6 +5859,229 @@ async fn scoped_workspace_orchestrator_status(
Ok(Json(workspace_orchestrator_response(&api, "observed")))
}
#[derive(Debug, Serialize, Deserialize)]
struct KnownWorkerRecord {
subject: RuntimeWorkerRef,
relation: String,
origin: String,
permissions: Vec<String>,
summary: WorkerSummary,
}
#[derive(Debug, Serialize, Deserialize)]
struct KnownWorkersResponse {
workspace_id: String,
items: Vec<KnownWorkerRecord>,
truncated: bool,
}
async fn list_known_workers(
State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedWorkspacePath>,
headers: HeaderMap,
) -> ApiResult<Json<KnownWorkersResponse>> {
validate_workspace_scope(&api, &path.workspace_id)?;
let source = authenticate_worker_mutation_source(&api, &path.workspace_id, &headers)?;
let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id);
let limit = api.config.max_records.clamp(1, 500);
let grants =
api.store
.list_active_worker_control_grants(&path.workspace_id, &controller, limit + 1)?;
let truncated = grants.len() > limit;
let mut items = Vec::with_capacity(grants.len().min(limit));
for grant in grants.into_iter().take(limit) {
let summary = api
.runtime
.worker(&grant.subject)
.map_err(|error| error.into_error())?;
items.push(KnownWorkerRecord {
subject: grant.subject,
relation: grant.relation,
origin: grant.origin,
permissions: grant.permissions,
summary,
});
}
Ok(Json(KnownWorkersResponse {
workspace_id: path.workspace_id,
items,
truncated,
}))
}
async fn spawn_known_worker(
State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedWorkspacePath>,
headers: HeaderMap,
Json(request): Json<CreateWorkspaceWorkerRequest>,
) -> ApiResult<Json<BrowserCreateWorkerResponse>> {
validate_workspace_scope(&api, &path.workspace_id)?;
let source = authenticate_worker_mutation_source(&api, &path.workspace_id, &headers)?;
let relation = if request.ticket_assignment.is_some() {
"assigned"
} else {
"spawned"
};
let operation_id = request
.control_operation_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| Error::InvalidInput("control_operation_id is required".to_string()))?
.to_string();
let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id);
if request.ticket_assignment.is_none()
&& let Some(existing) = api.store.get_worker_control_grant_by_operation(
&path.workspace_id,
&controller,
&operation_id,
)?
{
let worker = api
.runtime
.worker(&existing.subject)
.map_err(|error| error.into_error())?;
return Ok(Json(BrowserCreateWorkerResponse {
workspace_id: path.workspace_id,
console_href: format!(
"/w/{}/runtimes/{}/workers/{}/console",
encode_path_segment(&existing.workspace_id),
encode_path_segment(&existing.subject.runtime_id),
encode_path_segment(&existing.subject.worker_id),
),
worker_ref: existing.subject,
worker,
diagnostics: Vec::new(),
}));
}
let response = create_workspace_worker(State(api.clone()), headers, Json(request)).await?;
if let Err(error) = api
.store
.create_worker_control_grant(&WorkerControlGrantRecord {
workspace_id: path.workspace_id.clone(),
grant_id: new_id("wcg"),
controller,
subject: response.0.worker_ref.clone(),
relation: relation.to_string(),
origin: "worker_spawn".to_string(),
permissions: vec![
"send_input".to_string(),
"notify".to_string(),
"cancel".to_string(),
"stop".to_string(),
"restore".to_string(),
"remove".to_string(),
"observe".to_string(),
],
operation_id,
created_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true),
revoked_at: None,
})
{
let subject = response.0.worker_ref.clone();
let _ = api.runtime.delete_worker(&subject);
let _ = api
.store
.delete_worker_registry(&path.workspace_id, &subject);
return Err(ApiError::from(error));
}
Ok(response)
}
fn authorize_known_worker_permission(
api: &WorkspaceApi,
workspace_id: &str,
controller: &RuntimeWorkerRef,
subject: &RuntimeWorkerRef,
permission: &str,
) -> Result<()> {
let grant = api
.store
.get_active_worker_control_grant(workspace_id, controller, subject)?
.ok_or_else(|| Error::UnknownWorker {
worker: subject.clone(),
})?;
if !grant
.permissions
.iter()
.any(|candidate| candidate == permission)
{
return Err(Error::InvalidInput(format!(
"worker control permission `{permission}` was not granted"
)));
}
Ok(())
}
async fn send_known_worker_input(
State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedRuntimeWorkerPath>,
headers: HeaderMap,
Json(request): Json<WorkerInputRequest>,
) -> ApiResult<Json<WorkerInputResult>> {
validate_workspace_scope(&api, &path.workspace_id)?;
let source = authenticate_worker_mutation_source(&api, &path.workspace_id, &headers)?;
let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id);
let permission = match request.kind {
WorkerInputKind::Notify => "notify",
_ => "send_input",
};
authorize_known_worker_permission(
&api,
&path.workspace_id,
&controller,
&path.worker,
permission,
)?;
scoped_send_runtime_worker_input(State(api), AxumPath(path), Json(request)).await
}
async fn cancel_known_worker(
State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedRuntimeWorkerPath>,
headers: HeaderMap,
Json(request): Json<WorkerLifecycleRequest>,
) -> ApiResult<Json<WorkerLifecycleResult>> {
validate_workspace_scope(&api, &path.workspace_id)?;
let source = authenticate_worker_mutation_source(&api, &path.workspace_id, &headers)?;
let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id);
authorize_known_worker_permission(
&api,
&path.workspace_id,
&controller,
&path.worker,
"cancel",
)?;
scoped_cancel_runtime_worker(State(api), AxumPath(path), Json(request)).await
}
async fn stop_known_worker(
State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedRuntimeWorkerPath>,
headers: HeaderMap,
Json(request): Json<WorkerLifecycleRequest>,
) -> ApiResult<Json<WorkerLifecycleResult>> {
validate_workspace_scope(&api, &path.workspace_id)?;
let source = authenticate_worker_mutation_source(&api, &path.workspace_id, &headers)?;
let subject = path.worker.clone();
let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id);
authorize_known_worker_permission(&api, &path.workspace_id, &controller, &subject, "stop")?;
scoped_stop_runtime_worker(State(api), AxumPath(path), Json(request)).await
}
async fn restore_known_worker(
State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedRuntimeWorkerPath>,
headers: HeaderMap,
) -> ApiResult<Json<WorkerRestoreResponse>> {
validate_workspace_scope(&api, &path.workspace_id)?;
let source = authenticate_worker_mutation_source(&api, &path.workspace_id, &headers)?;
let subject = path.worker.clone();
let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id);
authorize_known_worker_permission(&api, &path.workspace_id, &controller, &subject, "restore")?;
scoped_restore_runtime_worker(State(api), AxumPath(path), Query(Default::default())).await
}
async fn scoped_list_worker_observation_sessions(
State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedWorkspacePath>,
@@ -5839,26 +6089,34 @@ async fn scoped_list_worker_observation_sessions(
) -> ApiResult<Json<serde_json::Value>> {
validate_workspace_scope(&api, &path.workspace_id)?;
let source = authenticate_worker_mutation_source(&api, &path.workspace_id, &headers)?;
authorize_workspace_orchestrator_observation(&api, &source)?;
let sessions = workers_response(api.clone())?
.items
let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id);
let sessions = api
.store
.list_active_worker_control_grants(&path.workspace_id, &controller, 100)?
.into_iter()
.filter(|worker| {
!matches!(
.filter(|grant| {
grant
.permissions
.iter()
.any(|permission| permission == "observe")
})
.filter_map(|grant| {
let worker = api.runtime.worker(&grant.subject).ok()?;
if matches!(
worker.state.as_str(),
"stopped" | "failed" | "rejected" | "disconnected"
) && (worker.worker.runtime_id != source.runtime_id
|| worker.worker.worker_id != source.worker_id)
})
.take(100)
.map(|worker| WorkerObservationSubject {
subject: WorkerObservationSubjectRef::RuntimeWorker {
runtime_id: worker.worker.runtime_id,
worker_id: worker.worker.worker_id,
},
display_name: worker.display_name,
relation: "workspace_orchestrator_grant".to_string(),
status: worker.state,
) {
return None;
}
Some(WorkerObservationSubject {
subject: WorkerObservationSubjectRef::RuntimeWorker {
runtime_id: grant.subject.runtime_id,
worker_id: grant.subject.worker_id,
},
display_name: worker.display_name,
relation: grant.relation,
status: worker.state,
})
})
.collect::<Vec<_>>();
Ok(Json(serde_json::json!({ "sessions": sessions })))
@@ -5872,7 +6130,6 @@ async fn scoped_capture_worker_observation_session(
) -> ApiResult<Json<serde_json::Value>> {
validate_workspace_scope(&api, &path.workspace_id)?;
let source = authenticate_worker_mutation_source(&api, &path.workspace_id, &headers)?;
authorize_workspace_orchestrator_observation(&api, &source)?;
let WorkerObservationSubjectRef::RuntimeWorker {
runtime_id,
worker_id,
@@ -5883,17 +6140,16 @@ async fn scoped_capture_worker_observation_session(
}));
};
let target = RuntimeWorkerRef::new(runtime_id, worker_id);
let granted = workers_response(api.clone())?
.items
.into_iter()
.any(|worker| {
worker.worker == target
&& !matches!(
worker.state.as_str(),
"stopped" | "failed" | "rejected" | "disconnected"
)
});
if !granted {
let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id);
authorize_known_worker_permission(&api, &path.workspace_id, &controller, &target, "observe")?;
let target_summary = api
.runtime
.worker(&target)
.map_err(|error| error.into_error())?;
if matches!(
target_summary.state.as_str(),
"stopped" | "failed" | "rejected" | "disconnected"
) {
return Err(ApiError::from(Error::UnknownWorker { worker: target }));
}
@@ -5927,25 +6183,6 @@ async fn scoped_capture_worker_observation_session(
})))
}
fn authorize_workspace_orchestrator_observation(
api: &WorkspaceApi,
source: &WorkerMutationSource,
) -> ApiResult<()> {
let Some(orchestrator) = find_workspace_orchestrator(api) else {
return Err(ApiError::from(Error::UnknownWorker {
worker: RuntimeWorkerRef::new(source.runtime_id.clone(), source.worker_id.clone()),
}));
};
if orchestrator.worker.runtime_id != source.runtime_id
|| orchestrator.worker.worker_id != source.worker_id
{
return Err(ApiError::from(Error::UnknownWorker {
worker: RuntimeWorkerRef::new(source.runtime_id.clone(), source.worker_id.clone()),
}));
}
Ok(())
}
async fn scoped_start_workspace_orchestrator(
State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedWorkspacePath>,
@@ -6028,6 +6265,14 @@ async fn scoped_start_workspace_orchestrator(
result.diagnostics,
));
}
let worker = result.worker.as_ref().expect("accepted Worker was checked");
record_worker_summary(
&api,
worker,
&worker.display_name,
Some("builtin:orchestrator".to_string()),
WorkerRegistryDisplayNamePolicy::UseProvided,
)?;
*api.orchestrator_attention_fingerprint
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner()) = None;
@@ -8446,6 +8691,7 @@ async fn create_workspace_worker(
ticket_assignment,
initial_submit,
working_directory,
control_operation_id: _,
} = request;
let config_state = api
.config_store
@@ -12273,6 +12519,7 @@ mod tests {
selector: "builtin:coder-review".to_string(),
}],
working_directory: None,
control_operation_id: None,
}),
)
.await
@@ -12312,6 +12559,7 @@ mod tests {
ticket_assignment: None,
initial_submit: Vec::new(),
working_directory: None,
control_operation_id: None,
}),
)
.await
@@ -12346,6 +12594,7 @@ mod tests {
ticket_assignment: None,
initial_submit: Vec::new(),
working_directory: None,
control_operation_id: None,
}),
)
.await
@@ -12447,6 +12696,7 @@ mod tests {
selector: "builtin:coder-review".to_string(),
}],
working_directory: None,
control_operation_id: None,
}),
)
.await
@@ -12745,6 +12995,7 @@ mod tests {
ticket_assignment: None,
initial_submit: Vec::new(),
working_directory: None,
control_operation_id: None,
}),
)
.await
@@ -12761,6 +13012,7 @@ mod tests {
ticket_assignment: None,
initial_submit: Vec::new(),
working_directory: None,
control_operation_id: None,
}),
)
.await
@@ -12783,6 +13035,21 @@ mod tests {
);
assert_ne!(dedicated.worker.worker_id, generic.worker_ref.worker_id);
api.store
.create_worker_control_grant(&WorkerControlGrantRecord {
workspace_id: workspace_id.clone(),
grant_id: "orchestrator-controls-generic".to_string(),
controller: dedicated.worker.clone(),
subject: generic.worker_ref.clone(),
relation: "spawned".to_string(),
origin: "test".to_string(),
permissions: vec!["observe".to_string()],
operation_id: "observe-generic".to_string(),
created_at: "2026-07-27T00:00:00Z".to_string(),
revoked_at: None,
})
.unwrap();
let mut observation_headers = HeaderMap::new();
observation_headers.insert(
"x-yoi-runtime-id",
@@ -12792,6 +13059,19 @@ mod tests {
"x-yoi-worker-id",
axum::http::HeaderValue::from_str(&dedicated.worker.worker_id).unwrap(),
);
let Json(known) = list_known_workers(
State(api.clone()),
AxumPath(ScopedWorkspacePath {
workspace_id: workspace_id.clone(),
}),
observation_headers.clone(),
)
.await
.unwrap();
assert_eq!(known.items.len(), 1);
assert_eq!(known.items[0].subject, generic.worker_ref);
assert_eq!(known.items[0].permissions, ["observe"]);
let Json(sessions) = scoped_list_worker_observation_sessions(
State(api.clone()),
AxumPath(ScopedWorkspacePath {
@@ -12836,7 +13116,7 @@ mod tests {
"x-yoi-worker-id",
axum::http::HeaderValue::from_str(&generic.worker_ref.worker_id).unwrap(),
);
let error = scoped_list_worker_observation_sessions(
let Json(unauthorized) = scoped_list_worker_observation_sessions(
State(api.clone()),
AxumPath(ScopedWorkspacePath {
workspace_id: workspace_id.clone(),
@@ -12844,8 +13124,8 @@ mod tests {
unauthorized_headers,
)
.await
.unwrap_err();
assert_eq!(error.into_response().status(), StatusCode::NOT_FOUND);
.unwrap();
assert!(unauthorized["sessions"].as_array().unwrap().is_empty());
let Json(existing) = scoped_start_workspace_orchestrator(
State(api.clone()),
@@ -14883,6 +15163,9 @@ mod tests {
)
.unwrap();
let target = spawned.worker.unwrap().worker;
let target_summary = api.runtime.worker(&target).unwrap();
sync_worker_observation(&api, &target_summary).unwrap();
seed_worker_control_grant(&api, &source, &target, "caller-guard-target");
let running_response = executor
.execute_async(
verified_source(),
@@ -14992,6 +15275,7 @@ mod tests {
.unwrap();
let summary = api.runtime.worker(&target).unwrap();
let record = sync_worker_observation(&api, &summary).unwrap();
seed_worker_control_grant(&api, &source, &target, "embedded-valid-proof");
let response = WorkspaceWorkerRemoveExecutor::new(&api)
.execute_async(
@@ -15208,12 +15492,12 @@ mod tests {
)
.await
.unwrap();
assert_eq!(route_response.status(), StatusCode::FORBIDDEN);
assert_eq!(route_response.status(), StatusCode::NOT_FOUND);
let route_body = axum::body::to_bytes(route_response.into_body(), usize::MAX)
.await
.unwrap();
let route_body = String::from_utf8(route_body.to_vec()).unwrap();
assert!(route_body.contains("orchestrator_required"));
assert!(route_body.contains("unknown_worker"));
assert!(!route_body.contains("source"));
assert!(!route_body.contains("proof"));
@@ -15262,6 +15546,28 @@ mod tests {
.unwrap();
}
fn seed_worker_control_grant(
api: &WorkspaceApi,
controller: &RuntimeWorkerRef,
subject: &RuntimeWorkerRef,
operation_id: &str,
) {
api.store
.create_worker_control_grant(&WorkerControlGrantRecord {
workspace_id: api.config.workspace_id.clone(),
grant_id: format!("grant-{operation_id}"),
controller: controller.clone(),
subject: subject.clone(),
relation: "spawned".to_string(),
origin: "test".to_string(),
permissions: vec!["remove".to_string()],
operation_id: operation_id.to_string(),
created_at: now_registry_timestamp(),
revoked_at: None,
})
.unwrap();
}
fn seed_cleanup_worker(
api: &WorkspaceApi,
runtime_worker_id: u64,