fix: expose canonical SubWorker control tools

This commit is contained in:
2026-08-30 20:25:17 +09:00
parent a9aa09636f
commit 44ba5fd6d4
9 changed files with 590 additions and 301 deletions
+4 -2
View File
@@ -884,8 +884,10 @@ where
.register_tools(tools::web_builtin_tools(web_config)); .register_tools(tools::web_builtin_tools(web_config));
} }
let worker_enabled = feature_config.worker.enabled;
let sub_worker_enabled = feature_config.sub_worker.enabled;
let mut feature_registry = FeatureRegistryBuilder::new(); let mut feature_registry = FeatureRegistryBuilder::new();
if feature_config.sub_worker.enabled { if sub_worker_enabled && !worker_enabled {
feature_registry.add_module( feature_registry.add_module(
crate::feature::builtin::manage_worker::sub_worker_control_feature( crate::feature::builtin::manage_worker::sub_worker_control_feature(
worker.workspace_client_handle(), worker.workspace_client_handle(),
@@ -970,7 +972,7 @@ where
feature_registry.add_module( feature_registry.add_module(
crate::feature::builtin::manage_worker::manage_worker_feature( crate::feature::builtin::manage_worker::manage_worker_feature(
workspace_client, workspace_client,
Some(spawned_registry.clone()), sub_worker_enabled.then(|| spawned_registry.clone()),
feature_config.worker.direct_spawn, feature_config.worker.direct_spawn,
), ),
); );
@@ -63,6 +63,7 @@ struct WorkspaceWorkerControlService {
client: Arc<dyn WorkspaceClient>, client: Arc<dyn WorkspaceClient>,
workspace_id: String, workspace_id: String,
registry: Option<Arc<SpawnedWorkerRegistry>>, registry: Option<Arc<SpawnedWorkerRegistry>>,
runtime_worker_control: bool,
} }
impl std::fmt::Debug for WorkspaceWorkerControlService { impl std::fmt::Debug for WorkspaceWorkerControlService {
@@ -169,6 +170,11 @@ impl WorkerControlService for WorkspaceWorkerControlService {
worker_id: &str, worker_id: &str,
reason: &str, reason: &str,
) -> Result<WorkspaceResponse, WorkspaceClientError> { ) -> Result<WorkspaceResponse, WorkspaceClientError> {
if !self.runtime_worker_control {
return Err(WorkspaceClientError::Unavailable(
"Runtime Worker control is not enabled for this Worker".to_string(),
));
}
self.client self.client
.execute_worker_remove(runtime_id, worker_id, reason) .execute_worker_remove(runtime_id, worker_id, reason)
} }
@@ -177,6 +183,11 @@ impl WorkerControlService for WorkspaceWorkerControlService {
&self, &self,
request: WorkspaceRequest, request: WorkspaceRequest,
) -> Result<WorkspaceResponse, WorkspaceClientError> { ) -> Result<WorkspaceResponse, WorkspaceClientError> {
if !self.runtime_worker_control {
return Err(WorkspaceClientError::Request(
"Runtime Worker control is not enabled for this Worker".to_string(),
));
}
self.client.execute(request) self.client.execute(request)
} }
@@ -204,6 +215,11 @@ impl WorkerControlService for WorkspaceWorkerControlService {
runtime_id, runtime_id,
worker_id, worker_id,
} => { } => {
if !self.runtime_worker_control {
return Err(WorkspaceClientError::Unavailable(
"Runtime Worker control is not enabled for this Worker".to_string(),
));
}
let response = self.client.execute(WorkspaceRequest::get(format!( let response = self.client.execute(WorkspaceRequest::get(format!(
"/api/w/{}/worker-control/workers", "/api/w/{}/worker-control/workers",
self.workspace_id self.workspace_id
@@ -348,6 +364,7 @@ pub fn manage_worker_feature(
client: client.clone(), client: client.clone(),
workspace_id, workspace_id,
registry, registry,
runtime_worker_control: true,
}); });
ManageWorkerFeature { ManageWorkerFeature {
client, client,
@@ -356,6 +373,12 @@ pub fn manage_worker_feature(
} }
} }
const SUB_WORKER_CONTROL_OPERATIONS: &[WorkerOperation] = &[
WorkerOperation::List,
WorkerOperation::SendInput,
WorkerOperation::Stop,
];
pub struct SubWorkerControlFeature { pub struct SubWorkerControlFeature {
client: Arc<dyn WorkspaceClient>, client: Arc<dyn WorkspaceClient>,
registry: Arc<SpawnedWorkerRegistry>, registry: Arc<SpawnedWorkerRegistry>,
@@ -369,13 +392,19 @@ impl SubWorkerControlFeature {
impl FeatureModule for SubWorkerControlFeature { impl FeatureModule for SubWorkerControlFeature {
fn descriptor(&self) -> FeatureDescriptor { fn descriptor(&self) -> FeatureDescriptor {
FeatureDescriptor::builtin("sub-worker-control", "SubWorker Control") let mut descriptor = FeatureDescriptor::builtin("sub-worker-control", "SubWorker Control")
.with_description("Parent-owned SubWorker control service provider") .with_description(
"Parent-owned SubWorker lifecycle through the canonical Worker control surface",
)
.with_provided_service(ServiceDeclaration::new( .with_provided_service(ServiceDeclaration::new(
ServiceId::builtin(WORKER_CONTROL_SERVICE_ID), ServiceId::builtin(WORKER_CONTROL_SERVICE_ID),
WORKER_LIFECYCLE_SERVICE_VERSION, WORKER_LIFECYCLE_SERVICE_VERSION,
"Parent-owned SubWorker discovery and control operations", "Parent-owned SubWorker discovery and control operations",
)) ));
for operation in SUB_WORKER_CONTROL_OPERATIONS {
descriptor = descriptor.with_tool(sub_worker_tool_declaration(*operation));
}
descriptor
} }
fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> { fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> {
@@ -383,6 +412,7 @@ impl FeatureModule for SubWorkerControlFeature {
workspace_id: self.client.workspace_id().unwrap_or_default().to_string(), workspace_id: self.client.workspace_id().unwrap_or_default().to_string(),
client: self.client.clone(), client: self.client.clone(),
registry: Some(self.registry.clone()), registry: Some(self.registry.clone()),
runtime_worker_control: false,
}); });
context.services().provide( context.services().provide(
ServiceDeclaration::new( ServiceDeclaration::new(
@@ -390,8 +420,16 @@ impl FeatureModule for SubWorkerControlFeature {
WORKER_LIFECYCLE_SERVICE_VERSION, WORKER_LIFECYCLE_SERVICE_VERSION,
"Parent-owned SubWorker discovery and control operations", "Parent-owned SubWorker discovery and control operations",
), ),
control, control.clone(),
) )?;
for operation in SUB_WORKER_CONTROL_OPERATIONS {
context.tools().register(worker_tool_contribution(
*operation,
control.clone(),
false,
))?;
}
Ok(())
} }
} }
@@ -418,10 +456,7 @@ impl FeatureModule for ManageWorkerFeature {
)); ));
for operation in WorkerOperation::ALL { for operation in WorkerOperation::ALL {
if operation != WorkerOperation::Spawn || self.direct_spawn { if operation != WorkerOperation::Spawn || self.direct_spawn {
descriptor = descriptor.with_tool(ToolDeclaration::new( descriptor = descriptor.with_tool(worker_tool_declaration(operation));
operation.tool_name(),
operation.description(),
));
} }
} }
descriptor descriptor
@@ -463,29 +498,11 @@ impl FeatureModule for ManageWorkerFeature {
if operation == WorkerOperation::Spawn && !self.direct_spawn { if operation == WorkerOperation::Spawn && !self.direct_spawn {
continue; continue;
} }
let definition = match operation { context.tools().register(worker_tool_contribution(
WorkerOperation::List => { operation,
definition::<WorkerListInput>(operation, self.control.clone()) self.control.clone(),
} true,
WorkerOperation::Spawn => { ))?;
definition::<WorkerSpawnInput>(operation, self.control.clone())
}
WorkerOperation::SendInput | WorkerOperation::Notify => {
definition::<WorkerMessageInput>(operation, self.control.clone())
}
WorkerOperation::Cancel | WorkerOperation::Stop => {
definition::<WorkerStopInput>(operation, self.control.clone())
}
WorkerOperation::Restore => {
definition::<WorkerTargetInput>(operation, self.control.clone())
}
WorkerOperation::Remove => {
definition::<WorkerRemoveInput>(operation, self.control.clone())
}
};
context
.tools()
.register(ToolContribution::new(operation.tool_name(), definition))?;
} }
Ok(()) Ok(())
} }
@@ -558,6 +575,34 @@ struct WorkerTargetInput {
subject: WorkerSubjectInput, subject: WorkerSubjectInput,
} }
#[derive(Debug, Deserialize, JsonSchema)]
#[serde(tag = "kind", rename_all = "snake_case")]
enum DirectSubWorkerSubjectInput {
SubWorker { name: String },
}
impl From<DirectSubWorkerSubjectInput> for WorkerSubjectInput {
fn from(subject: DirectSubWorkerSubjectInput) -> Self {
match subject {
DirectSubWorkerSubjectInput::SubWorker { name } => Self::SubWorker { name },
}
}
}
#[derive(Debug, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
struct DirectSubWorkerMessageInput {
subject: DirectSubWorkerSubjectInput,
content: String,
}
#[derive(Debug, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
struct DirectSubWorkerStopInput {
subject: DirectSubWorkerSubjectInput,
reason: Option<String>,
}
#[derive(Debug, Deserialize, JsonSchema)] #[derive(Debug, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)] #[serde(deny_unknown_fields)]
struct WorkerMessageInput { struct WorkerMessageInput {
@@ -583,6 +628,7 @@ struct WorkerRemoveInput {
struct WorkspaceWorkerTool { struct WorkspaceWorkerTool {
operation: WorkerOperation, operation: WorkerOperation,
control: Arc<dyn WorkerControlService>, control: Arc<dyn WorkerControlService>,
runtime_worker_control: bool,
} }
#[derive(Debug, Clone, Copy, PartialEq, Eq)] #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -654,14 +700,20 @@ impl Tool for WorkspaceWorkerTool {
let response = match self.operation { let response = match self.operation {
WorkerOperation::List => { WorkerOperation::List => {
parse::<WorkerListInput>(input_json, "WorkerList")?; parse::<WorkerListInput>(input_json, "WorkerList")?;
let response = self let response = if self.runtime_worker_control {
.control self.control
.execute_runtime(WorkspaceRequest::get(format!( .execute_runtime(WorkspaceRequest::get(format!(
"/api/w/{}/worker-control/workers", "/api/w/{}/worker-control/workers",
self.control.workspace_id() self.control.workspace_id()
))) )))
.await .await
.map_err(control_tool_error)?; .map_err(control_tool_error)?
} else {
WorkspaceResponse {
status: 200,
body: r#"{"items":[]}"#.to_string(),
}
};
self.with_subworkers(response)? self.with_subworkers(response)?
} }
WorkerOperation::Spawn => { WorkerOperation::Spawn => {
@@ -701,7 +753,18 @@ impl Tool for WorkspaceWorkerTool {
.map_err(control_tool_error)? .map_err(control_tool_error)?
} }
WorkerOperation::SendInput | WorkerOperation::Notify => { WorkerOperation::SendInput | WorkerOperation::Notify => {
let input = parse::<WorkerMessageInput>(input_json, self.operation.tool_name())?; let input = if self.runtime_worker_control {
parse::<WorkerMessageInput>(input_json, self.operation.tool_name())?
} else {
let input = parse::<DirectSubWorkerMessageInput>(
input_json,
self.operation.tool_name(),
)?;
WorkerMessageInput {
subject: input.subject.into(),
content: input.content,
}
};
let content = non_empty(input.content, "content")?; let content = non_empty(input.content, "content")?;
if content.len() > 16 * 1024 { if content.len() > 16 * 1024 {
return Err(ToolError::ExecutionFailed( return Err(ToolError::ExecutionFailed(
@@ -738,7 +801,16 @@ impl Tool for WorkspaceWorkerTool {
} }
} }
WorkerOperation::Cancel | WorkerOperation::Stop => { WorkerOperation::Cancel | WorkerOperation::Stop => {
let input = parse::<WorkerStopInput>(input_json, self.operation.tool_name())?; let input = if self.runtime_worker_control {
parse::<WorkerStopInput>(input_json, self.operation.tool_name())?
} else {
let input =
parse::<DirectSubWorkerStopInput>(input_json, self.operation.tool_name())?;
WorkerStopInput {
subject: input.subject.into(),
reason: input.reason,
}
};
match input.subject { match input.subject {
WorkerSubjectInput::SubWorker { name } => { WorkerSubjectInput::SubWorker { name } => {
if self.operation != WorkerOperation::Stop { if self.operation != WorkerOperation::Stop {
@@ -911,7 +983,7 @@ fn render_subworker_stop_summary(summary: &SubWorkerStopSummary) -> String {
.as_ref() .as_ref()
.map(|stat| format!("+{}/-{} Changes · ", stat.added, stat.deleted)) .map(|stat| format!("+{}/-{} Changes · ", stat.added, stat.deleted))
.unwrap_or_default(); .unwrap_or_default();
format!("SubWorkerStop - done\n {tools}\n {changes}{elapsed}",) format!("WorkerStop - done\n {tools}\n {changes}{elapsed}",)
} }
fn format_elapsed(elapsed_ms: u64) -> String { fn format_elapsed(elapsed_ms: u64) -> String {
@@ -925,9 +997,62 @@ fn format_elapsed(elapsed_ms: u64) -> String {
} }
} }
fn sub_worker_tool_declaration(operation: WorkerOperation) -> ToolDeclaration {
let description = match operation {
WorkerOperation::List => "List this Worker's direct SubWorkers.",
WorkerOperation::SendInput => "Send a new user turn to a direct SubWorker.",
WorkerOperation::Stop => {
"Stop a direct SubWorker and release all authority delegated to its child session."
}
_ => unreachable!("unsupported direct SubWorker operation"),
};
ToolDeclaration::new(operation.tool_name(), description)
}
fn worker_tool_declaration(operation: WorkerOperation) -> ToolDeclaration {
ToolDeclaration::new(operation.tool_name(), operation.description())
}
fn worker_tool_contribution(
operation: WorkerOperation,
control: Arc<dyn WorkerControlService>,
runtime_worker_control: bool,
) -> ToolContribution {
let definition = match operation {
WorkerOperation::List => {
definition::<WorkerListInput>(operation, control, runtime_worker_control)
}
WorkerOperation::Spawn => {
definition::<WorkerSpawnInput>(operation, control, runtime_worker_control)
}
WorkerOperation::SendInput | WorkerOperation::Notify => {
if runtime_worker_control {
definition::<WorkerMessageInput>(operation, control, true)
} else {
definition::<DirectSubWorkerMessageInput>(operation, control, false)
}
}
WorkerOperation::Cancel | WorkerOperation::Stop => {
if runtime_worker_control {
definition::<WorkerStopInput>(operation, control, true)
} else {
definition::<DirectSubWorkerStopInput>(operation, control, false)
}
}
WorkerOperation::Restore => {
definition::<WorkerTargetInput>(operation, control, runtime_worker_control)
}
WorkerOperation::Remove => {
definition::<WorkerRemoveInput>(operation, control, runtime_worker_control)
}
};
ToolContribution::new(operation.tool_name(), definition)
}
fn definition<I: JsonSchema + 'static>( fn definition<I: JsonSchema + 'static>(
operation: WorkerOperation, operation: WorkerOperation,
control: Arc<dyn WorkerControlService>, control: Arc<dyn WorkerControlService>,
runtime_worker_control: bool,
) -> ToolDefinition { ) -> ToolDefinition {
Arc::new(move || { Arc::new(move || {
let schema = schemars::schema_for!(I); let schema = schemars::schema_for!(I);
@@ -938,6 +1063,7 @@ fn definition<I: JsonSchema + 'static>(
let tool: Arc<dyn Tool> = Arc::new(WorkspaceWorkerTool { let tool: Arc<dyn Tool> = Arc::new(WorkspaceWorkerTool {
operation, operation,
control: control.clone(), control: control.clone(),
runtime_worker_control,
}); });
(meta, tool) (meta, tool)
}) })
@@ -1036,11 +1162,101 @@ mod tests {
} }
} }
#[derive(Debug, Default)]
struct RecordingSubWorkerControl {
sent: Mutex<Vec<(String, String)>>,
stopped: Mutex<Vec<String>>,
}
#[async_trait]
impl WorkerLifecycleService for RecordingSubWorkerControl {
async fn spawn(
&self,
_request: WorkerLifecycleSpawnRequest,
) -> Result<WorkspaceResponse, WorkspaceClientError> {
Err(WorkspaceClientError::Unavailable(
"Runtime Worker control is disabled in this test".to_string(),
))
}
}
#[async_trait]
impl WorkerControlService for RecordingSubWorkerControl {
fn workspace_id(&self) -> &str {
"workspace-test"
}
fn known_subworkers(&self) -> Vec<serde_json::Value> {
Vec::new()
}
async fn spawn_worker(
&self,
_request: WorkerLifecycleSpawnRequest,
) -> Result<WorkspaceResponse, WorkspaceClientError> {
Err(WorkspaceClientError::Unavailable(
"Runtime Worker control is disabled in this test".to_string(),
))
}
fn remove_runtime_worker(
&self,
_runtime_id: &str,
_worker_id: &str,
_reason: &str,
) -> Result<WorkspaceResponse, WorkspaceClientError> {
Err(WorkspaceClientError::Unavailable(
"Runtime Worker control is disabled in this test".to_string(),
))
}
async fn execute_runtime(
&self,
_request: WorkspaceRequest,
) -> Result<WorkspaceResponse, WorkspaceClientError> {
Err(WorkspaceClientError::Unavailable(
"Runtime Worker control is disabled in this test".to_string(),
))
}
async fn ensure_permission(
&self,
_subject: &crate::feature::builtin::WorkerObservationSubjectRef,
_permission: &str,
) -> 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> { fn test_control(client: Arc<RecordingWorkspaceClient>) -> Arc<dyn WorkerControlService> {
Arc::new(WorkspaceWorkerControlService { Arc::new(WorkspaceWorkerControlService {
client, client,
workspace_id: "workspace%2Ftest".to_string(), workspace_id: "workspace%2Ftest".to_string(),
registry: None, registry: None,
runtime_worker_control: true,
}) })
} }
@@ -1050,6 +1266,7 @@ mod tests {
let tool = WorkspaceWorkerTool { let tool = WorkspaceWorkerTool {
operation: WorkerOperation::Spawn, operation: WorkerOperation::Spawn,
control: test_control(client.clone()), control: test_control(client.clone()),
runtime_worker_control: true,
}; };
tool.execute( tool.execute(
&serde_json::json!({ &serde_json::json!({
@@ -1128,6 +1345,123 @@ mod tests {
assert!(report.services.providers().is_empty()); assert!(report.services.providers().is_empty());
} }
#[tokio::test]
async fn sub_worker_control_surface_lists_registry_without_workspace_authority() {
let client = Arc::new(RecordingWorkspaceClient::default());
let runtime_base = tempfile::tempdir().unwrap();
let runtime_dir = Arc::new(
crate::runtime::dir::RuntimeDir::create(runtime_base.path(), "sub-worker-control")
.await
.unwrap(),
);
let registry = SpawnedWorkerRegistry::new(runtime_dir);
let control: Arc<dyn WorkerControlService> = Arc::new(WorkspaceWorkerControlService {
client: client.clone(),
workspace_id: "workspace%2Ftest".to_string(),
registry: Some(registry),
runtime_worker_control: false,
});
let tool = WorkspaceWorkerTool {
operation: WorkerOperation::List,
control,
runtime_worker_control: false,
};
let output = tool
.execute(
"{}",
ToolExecutionContext::new("call-list", "batch-list", 0),
)
.await
.unwrap();
assert!(client.requests.lock().unwrap().is_empty());
let value: serde_json::Value =
serde_json::from_str(output.content.as_deref().unwrap()).unwrap();
assert_eq!(value["items"], serde_json::json!([]));
}
#[tokio::test]
async fn canonical_send_and_stop_tools_route_to_direct_subworker_control() {
let control = Arc::new(RecordingSubWorkerControl::default());
let send_tool = WorkspaceWorkerTool {
operation: WorkerOperation::SendInput,
control: control.clone(),
runtime_worker_control: false,
};
let stop_tool = WorkspaceWorkerTool {
operation: WorkerOperation::Stop,
control: control.clone(),
runtime_worker_control: false,
};
send_tool
.execute(
r#"{"subject":{"kind":"sub_worker","name":"reviewer"},"content":"continue"}"#,
ToolExecutionContext::new("call-send", "batch-control", 0),
)
.await
.unwrap();
stop_tool
.execute(
r#"{"subject":{"kind":"sub_worker","name":"reviewer"},"reason":"review completed"}"#,
ToolExecutionContext::new("call-stop", "batch-control", 1),
)
.await
.unwrap();
assert_eq!(
*control.sent.lock().unwrap(),
vec![("reviewer".to_string(), "continue".to_string())]
);
assert_eq!(
*control.stopped.lock().unwrap(),
vec!["reviewer".to_string()]
);
}
#[tokio::test]
async fn sub_worker_control_surface_rejects_runtime_subject_without_workspace_call() {
let client = Arc::new(RecordingWorkspaceClient::default());
let control: Arc<dyn WorkerControlService> = Arc::new(WorkspaceWorkerControlService {
client: client.clone(),
workspace_id: "workspace%2Ftest".to_string(),
registry: None,
runtime_worker_control: false,
});
let permission_error = control
.ensure_permission(
&crate::feature::builtin::WorkerObservationSubjectRef::RuntimeWorker {
runtime_id: "runtime-1".to_string(),
worker_id: "worker-1".to_string(),
},
"stop",
)
.await
.unwrap_err();
assert!(matches!(
permission_error,
WorkspaceClientError::Unavailable(_)
));
let tool = WorkspaceWorkerTool {
operation: WorkerOperation::Stop,
control,
runtime_worker_control: false,
};
let error = tool
.execute(
r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"worker-1"},"reason":"not authorized"}"#,
ToolExecutionContext::new("call-stop", "batch-stop", 0),
)
.await
.unwrap_err();
assert!(matches!(error, ToolError::InvalidArgument(_)));
assert!(client.requests.lock().unwrap().is_empty());
}
#[test] #[test]
fn worker_service_can_remain_enabled_without_direct_spawn_surface() { fn worker_service_can_remain_enabled_without_direct_spawn_surface() {
let client = Arc::new(RecordingWorkspaceClient::default()); let client = Arc::new(RecordingWorkspaceClient::default());
@@ -1146,7 +1480,7 @@ mod tests {
} }
#[test] #[test]
fn worker_tool_family_is_distinct_from_sub_worker_tools() { fn canonical_worker_tool_family_uses_one_subject_based_namespace() {
assert_eq!( assert_eq!(
WorkerOperation::ALL.map(WorkerOperation::tool_name), WorkerOperation::ALL.map(WorkerOperation::tool_name),
[ [
@@ -1254,6 +1588,7 @@ mod tests {
WorkspaceWorkerTool { WorkspaceWorkerTool {
operation, operation,
control: test_control(client.clone()), control: test_control(client.clone()),
runtime_worker_control: true,
} }
.execute( .execute(
&args.to_string(), &args.to_string(),
@@ -1278,6 +1613,7 @@ mod tests {
let tool = WorkspaceWorkerTool { let tool = WorkspaceWorkerTool {
operation: WorkerOperation::Remove, operation: WorkerOperation::Remove,
control: test_control(client.clone()), control: test_control(client.clone()),
runtime_worker_control: true,
}; };
tool.execute( tool.execute(
&serde_json::json!({ &serde_json::json!({
@@ -1327,6 +1663,7 @@ mod tests {
let tool = WorkspaceWorkerTool { let tool = WorkspaceWorkerTool {
operation: WorkerOperation::Remove, operation: WorkerOperation::Remove,
control: test_control(client.clone()), control: test_control(client.clone()),
runtime_worker_control: true,
}; };
for reason in [" ".to_string(), "x".repeat(513)] { for reason in [" ".to_string(), "x".repeat(513)] {
let _error = tool let _error = tool
@@ -1397,7 +1734,7 @@ mod tests {
assert_eq!( assert_eq!(
output.summary, output.summary,
"SubWorkerStop - done\n 26 Read, 5 Grep\n +215/-148 Changes · 1m 18s" "WorkerStop - done\n 26 Read, 5 Grep\n +215/-148 Changes · 1m 18s"
); );
assert_eq!( assert_eq!(
serde_json::from_str::<SubWorkerStopSummary>(output.content.as_deref().unwrap()) serde_json::from_str::<SubWorkerStopSummary>(output.content.as_deref().unwrap())
+18 -1
View File
@@ -901,6 +901,23 @@ mod tests {
.unwrap() .unwrap()
.contains("BOUNDARY_MARKER") .contains("BOUNDARY_MARKER")
); );
catalog.worker_orchestration_guidance_section().unwrap(); let orchestration = catalog.worker_orchestration_guidance_section().unwrap();
for name in [
"SubWorkerSpawn",
"WorkerList",
"WorkerSendInput",
"WorkerStop",
] {
assert!(
orchestration.contains(name),
"missing canonical tool {name}"
);
}
for alias in ["SubWorkerList", "SubWorkerSend", "SubWorkerStop"] {
assert!(
!orchestration.contains(alias),
"guidance referenced stale alias {alias}"
);
}
} }
} }
+39 -3
View File
@@ -181,12 +181,19 @@ impl ToolCapabilities {
"MemoryReadDocument" => capabilities.memory_read_document = true, "MemoryReadDocument" => capabilities.memory_read_document = true,
"MemoryUpdateDocument" => capabilities.memory_update_document = true, "MemoryUpdateDocument" => capabilities.memory_update_document = true,
"SubWorkerSpawn" => capabilities.sub_worker_spawn = true, "SubWorkerSpawn" => capabilities.sub_worker_spawn = true,
"SubWorkerSend" => capabilities.sub_worker_send = true,
"SubWorkerStop" => capabilities.sub_worker_stop = true,
"SubWorkerList" => capabilities.sub_worker_list = true,
_ => {} _ => {}
} }
} }
if capabilities.sub_worker_spawn {
for name in names {
match name.as_str() {
"WorkerSendInput" => capabilities.sub_worker_send = true,
"WorkerStop" => capabilities.sub_worker_stop = true,
"WorkerList" => capabilities.sub_worker_list = true,
_ => {}
}
}
}
capabilities capabilities
} }
@@ -316,6 +323,35 @@ fn append_trailing_section(
mod tests { mod tests {
use super::*; use super::*;
#[test]
fn sub_worker_capabilities_follow_the_registered_canonical_control_tools() {
let names = [
"SubWorkerSpawn",
"WorkerList",
"WorkerSendInput",
"WorkerStop",
]
.map(str::to_string);
let capabilities = ToolCapabilities::from_tool_names(&names);
assert!(capabilities.sub_worker_management());
assert!(capabilities.sub_worker_list);
assert!(capabilities.sub_worker_send);
assert!(capabilities.sub_worker_stop);
let stale_aliases = [
"SubWorkerSpawn",
"SubWorkerList",
"SubWorkerSend",
"SubWorkerStop",
]
.map(str::to_string);
let capabilities = ToolCapabilities::from_tool_names(&stale_aliases);
assert!(capabilities.sub_worker_spawn);
assert!(!capabilities.sub_worker_list);
assert!(!capabilities.sub_worker_send);
assert!(!capabilities.sub_worker_stop);
}
#[test] #[test]
fn rejects_legacy_prefix_relative_and_missing_names() { fn rejects_legacy_prefix_relative_and_missing_names() {
for reference in ["legacy/custom", "custom.md", "../custom", "missing"] { for reference in ["legacy/custom", "custom.md", "../custom", "missing"] {
+3 -210
View File
@@ -1,226 +1,19 @@
#![cfg_attr(not(test), allow(dead_code, unused_imports))] //! Socket communication retained for the legacy top-level Worker callback protocol.
//! Parent-facing tools for in-process Internal SubWorker sessions.
//! //!
//! Legacy direct-child tool constructors are test-only; production exposes the //! Direct Internal SubWorker lifecycle is exposed through `worker.control` and the
//! registry through the unified `worker.control` service and Worker tools. //! canonical Worker tools, not a second SubWorker-specific tool family.
//! 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.
use std::path::Path; use std::path::Path;
use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use async_trait::async_trait;
use protocol::stream::{JsonLineReader, JsonLineWriter}; use protocol::stream::{JsonLineReader, JsonLineWriter};
use protocol::{Event, Method}; use protocol::{Event, Method};
use serde::{Deserialize, Serialize};
use tokio::net::UnixStream; use tokio::net::UnixStream;
use crate::spawn::registry::SpawnedWorkerRegistry;
/// Timeout applied to each socket-level operation — connect, write, /// Timeout applied to each socket-level operation — connect, write,
/// read. Kept short so a stuck child doesn't block the spawner's turn. /// read. Kept short so a stuck child doesn't block the spawner's turn.
const SOCKET_OP_TIMEOUT: Duration = Duration::from_secs(5); const SOCKET_OP_TIMEOUT: Duration = Duration::from_secs(5);
// ---------------------------------------------------------------------------
// Shared input types
// ---------------------------------------------------------------------------
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct NameInput {
/// Name of a previously spawned SubWorker.
name: String,
}
#[derive(Debug, Deserialize, schemars::JsonSchema)]
#[serde(deny_unknown_fields)]
struct SubWorkerListInput {}
#[derive(Debug, Serialize)]
struct SubWorkerListItem {
name: String,
}
struct SubWorkerListTool {
registry: Arc<SpawnedWorkerRegistry>,
}
#[async_trait]
impl Tool for SubWorkerListTool {
async fn execute(
&self,
input_json: &str,
_ctx: agen::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
let _input: SubWorkerListInput = serde_json::from_str(input_json).map_err(|error| {
ToolError::InvalidArgument(format!("invalid SubWorkerList input: {error}"))
})?;
let items = self
.registry
.list_internal()
.into_iter()
.map(|record| SubWorkerListItem {
name: record.worker_name,
})
.collect::<Vec<_>>();
let count = items.len();
let content = serde_json::to_string_pretty(&serde_json::json!({ "sub_workers": items }))
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?;
Ok(ToolOutput {
summary: format!("listed {count} child SubWorker(s)"),
content: Some(content),
attachments: Vec::new(),
})
}
}
#[cfg(test)]
pub fn sub_worker_list_tool(registry: Arc<SpawnedWorkerRegistry>) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(SubWorkerListInput);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
let meta = ToolMeta::new("SubWorkerList")
.description("List child SubWorkers owned by this Worker. Peer Workers and general Runtime Workers are excluded.")
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(SubWorkerListTool {
registry: registry.clone(),
});
(meta, tool)
})
}
// ---------------------------------------------------------------------------
// SubWorkerSend
// ---------------------------------------------------------------------------
const SEND_TO_POD_DESCRIPTION: &str = "Send a text message to a previously spawned SubWorker. The SubWorker \
processes it as a user turn. Fails if the SubWorker is already executing a \
turn retry after it finishes. Does not wait for the turn to complete; \
use worker-observation tools to inspect its committed session.";
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct SubWorkerSendInput {
/// Target SubWorker name.
name: String,
/// Text delivered to the SubWorker as the next user message.
message: String,
}
struct SubWorkerSendTool {
registry: Arc<SpawnedWorkerRegistry>,
}
#[async_trait]
impl Tool for SubWorkerSendTool {
async fn execute(
&self,
input_json: &str,
_ctx: agen::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
let input: SubWorkerSendInput = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid SubWorkerSend input: {e}")))?;
if let Some(record) = self.registry.get_internal(&input.name) {
record.session.send(input.message).await.map_err(|error| {
ToolError::ExecutionFailed(format!("send to `{}`: {error}", input.name))
})?;
return Ok(ToolOutput {
summary: format!("sent message to `{}`", input.name),
content: None,
attachments: Vec::new(),
});
}
Err(unknown_worker_err(&input.name))
}
}
#[cfg(test)]
pub fn sub_worker_send_tool(registry: Arc<SpawnedWorkerRegistry>) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(SubWorkerSendInput);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
let meta = ToolMeta::new("SubWorkerSend")
.description(SEND_TO_POD_DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(SubWorkerSendTool {
registry: registry.clone(),
});
(meta, tool)
})
}
// ---------------------------------------------------------------------------
// SubWorkerStop
// ---------------------------------------------------------------------------
const STOP_POD_DESCRIPTION: &str = "Cancel and stop a spawned Internal SubWorker session, remove it from the parent's direct-child registry, and reclaim delegated Write scope.";
struct SubWorkerStopTool {
registry: Arc<SpawnedWorkerRegistry>,
}
#[async_trait]
impl Tool for SubWorkerStopTool {
async fn execute(
&self,
input_json: &str,
_ctx: agen::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
let input: NameInput = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid SubWorkerStop input: {e}")))?;
if let Some(summary) = self
.registry
.remove_internal(&input.name)
.await
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?
{
return Ok(ToolOutput {
summary: format!(
"SubWorkerStop - done\n {} tool kind{}\n {}ms",
summary.tool_counts.len(),
if summary.tool_counts.len() == 1 {
""
} else {
"s"
},
summary.elapsed_ms,
),
content: Some(
serde_json::to_string(&summary)
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?,
),
attachments: Vec::new(),
});
}
Err(unknown_worker_err(&input.name))
}
}
#[cfg(test)]
pub fn sub_worker_stop_tool(registry: Arc<SpawnedWorkerRegistry>) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(NameInput);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
let meta = ToolMeta::new("SubWorkerStop")
.description(STOP_POD_DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(SubWorkerStopTool {
registry: registry.clone(),
});
(meta, tool)
})
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
fn unknown_worker_err(name: &str) -> ToolError {
ToolError::InvalidArgument(format!("no spawned worker named `{name}`"))
}
/// Connect with a timeout, drain the server's connect-time snapshot, /// Connect with a timeout, drain the server's connect-time snapshot,
/// write one `Method` line, flush, and close. /// write one `Method` line, flush, and close.
/// ///
+16 -22
View File
@@ -1226,14 +1226,11 @@ extract_threshold = 4000
Err(tokio::sync::mpsc::error::TryRecvError::Empty) Err(tokio::sync::mpsc::error::TryRecvError::Empty)
)); ));
let context = agen::tool::ToolExecutionContext::direct();
let list = (crate::spawn::comm_tools::sub_worker_list_tool(registry.clone()))().1;
let listed = list.execute("{}", context.clone()).await.unwrap();
assert!( assert!(
listed registry
.content .list_internal()
.unwrap_or_default() .iter()
.contains("reviewer-child") .any(|record| record.worker_name == "reviewer-child")
); );
let observation = let observation =
@@ -1256,11 +1253,9 @@ extract_threshold = 4000
.contains("reviewed") .contains("reviewed")
); );
let send = (crate::spawn::comm_tools::sub_worker_send_tool(registry.clone()))().1; record
send.execute( .session
r#"{"name":"reviewer-child","message":"review follow-up"}"#, .send("review follow-up".to_string())
context.clone(),
)
.await .await
.unwrap(); .unwrap();
assert_eq!( assert_eq!(
@@ -1277,10 +1272,9 @@ extract_threshold = 4000
assert!(latest_capture.session.entries.len() > first_capture.session.entries.len()); assert!(latest_capture.session.entries.len() > first_capture.session.entries.len());
fail_requests.store(true, Ordering::SeqCst); fail_requests.store(true, Ordering::SeqCst);
send.execute( record
r#"{"name":"reviewer-child","message":"trigger terminal failure"}"#, .session
context.clone(), .send("trigger terminal failure".to_string())
)
.await .await
.unwrap(); .unwrap();
assert_eq!( assert_eq!(
@@ -1298,10 +1292,13 @@ extract_threshold = 4000
); );
assert!(registry.get_internal("reviewer-child").is_some()); assert!(registry.get_internal("reviewer-child").is_some());
let stop = (crate::spawn::comm_tools::sub_worker_stop_tool(registry.clone()))().1; assert!(
stop.execute(r#"{"name":"reviewer-child"}"#, context) registry
.remove_internal("reviewer-child")
.await .await
.unwrap(); .unwrap()
.is_some()
);
assert!(registry.get_internal("reviewer-child").is_none()); assert!(registry.get_internal("reviewer-child").is_none());
assert!(spawner_scope.snapshot().is_writable(&workspace_root)); assert!(spawner_scope.snapshot().is_writable(&workspace_root));
@@ -1315,9 +1312,6 @@ extract_threshold = 4000
.await .await
.unwrap(); .unwrap();
assert!(spawner_scope.snapshot().is_writable(&workspace_root)); assert!(spawner_scope.snapshot().is_writable(&workspace_root));
drop(list);
drop(send);
drop(stop);
drop(observation); drop(observation);
drop(tool); drop(tool);
drop(registry); drop(registry);
+119 -9
View File
@@ -18,7 +18,8 @@ use workdir::{
use worker::{ use worker::{
Event, Method, Worker, WorkerController, WorkerFilesystemAuthority, WorkerHandle, Event, Method, Worker, WorkerController, WorkerFilesystemAuthority, WorkerHandle,
WorkerManifest, WorkerStatus, WorkerWorkspaceContext, WorkerManifest, WorkerStatus, WorkerWorkspaceContext, WorkspaceClient, WorkspaceClientError,
WorkspaceRequest, WorkspaceResponse,
}; };
type TestStore = CombinedStore<FsStore, FsWorkerStore>; type TestStore = CombinedStore<FsStore, FsWorkerStore>;
@@ -179,9 +180,48 @@ async fn make_worker_with_pwd(
make_worker_with_pwd_and_manifest(client, MANIFEST_TOML).await make_worker_with_pwd_and_manifest(client, MANIFEST_TOML).await
} }
#[derive(Debug)]
struct NoopWorkspaceClient;
impl WorkspaceClient for NoopWorkspaceClient {
fn workspace_id(&self) -> Option<&str> {
Some("workspace-test")
}
fn kind(&self) -> &str {
"test-noop"
}
fn is_available(&self) -> bool {
true
}
fn execute(
&self,
_request: WorkspaceRequest,
) -> Result<WorkspaceResponse, WorkspaceClientError> {
Err(WorkspaceClientError::Unavailable(
"test client does not execute requests".to_string(),
))
}
}
async fn make_worker_with_pwd_and_manifest( async fn make_worker_with_pwd_and_manifest(
client: MockClient, client: MockClient,
manifest_toml: &str, manifest_toml: &str,
) -> (Worker<MockClient, TestStore>, std::path::PathBuf) {
make_worker_with_pwd_manifest_and_workspace_context(
client,
manifest_toml,
WorkerWorkspaceContext::local_filesystem(None),
)
.await
}
async fn make_worker_with_pwd_manifest_and_workspace_context(
client: MockClient,
manifest_toml: &str,
workspace_context: WorkerWorkspaceContext,
) -> (Worker<MockClient, TestStore>, std::path::PathBuf) { ) -> (Worker<MockClient, TestStore>, std::path::PathBuf) {
let manifest = WorkerManifest::from_toml(manifest_toml).unwrap(); let manifest = WorkerManifest::from_toml(manifest_toml).unwrap();
let store_tmp = tempfile::tempdir().unwrap(); let store_tmp = tempfile::tempdir().unwrap();
@@ -202,14 +242,7 @@ async fn make_worker_with_pwd_and_manifest(
let worker = let worker =
Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client); Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client);
let authority = WorkerFilesystemAuthority::local(pwd.clone(), pwd.clone()); let authority = WorkerFilesystemAuthority::local(pwd.clone(), pwd.clone());
let worker = Worker::new( let worker = Worker::new(manifest, worker, store, workspace_context, authority, scope)
manifest,
worker,
store,
WorkerWorkspaceContext::local_filesystem(None),
authority,
scope,
)
.await .await
.unwrap(); .unwrap();
(worker, pwd) (worker, pwd)
@@ -663,6 +696,83 @@ permission = "write"
"{} role SubWorker tool exposure mismatch: {names:?}", "{} role SubWorker tool exposure mismatch: {names:?}",
case.role case.role
); );
for control_tool in ["WorkerList", "WorkerSendInput", "WorkerStop"] {
assert_eq!(
names.iter().any(|name| name == control_tool),
case.sub_worker_enabled,
"{} role {control_tool} exposure mismatch: {names:?}",
case.role
);
}
for stale_alias in ["SubWorkerList", "SubWorkerSend", "SubWorkerStop"] {
assert!(
!names.iter().any(|name| name == stale_alias),
"{} role exposed stale alias {stale_alias}: {names:?}",
case.role
);
}
}
}
#[tokio::test]
async fn worker_and_sub_worker_features_install_one_canonical_control_surface() {
let manifest = r#"
[worker]
name = "combined-worker-control-feature-test"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[feature.worker]
enabled = true
direct_spawn = false
[feature.sub_worker]
enabled = true
[[scope.allow]]
target = "./"
permission = "write"
[[delegation_scope.allow]]
target = "/tmp"
permission = "write"
"#;
let client = MockClient::new(simple_text_events());
let client_for_assert = client.clone();
let worker = make_worker_with_pwd_manifest_and_workspace_context(
client,
manifest,
WorkerWorkspaceContext::with_client(None, Arc::new(NoopWorkspaceClient)),
)
.await
.0;
let handle = spawn_controller(worker).await;
handle.send(Method::run_text("Hello")).await.unwrap();
wait_for_status(&handle, WorkerStatus::Idle).await;
let request = wait_for_captured_request(&client_for_assert).await;
let names = request_tool_names(&request);
assert!(names.iter().any(|name| name == "SubWorkerSpawn"));
assert!(!names.iter().any(|name| name == "WorkerSpawn"));
for control_tool in ["WorkerList", "WorkerSendInput", "WorkerStop"] {
assert_eq!(
names
.iter()
.filter(|name| name.as_str() == control_tool)
.count(),
1,
"expected one {control_tool} contribution: {names:?}"
);
}
for stale_alias in ["SubWorkerList", "SubWorkerSend", "SubWorkerStop"] {
assert!(!names.iter().any(|name| name == stale_alias));
} }
} }
+1 -1
View File
@@ -44,4 +44,4 @@ Observation is read-only evidence access. It does not authorize Ticket, Memory,
## SubWorker output ## SubWorker output
SubWorkers no longer expose a separate output cursor tool. `SubWorkerList`, `SubWorkerSend`, and `SubWorkerStop` retain parent-owned lifecycle control, while committed child output is read through `worker-observation`. Turn-completion notifications carry no transcript and only tell the parent to inspect the authoritative committed session at a natural boundary. SubWorkers no longer expose a separate output cursor or lifecycle tool family. `WorkerList`, `WorkerSendInput`, and `WorkerStop` retain parent-owned lifecycle control through a `{ kind: "sub_worker", name }` subject, while committed child output is read through `worker-observation`. Turn-completion notifications carry no transcript and only tell the parent to inspect the authoritative committed session at a natural boundary.
@@ -2,7 +2,7 @@
--- ---
## SubWorker orchestration ## SubWorker orchestration
When SubWorker-management tools are available, SubWorker notifications are background signals for the parent Worker to handle at a natural stopping point. Do not ignore routine follow-up, but do not interrupt the current user request unnecessarily. When SubWorker-management tools are available, create direct children with `SubWorkerSpawn`, discover them with `WorkerList`, continue them with `WorkerSendInput`, and release their delegated authority with `WorkerStop`. Pass the exact `{ kind: "sub_worker", name }` subject returned by `WorkerList`; do not invent direct-only aliases. SubWorker notifications are background signals for the parent Worker to handle at a natural stopping point. Do not ignore routine follow-up, but do not interrupt the current user request unnecessarily.
The parent Worker does not need to keep a turn open or call tools solely to wait for a notification. Do not use `sleep` or polling loops just to wait for SubWorker output; if there is no useful immediate work, return control and handle the SubWorker when notified or when the user next asks. The parent Worker does not need to keep a turn open or call tools solely to wait for a notification. Do not use `sleep` or polling loops just to wait for SubWorker output; if there is no useful immediate work, return control and handle the SubWorker when notified or when the user next asks.