refactor: centralize worker control routing

crates/worker/src/feature/builtin/manage_worker.rs

@@ -28,16 +28,45 @@ 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";
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum WorkerControlSubject {
+    RuntimeWorker {
+        runtime_id: String,
+        worker_id: String,
+    },
+    SubWorker {
+        name: String,
+    },
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum WorkerControlInputKind {
+    User,
+    Notify,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum WorkerControlStopKind {
+    Cancel,
+    Stop,
+}
+
 #[async_trait]
 pub trait WorkerControlService: Send + Sync {
     fn workspace_id(&self) -> &str;
-    fn known_subworkers(&self) -> Vec<serde_json::Value>;
+    async fn list_workers(&self) -> Result<WorkspaceResponse, WorkspaceClientError>;
-    async fn send_subworker(
+    async fn send_input(
         &self,
-        name: &str,
+        subject: WorkerControlSubject,
         content: String,
+        kind: WorkerControlInputKind,
+    ) -> Result<WorkspaceResponse, WorkspaceClientError>;
+    async fn stop_worker(
+        &self,
+        subject: WorkerControlSubject,
+        kind: WorkerControlStopKind,
+        reason: Option<String>,
     ) -> Result<WorkspaceResponse, WorkspaceClientError>;
-    async fn stop_subworker(&self, name: &str) -> Result<WorkspaceResponse, WorkspaceClientError>;
     async fn spawn_worker(
         &self,
         request: WorkerLifecycleSpawnRequest,
@@ -66,22 +95,7 @@ struct WorkspaceWorkerControlService {
     runtime_worker_control: bool,
 }
 
-impl std::fmt::Debug for WorkspaceWorkerControlService {
+impl 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<serde_json::Value> {
         self.registry
             .as_ref()
@@ -105,51 +119,203 @@ impl WorkerControlService for WorkspaceWorkerControlService {
             })
             .unwrap_or_default()
     }
+}
 
-    async fn send_subworker(
+impl std::fmt::Debug for WorkspaceWorkerControlService {
-        &self,
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        name: &str,
+        formatter
-        content: String,
+            .debug_struct("WorkspaceWorkerControlService")
-    ) -> Result<WorkspaceResponse, WorkspaceClientError> {
+            .field("workspace_id", &self.workspace_id)
-        let record = self
+            .field("has_subworker_registry", &self.registry.is_some())
-            .registry
+            .finish_non_exhaustive()
-            .as_ref()
+    }
-            .and_then(|registry| registry.get_internal(name))
+}
+
+#[async_trait]
+impl WorkerControlService for WorkspaceWorkerControlService {
+    fn workspace_id(&self) -> &str {
+        &self.workspace_id
+    }
+
+    async fn list_workers(&self) -> Result<WorkspaceResponse, WorkspaceClientError> {
+        let response = if self.runtime_worker_control {
+            self.execute_runtime(WorkspaceRequest::get(format!(
+                "/api/w/{}/worker-control/workers",
+                self.workspace_id
+            )))
+            .await?
+        } else {
+            WorkspaceResponse {
+                status: 200,
+                body: serde_json::json!({ "items": [] }).to_string(),
+            }
+        };
+        if !response.is_success() {
+            return Ok(response);
+        }
+
+        let mut payload: serde_json::Value =
+            serde_json::from_str(&response.body).map_err(|error| {
+                WorkspaceClientError::Request(format!(
+                    "invalid Workspace control response: {error}"
+                ))
+            })?;
+        let items = payload
+            .get_mut("items")
+            .and_then(serde_json::Value::as_array_mut)
             .ok_or_else(|| {
                 WorkspaceClientError::Request(
-                    "unknown Worker or permission not granted".to_string(),
+                    "invalid Workspace control response: missing items".to_string(),
                 )
             })?;
-        record
+        let mut seen = std::collections::HashSet::new();
-            .session
+        items.retain(|item| {
-            .send(content)
+            item.get("subject")
-            .await
+                .map(serde_json::Value::to_string)
-            .map_err(|error| WorkspaceClientError::Request(error.to_string()))?;
+                .is_none_or(|subject| seen.insert(subject))
+        });
+        for worker in self.known_subworkers() {
+            let is_new = worker
+                .get("subject")
+                .map(serde_json::Value::to_string)
+                .is_none_or(|subject| seen.insert(subject));
+            if is_new {
+                items.push(worker);
+            }
+        }
         Ok(WorkspaceResponse {
-            status: 200,
+            status: response.status,
-            body: serde_json::json!({ "subject": { "kind": "sub_worker", "name": name } })
+            body: serde_json::to_string(&payload)
-                .to_string(),
+                .map_err(|error| WorkspaceClientError::Request(error.to_string()))?,
         })
     }
 
-    async fn stop_subworker(&self, name: &str) -> Result<WorkspaceResponse, WorkspaceClientError> {
+    async fn send_input(
-        let registry = self.registry.as_ref().ok_or_else(|| {
+        &self,
-            WorkspaceClientError::Request("unknown Worker or permission not granted".to_string())
+        subject: WorkerControlSubject,
-        })?;
+        content: String,
-        let summary = registry
+        kind: WorkerControlInputKind,
-            .remove_internal(name)
+    ) -> Result<WorkspaceResponse, WorkspaceClientError> {
-            .await
+        match subject {
-            .map_err(|error| WorkspaceClientError::Request(error.to_string()))?
+            WorkerControlSubject::RuntimeWorker {
-            .ok_or_else(|| {
+                runtime_id,
-                WorkspaceClientError::Request(
+                worker_id,
-                    "unknown Worker or permission not granted".to_string(),
+            } => {
+                let kind = match kind {
+                    WorkerControlInputKind::User => "user",
+                    WorkerControlInputKind::Notify => "notify",
+                };
+                self.execute_runtime(WorkspaceRequest::json(
+                    WorkspaceRequestMethod::Post,
+                    format!(
+                        "/api/w/{}/worker-control/workers/{runtime_id}/{worker_id}/input",
+                        self.workspace_id
+                    ),
+                    serde_json::json!({
+                        "kind": kind,
+                        "content": content,
+                    })
+                    .to_string(),
+                ))
+                .await
+            }
+            WorkerControlSubject::SubWorker { name } => {
+                if kind != WorkerControlInputKind::User {
+                    return Err(WorkspaceClientError::Request(
+                        "notify is not available for parent-owned subworkers".to_string(),
+                    ));
+                }
+                self.ensure_permission(
+                    &super::worker_observation::WorkerObservationSubjectRef::SubWorker {
+                        name: name.clone(),
+                    },
+                    "send_input",
                 )
-            })?;
+                .await?;
-        Ok(WorkspaceResponse {
+                let record = self
-            status: 200,
+                    .registry
-            body: serde_json::to_string(&summary)
+                    .as_ref()
-                .map_err(|error| WorkspaceClientError::Request(error.to_string()))?,
+                    .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_worker(
+        &self,
+        subject: WorkerControlSubject,
+        kind: WorkerControlStopKind,
+        reason: Option<String>,
+    ) -> Result<WorkspaceResponse, WorkspaceClientError> {
+        match subject {
+            WorkerControlSubject::RuntimeWorker {
+                runtime_id,
+                worker_id,
+            } => {
+                let action = match kind {
+                    WorkerControlStopKind::Cancel => "cancel",
+                    WorkerControlStopKind::Stop => "stop",
+                };
+                self.execute_runtime(WorkspaceRequest::json(
+                    WorkspaceRequestMethod::Post,
+                    format!(
+                        "/api/w/{}/worker-control/workers/{runtime_id}/{worker_id}/{action}",
+                        self.workspace_id
+                    ),
+                    serde_json::json!({ "reason": reason }).to_string(),
+                ))
+                .await
+            }
+            WorkerControlSubject::SubWorker { name } => {
+                if kind != WorkerControlStopKind::Stop {
+                    return Err(WorkspaceClientError::Request(
+                        "cancel is not available for parent-owned subworkers".to_string(),
+                    ));
+                }
+                self.ensure_permission(
+                    &super::worker_observation::WorkerObservationSubjectRef::SubWorker {
+                        name: name.clone(),
+                    },
+                    "stop",
+                )
+                .await?;
+                let registry = self.registry.as_ref().ok_or_else(|| {
+                    WorkspaceClientError::Request(
+                        "unknown Worker or permission not granted".to_string(),
+                    )
+                })?;
+                let summary = registry
+                    .remove_internal(&name)
+                    .await
+                    .map_err(|error| WorkspaceClientError::Request(error.to_string()))?
+                    .ok_or_else(|| {
+                        WorkspaceClientError::Request(
+                            "unknown Worker or permission not granted".to_string(),
+                        )
+                    })?;
+                Ok(WorkspaceResponse {
+                    status: 200,
+                    body: serde_json::to_string(&summary)
+                        .map_err(|error| WorkspaceClientError::Request(error.to_string()))?,
+                })
+            }
+        }
     }
 
     async fn spawn_worker(
@@ -345,10 +511,6 @@ impl std::fmt::Debug for ManageWorkerFeature {
     fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         formatter
             .debug_struct("ManageWorkerFeature")
-            .field(
-                "has_subworker_registry",
-                &!self.control.known_subworkers().is_empty(),
-            )
             .field("direct_spawn", &self.direct_spawn)
             .finish_non_exhaustive()
     }
@@ -700,21 +862,10 @@ impl Tool for WorkspaceWorkerTool {
         let response = match self.operation {
             WorkerOperation::List => {
                 parse::<WorkerListInput>(input_json, "WorkerList")?;
-                let response = if self.runtime_worker_control {
+                self.control
-                    self.control
+                    .list_workers()
-                        .execute_runtime(WorkspaceRequest::get(format!(
+                    .await
-                            "/api/w/{}/worker-control/workers",
+                    .map_err(control_tool_error)?
-                            self.control.workspace_id()
-                        )))
-                        .await
-                        .map_err(control_tool_error)?
-                } else {
-                    WorkspaceResponse {
-                        status: 200,
-                        body: r#"{"items":[]}"#.to_string(),
-                    }
-                };
-                self.with_subworkers(response)?
             }
             WorkerOperation::Spawn => {
                 let input = parse::<WorkerSpawnInput>(input_json, "WorkerSpawn")?;
@@ -771,34 +922,16 @@ impl Tool for WorkspaceWorkerTool {
                         "content must contain at most 16384 bytes".to_string(),
                     ));
                 }
-                match input.subject {
+                let subject = worker_control_subject(input.subject)?;
-                    WorkerSubjectInput::SubWorker { name } => {
+                let kind = if self.operation == WorkerOperation::Notify {
-                        if self.operation != WorkerOperation::SendInput {
+                    WorkerControlInputKind::Notify
-                            return Err(unsupported_subject(self.operation, "sub_worker"));
+                } else {
-                        }
+                    WorkerControlInputKind::User
-                        let subject = subworker_subject(&name)?;
+                };
-                        self.control
+                self.control
-                            .ensure_permission(&subject, "send_input")
+                    .send_input(subject, content, kind)
-                            .await
+                    .await
-                            .map_err(control_tool_error)?;
+                    .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 = if self.runtime_worker_control {
@@ -811,36 +944,16 @@ impl Tool for WorkspaceWorkerTool {
                         reason: input.reason,
                     }
                 };
-                match input.subject {
+                let subject = worker_control_subject(input.subject)?;
-                    WorkerSubjectInput::SubWorker { name } => {
+                let kind = if self.operation == WorkerOperation::Cancel {
-                        if self.operation != WorkerOperation::Stop {
+                    WorkerControlStopKind::Cancel
-                            return Err(unsupported_subject(self.operation, "sub_worker"));
+                } else {
-                        }
+                    WorkerControlStopKind::Stop
-                        let subject = subworker_subject(&name)?;
+                };
-                        self.control
+                self.control
-                            .ensure_permission(&subject, "stop")
+                    .stop_worker(subject, kind, input.reason)
-                            .await
+                    .await
-                            .map_err(control_tool_error)?;
+                    .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::<WorkerTargetInput>(input_json, "WorkerRestore")?;
@@ -875,31 +988,18 @@ impl Tool for WorkspaceWorkerTool {
     }
 }
 
-impl WorkspaceWorkerTool {
+fn worker_control_subject(subject: WorkerSubjectInput) -> Result<WorkerControlSubject, ToolError> {
-    fn with_subworkers(
+    match subject {
-        &self,
+        WorkerSubjectInput::RuntimeWorker {
-        mut response: WorkspaceResponse,
+            runtime_id,
-    ) -> Result<WorkspaceResponse, ToolError> {
+            worker_id,
-        if !response.is_success() {
+        } => Ok(WorkerControlSubject::RuntimeWorker {
-            return Ok(response);
+            runtime_id: authority_id(&runtime_id, "runtime_id")?,
-        }
+            worker_id: authority_id(&worker_id, "worker_id")?,
-        let mut body: serde_json::Value =
+        }),
-            serde_json::from_str(&response.body).map_err(|error| {
+        WorkerSubjectInput::SubWorker { name } => Ok(WorkerControlSubject::SubWorker {
-                ToolError::ExecutionFailed(format!("WorkerList returned invalid JSON: {error}"))
+            name: authority_id(&name, "name")?,
-            })?;
+        }),
-        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(),
-                )
-            })?;
-        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}"))
-        })?;
-        Ok(response)
     }
 }
 
@@ -919,16 +1019,6 @@ fn runtime_subject_ids(
     }
 }
 
-fn subworker_subject(
-    name: &str,
-) -> Result<super::worker_observation::WorkerObservationSubjectRef, ToolError> {
-    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}'",
@@ -1137,10 +1227,35 @@ mod tests {
             &self,
             request: WorkspaceRequest,
         ) -> Result<WorkspaceResponse, WorkspaceClientError> {
+            let is_worker_list = request.path.ends_with("/worker-control/workers");
             self.requests.lock().unwrap().push(request);
             Ok(WorkspaceResponse {
                 status: 200,
-                body: "{}".to_string(),
+                body: if is_worker_list {
+                    serde_json::json!({
+                        "items": [
+                            {
+                                "subject": {
+                                    "kind": "runtime_worker",
+                                    "runtime_id": "runtime-a",
+                                    "worker_id": "worker-a"
+                                },
+                                "summary": { "status": "idle" }
+                            },
+                            {
+                                "subject": {
+                                    "kind": "runtime_worker",
+                                    "runtime_id": "runtime-a",
+                                    "worker_id": "worker-a"
+                                },
+                                "summary": { "status": "idle" }
+                            }
+                        ]
+                    })
+                    .to_string()
+                } else {
+                    "{}".to_string()
+                },
             })
         }
 
@@ -1186,8 +1301,55 @@ mod tests {
             "workspace-test"
         }
 
-        fn known_subworkers(&self) -> Vec<serde_json::Value> {
+        async fn list_workers(&self) -> Result<WorkspaceResponse, WorkspaceClientError> {
-            Vec::new()
+            Ok(WorkspaceResponse {
+                status: 200,
+                body: serde_json::json!({ "items": [] }).to_string(),
+            })
+        }
+
+        async fn send_input(
+            &self,
+            subject: WorkerControlSubject,
+            content: String,
+            kind: WorkerControlInputKind,
+        ) -> Result<WorkspaceResponse, WorkspaceClientError> {
+            let WorkerControlSubject::SubWorker { name } = subject else {
+                return Err(WorkspaceClientError::Request(
+                    "expected subworker subject".to_string(),
+                ));
+            };
+            if kind != WorkerControlInputKind::User {
+                return Err(WorkspaceClientError::Request(
+                    "expected user input".to_string(),
+                ));
+            }
+            self.sent.lock().unwrap().push((name, content));
+            Ok(WorkspaceResponse {
+                status: 200,
+                body: serde_json::json!({ "status": "accepted" }).to_string(),
+            })
+        }
+
+        async fn stop_worker(
+            &self,
+            subject: WorkerControlSubject,
+            kind: WorkerControlStopKind,
+            _reason: Option<String>,
+        ) -> Result<WorkspaceResponse, WorkspaceClientError> {
+            let WorkerControlSubject::SubWorker { name } = subject else {
+                return Err(WorkspaceClientError::Request(
+                    "expected subworker subject".to_string(),
+                ));
+            };
+            if kind != WorkerControlStopKind::Stop {
+                return Err(WorkspaceClientError::Request("expected stop".to_string()));
+            }
+            self.stopped.lock().unwrap().push(name);
+            Ok(WorkspaceResponse {
+                status: 200,
+                body: serde_json::json!({ "status": "stopped" }).to_string(),
+            })
         }
 
         async fn spawn_worker(
@@ -1226,29 +1388,6 @@ mod tests {
         ) -> Result<(), WorkspaceClientError> {
             Ok(())
         }
-
-        async fn send_subworker(
-            &self,
-            name: &str,
-            content: String,
-        ) -> Result<WorkspaceResponse, WorkspaceClientError> {
-            self.sent.lock().unwrap().push((name.to_string(), content));
-            Ok(WorkspaceResponse {
-                status: 200,
-                body: serde_json::json!({ "status": "accepted" }).to_string(),
-            })
-        }
-
-        async fn stop_subworker(
-            &self,
-            name: &str,
-        ) -> Result<WorkspaceResponse, WorkspaceClientError> {
-            self.stopped.lock().unwrap().push(name.to_string());
-            Ok(WorkspaceResponse {
-                status: 200,
-                body: serde_json::json!({ "status": "stopped" }).to_string(),
-            })
-        }
     }
 
     fn test_control(client: Arc<RecordingWorkspaceClient>) -> Arc<dyn WorkerControlService> {
@@ -1345,6 +1484,73 @@ mod tests {
         assert!(report.services.providers().is_empty());
     }
 
+    #[tokio::test]
+    async fn control_service_deduplicates_worker_list_by_stable_subject() {
+        let client = Arc::new(RecordingWorkspaceClient::default());
+        let control = WorkspaceWorkerControlService {
+            client: client.clone(),
+            workspace_id: "workspace%2Ftest".to_string(),
+            registry: None,
+            runtime_worker_control: true,
+        };
+
+        let response = control.list_workers().await.unwrap();
+
+        let value: serde_json::Value = serde_json::from_str(&response.body).unwrap();
+        assert_eq!(value["items"].as_array().unwrap().len(), 1);
+        assert_eq!(
+            value["items"][0]["subject"],
+            serde_json::json!({
+                "kind": "runtime_worker",
+                "runtime_id": "runtime-a",
+                "worker_id": "worker-a"
+            })
+        );
+        assert_eq!(client.requests.lock().unwrap().len(), 1);
+    }
+
+    #[tokio::test]
+    async fn control_service_routes_by_subject_without_cross_provider_fallback() {
+        let client = Arc::new(RecordingWorkspaceClient::default());
+        let control = WorkspaceWorkerControlService {
+            client: client.clone(),
+            workspace_id: "workspace%2Ftest".to_string(),
+            registry: None,
+            runtime_worker_control: true,
+        };
+
+        control
+            .send_input(
+                WorkerControlSubject::RuntimeWorker {
+                    runtime_id: "runtime-1".to_string(),
+                    worker_id: "worker-1".to_string(),
+                },
+                "continue".to_string(),
+                WorkerControlInputKind::User,
+            )
+            .await
+            .unwrap();
+        let runtime_request_count = client.requests.lock().unwrap().len();
+        assert_eq!(runtime_request_count, 1);
+
+        let error = control
+            .send_input(
+                WorkerControlSubject::SubWorker {
+                    name: "missing-reviewer".to_string(),
+                },
+                "continue".to_string(),
+                WorkerControlInputKind::User,
+            )
+            .await
+            .unwrap_err();
+        assert!(matches!(error, WorkspaceClientError::Request(_)));
+        assert_eq!(
+            client.requests.lock().unwrap().len(),
+            runtime_request_count,
+            "a subworker failure must not fall through to Workspace control"
+        );
+    }
+
     #[tokio::test]
     async fn sub_worker_control_surface_lists_registry_without_workspace_authority() {
         let client = Arc::new(RecordingWorkspaceClient::default());