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

worker: add guarded WorkerRemove lifecycle

Browse Source
This commit is contained in:
Hare
2026-08-12 17:00:32 +09:00
parent ebe0f93744
commit 8ae930c5fc
12 changed files with 1565 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/feature/builtin/manage_worker.rs
+236 -74
View File
@@ -77,6 +77,11 @@ impl FeatureModule for ManageWorkerFeature {
self.client.clone(),
workspace_id.clone(),
),
WorkerOperation::Remove => definition::<WorkerRemoveInput>(
operation,
self.client.clone(),
workspace_id.clone(),
),
};
context
.tools()
@@ -149,6 +154,15 @@ struct WorkerStopInput {
reason: Option<String>,
}
#[derive(Debug, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
struct WorkerRemoveInput {
runtime_id: String,
worker_id: String,
expected_worker_revision: String,
reason: String,
}
struct WorkspaceWorkerTool {
operation: WorkerOperation,
client: Arc<dyn WorkspaceClient>,
@@ -161,10 +175,17 @@ enum WorkerOperation {
Spawn,
Stop,
Restore,
Remove,
}
impl WorkerOperation {
const ALL: [Self; 4] = [Self::List, Self::Spawn, Self::Stop, Self::Restore];
const ALL: [Self; 5] = [
Self::List,
Self::Spawn,
Self::Stop,
Self::Restore,
Self::Remove,
];
fn tool_name(self) -> &'static str {
match self {
@@ -172,6 +193,7 @@ impl WorkerOperation {
Self::Spawn => "WorkerSpawn",
Self::Stop => "WorkerStop",
Self::Restore => "WorkerRestore",
Self::Remove => "WorkerRemove",
}
}
@@ -187,6 +209,9 @@ impl WorkerOperation {
Self::Restore => {
"Restore a stopped Backend/Runtime Worker session in the current Workspace."
}
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."
}
}
}
}
@@ -198,82 +223,107 @@ impl Tool for WorkspaceWorkerTool {
input_json: &str,
ctx: ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
let request = match self.operation {
WorkerOperation::List => {
parse::<WorkerListInput>(input_json, "WorkerList")?;
WorkspaceRequest::get(format!("/api/w/{}/workers", self.workspace_id))
let response = match self.operation {
WorkerOperation::Remove => {
let input = parse::<WorkerRemoveInput>(input_json, "WorkerRemove")?;
let runtime_id = authority_id(&input.runtime_id, "runtime_id")?;
let worker_id = authority_id(&input.worker_id, "worker_id")?;
let expected_worker_revision =
non_empty(input.expected_worker_revision, "expected_worker_revision")?;
let reason = non_empty(input.reason, "reason")?;
if reason.len() > 512 {
return Err(ToolError::ExecutionFailed(
"reason must contain at most 512 bytes".to_string(),
));
}
self.client
.execute_worker_remove(
&runtime_id,
&worker_id,
&expected_worker_revision,
&reason,
)
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?
}
WorkerOperation::Spawn => {
let input = parse::<WorkerSpawnInput>(input_json, "WorkerSpawn")?;
let ticket_assignment = input
.ticket_id
.map(|ticket_id| {
let ticket_id = authority_id(&ticket_id, "ticket_id")?;
let call_id = non_empty(ctx.call_id.clone(), "tool call_id")?;
Ok::<_, ToolError>(WorkerSpawnTicketAssignmentRequest {
operation_id: format!("worker-spawn:{ticket_id}:{call_id}"),
ticket_id,
})
})
.transpose()?;
let request = WorkerSpawnRequest {
runtime_id: authority_id(&input.runtime_id, "runtime_id")?,
display_name: input
.display_name
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| "Workspace Worker".to_string()),
profile: non_empty(input.profile, "profile")?,
ticket_assignment,
initial_submit: input.initial_submit,
working_directory: WorkerWorkingDirectorySelection {
working_directory_id: authority_id(
&input.working_directory_id,
"working_directory_id",
)?,
relative_cwd: input
.relative_cwd
.map(|value| validate_relative_cwd(&value))
.transpose()?,
},
operation => {
let request = match operation {
WorkerOperation::List => {
parse::<WorkerListInput>(input_json, "WorkerList")?;
WorkspaceRequest::get(format!("/api/w/{}/workers", self.workspace_id))
}
WorkerOperation::Spawn => {
let input = parse::<WorkerSpawnInput>(input_json, "WorkerSpawn")?;
let ticket_assignment = input
.ticket_id
.map(|ticket_id| {
let ticket_id = authority_id(&ticket_id, "ticket_id")?;
let call_id = non_empty(ctx.call_id.clone(), "tool call_id")?;
Ok::<_, ToolError>(WorkerSpawnTicketAssignmentRequest {
operation_id: format!("worker-spawn:{ticket_id}:{call_id}"),
ticket_id,
})
})
.transpose()?;
let request = WorkerSpawnRequest {
runtime_id: authority_id(&input.runtime_id, "runtime_id")?,
display_name: input
.display_name
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| "Workspace Worker".to_string()),
profile: non_empty(input.profile, "profile")?,
ticket_assignment,
initial_submit: input.initial_submit,
working_directory: WorkerWorkingDirectorySelection {
working_directory_id: authority_id(
&input.working_directory_id,
"working_directory_id",
)?,
relative_cwd: input
.relative_cwd
.map(|value| validate_relative_cwd(&value))
.transpose()?,
},
};
WorkspaceRequest::json(
WorkspaceRequestMethod::Post,
format!("/api/w/{}/workers", self.workspace_id),
serde_json::to_string(&request)
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?,
)
}
WorkerOperation::Stop => {
let input = parse::<WorkerStopInput>(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/{}/runtimes/{runtime_id}/workers/{worker_id}/stop",
self.workspace_id
),
serde_json::json!({ "reason": input.reason }).to_string(),
)
}
WorkerOperation::Restore => {
let input = parse::<WorkerTargetInput>(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/{}/runtimes/{runtime_id}/workers/{worker_id}/restore",
self.workspace_id
),
"{}",
)
}
WorkerOperation::Remove => unreachable!("handled above"),
};
WorkspaceRequest::json(
WorkspaceRequestMethod::Post,
format!("/api/w/{}/workers", self.workspace_id),
serde_json::to_string(&request)
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?,
)
}
WorkerOperation::Stop => {
let input = parse::<WorkerStopInput>(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/{}/runtimes/{runtime_id}/workers/{worker_id}/stop",
self.workspace_id
),
serde_json::json!({ "reason": input.reason }).to_string(),
)
}
WorkerOperation::Restore => {
let input = parse::<WorkerTargetInput>(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/{}/runtimes/{runtime_id}/workers/{worker_id}/restore",
self.workspace_id
),
"{}",
)
self.client
.execute(request)
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?
}
};
let response = self
.client
.execute(request)
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?;
if !response.is_success() {
return Err(ToolError::ExecutionFailed(format!(
"Workspace Worker operation returned HTTP {}: {}",
@@ -356,6 +406,7 @@ mod tests {
#[derive(Debug, Default)]
struct RecordingWorkspaceClient {
requests: Mutex<Vec<WorkspaceRequest>>,
removals: Mutex<Vec<(String, String, String, String)>>,
}
impl WorkspaceClient for RecordingWorkspaceClient {
@@ -381,6 +432,25 @@ mod tests {
body: "{}".to_string(),
})
}
fn execute_worker_remove(
&self,
target_runtime_id: &str,
target_worker_id: &str,
expected_worker_revision: &str,
reason: &str,
) -> Result<WorkspaceResponse, WorkspaceClientError> {
self.removals.lock().unwrap().push((
target_runtime_id.to_string(),
target_worker_id.to_string(),
expected_worker_revision.to_string(),
reason.to_string(),
));
Ok(WorkspaceResponse {
status: 200,
body: r#"{"removed":true}"#.to_string(),
})
}
}
#[tokio::test]
@@ -433,7 +503,13 @@ mod tests {
fn worker_tool_family_is_distinct_from_sub_worker_tools() {
assert_eq!(
WorkerOperation::ALL.map(WorkerOperation::tool_name),
["WorkerList", "WorkerSpawn", "WorkerStop", "WorkerRestore"]
[
"WorkerList",
"WorkerSpawn",
"WorkerStop",
"WorkerRestore",
"WorkerRemove",
]
);
}
@@ -481,6 +557,92 @@ mod tests {
assert!(value.get("initial_text").is_none());
}
#[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(),
};
tool.execute(
&serde_json::json!({
"runtime_id": "runtime-1",
"worker_id": "worker-7",
"expected_worker_revision": "2026-08-11T20:00:00Z",
"reason": " retire completed Worker "
})
.to_string(),
ToolExecutionContext::new("call-remove", "batch-remove", 0),
)
.await
.unwrap();
assert_eq!(
client.removals.lock().unwrap().as_slice(),
[(
"runtime-1".to_string(),
"worker-7".to_string(),
"2026-08-11T20:00:00Z".to_string(),
"retire completed Worker".to_string(),
)]
);
let schema = serde_json::to_value(schemars::schema_for!(WorkerRemoveInput))
.unwrap()
.to_string();
for field in [
"runtime_id",
"worker_id",
"expected_worker_revision",
"reason",
] {
assert!(schema.contains(field));
}
for forbidden in ["proof", "actor", "workspace_id", "policy", "plan", "stage"] {
assert!(!schema.contains(forbidden), "schema leaked {forbidden}");
}
}
#[tokio::test]
async fn worker_remove_rejects_empty_oversized_and_unknown_authority_input() {
let client = Arc::new(RecordingWorkspaceClient::default());
let tool = WorkspaceWorkerTool {
operation: WorkerOperation::Remove,
client: client.clone(),
workspace_id: "workspace%2Ftest".to_string(),
};
for reason in [" ".to_string(), "x".repeat(513)] {
let _error = tool
.execute(
&serde_json::json!({
"runtime_id": "runtime-1",
"worker_id": "worker-7",
"expected_worker_revision": "revision-1",
"reason": reason,
})
.to_string(),
ToolExecutionContext::new("call-invalid", "batch-remove", 0),
)
.await
.unwrap_err();
}
let _error = tool
.execute(
&serde_json::json!({
"runtime_id": "runtime-1",
"worker_id": "worker-7",
"expected_worker_revision": "revision-1",
"reason": "retire",
"source_proof": "caller-controlled"
})
.to_string(),
ToolExecutionContext::new("call-spoof", "batch-remove", 0),
)
.await
.unwrap_err();
assert!(client.removals.lock().unwrap().is_empty());
}
#[test]
fn worker_inputs_reject_paths_and_parent_traversal() {
assert!(authority_id("https://runtime.example", "runtime_id").is_err());
crates/worker/src/prompt/catalog.rs
+15
View File
@@ -746,6 +746,21 @@ compact_system = "PREFIX\n{% include \"$yoi/internal/compact_system\" %}"
assert!(rendered.contains("bypass user/Ticket authorization"));
}
#[test]
fn orchestrator_role_prompt_fences_worker_remove_authority() {
let source = include_str!("../../../../resources/prompts/role/orchestrator.md");
assert!(source.contains("Use `WorkerRemove` only for a terminal or authoritatively reassigned non-internal Coder"));
assert!(source.contains("exact current `updated_at` value"));
assert!(source.contains("must have no current Ticket assignment"));
assert!(source.contains("pending notification, Reviewer handoff, legal hold, or pin"));
assert!(source.contains("After removal, reread the Worker catalog and attachment state"));
assert!(source.contains("attachment-close, and attachment-release conflicts"));
assert!(source.contains("preserves the Workdir materialization"));
assert!(!source.contains("source proof"));
assert!(!source.contains("provider handle"));
assert!(!source.contains("retention plan"));
}
#[test]
fn sub_worker_spawn_tool_description_renders_profile_block() {
let cat = PromptCatalog::builtins_only().unwrap();
crates/worker/src/worker.rs
+2
View File
@@ -230,6 +230,8 @@ pub trait WorkspaceClient: std::fmt::Debug + Send + Sync {
&self,
_target_runtime_id: &str,
_target_worker_id: &str,
_expected_worker_revision: &str,
_reason: &str,
) -> Result<WorkspaceResponse, WorkspaceClientError> {
Err(WorkspaceClientError::Unavailable(
"Runtime-owned WorkerRemove forwarding is unavailable".to_string(),