From f41ab0e277658135d11c9923bd4ecef0e9e008b1 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 16 Aug 2026 23:08:32 +0900 Subject: [PATCH 1/4] worker: persist known-worker control grants --- crates/worker/src/controller.rs | 4 +- .../src/feature/builtin/manage_worker.rs | 265 +++++++++- .../src/feature/builtin/worker_observation.rs | 146 ++---- crates/workspace-server/src/server.rs | 452 +++++++++++++++--- crates/workspace-server/src/store.rs | 423 +++++++++++++++- resources/profiles/base.dcdl | 2 +- .../prompts/common/worker-observation.md | 2 +- 7 files changed, 1082 insertions(+), 212 deletions(-) diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 73a96bd6..5e696c3b 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -21,7 +21,7 @@ use crate::shutdown_after_idle::{ ShutdownAfterIdleRequest, TicketIntakeReadyShutdownHook, is_ticket_intake_role, take_shutdown_request_after_status, }; -use crate::spawn::comm_tools::{sub_worker_list_tool, sub_worker_send_tool, sub_worker_stop_tool}; +use crate::spawn::comm_tools::{sub_worker_send_tool, sub_worker_stop_tool}; use crate::spawn::registry::SpawnedWorkerRegistry; use crate::spawn::tool::sub_worker_spawn_tool; use crate::worker::{ @@ -802,6 +802,7 @@ where feature_registry.add_module( crate::feature::builtin::manage_worker::manage_worker_feature( workspace_client, + Some(spawned_registry.clone()), feature_config.worker.direct_spawn, ), ); @@ -920,7 +921,6 @@ where scope_handle, prompts, )); - engine.register_tool(sub_worker_list_tool(spawned_registry.clone())); engine.register_tool(sub_worker_send_tool(spawned_registry.clone())); engine.register_tool(sub_worker_stop_tool(spawned_registry.clone())); observation_providers.push(Arc::new( diff --git a/crates/worker/src/feature/builtin/manage_worker.rs b/crates/worker/src/feature/builtin/manage_worker.rs index 7d6a6b2c..73aff8bc 100644 --- a/crates/worker/src/feature/builtin/manage_worker.rs +++ b/crates/worker/src/feature/builtin/manage_worker.rs @@ -8,6 +8,7 @@ use llm_engine::tool::{ }; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; +use uuid::Uuid; use protocol::Segment; @@ -15,6 +16,7 @@ use crate::feature::{ FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureModule, ServiceDeclaration, ServiceId, ToolContribution, ToolDeclaration, }; +use crate::spawn::registry::SpawnedWorkerRegistry; use crate::worker::{ WorkspaceClient, WorkspaceClientError, WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse, @@ -25,8 +27,16 @@ const FEATURE_NAME: &str = "Worker"; const FEATURE_DESCRIPTION: &str = "Workspace-authority tools for managing Workdir-bound Backend/Runtime Worker sessions."; pub const WORKER_LIFECYCLE_SERVICE_ID: &str = "worker.lifecycle"; +pub const WORKER_CONTROL_SERVICE_ID: &str = "worker.control"; const WORKER_LIFECYCLE_SERVICE_VERSION: &str = "1"; +pub trait WorkerControlService: Send + Sync {} + +#[derive(Debug)] +struct WorkspaceWorkerControlService; + +impl WorkerControlService for WorkspaceWorkerControlService {} + #[async_trait] pub trait WorkerLifecycleService: Send + Sync { async fn spawn( @@ -70,10 +80,15 @@ impl WorkerLifecycleService for WorkspaceWorkerLifecycleService { )); } }; + let control_operation_id = ticket_assignment + .as_ref() + .map(|assignment| assignment.operation_id.clone()) + .unwrap_or_else(|| format!("worker-spawn-{}", Uuid::now_v7())); let body = WorkerSpawnRequest { runtime_id: request.runtime_id, display_name: request.display_name, profile: request.profile, + control_operation_id, ticket_assignment, initial_submit: request.initial_submit, working_directory: WorkerWorkingDirectorySelection { @@ -83,25 +98,38 @@ impl WorkerLifecycleService for WorkspaceWorkerLifecycleService { }; self.client.execute(WorkspaceRequest::json( WorkspaceRequestMethod::Post, - format!("/api/w/{}/workers", self.workspace_id), + format!("/api/w/{}/worker-control/workers", self.workspace_id), serde_json::to_string(&body) .map_err(|error| WorkspaceClientError::Request(error.to_string()))?, )) } } -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct ManageWorkerFeature { client: Arc, + registry: Option>, direct_spawn: bool, } +impl std::fmt::Debug for ManageWorkerFeature { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ManageWorkerFeature") + .field("has_registry", &self.registry.is_some()) + .field("direct_spawn", &self.direct_spawn) + .finish_non_exhaustive() + } +} + pub fn manage_worker_feature( client: Arc, + registry: Option>, direct_spawn: bool, ) -> ManageWorkerFeature { ManageWorkerFeature { client, + registry, direct_spawn, } } @@ -114,6 +142,11 @@ impl FeatureModule for ManageWorkerFeature { ServiceId::builtin(WORKER_LIFECYCLE_SERVICE_ID), WORKER_LIFECYCLE_SERVICE_VERSION, "Workspace-authoritative Worker lifecycle operations", + )) + .with_provided_service(ServiceDeclaration::new( + ServiceId::builtin(WORKER_CONTROL_SERVICE_ID), + WORKER_LIFECYCLE_SERVICE_VERSION, + "Known-Worker discovery and permission-fenced control operations", )); for operation in WorkerOperation::ALL { if operation != WorkerOperation::Spawn || self.direct_spawn { @@ -150,6 +183,14 @@ impl FeatureModule for ManageWorkerFeature { ), lifecycle, )?; + context.services().provide( + ServiceDeclaration::new( + ServiceId::builtin(WORKER_CONTROL_SERVICE_ID), + WORKER_LIFECYCLE_SERVICE_VERSION, + "Known-Worker discovery and permission-fenced control operations", + ), + Arc::new(WorkspaceWorkerControlService) as Arc, + )?; for operation in WorkerOperation::ALL { if operation == WorkerOperation::Spawn && !self.direct_spawn { continue; @@ -159,26 +200,45 @@ impl FeatureModule for ManageWorkerFeature { operation, self.client.clone(), workspace_id.clone(), + self.registry.clone(), ), WorkerOperation::Spawn => definition::( operation, self.client.clone(), workspace_id.clone(), + self.registry.clone(), + ), + WorkerOperation::SendInput | WorkerOperation::Notify => { + definition::( + operation, + self.client.clone(), + workspace_id.clone(), + self.registry.clone(), + ) + } + WorkerOperation::Cancel => definition::( + operation, + self.client.clone(), + workspace_id.clone(), + self.registry.clone(), ), WorkerOperation::Stop => definition::( operation, self.client.clone(), workspace_id.clone(), + self.registry.clone(), ), WorkerOperation::Restore => definition::( operation, self.client.clone(), workspace_id.clone(), + self.registry.clone(), ), WorkerOperation::Remove => definition::( operation, self.client.clone(), workspace_id.clone(), + self.registry.clone(), ), }; context @@ -224,6 +284,7 @@ struct WorkerSpawnRequest { runtime_id: String, display_name: String, profile: String, + control_operation_id: String, #[serde(skip_serializing_if = "Option::is_none")] ticket_assignment: Option, initial_submit: Vec, @@ -244,6 +305,14 @@ struct WorkerTargetInput { worker_id: String, } +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +struct WorkerMessageInput { + runtime_id: String, + worker_id: String, + content: String, +} + #[derive(Debug, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] struct WorkerStopInput { @@ -266,21 +335,28 @@ struct WorkspaceWorkerTool { operation: WorkerOperation, client: Arc, workspace_id: String, + registry: Option>, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum WorkerOperation { List, Spawn, + SendInput, + Notify, + Cancel, Stop, Restore, Remove, } impl WorkerOperation { - const ALL: [Self; 5] = [ + const ALL: [Self; 8] = [ Self::List, Self::Spawn, + Self::SendInput, + Self::Notify, + Self::Cancel, Self::Stop, Self::Restore, Self::Remove, @@ -290,6 +366,9 @@ impl WorkerOperation { match self { Self::List => "WorkerList", Self::Spawn => "WorkerSpawn", + Self::SendInput => "WorkerSendInput", + Self::Notify => "WorkerNotify", + Self::Cancel => "WorkerCancel", Self::Stop => "WorkerStop", Self::Restore => "WorkerRestore", Self::Remove => "WorkerRemove", @@ -299,12 +378,15 @@ impl WorkerOperation { fn description(self) -> &'static str { match self { Self::List => { - "List Backend/Runtime Worker sessions in the current Workspace. SubWorkers are excluded." + "List only known Runtime Workers and direct SubWorkers granted to the current Worker." } Self::Spawn => { "Spawn a Backend/Runtime Worker session in an existing Workspace Workdir. The Workdir id is authority; filesystem paths and Runtime URLs are not accepted. `initial_submit` carries the normal typed user submission. After the Orchestrator has committed a Ticket to `inprogress`, set `ticket_id` with a Flow segment in `initial_submit` to atomically assign the new Coder Worker; the operation id is derived from the durable tool call rather than model input." } - Self::Stop => "Stop a Backend/Runtime Worker session in the current Workspace.", + Self::SendInput => "Send user input to a known Runtime Worker when allowed.", + Self::Notify => "Send an advisory notification to a known Runtime Worker when allowed.", + Self::Cancel => "Cancel the current turn of a known Runtime Worker when allowed.", + Self::Stop => "Stop a known Runtime Worker when allowed.", Self::Restore => { "Restore a stopped Backend/Runtime Worker session in the current Workspace." } @@ -348,7 +430,48 @@ impl Tool for WorkspaceWorkerTool { let request = match operation { WorkerOperation::List => { parse::(input_json, "WorkerList")?; - WorkspaceRequest::get(format!("/api/w/{}/workers", self.workspace_id)) + WorkspaceRequest::get(format!( + "/api/w/{}/worker-control/workers", + self.workspace_id + )) + } + WorkerOperation::SendInput | WorkerOperation::Notify => { + let tool_name = operation.tool_name(); + let input = parse::(input_json, tool_name)?; + let runtime_id = authority_id(&input.runtime_id, "runtime_id")?; + let worker_id = authority_id(&input.worker_id, "worker_id")?; + let content = non_empty(input.content, "content")?; + if content.len() > 16 * 1024 { + return Err(ToolError::ExecutionFailed( + "content must contain at most 16384 bytes".to_string(), + )); + } + WorkspaceRequest::json( + WorkspaceRequestMethod::Post, + format!( + "/api/w/{}/worker-control/workers/{runtime_id}/{worker_id}/input", + self.workspace_id + ), + serde_json::json!({ + "kind": if operation == WorkerOperation::Notify { "notify" } else { "user" }, + "content": content, + }) + .to_string(), + ) + } + WorkerOperation::Cancel => { + let input = parse::(input_json, "WorkerCancel")?; + let runtime_id = authority_id(&input.runtime_id, "runtime_id")?; + let worker_id = authority_id(&input.worker_id, "worker_id")?; + let reason = input.reason.unwrap_or_default(); + WorkspaceRequest::json( + WorkspaceRequestMethod::Post, + format!( + "/api/w/{}/worker-control/workers/{runtime_id}/{worker_id}/cancel", + self.workspace_id + ), + serde_json::json!({ "reason": reason }).to_string(), + ) } WorkerOperation::Spawn => { let input = parse::(input_json, "WorkerSpawn")?; @@ -398,7 +521,7 @@ impl Tool for WorkspaceWorkerTool { WorkspaceRequest::json( WorkspaceRequestMethod::Post, format!( - "/api/w/{}/runtimes/{runtime_id}/workers/{worker_id}/stop", + "/api/w/{}/worker-control/workers/{runtime_id}/{worker_id}/stop", self.workspace_id ), serde_json::json!({ "reason": input.reason }).to_string(), @@ -411,7 +534,7 @@ impl Tool for WorkspaceWorkerTool { WorkspaceRequest::json( WorkspaceRequestMethod::Post, format!( - "/api/w/{}/runtimes/{runtime_id}/workers/{worker_id}/restore", + "/api/w/{}/worker-control/workers/{runtime_id}/{worker_id}/restore", self.workspace_id ), "{}", @@ -424,10 +547,58 @@ impl Tool for WorkspaceWorkerTool { .map_err(|error| ToolError::ExecutionFailed(error.to_string()))? } }; + let response = if self.operation == WorkerOperation::List { + self.with_subworkers(response).await? + } else { + response + }; tool_output(self.operation, response) } } +impl WorkspaceWorkerTool { + async fn with_subworkers( + &self, + mut response: WorkspaceResponse, + ) -> Result { + let Some(registry) = &self.registry else { + return Ok(response); + }; + if !response.is_success() { + return Ok(response); + } + let mut body: serde_json::Value = + serde_json::from_str(&response.body).map_err(|error| { + ToolError::ExecutionFailed(format!("WorkerList returned invalid JSON: {error}")) + })?; + let items = body + .get_mut("items") + .and_then(serde_json::Value::as_array_mut) + .ok_or_else(|| { + ToolError::ExecutionFailed( + "WorkerList response did not contain an items array".to_string(), + ) + })?; + for internal in registry.list_internal() { + let child_name = internal.worker_name.clone(); + items.push(serde_json::json!({ + "subject": { "kind": "sub_worker", "name": child_name }, + "relation": "direct_child", + "origin": "sub_worker_spawn", + "permissions": ["send_input", "stop", "observe"], + "summary": { + "display_name": internal.worker_name, + "status": format!("{:?}", internal.session.status()).to_lowercase(), + } + })); + } + response.body = serde_json::to_string(&body).map_err(|error| { + ToolError::ExecutionFailed(format!("WorkerList could not encode its response: {error}")) + })?; + Ok(response) + } +} + fn tool_output( operation: WorkerOperation, response: WorkspaceResponse, @@ -449,6 +620,7 @@ fn definition( operation: WorkerOperation, client: Arc, workspace_id: String, + registry: Option>, ) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(I); @@ -460,6 +632,7 @@ fn definition( operation, client: client.clone(), workspace_id: workspace_id.clone(), + registry: registry.clone(), }); (meta, tool) }) @@ -567,6 +740,7 @@ mod tests { operation: WorkerOperation::Spawn, client: client.clone(), workspace_id: "workspace%2Ftest".to_string(), + registry: None, }; tool.execute( &serde_json::json!({ @@ -587,7 +761,10 @@ mod tests { let requests = client.requests.lock().unwrap(); assert_eq!(requests.len(), 1); - assert_eq!(requests[0].path, "/api/w/workspace%2Ftest/workers"); + assert_eq!( + requests[0].path, + "/api/w/workspace%2Ftest/worker-control/workers" + ); let body: serde_json::Value = serde_json::from_str(requests[0].body.as_deref().unwrap()).unwrap(); assert_eq!(body["initial_submit"][0]["kind"], "flow"); @@ -609,7 +786,7 @@ mod tests { #[test] fn worker_service_can_remain_enabled_without_direct_spawn_surface() { let client = Arc::new(RecordingWorkspaceClient::default()); - let descriptor = manage_worker_feature(client, false).descriptor(); + let descriptor = manage_worker_feature(client, None, false).descriptor(); let tools: Vec<_> = descriptor .tools .iter() @@ -618,8 +795,8 @@ mod tests { assert!(!tools.contains(&"WorkerSpawn")); assert!(tools.contains(&"WorkerList")); assert_eq!( - descriptor.provides_services[0].id, - ServiceId::builtin(WORKER_LIFECYCLE_SERVICE_ID) + descriptor.provides_services[1].id, + ServiceId::builtin(WORKER_CONTROL_SERVICE_ID) ); } @@ -630,6 +807,9 @@ mod tests { [ "WorkerList", "WorkerSpawn", + "WorkerSendInput", + "WorkerNotify", + "WorkerCancel", "WorkerStop", "WorkerRestore", "WorkerRemove", @@ -654,6 +834,7 @@ mod tests { runtime_id: "runtime-1".to_string(), display_name: "Coder".to_string(), profile: "builtin:coder".to_string(), + control_operation_id: "spawn-operation-1".to_string(), ticket_assignment: None, initial_submit: vec![ Segment::Flow { @@ -681,6 +862,64 @@ mod tests { assert!(value.get("initial_text").is_none()); } + #[tokio::test] + async fn worker_message_and_cancel_use_permission_fenced_control_routes() { + let client = Arc::new(RecordingWorkspaceClient::default()); + for (operation, args, expected_suffix, expected_kind) in [ + ( + WorkerOperation::SendInput, + serde_json::json!({ + "runtime_id": "runtime-1", + "worker_id": "worker-7", + "content": "continue", + }), + "/input", + Some("user"), + ), + ( + WorkerOperation::Notify, + serde_json::json!({ + "runtime_id": "runtime-1", + "worker_id": "worker-7", + "content": "review ready", + }), + "/input", + Some("notify"), + ), + ( + WorkerOperation::Cancel, + serde_json::json!({ + "runtime_id": "runtime-1", + "worker_id": "worker-7", + "reason": "superseded", + }), + "/cancel", + None, + ), + ] { + WorkspaceWorkerTool { + operation, + client: client.clone(), + workspace_id: "workspace%2Ftest".to_string(), + registry: None, + } + .execute( + &args.to_string(), + ToolExecutionContext::new("call-control", "batch-control", 0), + ) + .await + .unwrap(); + let request = client.requests.lock().unwrap().last().cloned().unwrap(); + assert!(request.path.ends_with(expected_suffix)); + assert!(request.path.contains("/worker-control/workers/")); + if let Some(expected_kind) = expected_kind { + let body: serde_json::Value = + serde_json::from_str(request.body.as_deref().unwrap()).unwrap(); + assert_eq!(body["kind"], expected_kind); + } + } + } + #[tokio::test] async fn worker_remove_forwards_only_target_revision_and_bounded_reason() { let client = Arc::new(RecordingWorkspaceClient::default()); @@ -688,6 +927,7 @@ mod tests { operation: WorkerOperation::Remove, client: client.clone(), workspace_id: "workspace%2Ftest".to_string(), + registry: None, }; tool.execute( &serde_json::json!({ @@ -734,6 +974,7 @@ mod tests { operation: WorkerOperation::Remove, client: client.clone(), workspace_id: "workspace%2Ftest".to_string(), + registry: None, }; for reason in [" ".to_string(), "x".repeat(513)] { let _error = tool diff --git a/crates/worker/src/feature/builtin/worker_observation.rs b/crates/worker/src/feature/builtin/worker_observation.rs index 376c31ce..cae10728 100644 --- a/crates/worker/src/feature/builtin/worker_observation.rs +++ b/crates/worker/src/feature/builtin/worker_observation.rs @@ -7,10 +7,11 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use session_store::collect_state; +use super::manage_worker::WORKER_CONTROL_SERVICE_ID; use crate::feature::{ FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureInstructionContribution, - FeatureInstructionDeclaration, FeatureInstructionId, FeatureModule, ToolContribution, - ToolDeclaration, + FeatureInstructionDeclaration, FeatureInstructionId, FeatureModule, ServiceId, + ServiceRequirement, ToolContribution, ToolDeclaration, }; use crate::session_capture::{ ReadDetail, ReadOptions, ReadSelector, ReferenceKind, SearchOptions, SessionCapture, @@ -111,11 +112,17 @@ impl WorkerObservationProvider for WorkspaceClientWorkerObservationProvider { async fn list_worker_sessions( &self, ) -> Result, WorkerObservationError> { + let workspace_id = self.client.workspace_id().ok_or_else(|| { + WorkerObservationError::Unavailable( + "Workspace observation requires a scoped Workspace client".to_string(), + ) + })?; let response = self .client - .execute(crate::worker::WorkspaceRequest::get( - "/worker-observation/sessions", - )) + .execute(crate::worker::WorkspaceRequest::get(format!( + "/api/w/{}/worker-observation/sessions", + workspace_id + ))) .map_err(workspace_client_error)?; let body = workspace_response_body(response)?; serde_json::from_str::(&body) @@ -129,11 +136,16 @@ impl WorkerObservationProvider for WorkspaceClientWorkerObservationProvider { ) -> Result { let body = serde_json::to_string(subject) .map_err(|error| WorkerObservationError::Unavailable(error.to_string()))?; + let workspace_id = self.client.workspace_id().ok_or_else(|| { + WorkerObservationError::Unavailable( + "Workspace observation requires a scoped Workspace client".to_string(), + ) + })?; let response = self .client .execute(crate::worker::WorkspaceRequest::json( crate::worker::WorkspaceRequestMethod::Post, - "/worker-observation/session", + format!("/api/w/{}/worker-observation/session", workspace_id), body, )) .map_err(workspace_client_error)?; @@ -190,11 +202,11 @@ impl FeatureModule for WorkerObservationFeature { .with_description( "Read-only exploration of explicitly granted active Worker sessions.", ) - .with_instruction(observation_instruction()) - .with_tool(ToolDeclaration::new( - "ListWorkerSessions", - "List bounded summaries of active Worker sessions granted to this Worker.", + .with_service_requirement(ServiceRequirement::required( + ServiceId::builtin(WORKER_CONTROL_SERVICE_ID), + "Worker observation extends the known-Worker control authority", )) + .with_instruction(observation_instruction()) .with_tool(ToolDeclaration::new( "ViewSessionOverview", "Show a sparse overview of the latest committed capture for one granted Worker session.", @@ -215,10 +227,6 @@ impl FeatureModule for WorkerObservationFeature { .register(FeatureInstructionContribution::new( observation_instruction(), ))?; - context.tools().register(ToolContribution::new( - "ListWorkerSessions", - list_definition(self.provider.clone()), - ))?; context.tools().register(ToolContribution::new( "ViewSessionOverview", overview_definition(self.provider.clone()), @@ -346,20 +354,6 @@ impl WorkerObservationProvider for SpawnedSubWorkerObservationProvider { } } -fn list_definition(provider: Arc) -> ToolDefinition { - Arc::new(move || { - let schema = serde_json::to_value(schemars::schema_for!(ListWorkerSessionsParams)) - .unwrap_or_else(|_| serde_json::json!({})); - let meta = ToolMeta::new("ListWorkerSessions") - .description("List active Worker sessions explicitly granted to this Worker.") - .input_schema(schema); - let tool: Arc = Arc::new(ListWorkerSessionsTool { - provider: provider.clone(), - }); - (meta, tool) - }) -} - fn overview_definition(provider: Arc) -> ToolDefinition { Arc::new(move || { let schema = serde_json::to_value(schemars::schema_for!(ViewSessionOverviewParams)) @@ -402,13 +396,6 @@ fn read_definition(provider: Arc) -> ToolDefiniti }) } -#[derive(Debug, Deserialize, JsonSchema)] -#[serde(deny_unknown_fields)] -struct ListWorkerSessionsParams { - #[serde(default)] - limit: Option, -} - #[derive(Debug, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] struct ViewSessionOverviewParams { @@ -454,43 +441,6 @@ fn default_read_mode() -> String { "compact".to_string() } -struct ListWorkerSessionsTool { - provider: Arc, -} - -#[async_trait] -impl Tool for ListWorkerSessionsTool { - async fn execute( - &self, - input_json: &str, - _context: llm_engine::tool::ToolExecutionContext, - ) -> Result { - let params: ListWorkerSessionsParams = parse_input("ListWorkerSessions", input_json)?; - let limit = bounded_limit(params.limit); - let mut subjects = self - .provider - .list_worker_sessions() - .await - .map_err(tool_error)?; - subjects.truncate(limit); - let sessions = subjects - .iter() - .map(|subject| { - serde_json::json!({ - "subject": bounded_subject(&subject.subject), - "display_name": truncate_text(&subject.display_name, 200), - "relation": truncate_text(&subject.relation, 64), - "status": truncate_text(&subject.status, 64), - }) - }) - .collect::>(); - json_output( - format!("Listed {} Worker session(s).", sessions.len()), - serde_json::json!({ "sessions": sessions }), - ) - } -} - struct ViewSessionOverviewTool { provider: Arc, } @@ -692,31 +642,6 @@ fn parse_tool_part(value: &str) -> Result { .ok_or_else(|| ToolError::InvalidArgument(format!("invalid tool_part {value:?}"))) } -fn bounded_subject(subject: &WorkerObservationSubjectRef) -> WorkerObservationSubjectRef { - match subject { - WorkerObservationSubjectRef::RuntimeWorker { - runtime_id, - worker_id, - } => WorkerObservationSubjectRef::RuntimeWorker { - runtime_id: truncate_text(runtime_id, 200), - worker_id: truncate_text(worker_id, 200), - }, - WorkerObservationSubjectRef::SubWorker { name } => WorkerObservationSubjectRef::SubWorker { - name: truncate_text(name, 200), - }, - } -} - -fn truncate_text(value: &str, max_chars: usize) -> String { - if value.chars().count() <= max_chars { - value.to_string() - } else { - let mut truncated = value.chars().take(max_chars).collect::(); - truncated.push('…'); - truncated - } -} - fn bounded_limit(limit: Option) -> usize { limit.unwrap_or(DEFAULT_PAGE_LIMIT).clamp(1, MAX_PAGE_LIMIT) } @@ -801,7 +726,7 @@ mod tests { let catalog = crate::PromptCatalog::builtins_only().unwrap(); let source = &catalog.projection().templates["common.worker_observation"]; for token in [ - "ListWorkerSessions", + "WorkerList", "ViewSessionOverview", "SearchSessionEntries", "ReadSessionEntry", @@ -812,7 +737,7 @@ mod tests { } #[test] - fn worker_observation_installs_without_session_explore_or_memory_extract() { + fn worker_observation_requires_worker_control_service() { let provider = Arc::new(FakeProvider { captures: Mutex::new(Vec::new()), }); @@ -821,15 +746,15 @@ mod tests { let report = FeatureRegistryBuilder::new() .with_module(WorkerObservationFeature::new(provider)) .install_into_pending(&mut pending_tools, &mut hook_builder); - assert!(report.reports[0].installed); + assert!(!report.reports[0].installed); + assert!(report.installed_tool_names().is_empty()); + let descriptor = WorkerObservationFeature::new(Arc::new(FakeProvider { + captures: Mutex::new(Vec::new()), + })) + .descriptor(); assert_eq!( - report.installed_tool_names(), - [ - "ListWorkerSessions", - "ViewSessionOverview", - "SearchSessionEntries", - "ReadSessionEntry", - ] + descriptor.requires_services[0].id, + ServiceId::builtin(WORKER_CONTROL_SERVICE_ID) ); } @@ -838,13 +763,6 @@ mod tests { let provider = Arc::new(FakeProvider { captures: Mutex::new(vec![message("u1", Role::User, "first")]), }); - let list = list_definition(provider.clone())().1; - let listed = list - .execute("{}", llm_engine::tool::ToolExecutionContext::direct()) - .await - .unwrap(); - assert!(listed.content.unwrap().contains("granted")); - let read = read_definition(provider.clone())().1; let hidden = read .execute( diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 664a8866..b77e82b4 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -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, #[serde(default)] pub working_directory: Option, + /// Backend idempotency key used only for authenticated Worker-owned spawn/control. + #[serde(default)] + pub control_operation_id: Option, } #[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, + summary: WorkerSummary, +} + +#[derive(Debug, Serialize, Deserialize)] +struct KnownWorkersResponse { + workspace_id: String, + items: Vec, + truncated: bool, +} + +async fn list_known_workers( + State(api): State, + AxumPath(path): AxumPath, + headers: HeaderMap, +) -> ApiResult> { + 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, + AxumPath(path): AxumPath, + headers: HeaderMap, + Json(request): Json, +) -> ApiResult> { + 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, + AxumPath(path): AxumPath, + headers: HeaderMap, + Json(request): Json, +) -> ApiResult> { + 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, + AxumPath(path): AxumPath, + headers: HeaderMap, + Json(request): Json, +) -> ApiResult> { + 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, + AxumPath(path): AxumPath, + headers: HeaderMap, + Json(request): Json, +) -> ApiResult> { + 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, + AxumPath(path): AxumPath, + headers: HeaderMap, +) -> ApiResult> { + 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, AxumPath(path): AxumPath, @@ -5839,26 +6089,34 @@ async fn scoped_list_worker_observation_sessions( ) -> ApiResult> { 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::>(); Ok(Json(serde_json::json!({ "sessions": sessions }))) @@ -5872,7 +6130,6 @@ async fn scoped_capture_worker_observation_session( ) -> ApiResult> { 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, AxumPath(path): AxumPath, @@ -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, diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index 82df8bc3..ed298ead 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -181,6 +181,11 @@ const MIGRATIONS: &[Migration] = &[ name: "persist Workspace config schema contribution bundles", apply: persist_workspace_config_schema_bundles, }, + Migration { + version: 33, + name: "create durable Runtime Worker control grants", + apply: create_worker_control_grant_authority, + }, ]; struct Migration { @@ -325,6 +330,23 @@ pub struct WorkerRegistryRecord { pub updated_at: String, } +/// Durable authority describing which Runtime Worker another Runtime Worker may +/// discover and control. Revoked grants remain as audit evidence but are never +/// returned by active-grant queries. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct WorkerControlGrantRecord { + pub workspace_id: String, + pub grant_id: String, + pub controller: RuntimeWorkerRef, + pub subject: RuntimeWorkerRef, + pub relation: String, + pub origin: String, + pub permissions: Vec, + pub operation_id: String, + pub created_at: String, + pub revoked_at: Option, +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct TicketWorkerAssignmentRecord { pub workspace_id: String, @@ -732,6 +754,35 @@ pub trait ControlPlaneStore: Send + Sync { fn delete_worker_registry(&self, workspace_id: &str, worker: &RuntimeWorkerRef) -> Result; + fn create_worker_control_grant( + &self, + record: &WorkerControlGrantRecord, + ) -> Result; + fn get_worker_control_grant_by_operation( + &self, + workspace_id: &str, + controller: &RuntimeWorkerRef, + operation_id: &str, + ) -> Result>; + fn get_active_worker_control_grant( + &self, + workspace_id: &str, + controller: &RuntimeWorkerRef, + subject: &RuntimeWorkerRef, + ) -> Result>; + fn list_active_worker_control_grants( + &self, + workspace_id: &str, + controller: &RuntimeWorkerRef, + limit: usize, + ) -> Result>; + fn revoke_worker_control_grant( + &self, + workspace_id: &str, + grant_id: &str, + revoked_at: &str, + ) -> Result; + fn get_ticket_assignment_operation( &self, workspace_id: &str, @@ -2321,6 +2372,157 @@ impl ControlPlaneStore for SqliteWorkspaceStore { }) } + fn create_worker_control_grant( + &self, + record: &WorkerControlGrantRecord, + ) -> Result { + self.with_conn(|conn| { + let permissions_json = serde_json::to_string(&record.permissions) + .map_err(|error| Error::Store(error.to_string()))?; + conn.execute( + r#"INSERT INTO worker_control_grants ( + workspace_id, grant_id, + controller_runtime_id, controller_worker_id, + subject_runtime_id, subject_worker_id, + relation, origin, permissions_json, operation_id, created_at, revoked_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12) + ON CONFLICT ( + workspace_id, + controller_runtime_id, + controller_worker_id, + operation_id + ) DO NOTHING"#, + params![ + record.workspace_id, + record.grant_id, + record.controller.runtime_id, + record.controller.worker_id, + record.subject.runtime_id, + record.subject.worker_id, + record.relation, + record.origin, + permissions_json, + record.operation_id, + record.created_at, + record.revoked_at, + ], + )?; + + let persisted = read_worker_control_grant_by_operation( + conn, + record.workspace_id.as_str(), + &record.controller, + record.operation_id.as_str(), + )? + .ok_or_else(|| Error::Store("worker control grant was not persisted".to_string()))?; + if persisted.subject != record.subject + || persisted.relation != record.relation + || persisted.origin != record.origin + || persisted.permissions != record.permissions + { + return Err(Error::InvalidInput(format!( + "worker control operation `{}` was already used with different input", + record.operation_id + ))); + } + Ok(persisted) + }) + } + + fn get_worker_control_grant_by_operation( + &self, + workspace_id: &str, + controller: &RuntimeWorkerRef, + operation_id: &str, + ) -> Result> { + self.with_conn(|conn| { + read_worker_control_grant_by_operation(conn, workspace_id, controller, operation_id) + }) + } + + fn get_active_worker_control_grant( + &self, + workspace_id: &str, + controller: &RuntimeWorkerRef, + subject: &RuntimeWorkerRef, + ) -> Result> { + self.with_conn(|conn| { + conn.query_row( + r#"SELECT workspace_id, grant_id, + controller_runtime_id, controller_worker_id, + subject_runtime_id, subject_worker_id, + relation, origin, permissions_json, operation_id, created_at, revoked_at + FROM worker_control_grants + WHERE workspace_id = ?1 + AND controller_runtime_id = ?2 AND controller_worker_id = ?3 + AND subject_runtime_id = ?4 AND subject_worker_id = ?5 + AND revoked_at IS NULL + ORDER BY created_at DESC + LIMIT 1"#, + params![ + workspace_id, + controller.runtime_id, + controller.worker_id, + subject.runtime_id, + subject.worker_id, + ], + read_worker_control_grant_record, + ) + .optional() + .map_err(Error::from) + }) + } + + fn list_active_worker_control_grants( + &self, + workspace_id: &str, + controller: &RuntimeWorkerRef, + limit: usize, + ) -> Result> { + self.with_conn(|conn| { + let mut stmt = conn.prepare( + r#"SELECT workspace_id, grant_id, + controller_runtime_id, controller_worker_id, + subject_runtime_id, subject_worker_id, + relation, origin, permissions_json, operation_id, created_at, revoked_at + FROM worker_control_grants + WHERE workspace_id = ?1 + AND controller_runtime_id = ?2 AND controller_worker_id = ?3 + AND revoked_at IS NULL + ORDER BY created_at ASC, grant_id ASC + LIMIT ?4"#, + )?; + let rows = stmt.query_map( + params![ + workspace_id, + controller.runtime_id, + controller.worker_id, + limit as i64, + ], + read_worker_control_grant_record, + )?; + rows.collect::, _>>() + .map_err(Error::from) + }) + } + + fn revoke_worker_control_grant( + &self, + workspace_id: &str, + grant_id: &str, + revoked_at: &str, + ) -> Result { + self.with_conn(|conn| { + let changed = conn.execute( + r#"UPDATE worker_control_grants + SET revoked_at = ?3 + WHERE workspace_id = ?1 AND grant_id = ?2 AND revoked_at IS NULL"#, + params![workspace_id, grant_id, revoked_at], + )?; + Ok(changed > 0) + }) + } + fn get_ticket_assignment_operation( &self, workspace_id: &str, @@ -3627,6 +3829,57 @@ fn read_worker_registry_record(row: &rusqlite::Row<'_>) -> rusqlite::Result, +) -> rusqlite::Result { + let permissions_json: String = row.get(8)?; + let permissions = serde_json::from_str(&permissions_json).map_err(|error| { + rusqlite::Error::FromSqlConversionFailure(8, rusqlite::types::Type::Text, Box::new(error)) + })?; + Ok(WorkerControlGrantRecord { + workspace_id: row.get(0)?, + grant_id: row.get(1)?, + controller: RuntimeWorkerRef::new( + row.get::<_, String>(2)?, + row.get::<_, u64>(3)?.to_string(), + ), + subject: RuntimeWorkerRef::new(row.get::<_, String>(4)?, row.get::<_, u64>(5)?.to_string()), + relation: row.get(6)?, + origin: row.get(7)?, + permissions, + operation_id: row.get(9)?, + created_at: row.get(10)?, + revoked_at: row.get(11)?, + }) +} + +fn read_worker_control_grant_by_operation( + conn: &Connection, + workspace_id: &str, + controller: &RuntimeWorkerRef, + operation_id: &str, +) -> Result> { + conn.query_row( + r#"SELECT workspace_id, grant_id, + controller_runtime_id, controller_worker_id, + subject_runtime_id, subject_worker_id, + relation, origin, permissions_json, operation_id, created_at, revoked_at + FROM worker_control_grants + WHERE workspace_id = ?1 + AND controller_runtime_id = ?2 AND controller_worker_id = ?3 + AND operation_id = ?4"#, + params![ + workspace_id, + controller.runtime_id, + controller.worker_id, + operation_id, + ], + read_worker_control_grant_record, + ) + .optional() + .map_err(Error::from) +} + fn current_ticket_worker_assignment_select_sql() -> String { "SELECT a.workspace_id, a.ticket_id, a.assignment_id, a.runtime_id, a.worker_id, \ a.assigned_by, a.assigned_at \ @@ -4639,6 +4892,58 @@ pub(crate) fn persist_workspace_config_schema_bundles(conn: &Connection) -> Resu Ok(()) } +fn create_worker_control_grant_authority(conn: &Connection) -> Result<()> { + conn.execute_batch( + r#" + CREATE TABLE worker_control_grants ( + workspace_id TEXT NOT NULL, + grant_id TEXT NOT NULL, + controller_runtime_id TEXT NOT NULL, + controller_worker_id INTEGER NOT NULL, + subject_runtime_id TEXT NOT NULL, + subject_worker_id INTEGER NOT NULL, + relation TEXT NOT NULL, + origin TEXT NOT NULL, + permissions_json TEXT NOT NULL, + operation_id TEXT NOT NULL, + created_at TEXT NOT NULL, + revoked_at TEXT, + PRIMARY KEY (workspace_id, grant_id), + UNIQUE ( + workspace_id, + controller_runtime_id, + controller_worker_id, + operation_id + ), + FOREIGN KEY (workspace_id, controller_runtime_id, controller_worker_id) + REFERENCES worker_registry (workspace_id, runtime_id, runtime_worker_id) + ON DELETE CASCADE, + FOREIGN KEY (workspace_id, subject_runtime_id, subject_worker_id) + REFERENCES worker_registry (workspace_id, runtime_id, runtime_worker_id) + ON DELETE CASCADE + ); + + CREATE INDEX idx_worker_control_grants_controller_active + ON worker_control_grants ( + workspace_id, + controller_runtime_id, + controller_worker_id, + revoked_at, + created_at + ); + + CREATE INDEX idx_worker_control_grants_subject_active + ON worker_control_grants ( + workspace_id, + subject_runtime_id, + subject_worker_id, + revoked_at + ); + "#, + )?; + Ok(()) +} + fn create_worker_mutation_source_proof_replay_guard(conn: &Connection) -> Result<()> { conn.execute_batch( r#" @@ -5312,7 +5617,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT); apply_migrations(&conn).unwrap(); - assert_eq!(current_schema_version(&conn).unwrap(), 32); + assert_eq!(current_schema_version(&conn).unwrap(), 33); assert!(table_exists(&conn, "worker_workdir_attachment_reservations").unwrap()); } @@ -5345,7 +5650,7 @@ CREATE TABLE flow_events (event_id TEXT PRIMARY KEY); apply_migrations(&conn).unwrap(); - assert_eq!(current_schema_version(&conn).unwrap(), 32); + assert_eq!(current_schema_version(&conn).unwrap(), 33); assert!(table_exists(&conn, "flow_sources").unwrap()); assert!(table_exists(&conn, "flow_source_revisions").unwrap()); assert!(!table_exists(&conn, "flow_instances").unwrap()); @@ -5412,7 +5717,7 @@ INSERT INTO worker_workdir_attachment_reservations ( apply_migrations(&conn).unwrap(); - assert_eq!(current_schema_version(&conn).unwrap(), 32); + assert_eq!(current_schema_version(&conn).unwrap(), 33); let repositories_sql: String = conn .query_row( "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'repositories'", @@ -5592,7 +5897,7 @@ INSERT INTO workdir_registry ( let db = dir.path().join("control-plane.sqlite"); let store = SqliteWorkspaceStore::open(&db).unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 32); + assert_eq!(store.schema_version().await.unwrap(), 33); assert!( !store .with_conn(|conn| table_exists(conn, "worker_workspace_credentials")) @@ -5609,7 +5914,7 @@ INSERT INTO workdir_registry ( store.upsert_workspace(&record).await.unwrap(); let reopened = SqliteWorkspaceStore::open(&db).unwrap(); - assert_eq!(reopened.schema_version().await.unwrap(), 32); + assert_eq!(reopened.schema_version().await.unwrap(), 33); assert_eq!( reopened.get_workspace("local-dev").await.unwrap(), Some(record) @@ -6156,7 +6461,7 @@ INSERT INTO workdir_registry ( .unwrap(); let store = SqliteWorkspaceStore::from_connection(conn).unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 32); + assert_eq!(store.schema_version().await.unwrap(), 33); store .with_conn(|conn| { @@ -6345,7 +6650,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(), 32); + assert_eq!(store.schema_version().await.unwrap(), 33); let workspace = WorkspaceRecord { workspace_id: "local-dev".to_string(), owner_account_id: None, @@ -6411,7 +6716,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(), 32); + assert_eq!(store.schema_version().await.unwrap(), 33); let workspace = WorkspaceRecord { workspace_id: "local-dev".to_string(), owner_account_id: None, @@ -6671,10 +6976,110 @@ CREATE TABLE ticket_assignment_operations ( ); } + #[tokio::test] + async fn worker_control_grants_are_idempotent_scoped_and_revocable() { + let store = SqliteWorkspaceStore::in_memory().unwrap(); + store + .upsert_workspace(&WorkspaceRecord { + workspace_id: "workspace-control".to_string(), + owner_account_id: None, + display_name: "Control grants".to_string(), + state: "active".to_string(), + created_at: "2026-07-27T00:00:00Z".to_string(), + updated_at: "2026-07-27T00:00:00Z".to_string(), + }) + .await + .unwrap(); + let worker_record = |worker_id: &str, display_name: &str| WorkerRegistryRecord { + workspace_id: "workspace-control".to_string(), + worker: RuntimeWorkerRef::new("runtime-a", worker_id), + display_name: display_name.to_string(), + profile: None, + retention_state: "normal".to_string(), + transcript_ref: None, + session_ref: None, + summary_ref: None, + diagnostics_ref: None, + created_at: "2026-07-27T00:00:00Z".to_string(), + updated_at: "2026-07-27T00:00:00Z".to_string(), + }; + let controller_record = worker_record("1", "Controller"); + let subject_record = worker_record("2", "Subject"); + store.upsert_worker_registry(&controller_record).unwrap(); + store.upsert_worker_registry(&subject_record).unwrap(); + + let grant = WorkerControlGrantRecord { + workspace_id: "workspace-control".to_string(), + grant_id: "grant-1".to_string(), + controller: controller_record.worker.clone(), + subject: subject_record.worker.clone(), + relation: "spawned".to_string(), + origin: "worker_spawn".to_string(), + permissions: vec![ + "observe".to_string(), + "send_input".to_string(), + "stop".to_string(), + ], + operation_id: "spawn-op-1".to_string(), + created_at: "2026-07-27T00:00:01Z".to_string(), + revoked_at: None, + }; + assert_eq!(store.create_worker_control_grant(&grant).unwrap(), grant); + assert_eq!(store.create_worker_control_grant(&grant).unwrap(), grant); + assert_eq!( + store + .list_active_worker_control_grants( + "workspace-control", + &controller_record.worker, + 10, + ) + .unwrap(), + vec![grant.clone()] + ); + assert_eq!( + store + .get_active_worker_control_grant( + "workspace-control", + &controller_record.worker, + &subject_record.worker, + ) + .unwrap(), + Some(grant.clone()) + ); + + let conflicting_replay = WorkerControlGrantRecord { + subject: controller_record.worker.clone(), + ..grant.clone() + }; + assert!(matches!( + store.create_worker_control_grant(&conflicting_replay), + Err(Error::InvalidInput(_)) + )); + assert!( + store + .revoke_worker_control_grant( + "workspace-control", + &grant.grant_id, + "2026-07-27T00:00:02Z", + ) + .unwrap() + ); + assert!( + store + .list_active_worker_control_grants( + "workspace-control", + &controller_record.worker, + 10, + ) + .unwrap() + .is_empty() + ); + } + #[tokio::test] async fn account_and_login_records_round_trip() { let store = SqliteWorkspaceStore::in_memory().unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 32); + assert_eq!(store.schema_version().await.unwrap(), 33); let now = "2026-07-22T00:00:00Z".to_string(); let account = AccountRecord { account_id: "acct-user-alice".to_string(), diff --git a/resources/profiles/base.dcdl b/resources/profiles/base.dcdl index 480be986..f8ab343c 100644 --- a/resources/profiles/base.dcdl +++ b/resources/profiles/base.dcdl @@ -27,7 +27,7 @@ feature = { web = { enabled = true; }; image = { enabled = true; }; sub_worker = { enabled = true; }; - worker = { enabled = false; }; + worker = { enabled = true; }; objective = { enabled = true; }; ticket = { enabled = true; authoring = true; thread = true; }; }; diff --git a/resources/prompts/common/worker-observation.md b/resources/prompts/common/worker-observation.md index 5040fe3c..4b43f651 100644 --- a/resources/prompts/common/worker-observation.md +++ b/resources/prompts/common/worker-observation.md @@ -2,7 +2,7 @@ Worker-session tools are a read-only exploration surface over host-granted active Worker sessions. -- Use `ListWorkerSessions` to discover only the sessions already granted to you. Reuse the returned structured `subject` exactly: Runtime peers use `{ kind: "runtime_worker", runtime_id, worker_id }`, while parent-owned children use `{ kind: "sub_worker", name }`. Do not guess subject identifiers. +- Use `WorkerList` to discover known Workers and reuse its returned structured `subject` exactly. Runtime Workers use `{ kind: "runtime_worker", runtime_id, worker_id }`, while parent-owned children use `{ kind: "sub_worker", name }`. Do not guess subject identifiers. - Use `ViewSessionOverview` for sparse orientation, `SearchSessionEntries` for bounded range/filter queries, and `ReadSessionEntry` for one bounded entry. - `SessionEntryRef` is the common entry identity across overview, search, reads, and evidence conversion. Reuse returned `E...` values; never invent them. - Every operation rereads the latest committed capture. Existing references remain stable when entries append. From 8be2cfd2a33390fbe054a287fd05e599388156bc Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 17 Aug 2026 00:18:49 +0900 Subject: [PATCH 2/4] worker: unify granted Worker control --- crates/worker/src/controller.rs | 3 - .../src/feature/builtin/manage_worker.rs | 846 +++++++++++++----- .../src/feature/builtin/worker_observation.rs | 53 +- crates/worker/src/spawn/comm_tools.rs | 8 +- crates/workspace-server/src/hosts.rs | 36 + crates/workspace-server/src/server.rs | 598 +++++++++++-- crates/workspace-server/src/store.rs | 56 +- resources/profiles/base.dcdl | 4 +- resources/profiles/coder.dcdl | 2 +- 9 files changed, 1304 insertions(+), 302 deletions(-) diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 5e696c3b..ebb07952 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -21,7 +21,6 @@ use crate::shutdown_after_idle::{ ShutdownAfterIdleRequest, TicketIntakeReadyShutdownHook, is_ticket_intake_role, take_shutdown_request_after_status, }; -use crate::spawn::comm_tools::{sub_worker_send_tool, sub_worker_stop_tool}; use crate::spawn::registry::SpawnedWorkerRegistry; use crate::spawn::tool::sub_worker_spawn_tool; use crate::worker::{ @@ -921,8 +920,6 @@ where scope_handle, prompts, )); - engine.register_tool(sub_worker_send_tool(spawned_registry.clone())); - engine.register_tool(sub_worker_stop_tool(spawned_registry.clone())); observation_providers.push(Arc::new( crate::feature::builtin::worker_observation::SpawnedSubWorkerObservationProvider::new( spawned_registry, diff --git a/crates/worker/src/feature/builtin/manage_worker.rs b/crates/worker/src/feature/builtin/manage_worker.rs index 73aff8bc..a66a531c 100644 --- a/crates/worker/src/feature/builtin/manage_worker.rs +++ b/crates/worker/src/feature/builtin/manage_worker.rs @@ -30,12 +30,224 @@ pub const WORKER_LIFECYCLE_SERVICE_ID: &str = "worker.lifecycle"; pub const WORKER_CONTROL_SERVICE_ID: &str = "worker.control"; const WORKER_LIFECYCLE_SERVICE_VERSION: &str = "1"; -pub trait WorkerControlService: Send + Sync {} +#[async_trait] +pub trait WorkerControlService: Send + Sync { + fn workspace_id(&self) -> &str; + fn known_subworkers(&self) -> Vec; + async fn send_subworker( + &self, + name: &str, + content: String, + ) -> Result; + async fn stop_subworker(&self, name: &str) -> Result; + async fn spawn_worker( + &self, + request: WorkerLifecycleSpawnRequest, + ) -> Result; + fn remove_runtime_worker( + &self, + runtime_id: &str, + worker_id: &str, + expected_worker_revision: &str, + reason: &str, + ) -> Result; + async fn execute_runtime( + &self, + request: WorkspaceRequest, + ) -> Result; + async fn ensure_permission( + &self, + subject: &super::worker_observation::WorkerObservationSubjectRef, + permission: &str, + ) -> Result<(), WorkspaceClientError>; +} -#[derive(Debug)] -struct WorkspaceWorkerControlService; +struct WorkspaceWorkerControlService { + client: Arc, + workspace_id: String, + registry: Option>, +} -impl WorkerControlService for WorkspaceWorkerControlService {} +impl std::fmt::Debug for WorkspaceWorkerControlService { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkspaceWorkerControlService") + .field("workspace_id", &self.workspace_id) + .field("has_subworker_registry", &self.registry.is_some()) + .finish_non_exhaustive() + } +} + +#[async_trait] +impl WorkerControlService for WorkspaceWorkerControlService { + fn workspace_id(&self) -> &str { + &self.workspace_id + } + + fn known_subworkers(&self) -> Vec { + self.registry + .as_ref() + .map(|registry| { + registry + .list_internal() + .into_iter() + .map(|internal| { + serde_json::json!({ + "subject": { "kind": "sub_worker", "name": internal.worker_name }, + "relation": "direct_child", + "origin": "sub_worker_spawn", + "permissions": ["send_input", "stop", "observe"], + "summary": { + "display_name": internal.worker_name, + "status": format!("{:?}", internal.session.status()).to_lowercase(), + } + }) + }) + .collect() + }) + .unwrap_or_default() + } + + async fn send_subworker( + &self, + name: &str, + content: String, + ) -> Result { + let record = self + .registry + .as_ref() + .and_then(|registry| registry.get_internal(name)) + .ok_or_else(|| { + WorkspaceClientError::Request( + "unknown Worker or permission not granted".to_string(), + ) + })?; + record + .session + .send(content) + .await + .map_err(|error| WorkspaceClientError::Request(error.to_string()))?; + Ok(WorkspaceResponse { + status: 200, + body: serde_json::json!({ "subject": { "kind": "sub_worker", "name": name } }) + .to_string(), + }) + } + + async fn stop_subworker(&self, name: &str) -> Result { + let registry = self.registry.as_ref().ok_or_else(|| { + WorkspaceClientError::Request("unknown Worker or permission not granted".to_string()) + })?; + registry + .remove_internal(name) + .await + .map_err(|error| WorkspaceClientError::Request(error.to_string()))?; + Ok(WorkspaceResponse { + status: 200, + body: serde_json::json!({ "subject": { "kind": "sub_worker", "name": name } }) + .to_string(), + }) + } + + async fn spawn_worker( + &self, + request: WorkerLifecycleSpawnRequest, + ) -> Result { + WorkspaceWorkerLifecycleService { + client: self.client.clone(), + workspace_id: self.workspace_id.clone(), + } + .spawn(request) + .await + } + + fn remove_runtime_worker( + &self, + runtime_id: &str, + worker_id: &str, + expected_worker_revision: &str, + reason: &str, + ) -> Result { + self.client + .execute_worker_remove(runtime_id, worker_id, expected_worker_revision, reason) + } + + async fn execute_runtime( + &self, + request: WorkspaceRequest, + ) -> Result { + self.client.execute(request) + } + + async fn ensure_permission( + &self, + subject: &super::worker_observation::WorkerObservationSubjectRef, + permission: &str, + ) -> Result<(), WorkspaceClientError> { + match subject { + super::worker_observation::WorkerObservationSubjectRef::SubWorker { name } => { + let known = self + .registry + .as_ref() + .and_then(|registry| registry.get_internal(name)) + .is_some(); + if known && matches!(permission, "send_input" | "stop" | "observe") { + Ok(()) + } else { + Err(WorkspaceClientError::Request( + "unknown Worker or permission not granted".to_string(), + )) + } + } + super::worker_observation::WorkerObservationSubjectRef::RuntimeWorker { + runtime_id, + worker_id, + } => { + let response = self.client.execute(WorkspaceRequest::get(format!( + "/api/w/{}/worker-control/workers", + self.workspace_id + )))?; + if !response.is_success() { + return Err(WorkspaceClientError::Request(format!( + "Workspace control request returned {}: {}", + response.status, response.body + ))); + } + let body: serde_json::Value = + serde_json::from_str(&response.body).map_err(|error| { + WorkspaceClientError::Request(format!( + "invalid Workspace control response: {error}" + )) + })?; + let granted = body + .get("items") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .any(|item| { + item.get("subject") + .and_then(|value| value.get("runtime_id")) + == Some(&serde_json::Value::String(runtime_id.clone())) + && item.get("subject").and_then(|value| value.get("worker_id")) + == Some(&serde_json::Value::String(worker_id.clone())) + && item + .get("permissions") + .and_then(serde_json::Value::as_array) + .is_some_and(|permissions| { + permissions.iter().any(|candidate| candidate == permission) + }) + }); + if granted { + Ok(()) + } else { + Err(WorkspaceClientError::Request( + "unknown Worker or permission not granted".to_string(), + )) + } + } + } + } +} #[async_trait] pub trait WorkerLifecycleService: Send + Sync { @@ -108,7 +320,7 @@ impl WorkerLifecycleService for WorkspaceWorkerLifecycleService { #[derive(Clone)] pub struct ManageWorkerFeature { client: Arc, - registry: Option>, + control: Arc, direct_spawn: bool, } @@ -116,7 +328,10 @@ impl std::fmt::Debug for ManageWorkerFeature { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter .debug_struct("ManageWorkerFeature") - .field("has_registry", &self.registry.is_some()) + .field( + "has_subworker_registry", + &!self.control.known_subworkers().is_empty(), + ) .field("direct_spawn", &self.direct_spawn) .finish_non_exhaustive() } @@ -127,9 +342,15 @@ pub fn manage_worker_feature( registry: Option>, direct_spawn: bool, ) -> ManageWorkerFeature { + let workspace_id = client.workspace_id().unwrap_or_default().to_string(); + let control: Arc = Arc::new(WorkspaceWorkerControlService { + client: client.clone(), + workspace_id, + registry, + }); ManageWorkerFeature { client, - registry, + control, direct_spawn, } } @@ -189,57 +410,37 @@ impl FeatureModule for ManageWorkerFeature { WORKER_LIFECYCLE_SERVICE_VERSION, "Known-Worker discovery and permission-fenced control operations", ), - Arc::new(WorkspaceWorkerControlService) as Arc, + self.control.clone(), )?; for operation in WorkerOperation::ALL { if operation == WorkerOperation::Spawn && !self.direct_spawn { continue; } let definition = match operation { - WorkerOperation::List => definition::( - operation, - self.client.clone(), - workspace_id.clone(), - self.registry.clone(), - ), - WorkerOperation::Spawn => definition::( - operation, - self.client.clone(), - workspace_id.clone(), - self.registry.clone(), - ), + WorkerOperation::List => { + definition::(operation, self.control.clone()) + } + WorkerOperation::Spawn => { + definition::(operation, self.control.clone()) + } WorkerOperation::SendInput | WorkerOperation::Notify => { - definition::( - operation, - self.client.clone(), - workspace_id.clone(), - self.registry.clone(), - ) + definition::(operation, self.control.clone()) + } + WorkerOperation::Cancel | WorkerOperation::Stop => { + definition::(operation, self.control.clone()) + } + WorkerOperation::Restore => { + definition::(operation, self.control.clone()) + } + WorkerOperation::Remove => { + definition::(operation, self.control.clone()) + } + WorkerOperation::Share | WorkerOperation::Transfer => { + definition::(operation, self.control.clone()) + } + WorkerOperation::Revoke => { + definition::(operation, self.control.clone()) } - WorkerOperation::Cancel => definition::( - operation, - self.client.clone(), - workspace_id.clone(), - self.registry.clone(), - ), - WorkerOperation::Stop => definition::( - operation, - self.client.clone(), - workspace_id.clone(), - self.registry.clone(), - ), - WorkerOperation::Restore => definition::( - operation, - self.client.clone(), - workspace_id.clone(), - self.registry.clone(), - ), - WorkerOperation::Remove => definition::( - operation, - self.client.clone(), - workspace_id.clone(), - self.registry.clone(), - ), }; context .tools() @@ -298,26 +499,35 @@ struct WorkerWorkingDirectorySelection { relative_cwd: Option, } +#[derive(Debug, Clone, Deserialize, JsonSchema)] +#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] +enum WorkerSubjectInput { + RuntimeWorker { + runtime_id: String, + worker_id: String, + }, + SubWorker { + name: String, + }, +} + #[derive(Debug, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] struct WorkerTargetInput { - runtime_id: String, - worker_id: String, + subject: WorkerSubjectInput, } #[derive(Debug, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] struct WorkerMessageInput { - runtime_id: String, - worker_id: String, + subject: WorkerSubjectInput, content: String, } #[derive(Debug, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] struct WorkerStopInput { - runtime_id: String, - worker_id: String, + subject: WorkerSubjectInput, #[serde(default)] reason: Option, } @@ -325,17 +535,27 @@ struct WorkerStopInput { #[derive(Debug, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] struct WorkerRemoveInput { - runtime_id: String, - worker_id: String, + subject: WorkerSubjectInput, expected_worker_revision: String, reason: String, } +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +struct WorkerRevokeInput { + grant_id: String, +} + +#[derive(Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +struct WorkerDelegateInput { + grant_id: String, + target_controller: WorkerSubjectInput, +} + struct WorkspaceWorkerTool { operation: WorkerOperation, - client: Arc, - workspace_id: String, - registry: Option>, + control: Arc, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -348,10 +568,13 @@ enum WorkerOperation { Stop, Restore, Remove, + Share, + Transfer, + Revoke, } impl WorkerOperation { - const ALL: [Self; 8] = [ + const ALL: [Self; 11] = [ Self::List, Self::Spawn, Self::SendInput, @@ -360,6 +583,9 @@ impl WorkerOperation { Self::Stop, Self::Restore, Self::Remove, + Self::Share, + Self::Transfer, + Self::Revoke, ]; fn tool_name(self) -> &'static str { @@ -372,6 +598,9 @@ impl WorkerOperation { Self::Stop => "WorkerStop", Self::Restore => "WorkerRestore", Self::Remove => "WorkerRemove", + Self::Share => "WorkerShare", + Self::Transfer => "WorkerTransfer", + Self::Revoke => "WorkerRevoke", } } @@ -393,6 +622,11 @@ impl WorkerOperation { Self::Remove => { "Remove an eligible stopped, unassigned, non-internal Worker. Supply the current Worker revision and a bounded reason; Backend validation and retention are authoritative." } + Self::Share => "Share one controlled Runtime Worker with another known Runtime Worker.", + Self::Transfer => { + "Transfer one controlled Runtime Worker to another known Runtime Worker." + } + Self::Revoke => "Revoke one durable Runtime Worker control grant owned by this Worker.", } } } @@ -405,10 +639,142 @@ impl Tool for WorkspaceWorkerTool { ctx: ToolExecutionContext, ) -> Result { let response = match self.operation { + WorkerOperation::List => { + parse::(input_json, "WorkerList")?; + let response = self + .control + .execute_runtime(WorkspaceRequest::get(format!( + "/api/w/{}/worker-control/workers", + self.control.workspace_id() + ))) + .await + .map_err(control_tool_error)?; + self.with_subworkers(response)? + } + WorkerOperation::Spawn => { + let input = parse::(input_json, "WorkerSpawn")?; + let ticket_id = input + .ticket_id + .map(|ticket_id| authority_id(&ticket_id, "ticket_id")) + .transpose()?; + let operation_id = ticket_id + .as_ref() + .map(|ticket_id| { + let call_id = non_empty(ctx.call_id.clone(), "tool call_id")?; + Ok::<_, ToolError>(format!("worker-spawn:{ticket_id}:{call_id}")) + }) + .transpose()?; + self.control + .spawn_worker(WorkerLifecycleSpawnRequest { + runtime_id: authority_id(&input.runtime_id, "runtime_id")?, + working_directory_id: authority_id( + &input.working_directory_id, + "working_directory_id", + )?, + relative_cwd: input + .relative_cwd + .map(|value| validate_relative_cwd(&value)) + .transpose()?, + profile: non_empty(input.profile, "profile")?, + ticket_id, + operation_id, + display_name: input + .display_name + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| "Workspace Worker".to_string()), + initial_submit: input.initial_submit, + }) + .await + .map_err(control_tool_error)? + } + WorkerOperation::SendInput | WorkerOperation::Notify => { + let input = parse::(input_json, self.operation.tool_name())?; + let content = non_empty(input.content, "content")?; + if content.len() > 16 * 1024 { + return Err(ToolError::ExecutionFailed( + "content must contain at most 16384 bytes".to_string(), + )); + } + match input.subject { + WorkerSubjectInput::SubWorker { name } => { + if self.operation != WorkerOperation::SendInput { + return Err(unsupported_subject(self.operation, "sub_worker")); + } + let subject = subworker_subject(&name)?; + self.control + .ensure_permission(&subject, "send_input") + .await + .map_err(control_tool_error)?; + self.control + .send_subworker(&name, content) + .await + .map_err(control_tool_error)? + } + subject @ WorkerSubjectInput::RuntimeWorker { .. } => { + let (runtime_id, worker_id) = + runtime_subject_ids(&subject, self.operation)?; + self.control.execute_runtime(WorkspaceRequest::json( + WorkspaceRequestMethod::Post, + format!("/api/w/{}/worker-control/workers/{runtime_id}/{worker_id}/input", self.control.workspace_id()), + serde_json::json!({ + "kind": if self.operation == WorkerOperation::Notify { "notify" } else { "user" }, + "content": content, + }).to_string(), + )).await.map_err(control_tool_error)? + } + } + } + WorkerOperation::Cancel | WorkerOperation::Stop => { + let input = parse::(input_json, self.operation.tool_name())?; + match input.subject { + WorkerSubjectInput::SubWorker { name } => { + if self.operation != WorkerOperation::Stop { + return Err(unsupported_subject(self.operation, "sub_worker")); + } + let subject = subworker_subject(&name)?; + self.control + .ensure_permission(&subject, "stop") + .await + .map_err(control_tool_error)?; + self.control + .stop_subworker(&name) + .await + .map_err(control_tool_error)? + } + subject @ WorkerSubjectInput::RuntimeWorker { .. } => { + let (runtime_id, worker_id) = + runtime_subject_ids(&subject, self.operation)?; + let action = if self.operation == WorkerOperation::Cancel { + "cancel" + } else { + "stop" + }; + self.control.execute_runtime(WorkspaceRequest::json( + WorkspaceRequestMethod::Post, + format!("/api/w/{}/worker-control/workers/{runtime_id}/{worker_id}/{action}", self.control.workspace_id()), + serde_json::json!({ "reason": input.reason }).to_string(), + )).await.map_err(control_tool_error)? + } + } + } + WorkerOperation::Restore => { + let input = parse::(input_json, "WorkerRestore")?; + let (runtime_id, worker_id) = runtime_subject_ids(&input.subject, self.operation)?; + self.control + .execute_runtime(WorkspaceRequest::json( + WorkspaceRequestMethod::Post, + format!( + "/api/w/{}/worker-control/workers/{runtime_id}/{worker_id}/restore", + self.control.workspace_id() + ), + "{}", + )) + .await + .map_err(control_tool_error)? + } WorkerOperation::Remove => { let input = parse::(input_json, "WorkerRemove")?; - let runtime_id = authority_id(&input.runtime_id, "runtime_id")?; - let worker_id = authority_id(&input.worker_id, "worker_id")?; + let (runtime_id, worker_id) = runtime_subject_ids(&input.subject, self.operation)?; let expected_worker_revision = non_empty(input.expected_worker_revision, "expected_worker_revision")?; let reason = non_empty(input.reason, "reason")?; @@ -417,153 +783,78 @@ impl Tool for WorkspaceWorkerTool { "reason must contain at most 512 bytes".to_string(), )); } - self.client - .execute_worker_remove( + self.control + .remove_runtime_worker( &runtime_id, &worker_id, &expected_worker_revision, &reason, ) - .map_err(|error| ToolError::ExecutionFailed(error.to_string()))? + .map_err(control_tool_error)? } - operation => { - let request = match operation { - WorkerOperation::List => { - parse::(input_json, "WorkerList")?; - WorkspaceRequest::get(format!( - "/api/w/{}/worker-control/workers", - self.workspace_id - )) - } - WorkerOperation::SendInput | WorkerOperation::Notify => { - let tool_name = operation.tool_name(); - let input = parse::(input_json, tool_name)?; - let runtime_id = authority_id(&input.runtime_id, "runtime_id")?; - let worker_id = authority_id(&input.worker_id, "worker_id")?; - let content = non_empty(input.content, "content")?; - if content.len() > 16 * 1024 { - return Err(ToolError::ExecutionFailed( - "content must contain at most 16384 bytes".to_string(), - )); - } - WorkspaceRequest::json( - WorkspaceRequestMethod::Post, - format!( - "/api/w/{}/worker-control/workers/{runtime_id}/{worker_id}/input", - self.workspace_id - ), - serde_json::json!({ - "kind": if operation == WorkerOperation::Notify { "notify" } else { "user" }, - "content": content, - }) - .to_string(), - ) - } - WorkerOperation::Cancel => { - let input = parse::(input_json, "WorkerCancel")?; - let runtime_id = authority_id(&input.runtime_id, "runtime_id")?; - let worker_id = authority_id(&input.worker_id, "worker_id")?; - let reason = input.reason.unwrap_or_default(); - WorkspaceRequest::json( - WorkspaceRequestMethod::Post, - format!( - "/api/w/{}/worker-control/workers/{runtime_id}/{worker_id}/cancel", - self.workspace_id - ), - serde_json::json!({ "reason": reason }).to_string(), - ) - } - WorkerOperation::Spawn => { - let input = parse::(input_json, "WorkerSpawn")?; - let ticket_id = input - .ticket_id - .map(|ticket_id| authority_id(&ticket_id, "ticket_id")) - .transpose()?; - let operation_id = ticket_id - .as_ref() - .map(|ticket_id| { - let call_id = non_empty(ctx.call_id.clone(), "tool call_id")?; - Ok::<_, ToolError>(format!("worker-spawn:{ticket_id}:{call_id}")) - }) - .transpose()?; - let lifecycle = WorkspaceWorkerLifecycleService { - client: self.client.clone(), - workspace_id: self.workspace_id.clone(), - }; - let response = lifecycle - .spawn(WorkerLifecycleSpawnRequest { - runtime_id: authority_id(&input.runtime_id, "runtime_id")?, - working_directory_id: authority_id( - &input.working_directory_id, - "working_directory_id", - )?, - relative_cwd: input - .relative_cwd - .map(|value| validate_relative_cwd(&value)) - .transpose()?, - profile: non_empty(input.profile, "profile")?, - ticket_id, - operation_id, - display_name: input - .display_name - .filter(|value| !value.trim().is_empty()) - .unwrap_or_else(|| "Workspace Worker".to_string()), - initial_submit: input.initial_submit, - }) - .await - .map_err(|error| ToolError::ExecutionFailed(error.to_string()))?; - return tool_output(self.operation, response); - } - WorkerOperation::Stop => { - let input = parse::(input_json, "WorkerStop")?; - let runtime_id = authority_id(&input.runtime_id, "runtime_id")?; - let worker_id = authority_id(&input.worker_id, "worker_id")?; - WorkspaceRequest::json( - WorkspaceRequestMethod::Post, - format!( - "/api/w/{}/worker-control/workers/{runtime_id}/{worker_id}/stop", - self.workspace_id - ), - serde_json::json!({ "reason": input.reason }).to_string(), - ) - } - WorkerOperation::Restore => { - let input = parse::(input_json, "WorkerRestore")?; - let runtime_id = authority_id(&input.runtime_id, "runtime_id")?; - let worker_id = authority_id(&input.worker_id, "worker_id")?; - WorkspaceRequest::json( - WorkspaceRequestMethod::Post, - format!( - "/api/w/{}/worker-control/workers/{runtime_id}/{worker_id}/restore", - self.workspace_id - ), - "{}", - ) - } - WorkerOperation::Remove => unreachable!("handled above"), + WorkerOperation::Share | WorkerOperation::Transfer => { + let input = parse::(input_json, self.operation.tool_name())?; + let grant_id = authority_id(&input.grant_id, "grant_id")?; + let (runtime_id, worker_id) = + runtime_subject_ids(&input.target_controller, self.operation)?; + let operation_id = format!( + "worker-control-{}:{}", + if self.operation == WorkerOperation::Transfer { + "transfer" + } else { + "share" + }, + non_empty(ctx.call_id.clone(), "tool call_id")? + ); + let action = if self.operation == WorkerOperation::Transfer { + "transfer" + } else { + "share" }; - self.client - .execute(request) - .map_err(|error| ToolError::ExecutionFailed(error.to_string()))? + self.control + .execute_runtime(WorkspaceRequest::json( + WorkspaceRequestMethod::Post, + format!( + "/api/w/{}/worker-control/grants/{grant_id}/{action}", + self.control.workspace_id() + ), + serde_json::json!({ + "target_controller": { + "runtime_id": runtime_id, + "worker_id": worker_id, + }, + "operation_id": operation_id, + }) + .to_string(), + )) + .await + .map_err(control_tool_error)? + } + WorkerOperation::Revoke => { + let input = parse::(input_json, "WorkerRevoke")?; + let grant_id = authority_id(&input.grant_id, "grant_id")?; + self.control + .execute_runtime(WorkspaceRequest::json( + WorkspaceRequestMethod::Post, + format!( + "/api/w/{}/worker-control/grants/{grant_id}/revoke", + self.control.workspace_id() + ), + "{}", + )) + .await + .map_err(control_tool_error)? } - }; - let response = if self.operation == WorkerOperation::List { - self.with_subworkers(response).await? - } else { - response }; tool_output(self.operation, response) } } impl WorkspaceWorkerTool { - async fn with_subworkers( + fn with_subworkers( &self, mut response: WorkspaceResponse, ) -> Result { - let Some(registry) = &self.registry else { - return Ok(response); - }; if !response.is_success() { return Ok(response); } @@ -579,19 +870,7 @@ impl WorkspaceWorkerTool { "WorkerList response did not contain an items array".to_string(), ) })?; - for internal in registry.list_internal() { - let child_name = internal.worker_name.clone(); - items.push(serde_json::json!({ - "subject": { "kind": "sub_worker", "name": child_name }, - "relation": "direct_child", - "origin": "sub_worker_spawn", - "permissions": ["send_input", "stop", "observe"], - "summary": { - "display_name": internal.worker_name, - "status": format!("{:?}", internal.session.status()).to_lowercase(), - } - })); - } + items.extend(self.control.known_subworkers()); response.body = serde_json::to_string(&body).map_err(|error| { ToolError::ExecutionFailed(format!("WorkerList could not encode its response: {error}")) })?; @@ -599,6 +878,43 @@ impl WorkspaceWorkerTool { } } +fn runtime_subject_ids( + subject: &WorkerSubjectInput, + operation: WorkerOperation, +) -> Result<(String, String), ToolError> { + match subject { + WorkerSubjectInput::RuntimeWorker { + runtime_id, + worker_id, + } => Ok(( + authority_id(runtime_id, "runtime_id")?, + authority_id(worker_id, "worker_id")?, + )), + WorkerSubjectInput::SubWorker { .. } => Err(unsupported_subject(operation, "sub_worker")), + } +} + +fn subworker_subject( + name: &str, +) -> Result { + Ok( + super::worker_observation::WorkerObservationSubjectRef::SubWorker { + name: authority_id(name, "name")?, + }, + ) +} + +fn unsupported_subject(operation: WorkerOperation, kind: &str) -> ToolError { + ToolError::InvalidArgument(format!( + "{} does not support subject kind '{kind}'", + operation.tool_name() + )) +} + +fn control_tool_error(error: WorkspaceClientError) -> ToolError { + ToolError::ExecutionFailed(error.to_string()) +} + fn tool_output( operation: WorkerOperation, response: WorkspaceResponse, @@ -618,9 +934,7 @@ fn tool_output( fn definition( operation: WorkerOperation, - client: Arc, - workspace_id: String, - registry: Option>, + control: Arc, ) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(I); @@ -630,9 +944,7 @@ fn definition( .input_schema(schema_value); let tool: Arc = Arc::new(WorkspaceWorkerTool { operation, - client: client.clone(), - workspace_id: workspace_id.clone(), - registry: registry.clone(), + control: control.clone(), }); (meta, tool) }) @@ -733,14 +1045,20 @@ mod tests { } } + fn test_control(client: Arc) -> Arc { + Arc::new(WorkspaceWorkerControlService { + client, + workspace_id: "workspace%2Ftest".to_string(), + registry: None, + }) + } + #[tokio::test] async fn worker_spawn_forwards_typed_initial_submit_to_workspace_api() { let client = Arc::new(RecordingWorkspaceClient::default()); let tool = WorkspaceWorkerTool { operation: WorkerOperation::Spawn, - client: client.clone(), - workspace_id: "workspace%2Ftest".to_string(), - registry: None, + control: test_control(client.clone()), }; tool.execute( &serde_json::json!({ @@ -813,6 +1131,9 @@ mod tests { "WorkerStop", "WorkerRestore", "WorkerRemove", + "WorkerShare", + "WorkerTransfer", + "WorkerRevoke", ] ); } @@ -869,8 +1190,11 @@ mod tests { ( WorkerOperation::SendInput, serde_json::json!({ - "runtime_id": "runtime-1", - "worker_id": "worker-7", + "subject": { + "kind": "runtime_worker", + "runtime_id": "runtime-1", + "worker_id": "worker-7", + }, "content": "continue", }), "/input", @@ -879,8 +1203,11 @@ mod tests { ( WorkerOperation::Notify, serde_json::json!({ - "runtime_id": "runtime-1", - "worker_id": "worker-7", + "subject": { + "kind": "runtime_worker", + "runtime_id": "runtime-1", + "worker_id": "worker-7", + }, "content": "review ready", }), "/input", @@ -889,8 +1216,11 @@ mod tests { ( WorkerOperation::Cancel, serde_json::json!({ - "runtime_id": "runtime-1", - "worker_id": "worker-7", + "subject": { + "kind": "runtime_worker", + "runtime_id": "runtime-1", + "worker_id": "worker-7", + }, "reason": "superseded", }), "/cancel", @@ -899,9 +1229,7 @@ mod tests { ] { WorkspaceWorkerTool { operation, - client: client.clone(), - workspace_id: "workspace%2Ftest".to_string(), - registry: None, + control: test_control(client.clone()), } .execute( &args.to_string(), @@ -920,19 +1248,59 @@ mod tests { } } + #[tokio::test] + async fn worker_share_and_transfer_use_typed_runtime_subjects_and_operation_ids() { + let client = Arc::new(RecordingWorkspaceClient::default()); + for (operation, action) in [ + (WorkerOperation::Share, "share"), + (WorkerOperation::Transfer, "transfer"), + ] { + WorkspaceWorkerTool { + operation, + control: test_control(client.clone()), + } + .execute( + &serde_json::json!({ + "grant_id": "grant-1", + "target_controller": { + "kind": "runtime_worker", + "runtime_id": "runtime-2", + "worker_id": "worker-9", + }, + }) + .to_string(), + ToolExecutionContext::new("call-delegate", "batch-delegate", 0), + ) + .await + .unwrap(); + let request = client.requests.lock().unwrap().last().cloned().unwrap(); + assert!(request.path.ends_with(&format!("/grant-1/{action}"))); + let body: serde_json::Value = + serde_json::from_str(request.body.as_deref().unwrap()).unwrap(); + assert_eq!(body["target_controller"]["runtime_id"], "runtime-2"); + assert!( + body["operation_id"] + .as_str() + .unwrap() + .contains("call-delegate") + ); + } + } + #[tokio::test] async fn worker_remove_forwards_only_target_revision_and_bounded_reason() { let client = Arc::new(RecordingWorkspaceClient::default()); let tool = WorkspaceWorkerTool { operation: WorkerOperation::Remove, - client: client.clone(), - workspace_id: "workspace%2Ftest".to_string(), - registry: None, + control: test_control(client.clone()), }; tool.execute( &serde_json::json!({ - "runtime_id": "runtime-1", - "worker_id": "worker-7", + "subject": { + "kind": "runtime_worker", + "runtime_id": "runtime-1", + "worker_id": "worker-7", + }, "expected_worker_revision": "2026-08-11T20:00:00Z", "reason": " retire completed Worker " }) @@ -972,16 +1340,17 @@ mod tests { let client = Arc::new(RecordingWorkspaceClient::default()); let tool = WorkspaceWorkerTool { operation: WorkerOperation::Remove, - client: client.clone(), - workspace_id: "workspace%2Ftest".to_string(), - registry: None, + control: test_control(client.clone()), }; for reason in [" ".to_string(), "x".repeat(513)] { let _error = tool .execute( &serde_json::json!({ - "runtime_id": "runtime-1", - "worker_id": "worker-7", + "subject": { + "kind": "runtime_worker", + "runtime_id": "runtime-1", + "worker_id": "worker-7", + }, "expected_worker_revision": "revision-1", "reason": reason, }) @@ -994,8 +1363,11 @@ mod tests { let _error = tool .execute( &serde_json::json!({ - "runtime_id": "runtime-1", - "worker_id": "worker-7", + "subject": { + "kind": "runtime_worker", + "runtime_id": "runtime-1", + "worker_id": "worker-7", + }, "expected_worker_revision": "revision-1", "reason": "retire", "source_proof": "caller-controlled" diff --git a/crates/worker/src/feature/builtin/worker_observation.rs b/crates/worker/src/feature/builtin/worker_observation.rs index cae10728..e9e65f22 100644 --- a/crates/worker/src/feature/builtin/worker_observation.rs +++ b/crates/worker/src/feature/builtin/worker_observation.rs @@ -7,7 +7,7 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use session_store::collect_state; -use super::manage_worker::WORKER_CONTROL_SERVICE_ID; +use super::manage_worker::{WORKER_CONTROL_SERVICE_ID, WorkerControlService}; use crate::feature::{ FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureInstructionContribution, FeatureInstructionDeclaration, FeatureInstructionId, FeatureModule, ServiceId, @@ -186,6 +186,43 @@ fn workspace_client_error(error: crate::worker::WorkspaceClientError) -> WorkerO } #[derive(Clone)] +struct ControlAuthorizedObservationProvider { + control: Arc, + inner: Arc, +} + +#[async_trait] +impl WorkerObservationProvider for ControlAuthorizedObservationProvider { + async fn list_worker_sessions( + &self, + ) -> Result, WorkerObservationError> { + let candidates = self.inner.list_worker_sessions().await?; + let mut granted = Vec::new(); + for candidate in candidates { + if self + .control + .ensure_permission(&candidate.subject, "observe") + .await + .is_ok() + { + granted.push(candidate); + } + } + Ok(granted) + } + + async fn capture_worker_session( + &self, + subject: &WorkerObservationSubjectRef, + ) -> Result { + self.control + .ensure_permission(subject, "observe") + .await + .map_err(|_| WorkerObservationError::NotFound)?; + self.inner.capture_worker_session(subject).await + } +} + pub struct WorkerObservationFeature { provider: Arc, } @@ -227,17 +264,25 @@ impl FeatureModule for WorkerObservationFeature { .register(FeatureInstructionContribution::new( observation_instruction(), ))?; + let control = context + .services() + .require::(&ServiceId::builtin(WORKER_CONTROL_SERVICE_ID))?; + let provider: Arc = + Arc::new(ControlAuthorizedObservationProvider { + control, + inner: self.provider.clone(), + }); context.tools().register(ToolContribution::new( "ViewSessionOverview", - overview_definition(self.provider.clone()), + overview_definition(provider.clone()), ))?; context.tools().register(ToolContribution::new( "SearchSessionEntries", - search_definition(self.provider.clone()), + search_definition(provider.clone()), ))?; context.tools().register(ToolContribution::new( "ReadSessionEntry", - read_definition(self.provider.clone()), + read_definition(provider), ))?; Ok(()) } diff --git a/crates/worker/src/spawn/comm_tools.rs b/crates/worker/src/spawn/comm_tools.rs index 458719b7..763ae531 100644 --- a/crates/worker/src/spawn/comm_tools.rs +++ b/crates/worker/src/spawn/comm_tools.rs @@ -1,6 +1,9 @@ +#![cfg_attr(not(test), allow(dead_code, unused_imports))] + //! Parent-facing tools for in-process Internal SubWorker sessions. //! -//! All four tools share the same parent-owned `SpawnedWorkerRegistry` of typed session handles. +//! Legacy direct-child tool constructors are test-only; production exposes the +//! registry through the unified `worker.control` service and Worker tools. //! There is no Runtime catalog lookup or child socket transport, so a Worker can operate only on //! its direct Internal children. The socket helper at the bottom remains solely for the legacy //! top-level Worker callback protocol and is not part of SubWorker communication. @@ -74,6 +77,7 @@ impl Tool for SubWorkerListTool { } } +#[cfg(test)] pub fn sub_worker_list_tool(registry: Arc) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(SubWorkerListInput); @@ -132,6 +136,7 @@ impl Tool for SubWorkerSendTool { } } +#[cfg(test)] pub fn sub_worker_send_tool(registry: Arc) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(SubWorkerSendInput); @@ -186,6 +191,7 @@ impl Tool for SubWorkerStopTool { } } +#[cfg(test)] pub fn sub_worker_stop_tool(registry: Arc) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(NameInput); diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index 847812d1..65cb9e26 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -342,6 +342,12 @@ pub struct WorkerTicketAssignmentRequest { pub(crate) fn worker_spawn_idempotency( request: &WorkerSpawnRequest, ) -> Result, String> { + if let Some(operation) = request.resolved_control_operation.as_ref() { + return Ok(Some(( + operation.operation_id.clone(), + operation.input_fingerprint.clone(), + ))); + } let Some(assignment) = request.ticket_assignment.as_ref() else { return Ok(None); }; @@ -353,6 +359,12 @@ pub(crate) fn worker_spawn_idempotency( ))) } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct WorkerControlOperation { + pub operation_id: String, + pub input_fingerprint: String, +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[serde(deny_unknown_fields)] pub struct WorkerSpawnRequest { @@ -384,6 +396,9 @@ pub struct WorkerSpawnRequest { /// Backend-authored peer-session grants. Browser/model input cannot set this field. #[serde(skip, default)] pub resolved_worker_observation_grants: Vec, + /// Trusted Backend operation identity used to make Worker-owned spawns replay-safe. + #[serde(skip, default)] + pub resolved_control_operation: Option, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] @@ -4801,10 +4816,28 @@ mod tests { resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), + resolved_control_operation: None, resolved_workspace_api: Some(test_workspace_api()), } } + #[test] + fn trusted_control_operation_is_runtime_spawn_idempotency_authority() { + let mut request = embedded_spawn_request(); + request.resolved_control_operation = Some(WorkerControlOperation { + operation_id: "control-op-1".to_string(), + input_fingerprint: "sha256:control-input".to_string(), + }); + + assert_eq!( + worker_spawn_idempotency(&request).unwrap(), + Some(( + "control-op-1".to_string(), + "sha256:control-input".to_string(), + )) + ); + } + #[test] fn spawn_config_bundle_ref_preserves_bundle_identity() { let mut request = embedded_spawn_request(); @@ -5016,6 +5049,7 @@ mod tests { resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), + resolved_control_operation: None, resolved_workspace_api: Some(test_workspace_api()), }, ) @@ -5113,6 +5147,7 @@ mod tests { resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), + resolved_control_operation: None, resolved_workspace_api: Some(test_workspace_api()), }, ) @@ -5149,6 +5184,7 @@ mod tests { resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), + resolved_control_operation: None, resolved_workspace_api: Some(test_workspace_api()), }, ) diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index b77e82b4..5e856c87 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -75,10 +75,10 @@ use crate::hosts::{ EmbeddedWorkerRuntime, HostSummary, RemoteRuntimeConfig, RemoteWorkerRuntime, RuntimeDiagnostic, RuntimeRegistry, RuntimeRegistryError, RuntimeRegistryUnregisterResult, RuntimeSummary, TicketWorkerRole, WorkerCapabilitySummary, WorkerCompletionsRequest, - WorkerCompletionsResult, WorkerImplementationSummary, WorkerInputKind, WorkerInputRequest, - WorkerInputResult, WorkerLifecycleRequest, WorkerLifecycleResult, WorkerOperationState, - WorkerRestoreResult, WorkerSpawnAcceptanceRequirement, WorkerSpawnIntent, WorkerSpawnRequest, - WorkerSpawnResult, WorkerSpawnWorkingDirectoryRequest, WorkerSummary, + WorkerCompletionsResult, WorkerControlOperation, WorkerImplementationSummary, WorkerInputKind, + WorkerInputRequest, WorkerInputResult, WorkerLifecycleRequest, WorkerLifecycleResult, + WorkerOperationState, WorkerRestoreResult, WorkerSpawnAcceptanceRequirement, WorkerSpawnIntent, + WorkerSpawnRequest, WorkerSpawnResult, WorkerSpawnWorkingDirectoryRequest, WorkerSummary, WorkerTicketAssignmentRequest, WorkerWorkspaceSummary, }; use crate::identity::WorkspaceIdentity; @@ -264,6 +264,7 @@ pub struct WorkspaceApi { workdir_sessions: Arc>>, workdir_session_locks: Arc>>>>, worker_remove_locks: Arc>>>>, + worker_control_locks: Arc>>>>, } #[derive(Clone)] @@ -274,6 +275,7 @@ struct WorkspaceWorkerRemoveExecutor { workdir_sessions: Arc>>, workdir_session_locks: Arc>>>>, worker_remove_locks: Arc>>>>, + worker_control_locks: Arc>>>>, } impl WorkspaceWorkerRemoveExecutor { @@ -285,6 +287,7 @@ impl WorkspaceWorkerRemoveExecutor { workdir_sessions: api.workdir_sessions.clone(), workdir_session_locks: api.workdir_session_locks.clone(), worker_remove_locks: api.worker_remove_locks.clone(), + worker_control_locks: api.worker_control_locks.clone(), } } @@ -347,17 +350,46 @@ impl WorkspaceWorkerRemoveExecutor { )); } let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id); - let granted = self + let grant = 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| { + .filter(|grant| { grant .permissions .iter() .any(|permission| permission == "remove") }); - if !granted { + let Some(grant) = grant else { + return Ok(worker_remove_error_response( + StatusCode::NOT_FOUND, + "unknown_worker", + "The target Worker is not known to the current Worker", + )); + }; + let control_lock = { + let mut locks = self + .worker_control_locks + .lock() + .map_err(|_| "Worker control lock registry was poisoned".to_string())?; + locks + .entry(grant.grant_id.clone()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone() + }; + let _control_guard = control_lock.lock().await; + let still_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(|current| { + current.grant_id == grant.grant_id + && current + .permissions + .iter() + .any(|permission| permission == "remove") + }); + if !still_granted { return Ok(worker_remove_error_response( StatusCode::NOT_FOUND, "unknown_worker", @@ -787,6 +819,7 @@ impl WorkspaceApi { workdir_sessions: Arc::new(Mutex::new(HashMap::new())), workdir_session_locks: Arc::new(Mutex::new(HashMap::new())), worker_remove_locks: Arc::new(Mutex::new(HashMap::new())), + worker_control_locks: Arc::new(Mutex::new(HashMap::new())), }; if let Some(dispatcher) = worker_remove_dispatcher { dispatcher @@ -1443,6 +1476,18 @@ pub fn build_router(api: WorkspaceApi) -> Router { "/api/w/{workspace_id}/worker-control/workers", get(list_known_workers).post(spawn_known_worker), ) + .route( + "/api/w/{workspace_id}/worker-control/grants/{grant_id}/share", + post(share_worker_control_grant), + ) + .route( + "/api/w/{workspace_id}/worker-control/grants/{grant_id}/transfer", + post(transfer_worker_control_grant), + ) + .route( + "/api/w/{workspace_id}/worker-control/grants/{grant_id}/revoke", + post(revoke_worker_control_grant), + ) .route( "/api/w/{workspace_id}/worker-control/workers/{runtime_id}/{worker_id}/input", post(send_known_worker_input), @@ -2052,7 +2097,7 @@ pub struct BrowserWorkingDirectoryCreateRequest { pub selector: Option, } -#[derive(Debug, Deserialize)] +#[derive(Debug, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct BrowserWorkerWorkingDirectorySelection { pub working_directory_id: String, @@ -2070,14 +2115,14 @@ pub struct BrowserWorkspaceOrchestratorResponse { pub diagnostics: Vec, } -#[derive(Debug, Deserialize)] +#[derive(Debug, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct CreateWorkspaceWorkerTicketAssignmentRequest { pub ticket_id: String, pub operation_id: String, } -#[derive(Debug, Deserialize)] +#[derive(Debug, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct CreateWorkspaceWorkerRequest { pub runtime_id: String, @@ -2093,6 +2138,9 @@ pub struct CreateWorkspaceWorkerRequest { /// Backend idempotency key used only for authenticated Worker-owned spawn/control. #[serde(default)] pub control_operation_id: Option, + /// Trusted resolution populated only by the authenticated worker-control handler. + #[serde(skip, default)] + pub resolved_control_operation: Option, } #[derive(Debug, Serialize, Deserialize)] @@ -2209,6 +2257,12 @@ struct ScopedWorkspacePath { workspace_id: String, } +#[derive(Debug, Deserialize)] +struct ScopedWorkerControlGrantPath { + workspace_id: String, + grant_id: String, +} + #[derive(Debug, Deserialize)] struct ScopedFlowPath { workspace_id: String, @@ -5274,6 +5328,7 @@ fn start_memory_staging_consolidation( resolved_config_bundle, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), + resolved_control_operation: None, resolved_workspace_api: None, }, )?; @@ -5861,6 +5916,7 @@ async fn scoped_workspace_orchestrator_status( #[derive(Debug, Serialize, Deserialize)] struct KnownWorkerRecord { + grant_id: String, subject: RuntimeWorkerRef, relation: String, origin: String, @@ -5895,6 +5951,7 @@ async fn list_known_workers( .worker(&grant.subject) .map_err(|error| error.into_error())?; items.push(KnownWorkerRecord { + grant_id: grant.grant_id, subject: grant.subject, relation: grant.relation, origin: grant.origin, @@ -5913,7 +5970,7 @@ async fn spawn_known_worker( State(api): State, AxumPath(path): AxumPath, headers: HeaderMap, - Json(request): Json, + Json(mut request): Json, ) -> ApiResult> { validate_workspace_scope(&api, &path.workspace_id)?; let source = authenticate_worker_mutation_source(&api, &path.workspace_id, &headers)?; @@ -5929,31 +5986,20 @@ async fn spawn_known_worker( .filter(|value| !value.is_empty()) .ok_or_else(|| Error::InvalidInput("control_operation_id is required".to_string()))? .to_string(); + let fingerprint_input = serde_json::to_vec(&request) + .map_err(|error| Error::InvalidInput(format!("invalid Worker spawn input: {error}")))?; + let input_fingerprint = format!( + "sha256:{}", + Sha256::digest(&fingerprint_input) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::() + ); + request.resolved_control_operation = Some(WorkerControlOperation { + operation_id: operation_id.clone(), + input_fingerprint, + }); 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 @@ -5971,6 +6017,8 @@ async fn spawn_known_worker( "stop".to_string(), "restore".to_string(), "remove".to_string(), + "share".to_string(), + "transfer".to_string(), "observe".to_string(), ], operation_id, @@ -5978,23 +6026,34 @@ async fn spawn_known_worker( 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); + // Runtime creation is idempotent under the trusted control operation. + // Preserve the unacknowledged Worker/assignment so a retry converges on + // the same subject and can finish grant persistence without creating a + // second Worker or leaving assignment cleanup races. return Err(ApiError::from(error)); } Ok(response) } +fn worker_control_lock(api: &WorkspaceApi, grant_id: &str) -> Arc> { + let mut locks = api + .worker_control_locks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + Arc::clone( + locks + .entry(grant_id.to_string()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))), + ) +} + fn authorize_known_worker_permission( api: &WorkspaceApi, workspace_id: &str, controller: &RuntimeWorkerRef, subject: &RuntimeWorkerRef, permission: &str, -) -> Result<()> { +) -> Result { let grant = api .store .get_active_worker_control_grant(workspace_id, controller, subject)? @@ -6010,7 +6069,184 @@ fn authorize_known_worker_permission( "worker control permission `{permission}` was not granted" ))); } - Ok(()) + Ok(grant) +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct DelegateWorkerControlGrantRequest { + target_controller: RuntimeWorkerRef, + operation_id: String, +} + +async fn share_worker_control_grant( + State(api): State, + AxumPath(path): AxumPath, + headers: HeaderMap, + Json(request): Json, +) -> ApiResult> { + delegate_worker_control_grant(api, path, headers, request, false).await +} + +async fn transfer_worker_control_grant( + State(api): State, + AxumPath(path): AxumPath, + headers: HeaderMap, + Json(request): Json, +) -> ApiResult> { + delegate_worker_control_grant(api, path, headers, request, true).await +} + +async fn delegate_worker_control_grant( + api: WorkspaceApi, + path: ScopedWorkerControlGrantPath, + headers: HeaderMap, + request: DelegateWorkerControlGrantRequest, + transfer: bool, +) -> ApiResult> { + 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 = if transfer { "transfer" } else { "share" }; + let grant = api + .store + .get_worker_control_grant(&path.workspace_id, &path.grant_id)? + .filter(|grant| { + grant.controller == controller + && grant.permissions.iter().any(|value| value == permission) + }) + .ok_or_else(|| Error::UnknownWorker { + worker: controller.clone(), + })?; + if request.target_controller == controller { + return Err(ApiError::from(Error::InvalidInput( + "target_controller must differ from the current Worker".to_string(), + ))); + } + api.store + .get_worker_registry(&path.workspace_id, &request.target_controller)? + .ok_or_else(|| Error::UnknownWorker { + worker: request.target_controller.clone(), + })?; + let operation_id = request.operation_id.trim(); + if operation_id.is_empty() || operation_id.len() > 200 { + return Err(ApiError::from(Error::InvalidInput( + "operation_id must contain 1..=200 bytes".to_string(), + ))); + } + let expected_origin = format!("worker_control_{permission}:{}", grant.grant_id); + if let Some(existing) = api.store.get_worker_control_grant_by_operation( + &path.workspace_id, + &request.target_controller, + operation_id, + )? { + if existing.subject == grant.subject && existing.origin == expected_origin { + if transfer && grant.revoked_at.is_none() { + let lock = worker_control_lock(&api, &grant.grant_id); + let _guard = lock.lock().await; + let now = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + if api + .store + .get_worker_control_grant(&path.workspace_id, &grant.grant_id)? + .is_some_and(|current| current.revoked_at.is_none()) + { + api.store.revoke_worker_control_grant( + &path.workspace_id, + &grant.grant_id, + &now, + )?; + } + } + return Ok(Json(existing)); + } + return Err(ApiError::from(Error::InvalidInput( + "operation_id was already used with different grant input".to_string(), + ))); + } + if grant.revoked_at.is_some() { + return Err(ApiError::from(Error::UnknownWorker { + worker: grant.subject, + })); + } + let lock = worker_control_lock(&api, &grant.grant_id); + let _guard = lock.lock().await; + let current = api + .store + .get_worker_control_grant(&path.workspace_id, &path.grant_id)? + .filter(|candidate| { + candidate.controller == controller + && candidate.revoked_at.is_none() + && candidate + .permissions + .iter() + .any(|value| value == permission) + }) + .ok_or_else(|| Error::UnknownWorker { + worker: grant.subject.clone(), + })?; + let now = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + let delegated = api + .store + .create_worker_control_grant(&WorkerControlGrantRecord { + workspace_id: path.workspace_id.clone(), + grant_id: new_id("wcg"), + controller: request.target_controller, + subject: current.subject.clone(), + relation: if transfer { "transferred" } else { "shared" }.to_string(), + origin: expected_origin, + permissions: current.permissions.clone(), + operation_id: operation_id.to_string(), + created_at: now.clone(), + revoked_at: None, + })?; + if transfer + && !api + .store + .revoke_worker_control_grant(&path.workspace_id, ¤t.grant_id, &now)? + { + return Err(ApiError::from(Error::UnknownWorker { + worker: current.subject, + })); + } + Ok(Json(delegated)) +} + +async fn revoke_worker_control_grant( + State(api): State, + AxumPath(path): AxumPath, + headers: HeaderMap, +) -> ApiResult> { + 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 grant = api + .store + .get_worker_control_grant(&path.workspace_id, &path.grant_id)? + .filter(|grant| grant.controller == controller && grant.revoked_at.is_none()) + .ok_or_else(|| Error::UnknownWorker { + worker: controller.clone(), + })?; + let lock = worker_control_lock(&api, &grant.grant_id); + let _guard = lock.lock().await; + let current = api + .store + .get_worker_control_grant(&path.workspace_id, &path.grant_id)? + .filter(|candidate| candidate.controller == controller && candidate.revoked_at.is_none()) + .ok_or_else(|| Error::UnknownWorker { + worker: grant.subject.clone(), + })?; + let revoked_at = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + if !api + .store + .revoke_worker_control_grant(&path.workspace_id, &path.grant_id, &revoked_at)? + { + return Err(ApiError::from(Error::UnknownWorker { + worker: current.subject, + })); + } + let mut revoked = current; + revoked.revoked_at = Some(revoked_at); + Ok(Json(revoked)) } async fn send_known_worker_input( @@ -6026,6 +6262,15 @@ async fn send_known_worker_input( WorkerInputKind::Notify => "notify", _ => "send_input", }; + let grant = authorize_known_worker_permission( + &api, + &path.workspace_id, + &controller, + &path.worker, + permission, + )?; + let lock = worker_control_lock(&api, &grant.grant_id); + let _guard = lock.lock().await; authorize_known_worker_permission( &api, &path.workspace_id, @@ -6045,6 +6290,15 @@ async fn cancel_known_worker( 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 grant = authorize_known_worker_permission( + &api, + &path.workspace_id, + &controller, + &path.worker, + "cancel", + )?; + let lock = worker_control_lock(&api, &grant.grant_id); + let _guard = lock.lock().await; authorize_known_worker_permission( &api, &path.workspace_id, @@ -6065,6 +6319,10 @@ async fn stop_known_worker( 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); + let grant = + authorize_known_worker_permission(&api, &path.workspace_id, &controller, &subject, "stop")?; + let lock = worker_control_lock(&api, &grant.grant_id); + let _guard = lock.lock().await; authorize_known_worker_permission(&api, &path.workspace_id, &controller, &subject, "stop")?; scoped_stop_runtime_worker(State(api), AxumPath(path), Json(request)).await } @@ -6078,6 +6336,15 @@ async fn restore_known_worker( 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); + let grant = authorize_known_worker_permission( + &api, + &path.workspace_id, + &controller, + &subject, + "restore", + )?; + let lock = worker_control_lock(&api, &grant.grant_id); + let _guard = lock.lock().await; authorize_known_worker_permission(&api, &path.workspace_id, &controller, &subject, "restore")?; scoped_restore_runtime_worker(State(api), AxumPath(path), Query(Default::default())).await } @@ -6141,6 +6408,15 @@ async fn scoped_capture_worker_observation_session( }; let target = RuntimeWorkerRef::new(runtime_id, worker_id); let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id); + let grant = authorize_known_worker_permission( + &api, + &path.workspace_id, + &controller, + &target, + "observe", + )?; + let lock = worker_control_lock(&api, &grant.grant_id); + let _guard = lock.lock().await; authorize_known_worker_permission(&api, &path.workspace_id, &controller, &target, "observe")?; let target_summary = api .runtime @@ -6242,16 +6518,8 @@ async fn scoped_start_workspace_orchestrator( resolved_working_directory: None, resolved_config_bundle: None, resolved_worker_observation_enabled: true, - resolved_worker_observation_grants: workers_response(api.clone()) - .map(|response| { - response - .items - .into_iter() - .take(100) - .map(|worker| worker.worker) - .collect() - }) - .unwrap_or_default(), + resolved_worker_observation_grants: Vec::new(), + resolved_control_operation: None, resolved_workspace_api: None, }, )?; @@ -8692,6 +8960,7 @@ async fn create_workspace_worker( initial_submit, working_directory, control_operation_id: _, + resolved_control_operation, } = request; let config_state = api .config_store @@ -8771,6 +9040,7 @@ async fn create_workspace_worker( resolved_config_bundle, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), + resolved_control_operation, resolved_workspace_api: None, }; validate_ticket_assignment_spawn(&api, &runtime_id, &request)?; @@ -12322,6 +12592,7 @@ mod tests { resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), + resolved_control_operation: None, resolved_workspace_api: None, }; assert!( @@ -12351,6 +12622,7 @@ mod tests { resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), + resolved_control_operation: None, resolved_workspace_api: None, }; assert!( @@ -12474,6 +12746,7 @@ mod tests { resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), + resolved_control_operation: None, resolved_workspace_api: None, }; @@ -12520,6 +12793,7 @@ mod tests { }], working_directory: None, control_operation_id: None, + resolved_control_operation: None, }), ) .await @@ -12560,6 +12834,7 @@ mod tests { initial_submit: Vec::new(), working_directory: None, control_operation_id: None, + resolved_control_operation: None, }), ) .await @@ -12595,6 +12870,7 @@ mod tests { initial_submit: Vec::new(), working_directory: None, control_operation_id: None, + resolved_control_operation: None, }), ) .await @@ -12697,6 +12973,7 @@ mod tests { }], working_directory: None, control_operation_id: None, + resolved_control_operation: None, }), ) .await @@ -12978,6 +13255,104 @@ mod tests { assert!(retried.online); } + #[tokio::test] + async fn worker_control_spawn_retry_converges_on_one_worker_and_one_grant() { + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let api = test_api(workspace.path()).await; + let workspace_id = api.config.workspace_id.clone(); + let Json(controller_worker) = create_workspace_worker( + State(api.clone()), + HeaderMap::new(), + Json(CreateWorkspaceWorkerRequest { + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + display_name: "Control caller".to_string(), + profile: None, + ticket_assignment: None, + initial_submit: Vec::new(), + working_directory: None, + control_operation_id: None, + resolved_control_operation: None, + }), + ) + .await + .unwrap(); + let controller = controller_worker.worker_ref; + let mut headers = HeaderMap::new(); + headers.insert( + "x-yoi-runtime-id", + axum::http::HeaderValue::from_str(&controller.runtime_id).unwrap(), + ); + headers.insert( + "x-yoi-worker-id", + axum::http::HeaderValue::from_str(&controller.worker_id).unwrap(), + ); + let request = || CreateWorkspaceWorkerRequest { + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + display_name: "Idempotent controlled child".to_string(), + profile: None, + ticket_assignment: None, + initial_submit: Vec::new(), + working_directory: None, + control_operation_id: Some("control-spawn-retry".to_string()), + resolved_control_operation: None, + }; + + let Json(first) = spawn_known_worker( + State(api.clone()), + AxumPath(ScopedWorkspacePath { + workspace_id: workspace_id.clone(), + }), + headers.clone(), + Json(request()), + ) + .await + .unwrap(); + let Json(retried) = spawn_known_worker( + State(api.clone()), + AxumPath(ScopedWorkspacePath { + workspace_id: workspace_id.clone(), + }), + headers, + Json(request()), + ) + .await + .unwrap(); + + assert_eq!(retried.worker_ref, first.worker_ref); + let mut conflicting_request = request(); + conflicting_request.display_name = "Different controlled child".to_string(); + let conflict = spawn_known_worker( + State(api.clone()), + AxumPath(ScopedWorkspacePath { + workspace_id: workspace_id.clone(), + }), + { + let mut headers = HeaderMap::new(); + headers.insert( + "x-yoi-runtime-id", + axum::http::HeaderValue::from_str(&controller.runtime_id).unwrap(), + ); + headers.insert( + "x-yoi-worker-id", + axum::http::HeaderValue::from_str(&controller.worker_id).unwrap(), + ); + headers + }, + Json(conflicting_request), + ) + .await + .unwrap_err(); + assert_eq!(conflict.into_response().status(), StatusCode::BAD_GATEWAY); + let grants = api + .store + .list_active_worker_control_grants(&workspace_id, &controller, 10) + .unwrap(); + assert_eq!(grants.len(), 1); + assert_eq!(grants[0].subject, first.worker_ref); + assert_eq!(grants[0].operation_id, "control-spawn-retry"); + } + #[tokio::test] async fn explicit_orchestrator_launch_marks_only_the_dedicated_worker() { let workspace = tempfile::tempdir().unwrap(); @@ -12996,6 +13371,7 @@ mod tests { initial_submit: Vec::new(), working_directory: None, control_operation_id: None, + resolved_control_operation: None, }), ) .await @@ -13013,6 +13389,7 @@ mod tests { initial_submit: Vec::new(), working_directory: None, control_operation_id: None, + resolved_control_operation: None, }), ) .await @@ -13069,6 +13446,7 @@ mod tests { .await .unwrap(); assert_eq!(known.items.len(), 1); + assert_eq!(known.items[0].grant_id, "orchestrator-controls-generic"); assert_eq!(known.items[0].subject, generic.worker_ref); assert_eq!(known.items[0].permissions, ["observe"]); @@ -13097,7 +13475,7 @@ mod tests { AxumPath(ScopedWorkspacePath { workspace_id: workspace_id.clone(), }), - observation_headers, + observation_headers.clone(), Json(WorkerObservationSubjectRef::RuntimeWorker { runtime_id: generic.worker_ref.runtime_id.clone(), worker_id: generic.worker_ref.worker_id.clone(), @@ -13107,6 +13485,35 @@ mod tests { .unwrap(); assert!(capture["entries"].is_array()); + let Json(revoked) = revoke_worker_control_grant( + State(api.clone()), + AxumPath(ScopedWorkerControlGrantPath { + workspace_id: workspace_id.clone(), + grant_id: "orchestrator-controls-generic".to_string(), + }), + observation_headers.clone(), + ) + .await + .unwrap(); + assert!(revoked.revoked_at.is_some()); + let revoked_capture = scoped_capture_worker_observation_session( + State(api.clone()), + AxumPath(ScopedWorkspacePath { + workspace_id: workspace_id.clone(), + }), + observation_headers.clone(), + Json(WorkerObservationSubjectRef::RuntimeWorker { + runtime_id: generic.worker_ref.runtime_id.clone(), + worker_id: generic.worker_ref.worker_id.clone(), + }), + ) + .await + .unwrap_err(); + assert_eq!( + revoked_capture.into_response().status(), + StatusCode::NOT_FOUND + ); + let mut unauthorized_headers = HeaderMap::new(); unauthorized_headers.insert( "x-yoi-runtime-id", @@ -13127,6 +13534,80 @@ mod tests { .unwrap(); assert!(unauthorized["sessions"].as_array().unwrap().is_empty()); + for (grant_id, permission) in [ + ("orchestrator-share-source", "share"), + ("orchestrator-transfer-source", "transfer"), + ] { + api.store + .create_worker_control_grant(&WorkerControlGrantRecord { + workspace_id: workspace_id.clone(), + grant_id: grant_id.to_string(), + controller: dedicated.worker.clone(), + subject: generic.worker_ref.clone(), + relation: "spawned".to_string(), + origin: "test-delegation".to_string(), + permissions: vec!["observe".to_string(), permission.to_string()], + operation_id: format!("seed-{permission}"), + created_at: "2026-07-27T00:00:01Z".to_string(), + revoked_at: None, + }) + .unwrap(); + } + let target_controller = generic.worker_ref.clone(); + let Json(shared) = share_worker_control_grant( + State(api.clone()), + AxumPath(ScopedWorkerControlGrantPath { + workspace_id: workspace_id.clone(), + grant_id: "orchestrator-share-source".to_string(), + }), + observation_headers.clone(), + Json(DelegateWorkerControlGrantRequest { + target_controller: target_controller.clone(), + operation_id: "share-operation".to_string(), + }), + ) + .await + .unwrap(); + assert_eq!(shared.controller, target_controller); + assert_eq!(shared.relation, "shared"); + let Json(transferred) = transfer_worker_control_grant( + State(api.clone()), + AxumPath(ScopedWorkerControlGrantPath { + workspace_id: workspace_id.clone(), + grant_id: "orchestrator-transfer-source".to_string(), + }), + observation_headers.clone(), + Json(DelegateWorkerControlGrantRequest { + target_controller: generic.worker_ref.clone(), + operation_id: "transfer-operation".to_string(), + }), + ) + .await + .unwrap(); + let Json(transfer_replay) = transfer_worker_control_grant( + State(api.clone()), + AxumPath(ScopedWorkerControlGrantPath { + workspace_id: workspace_id.clone(), + grant_id: "orchestrator-transfer-source".to_string(), + }), + observation_headers, + Json(DelegateWorkerControlGrantRequest { + target_controller: generic.worker_ref.clone(), + operation_id: "transfer-operation".to_string(), + }), + ) + .await + .unwrap(); + assert_eq!(transfer_replay.grant_id, transferred.grant_id); + assert!( + api.store + .get_worker_control_grant(&workspace_id, "orchestrator-transfer-source") + .unwrap() + .unwrap() + .revoked_at + .is_some() + ); + let Json(existing) = scoped_start_workspace_orchestrator( State(api.clone()), AxumPath(ScopedWorkspacePath { @@ -13656,6 +14137,7 @@ mod tests { resolved_workspace_api: Some(test_worker_workspace_api( EMBEDDED_WORKER_RUNTIME_ID, )), + resolved_control_operation: None, }, ) .unwrap(); @@ -13804,6 +14286,7 @@ mod tests { resolved_workspace_api: Some(test_worker_workspace_api( EMBEDDED_WORKER_RUNTIME_ID, )), + resolved_control_operation: None, }, ) .unwrap() @@ -14003,6 +14486,7 @@ mod tests { resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), resolved_workspace_api: Some(test_worker_workspace_api(EMBEDDED_WORKER_RUNTIME_ID)), + resolved_control_operation: None, }; let source_worker = api .runtime @@ -14221,6 +14705,7 @@ mod tests { resolved_workspace_api: Some(test_worker_workspace_api( EMBEDDED_WORKER_RUNTIME_ID, )), + resolved_control_operation: None, }, ) .unwrap() @@ -14362,6 +14847,7 @@ mod tests { resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), + resolved_control_operation: None, resolved_workspace_api: None, }; let Json(first) = scoped_create_runtime_worker( @@ -14508,6 +14994,7 @@ mod tests { ticket_id: second_ticket.id.clone(), operation_id: "pending-spawn-operation".to_string(), }), + resolved_control_operation: None, ..request }; pending_request.resolved_workspace_api = @@ -14602,6 +15089,7 @@ mod tests { resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), + resolved_control_operation: None, resolved_workspace_api: None, }; let Json(created) = scoped_create_runtime_worker( @@ -15159,6 +15647,7 @@ mod tests { resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), resolved_workspace_api: None, + resolved_control_operation: None, }, ) .unwrap(); @@ -15239,6 +15728,7 @@ mod tests { resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), resolved_workspace_api: None, + resolved_control_operation: None, }, ) .unwrap(); @@ -17276,6 +17766,7 @@ mod tests { resolved_workspace_api: Some(test_worker_workspace_api( "embedded-worker-runtime", )), + resolved_control_operation: None, }, ) .expect("spawn worker"); @@ -17792,6 +18283,7 @@ mod tests { resolved_config_bundle: Some(runtime_test_bundle()), resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), + resolved_control_operation: None, resolved_workspace_api: None, }; let spawned = api diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index ed298ead..f9dd91cd 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -758,6 +758,11 @@ pub trait ControlPlaneStore: Send + Sync { &self, record: &WorkerControlGrantRecord, ) -> Result; + fn get_worker_control_grant( + &self, + workspace_id: &str, + grant_id: &str, + ) -> Result>; fn get_worker_control_grant_by_operation( &self, workspace_id: &str, @@ -2429,6 +2434,27 @@ impl ControlPlaneStore for SqliteWorkspaceStore { }) } + fn get_worker_control_grant( + &self, + workspace_id: &str, + grant_id: &str, + ) -> Result> { + self.with_conn(|conn| { + conn.query_row( + r#"SELECT workspace_id, grant_id, + controller_runtime_id, controller_worker_id, + subject_runtime_id, subject_worker_id, + relation, origin, permissions_json, operation_id, created_at, revoked_at + FROM worker_control_grants + WHERE workspace_id = ?1 AND grant_id = ?2"#, + params![workspace_id, grant_id], + read_worker_control_grant_record, + ) + .optional() + .map_err(Error::from) + }) + } + fn get_worker_control_grant_by_operation( &self, workspace_id: &str, @@ -6978,7 +7004,9 @@ CREATE TABLE ticket_assignment_operations ( #[tokio::test] async fn worker_control_grants_are_idempotent_scoped_and_revocable() { - let store = SqliteWorkspaceStore::in_memory().unwrap(); + let dir = tempfile::tempdir().unwrap(); + let database = dir.path().join("control-grants.db"); + let store = SqliteWorkspaceStore::open(&database).unwrap(); store .upsert_workspace(&WorkspaceRecord { workspace_id: "workspace-control".to_string(), @@ -7047,6 +7075,20 @@ CREATE TABLE ticket_assignment_operations ( Some(grant.clone()) ); + drop(store); + let store = SqliteWorkspaceStore::open(&database).unwrap(); + assert_eq!( + store + .list_active_worker_control_grants( + "workspace-control", + &controller_record.worker, + 10, + ) + .unwrap(), + vec![grant.clone()], + "known Runtime Worker grants survive Backend restart" + ); + let conflicting_replay = WorkerControlGrantRecord { subject: controller_record.worker.clone(), ..grant.clone() @@ -7074,6 +7116,18 @@ CREATE TABLE ticket_assignment_operations ( .unwrap() .is_empty() ); + assert!( + store + .delete_worker_registry("workspace-control", &subject_record.worker) + .unwrap() + ); + assert!( + store + .get_worker_control_grant("workspace-control", &grant.grant_id) + .unwrap() + .is_none(), + "deleting a subject Worker cascades its durable control grants" + ); } #[tokio::test] diff --git a/resources/profiles/base.dcdl b/resources/profiles/base.dcdl index f8ab343c..04148c7b 100644 --- a/resources/profiles/base.dcdl +++ b/resources/profiles/base.dcdl @@ -26,8 +26,8 @@ feature = { memory = { enabled = true; }; web = { enabled = true; }; image = { enabled = true; }; - sub_worker = { enabled = true; }; - worker = { enabled = true; }; + sub_worker = { enabled = false; }; + worker = { enabled = false; }; objective = { enabled = true; }; ticket = { enabled = true; authoring = true; thread = true; }; }; diff --git a/resources/profiles/coder.dcdl b/resources/profiles/coder.dcdl index 74f3795b..27a85776 100644 --- a/resources/profiles/coder.dcdl +++ b/resources/profiles/coder.dcdl @@ -10,7 +10,7 @@ import "./base.dcdl" // { web = { enabled = true; }; sub_worker = { enabled = true; }; flow = { enabled = true; }; - worker = { enabled = false; }; + worker = { enabled = true; }; ticket = { enabled = true; thread = true; }; }; } From 97828f8bd58bfc2227e36db53e0370d5d2ae5b5a Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 17 Aug 2026 00:38:48 +0900 Subject: [PATCH 3/4] server: fence Worker control operation identity --- crates/workspace-server/src/server.rs | 266 ++++++++++++++++++++++---- crates/workspace-server/src/store.rs | 216 ++++++++++++++++++++- 2 files changed, 434 insertions(+), 48 deletions(-) diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 5e856c87..570508da 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -103,8 +103,9 @@ use crate::skills; use crate::store::{ AccountRecord, ApiTokenRecord, AuthChallengeRecord, BrowserSessionRecord, ControlPlaneStore, DeviceLoginFlowRecord, FlowSourceRecord, PasskeyCredentialRecord, RepositoryRecord, - TicketWorkerAssignmentRecord, UserRecord, WorkdirRegistryRecord, WorkerControlGrantRecord, - WorkerRegistryRecord, WorkerWorkdirLinkRecord, WorkspaceRecord, + TicketWorkerAssignmentRecord, UserRecord, WorkdirRegistryRecord, + WorkerControlDelegationOperationRecord, WorkerControlGrantRecord, WorkerRegistryRecord, + WorkerWorkdirLinkRecord, WorkspaceRecord, }; use crate::{Error, Result}; use worker_runtime::catalog::{ @@ -5966,6 +5967,13 @@ async fn list_known_workers( })) } +fn scoped_worker_control_operation_id(controller: &RuntimeWorkerRef, operation_id: &str) -> String { + format!( + "worker-control:{}:{}:{operation_id}", + controller.runtime_id, controller.worker_id + ) +} + async fn spawn_known_worker( State(api): State, AxumPath(path): AxumPath, @@ -5986,8 +5994,12 @@ async fn spawn_known_worker( .filter(|value| !value.is_empty()) .ok_or_else(|| Error::InvalidInput("control_operation_id is required".to_string()))? .to_string(); - let fingerprint_input = serde_json::to_vec(&request) - .map_err(|error| Error::InvalidInput(format!("invalid Worker spawn input: {error}")))?; + let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id); + let fingerprint_input = serde_json::to_vec(&serde_json::json!({ + "controller": &controller, + "request": &request, + })) + .map_err(|error| Error::InvalidInput(format!("invalid Worker spawn input: {error}")))?; let input_fingerprint = format!( "sha256:{}", Sha256::digest(&fingerprint_input) @@ -5996,10 +6008,9 @@ async fn spawn_known_worker( .collect::() ); request.resolved_control_operation = Some(WorkerControlOperation { - operation_id: operation_id.clone(), + operation_id: scoped_worker_control_operation_id(&controller, &operation_id), input_fingerprint, }); - let controller = RuntimeWorkerRef::new(&source.runtime_id, &source.worker_id); let response = create_workspace_worker(State(api.clone()), headers, Json(request)).await?; if let Err(error) = api .store @@ -6123,51 +6134,84 @@ async fn delegate_worker_control_grant( "target_controller must differ from the current Worker".to_string(), ))); } - api.store - .get_worker_registry(&path.workspace_id, &request.target_controller)? - .ok_or_else(|| Error::UnknownWorker { - worker: request.target_controller.clone(), - })?; let operation_id = request.operation_id.trim(); if operation_id.is_empty() || operation_id.len() > 200 { return Err(ApiError::from(Error::InvalidInput( "operation_id must contain 1..=200 bytes".to_string(), ))); } - let expected_origin = format!("worker_control_{permission}:{}", grant.grant_id); - if let Some(existing) = api.store.get_worker_control_grant_by_operation( - &path.workspace_id, - &request.target_controller, - operation_id, - )? { - if existing.subject == grant.subject && existing.origin == expected_origin { - if transfer && grant.revoked_at.is_none() { - let lock = worker_control_lock(&api, &grant.grant_id); - let _guard = lock.lock().await; - let now = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); - if api - .store - .get_worker_control_grant(&path.workspace_id, &grant.grant_id)? - .is_some_and(|current| current.revoked_at.is_none()) - { - api.store.revoke_worker_control_grant( - &path.workspace_id, - &grant.grant_id, - &now, - )?; - } + let operation_input = serde_json::json!({ + "source_controller": &controller, + "source_grant_id": &grant.grant_id, + "action": permission, + "target_controller": &request.target_controller, + "subject": &grant.subject, + "permissions": &grant.permissions, + }); + let operation_bytes = serde_json::to_vec(&operation_input).map_err(|error| { + Error::InvalidInput(format!("invalid Worker delegation input: {error}")) + })?; + let input_fingerprint = format!( + "sha256:{}", + Sha256::digest(&operation_bytes) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::() + ); + let now = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + let operation = api.store.reserve_worker_control_delegation_operation( + &WorkerControlDelegationOperationRecord { + workspace_id: path.workspace_id.clone(), + source_controller: controller.clone(), + source_grant_id: grant.grant_id.clone(), + operation_id: operation_id.to_string(), + input_fingerprint, + delegated_grant_id: None, + created_at: now.clone(), + completed_at: None, + }, + )?; + if let Some(delegated_grant_id) = operation.delegated_grant_id.as_deref() { + let delegated = api + .store + .get_worker_control_grant(&path.workspace_id, delegated_grant_id)? + .ok_or_else(|| { + Error::Store("completed Worker delegation references a missing grant".to_string()) + })?; + if transfer && grant.revoked_at.is_none() { + let lock = worker_control_lock(&api, &grant.grant_id); + let _guard = lock.lock().await; + if api + .store + .get_worker_control_grant(&path.workspace_id, &grant.grant_id)? + .is_some_and(|current| current.revoked_at.is_none()) + { + api.store + .revoke_worker_control_grant(&path.workspace_id, &grant.grant_id, &now)?; } - return Ok(Json(existing)); } - return Err(ApiError::from(Error::InvalidInput( - "operation_id was already used with different grant input".to_string(), - ))); + return Ok(Json(delegated)); } if grant.revoked_at.is_some() { return Err(ApiError::from(Error::UnknownWorker { worker: grant.subject, })); } + api.store + .get_active_worker_control_grant( + &path.workspace_id, + &controller, + &request.target_controller, + )? + .ok_or_else(|| Error::UnknownWorker { + worker: request.target_controller.clone(), + })?; + api.store + .get_worker_registry(&path.workspace_id, &request.target_controller)? + .ok_or_else(|| Error::UnknownWorker { + worker: request.target_controller.clone(), + })?; + let lock = worker_control_lock(&api, &grant.grant_id); let _guard = lock.lock().await; let current = api @@ -6184,7 +6228,20 @@ async fn delegate_worker_control_grant( .ok_or_else(|| Error::UnknownWorker { worker: grant.subject.clone(), })?; - let now = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + api.store + .get_active_worker_control_grant( + &path.workspace_id, + &controller, + &request.target_controller, + )? + .ok_or_else(|| Error::UnknownWorker { + worker: request.target_controller.clone(), + })?; + let delegated_operation_id = format!( + "worker-control-delegate:{}:{}:{}:{}:{}", + controller.runtime_id, controller.worker_id, current.grant_id, permission, operation_id + ); + let expected_origin = format!("worker_control_{permission}:{}", current.grant_id); let delegated = api .store .create_worker_control_grant(&WorkerControlGrantRecord { @@ -6195,7 +6252,7 @@ async fn delegate_worker_control_grant( relation: if transfer { "transferred" } else { "shared" }.to_string(), origin: expected_origin, permissions: current.permissions.clone(), - operation_id: operation_id.to_string(), + operation_id: delegated_operation_id, created_at: now.clone(), revoked_at: None, })?; @@ -6208,6 +6265,13 @@ async fn delegate_worker_control_grant( worker: current.subject, })); } + api.store.complete_worker_control_delegation_operation( + &path.workspace_id, + &controller, + operation_id, + &delegated.grant_id, + &now, + )?; Ok(Json(delegated)) } @@ -13278,6 +13342,14 @@ mod tests { .await .unwrap(); let controller = controller_worker.worker_ref; + assert_ne!( + scoped_worker_control_operation_id(&controller, "same-operation"), + scoped_worker_control_operation_id( + &RuntimeWorkerRef::new(&controller.runtime_id, "different-controller"), + "same-operation", + ), + "Runtime idempotency keys are scoped to the authenticated controller" + ); let mut headers = HeaderMap::new(); headers.insert( "x-yoi-runtime-id", @@ -13570,6 +13642,122 @@ mod tests { .unwrap(); assert_eq!(shared.controller, target_controller); assert_eq!(shared.relation, "shared"); + + let alternate_target = RuntimeWorkerRef::new(EMBEDDED_WORKER_RUNTIME_ID, "1000"); + let now = now_registry_timestamp(); + api.store + .upsert_worker_registry(&WorkerRegistryRecord { + workspace_id: workspace_id.clone(), + worker: alternate_target.clone(), + display_name: "Alternate known target".to_string(), + profile: None, + retention_state: "normal".to_string(), + transcript_ref: None, + session_ref: None, + summary_ref: None, + diagnostics_ref: None, + created_at: now.clone(), + updated_at: now.clone(), + }) + .unwrap(); + let registry_only_target = RuntimeWorkerRef::new(EMBEDDED_WORKER_RUNTIME_ID, "1001"); + api.store + .upsert_worker_registry(&WorkerRegistryRecord { + workspace_id: workspace_id.clone(), + worker: registry_only_target.clone(), + display_name: "Registry-only target".to_string(), + profile: None, + retention_state: "normal".to_string(), + transcript_ref: None, + session_ref: None, + summary_ref: None, + diagnostics_ref: None, + created_at: now.clone(), + updated_at: now.clone(), + }) + .unwrap(); + let unknown_target = share_worker_control_grant( + State(api.clone()), + AxumPath(ScopedWorkerControlGrantPath { + workspace_id: workspace_id.clone(), + grant_id: "orchestrator-share-source".to_string(), + }), + observation_headers.clone(), + Json(DelegateWorkerControlGrantRequest { + target_controller: registry_only_target, + operation_id: "share-registry-only".to_string(), + }), + ) + .await + .unwrap_err(); + assert_eq!( + unknown_target.into_response().status(), + StatusCode::NOT_FOUND + ); + for (grant_id, operation_id, subject) in [ + ( + "orchestrator-knows-alternate", + "seed-known-alternate", + alternate_target.clone(), + ), + ( + "orchestrator-second-share-source", + "seed-second-share", + generic.worker_ref.clone(), + ), + ] { + api.store + .create_worker_control_grant(&WorkerControlGrantRecord { + workspace_id: workspace_id.clone(), + grant_id: grant_id.to_string(), + controller: dedicated.worker.clone(), + subject, + relation: "spawned".to_string(), + origin: "test-delegation-conflict".to_string(), + permissions: vec!["share".to_string()], + operation_id: operation_id.to_string(), + created_at: now.clone(), + revoked_at: None, + }) + .unwrap(); + } + let changed_target = share_worker_control_grant( + State(api.clone()), + AxumPath(ScopedWorkerControlGrantPath { + workspace_id: workspace_id.clone(), + grant_id: "orchestrator-share-source".to_string(), + }), + observation_headers.clone(), + Json(DelegateWorkerControlGrantRequest { + target_controller: alternate_target, + operation_id: "share-operation".to_string(), + }), + ) + .await + .unwrap_err(); + assert_eq!( + changed_target.into_response().status(), + StatusCode::BAD_REQUEST + ); + let changed_source_grant = share_worker_control_grant( + State(api.clone()), + AxumPath(ScopedWorkerControlGrantPath { + workspace_id: workspace_id.clone(), + grant_id: "orchestrator-second-share-source".to_string(), + }), + observation_headers.clone(), + Json(DelegateWorkerControlGrantRequest { + target_controller: generic.worker_ref.clone(), + operation_id: "share-operation".to_string(), + }), + ) + .await + .unwrap_err(); + assert_eq!( + changed_source_grant.into_response().status(), + StatusCode::BAD_REQUEST + ); + let Json(transferred) = transfer_worker_control_grant( State(api.clone()), AxumPath(ScopedWorkerControlGrantPath { diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index f9dd91cd..eb89c8f8 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -186,6 +186,11 @@ const MIGRATIONS: &[Migration] = &[ name: "create durable Runtime Worker control grants", apply: create_worker_control_grant_authority, }, + Migration { + version: 34, + name: "create Worker control delegation operation authority", + apply: create_worker_control_delegation_operation_authority, + }, ]; struct Migration { @@ -347,6 +352,18 @@ pub struct WorkerControlGrantRecord { pub revoked_at: Option, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct WorkerControlDelegationOperationRecord { + pub workspace_id: String, + pub source_controller: RuntimeWorkerRef, + pub source_grant_id: String, + pub operation_id: String, + pub input_fingerprint: String, + pub delegated_grant_id: Option, + pub created_at: String, + pub completed_at: Option, +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct TicketWorkerAssignmentRecord { pub workspace_id: String, @@ -787,6 +804,18 @@ pub trait ControlPlaneStore: Send + Sync { grant_id: &str, revoked_at: &str, ) -> Result; + fn reserve_worker_control_delegation_operation( + &self, + record: &WorkerControlDelegationOperationRecord, + ) -> Result; + fn complete_worker_control_delegation_operation( + &self, + workspace_id: &str, + source_controller: &RuntimeWorkerRef, + operation_id: &str, + delegated_grant_id: &str, + completed_at: &str, + ) -> Result; fn get_ticket_assignment_operation( &self, @@ -2549,6 +2578,95 @@ impl ControlPlaneStore for SqliteWorkspaceStore { }) } + fn reserve_worker_control_delegation_operation( + &self, + record: &WorkerControlDelegationOperationRecord, + ) -> Result { + self.with_conn(|conn| { + conn.execute( + r#"INSERT INTO worker_control_delegation_operations ( + workspace_id, source_controller_runtime_id, source_controller_worker_id, + source_grant_id, operation_id, input_fingerprint, + delegated_grant_id, created_at, completed_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9) + ON CONFLICT ( + workspace_id, source_controller_runtime_id, + source_controller_worker_id, operation_id + ) DO NOTHING"#, + params![ + record.workspace_id, + record.source_controller.runtime_id, + record.source_controller.worker_id, + record.source_grant_id, + record.operation_id, + record.input_fingerprint, + record.delegated_grant_id, + record.created_at, + record.completed_at, + ], + )?; + let persisted = read_worker_control_delegation_operation_by_key( + conn, + &record.workspace_id, + &record.source_controller, + &record.operation_id, + )? + .ok_or_else(|| { + Error::Store("worker control delegation operation was not persisted".to_string()) + })?; + if persisted.source_grant_id != record.source_grant_id + || persisted.input_fingerprint != record.input_fingerprint + { + return Err(Error::InvalidInput(format!( + "worker control delegation operation `{}` was already used with different input", + record.operation_id + ))); + } + Ok(persisted) + }) + } + + fn complete_worker_control_delegation_operation( + &self, + workspace_id: &str, + source_controller: &RuntimeWorkerRef, + operation_id: &str, + delegated_grant_id: &str, + completed_at: &str, + ) -> Result { + self.with_conn(|conn| { + conn.execute( + r#"UPDATE worker_control_delegation_operations + SET delegated_grant_id = ?5, completed_at = ?6 + WHERE workspace_id = ?1 + AND source_controller_runtime_id = ?2 + AND source_controller_worker_id = ?3 + AND operation_id = ?4 + AND (delegated_grant_id IS NULL OR delegated_grant_id = ?5)"#, + params![ + workspace_id, + source_controller.runtime_id, + source_controller.worker_id, + operation_id, + delegated_grant_id, + completed_at, + ], + )?; + read_worker_control_delegation_operation_by_key( + conn, + workspace_id, + source_controller, + operation_id, + )? + .filter(|record| record.delegated_grant_id.as_deref() == Some(delegated_grant_id)) + .ok_or_else(|| { + Error::InvalidInput(format!( + "worker control delegation operation `{operation_id}` completed with a different grant" + )) + }) + }) + } + fn get_ticket_assignment_operation( &self, workspace_id: &str, @@ -3906,6 +4024,52 @@ fn read_worker_control_grant_by_operation( .map_err(Error::from) } +fn read_worker_control_delegation_operation_record( + row: &rusqlite::Row<'_>, +) -> rusqlite::Result { + Ok(WorkerControlDelegationOperationRecord { + workspace_id: row.get(0)?, + source_controller: RuntimeWorkerRef::new( + row.get::<_, String>(1)?, + row.get::<_, u64>(2)?.to_string(), + ), + source_grant_id: row.get(3)?, + operation_id: row.get(4)?, + input_fingerprint: row.get(5)?, + delegated_grant_id: row.get(6)?, + created_at: row.get(7)?, + completed_at: row.get(8)?, + }) +} + +fn read_worker_control_delegation_operation_by_key( + conn: &Connection, + workspace_id: &str, + source_controller: &RuntimeWorkerRef, + operation_id: &str, +) -> Result> { + conn.query_row( + r#"SELECT workspace_id, + source_controller_runtime_id, source_controller_worker_id, + source_grant_id, operation_id, input_fingerprint, + delegated_grant_id, created_at, completed_at + FROM worker_control_delegation_operations + WHERE workspace_id = ?1 + AND source_controller_runtime_id = ?2 + AND source_controller_worker_id = ?3 + AND operation_id = ?4"#, + params![ + workspace_id, + source_controller.runtime_id, + source_controller.worker_id, + operation_id, + ], + read_worker_control_delegation_operation_record, + ) + .optional() + .map_err(Error::from) +} + fn current_ticket_worker_assignment_select_sql() -> String { "SELECT a.workspace_id, a.ticket_id, a.assignment_id, a.runtime_id, a.worker_id, \ a.assigned_by, a.assigned_at \ @@ -4970,6 +5134,40 @@ fn create_worker_control_grant_authority(conn: &Connection) -> Result<()> { Ok(()) } +fn create_worker_control_delegation_operation_authority(conn: &Connection) -> Result<()> { + conn.execute_batch( + r#" + CREATE TABLE worker_control_delegation_operations ( + workspace_id TEXT NOT NULL, + source_controller_runtime_id TEXT NOT NULL, + source_controller_worker_id INTEGER NOT NULL, + source_grant_id TEXT NOT NULL, + operation_id TEXT NOT NULL, + input_fingerprint TEXT NOT NULL, + delegated_grant_id TEXT, + created_at TEXT NOT NULL, + completed_at TEXT, + PRIMARY KEY ( + workspace_id, + source_controller_runtime_id, + source_controller_worker_id, + operation_id + ), + FOREIGN KEY (workspace_id, source_controller_runtime_id, source_controller_worker_id) + REFERENCES worker_registry (workspace_id, runtime_id, runtime_worker_id) + ON DELETE CASCADE, + FOREIGN KEY (workspace_id, source_grant_id) + REFERENCES worker_control_grants (workspace_id, grant_id) + ON DELETE CASCADE, + FOREIGN KEY (workspace_id, delegated_grant_id) + REFERENCES worker_control_grants (workspace_id, grant_id) + ON DELETE SET NULL + ); + "#, + )?; + Ok(()) +} + fn create_worker_mutation_source_proof_replay_guard(conn: &Connection) -> Result<()> { conn.execute_batch( r#" @@ -5643,7 +5841,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT); apply_migrations(&conn).unwrap(); - assert_eq!(current_schema_version(&conn).unwrap(), 33); + assert_eq!(current_schema_version(&conn).unwrap(), 34); assert!(table_exists(&conn, "worker_workdir_attachment_reservations").unwrap()); } @@ -5676,7 +5874,7 @@ CREATE TABLE flow_events (event_id TEXT PRIMARY KEY); apply_migrations(&conn).unwrap(); - assert_eq!(current_schema_version(&conn).unwrap(), 33); + assert_eq!(current_schema_version(&conn).unwrap(), 34); assert!(table_exists(&conn, "flow_sources").unwrap()); assert!(table_exists(&conn, "flow_source_revisions").unwrap()); assert!(!table_exists(&conn, "flow_instances").unwrap()); @@ -5743,7 +5941,7 @@ INSERT INTO worker_workdir_attachment_reservations ( apply_migrations(&conn).unwrap(); - assert_eq!(current_schema_version(&conn).unwrap(), 33); + assert_eq!(current_schema_version(&conn).unwrap(), 34); let repositories_sql: String = conn .query_row( "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'repositories'", @@ -5923,7 +6121,7 @@ INSERT INTO workdir_registry ( let db = dir.path().join("control-plane.sqlite"); let store = SqliteWorkspaceStore::open(&db).unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 33); + assert_eq!(store.schema_version().await.unwrap(), 34); assert!( !store .with_conn(|conn| table_exists(conn, "worker_workspace_credentials")) @@ -5940,7 +6138,7 @@ INSERT INTO workdir_registry ( store.upsert_workspace(&record).await.unwrap(); let reopened = SqliteWorkspaceStore::open(&db).unwrap(); - assert_eq!(reopened.schema_version().await.unwrap(), 33); + assert_eq!(reopened.schema_version().await.unwrap(), 34); assert_eq!( reopened.get_workspace("local-dev").await.unwrap(), Some(record) @@ -6487,7 +6685,7 @@ INSERT INTO workdir_registry ( .unwrap(); let store = SqliteWorkspaceStore::from_connection(conn).unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 33); + assert_eq!(store.schema_version().await.unwrap(), 34); store .with_conn(|conn| { @@ -6676,7 +6874,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(), 33); + assert_eq!(store.schema_version().await.unwrap(), 34); let workspace = WorkspaceRecord { workspace_id: "local-dev".to_string(), owner_account_id: None, @@ -6742,7 +6940,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(), 33); + assert_eq!(store.schema_version().await.unwrap(), 34); let workspace = WorkspaceRecord { workspace_id: "local-dev".to_string(), owner_account_id: None, @@ -7133,7 +7331,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(), 33); + assert_eq!(store.schema_version().await.unwrap(), 34); let now = "2026-07-22T00:00:00Z".to_string(); let account = AccountRecord { account_id: "acct-user-alice".to_string(), From eb0dd67d165b834787b74f6a58fd45aaa5a9607f Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 17 Aug 2026 00:48:37 +0900 Subject: [PATCH 4/4] server: recover completed Worker transfers --- crates/workspace-server/src/server.rs | 138 +++++++++++++++++++++----- 1 file changed, 113 insertions(+), 25 deletions(-) diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 570508da..efe0b2ae 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -6108,6 +6108,32 @@ async fn transfer_worker_control_grant( delegate_worker_control_grant(api, path, headers, request, true).await } +fn worker_control_delegation_input_fingerprint( + controller: &RuntimeWorkerRef, + grant: &WorkerControlGrantRecord, + action: &str, + target_controller: &RuntimeWorkerRef, +) -> Result { + let operation_input = serde_json::json!({ + "source_controller": controller, + "source_grant_id": &grant.grant_id, + "action": action, + "target_controller": target_controller, + "subject": &grant.subject, + "permissions": &grant.permissions, + }); + let operation_bytes = serde_json::to_vec(&operation_input).map_err(|error| { + Error::InvalidInput(format!("invalid Worker delegation input: {error}")) + })?; + Ok(format!( + "sha256:{}", + Sha256::digest(&operation_bytes) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::() + )) +} + async fn delegate_worker_control_grant( api: WorkspaceApi, path: ScopedWorkerControlGrantPath, @@ -6140,24 +6166,12 @@ async fn delegate_worker_control_grant( "operation_id must contain 1..=200 bytes".to_string(), ))); } - let operation_input = serde_json::json!({ - "source_controller": &controller, - "source_grant_id": &grant.grant_id, - "action": permission, - "target_controller": &request.target_controller, - "subject": &grant.subject, - "permissions": &grant.permissions, - }); - let operation_bytes = serde_json::to_vec(&operation_input).map_err(|error| { - Error::InvalidInput(format!("invalid Worker delegation input: {error}")) - })?; - let input_fingerprint = format!( - "sha256:{}", - Sha256::digest(&operation_bytes) - .iter() - .map(|byte| format!("{byte:02x}")) - .collect::() - ); + let input_fingerprint = worker_control_delegation_input_fingerprint( + &controller, + &grant, + permission, + &request.target_controller, + )?; let now = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); let operation = api.store.reserve_worker_control_delegation_operation( &WorkerControlDelegationOperationRecord { @@ -6256,6 +6270,13 @@ async fn delegate_worker_control_grant( created_at: now.clone(), revoked_at: None, })?; + api.store.complete_worker_control_delegation_operation( + &path.workspace_id, + &controller, + operation_id, + &delegated.grant_id, + &now, + )?; if transfer && !api .store @@ -6265,13 +6286,6 @@ async fn delegate_worker_control_grant( worker: current.subject, })); } - api.store.complete_worker_control_delegation_operation( - &path.workspace_id, - &controller, - operation_id, - &delegated.grant_id, - &now, - )?; Ok(Json(delegated)) } @@ -13758,6 +13772,71 @@ mod tests { StatusCode::BAD_REQUEST ); + let recovery_source = api + .store + .get_worker_control_grant(&workspace_id, "orchestrator-transfer-source") + .unwrap() + .unwrap(); + let recovery_operation_id = "transfer-operation"; + let recovery_fingerprint = worker_control_delegation_input_fingerprint( + &dedicated.worker, + &recovery_source, + "transfer", + &generic.worker_ref, + ) + .unwrap(); + let recovery_now = now_registry_timestamp(); + api.store + .reserve_worker_control_delegation_operation(&WorkerControlDelegationOperationRecord { + workspace_id: workspace_id.clone(), + source_controller: dedicated.worker.clone(), + source_grant_id: recovery_source.grant_id.clone(), + operation_id: recovery_operation_id.to_string(), + input_fingerprint: recovery_fingerprint, + delegated_grant_id: None, + created_at: recovery_now.clone(), + completed_at: None, + }) + .unwrap(); + let precompleted_transfer = api + .store + .create_worker_control_grant(&WorkerControlGrantRecord { + workspace_id: workspace_id.clone(), + grant_id: "precompleted-transfer-grant".to_string(), + controller: generic.worker_ref.clone(), + subject: recovery_source.subject.clone(), + relation: "transferred".to_string(), + origin: format!("worker_control_transfer:{}", recovery_source.grant_id), + permissions: recovery_source.permissions.clone(), + operation_id: format!( + "worker-control-delegate:{}:{}:{}:transfer:{}", + dedicated.worker.runtime_id, + dedicated.worker.worker_id, + recovery_source.grant_id, + recovery_operation_id, + ), + created_at: recovery_now.clone(), + revoked_at: None, + }) + .unwrap(); + api.store + .complete_worker_control_delegation_operation( + &workspace_id, + &dedicated.worker, + recovery_operation_id, + &precompleted_transfer.grant_id, + &recovery_now, + ) + .unwrap(); + assert!( + api.store + .get_worker_control_grant(&workspace_id, &recovery_source.grant_id) + .unwrap() + .unwrap() + .revoked_at + .is_none() + ); + let Json(transferred) = transfer_worker_control_grant( State(api.clone()), AxumPath(ScopedWorkerControlGrantPath { @@ -13772,6 +13851,15 @@ mod tests { ) .await .unwrap(); + assert_eq!(transferred.grant_id, precompleted_transfer.grant_id); + assert!( + api.store + .get_worker_control_grant(&workspace_id, &recovery_source.grant_id) + .unwrap() + .unwrap() + .revoked_at + .is_some() + ); let Json(transfer_replay) = transfer_worker_control_grant( State(api.clone()), AxumPath(ScopedWorkerControlGrantPath {