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

chore: merge develop into worker state snapshot

Browse Source
This commit is contained in:
Hare
2026-09-06 07:17:40 +09:00
parent e8b9adcde4 668a9062b3
commit 4269ebec04
36 changed files with 6350 additions and 2472 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+52 -11
View File
@@ -705,6 +705,7 @@ impl WorkerController {
runtime_base.to_path_buf(),
spawned_registry.clone(),
Some(method_tx.downgrade()),
None,
)
.await?;
if let Some(session) = fs_for_view.as_ref() {
@@ -1116,6 +1117,7 @@ pub(crate) async fn register_worker_tools<C, St>(
runtime_base: PathBuf,
spawned_registry: Arc<SpawnedWorkerRegistry>,
parent_method_tx: Option<mpsc::WeakSender<Method>>,
inherited_workdir_tool_broker: Option<workdir::WorkdirToolBroker>,
) -> std::io::Result<Option<workdir::WorkdirSessionHandle>>
where
C: LlmClient + Clone + 'static,
@@ -1124,21 +1126,26 @@ where
// Worker-immutable snapshots taken before the mutable worker borrow
// below so the worker borrow doesn't conflict with reads on `worker`.
let feature_config = worker.manifest().feature.clone();
let mut workdir_tool_broker = inherited_workdir_tool_broker;
if feature_config.manage_workdir.enabled && worker.workdir_session().is_none() {
let workspace_client = worker.workspace_client_handle();
worker.bind_workdir_session(Some(workdir::delegation_capable_session(
let broker = workdir::WorkdirToolBroker::new(
crate::feature::builtin::manage_workdir::WorkspaceAttachedWorkdirSession::handle(
workspace_client,
),
)));
}
if feature_config.sub_worker.enabled
);
worker.bind_workdir_session(Some(broker.tool_session()));
workdir_tool_broker = Some(broker);
} else if workdir_tool_broker.is_none()
&& let Some(existing) = worker.workdir_session().cloned()
&& !existing.is_delegation_capable()
{
worker.bind_workdir_session(Some(workdir::delegation_capable_session(existing)));
let broker = workdir::WorkdirToolBroker::new(existing);
worker.bind_workdir_session(Some(broker.tool_session()));
workdir_tool_broker = Some(broker);
}
let worker_workdir = worker.workdir_session().cloned();
let worker_workdir = workdir_tool_broker
.as_ref()
.map(workdir::WorkdirToolBroker::tool_session);
let local_filesystem = worker.local_working_directory().cloned();
let local_workspace_root = local_filesystem.as_ref().map(|local| local.root.clone());
let task_feature = worker.task_feature();
@@ -1305,8 +1312,17 @@ where
"manage Workdir tools require Backend Workspace API authority",
));
}
let shutdown_registry = spawned_registry.clone();
let reopen_registry = spawned_registry.clone();
feature_registry.add_module(
crate::feature::builtin::manage_workdir::manage_workdir_feature(workspace_client),
crate::feature::builtin::manage_workdir::ManageWorkdirFeature::with_child_lifecycle(
workspace_client,
Arc::new(move || {
let child_registry = shutdown_registry.clone();
Box::pin(async move { child_registry.shutdown_internal().await })
}),
Arc::new(move || reopen_registry.reopen_internal()),
),
);
}
if feature_config.workspace_worker_discovery.enabled {
@@ -1368,7 +1384,6 @@ where
}
let host_worker_observation_provider = worker.worker_observation_provider();
let source_workdir_session = worker.workdir_session().cloned();
{
let workspace_client = worker.workspace_client_handle();
let engine = worker.engine_mut();
@@ -1410,7 +1425,7 @@ where
runtime_base.clone(),
bash_output_dir.clone(),
spawner_workspace_root,
source_workdir_session,
workdir_tool_broker,
spawned_registry.clone(),
spawner_manifest,
prompts,
@@ -2292,7 +2307,16 @@ async fn controller_loop<C, St>(
// Memory/Workdir teardown so they cannot observe a partially closed Worker.
worker.stop_feature_runtime("controller shutdown").await;
if let Some(session) = worker.workdir_session()
let child_cleanup_succeeded = match spawned_registry.shutdown_internal().await {
Ok(()) => true,
Err(error) => {
tracing::warn!(%error, "Internal SubWorker cleanup failed before Workdir shutdown");
false
}
};
if child_cleanup_succeeded
&& let Some(session) = worker.workdir_session()
&& let Err(error) = session.close().await
{
tracing::warn!(%error, "Workdir session close failed");
@@ -3702,4 +3726,21 @@ mod tests {
.is_ok()
);
}
#[test]
fn controller_shutdown_orders_child_cleanup_before_workdir_close() {
let source = include_str!("controller.rs");
let shutdown_start = source
.rfind("worker.stop_feature_runtime(\"controller shutdown\")")
.expect("controller shutdown block");
let shutdown = &source[shutdown_start..];
let children = shutdown
.find("spawned_registry.shutdown_internal().await")
.expect("Internal SubWorker cleanup");
let workdir = shutdown
.find("session.close().await")
.expect("parent Workdir close");
assert!(children < workdir);
assert!(shutdown.contains("if child_cleanup_succeeded"));
}
}
crates/worker/src/feature/builtin/manage_workdir.rs
+201 -197
View File
@@ -5,6 +5,8 @@
//! endpoints, credentials, materializer handles, and operation sessions stay
//! behind [`WorkspaceClient`].
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use agen::tool::{Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput};
@@ -12,7 +14,7 @@ use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::json;
use workdir::http::{WorkdirSessionOperation, WorkdirSessionOperationResult};
use workdir::workspace::{WorkspaceWorkdirSessionFence, WorkspaceWorkdirSessionOperationRequest};
use workdir::workspace::WorkspaceWorkdirSessionOperationRequest;
use workdir::{
CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandStatus, EditRequest,
EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult,
@@ -52,16 +54,48 @@ const LIST_DESCRIPTION: &str = "List persistent Workdirs in the current Workspac
const CREATE_DESCRIPTION: &str = "Materialize a persistent Workdir on a selected Runtime from a Workspace repository and optional selector. This does not change this Worker's attachment; use WorkdirAttach explicitly after creation.";
const ATTACH_DESCRIPTION: &str = "Attach this Worker to one existing Workdir. The Backend enforces one active Workdir per Worker and one active Worker per Workdir, then opens an ephemeral operation session.";
const DETACH_DESCRIPTION: &str = "Detach this Worker from its active Workdir and release Workdir occupancy. Any ephemeral operation session is closed.";
pub(crate) type BeforeWorkdirRelease =
Arc<dyn Fn() -> Pin<Box<dyn Future<Output = std::io::Result<()>> + Send>> + Send + Sync>;
pub(crate) type AfterWorkdirAttach = Arc<dyn Fn() + Send + Sync>;
const DELETE_DESCRIPTION: &str = "Request removal of one persistent Workdir by id through durable Backend Workspace authority. The input includes only the Workdir id and a bounded reason. The result reports removed, retained, or attention_required without exposing operation-table or provider internals.";
#[derive(Clone, Debug)]
#[derive(Clone)]
pub struct ManageWorkdirFeature {
client: Arc<dyn WorkspaceClient>,
before_workdir_release: Option<BeforeWorkdirRelease>,
after_workdir_attach: Option<AfterWorkdirAttach>,
}
impl std::fmt::Debug for ManageWorkdirFeature {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("ManageWorkdirFeature")
.field("client_kind", &self.client.kind())
.field("release_guard", &self.before_workdir_release.is_some())
.finish()
}
}
impl ManageWorkdirFeature {
pub fn new(client: Arc<dyn WorkspaceClient>) -> Self {
Self { client }
Self {
client,
before_workdir_release: None,
after_workdir_attach: None,
}
}
pub(crate) fn with_child_lifecycle(
client: Arc<dyn WorkspaceClient>,
before_workdir_release: BeforeWorkdirRelease,
after_workdir_attach: AfterWorkdirAttach,
) -> Self {
Self {
client,
before_workdir_release: Some(before_workdir_release),
after_workdir_attach: Some(after_workdir_attach),
}
}
}
@@ -81,7 +115,10 @@ impl FeatureModule for ManageWorkdirFeature {
}
fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> {
let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone());
let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone()).with_child_lifecycle(
self.before_workdir_release.clone(),
self.after_workdir_attach.clone(),
);
for (name, definition) in [
(
LIST_TOOL,
@@ -142,9 +179,21 @@ impl FeatureModule for ManageWorkdirFeature {
}
}
#[derive(Clone, Debug)]
#[derive(Clone)]
struct WorkspaceHttpWorkdirBackend {
client: Arc<dyn WorkspaceClient>,
before_workdir_release: Option<BeforeWorkdirRelease>,
after_workdir_attach: Option<AfterWorkdirAttach>,
}
impl std::fmt::Debug for WorkspaceHttpWorkdirBackend {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("WorkspaceHttpWorkdirBackend")
.field("client_kind", &self.client.kind())
.field("release_guard", &self.before_workdir_release.is_some())
.finish()
}
}
/// Worker-local Workdir handle whose operation authority remains in the Workspace Backend.
@@ -156,8 +205,6 @@ struct WorkspaceHttpWorkdirBackend {
pub struct WorkspaceAttachedWorkdirSession {
client: Arc<dyn WorkspaceClient>,
workdir: Workdir,
expected_session_fence: Option<String>,
delegations: Vec<workdir::WorkdirDelegationRequest>,
}
impl WorkspaceAttachedWorkdirSession {
@@ -165,8 +212,6 @@ impl WorkspaceAttachedWorkdirSession {
Arc::new(Self {
client,
workdir: Workdir::new("workspace-attachment"),
expected_session_fence: None,
delegations: Vec::new(),
})
}
@@ -183,16 +228,13 @@ impl WorkspaceAttachedWorkdirSession {
"/api/w/{}/workers/self/workdir-session/operations",
encode_path_segment(workspace_id)
),
serde_json::to_string(&WorkspaceWorkdirSessionOperationRequest {
expected_session_fence: self.expected_session_fence.clone(),
delegations: self.delegations.clone(),
operation,
})
.map_err(|error| {
WorkdirError::Transport(format!(
"failed to encode Workspace Workdir operation: {error}"
))
})?,
serde_json::to_string(&WorkspaceWorkdirSessionOperationRequest { operation }).map_err(
|error| {
WorkdirError::Transport(format!(
"failed to encode Workspace Workdir operation: {error}"
))
},
)?,
);
let response = self
.client
@@ -241,59 +283,6 @@ impl WorkdirSession for WorkspaceAttachedWorkdirSession {
WorkdirSessionCapabilities::ALL
}
fn transports_delegation_context(&self) -> bool {
true
}
async fn capture_delegation_source(
&self,
request: &workdir::WorkdirDelegationRequest,
) -> Result<WorkdirSessionHandle, WorkdirError> {
let expected_session_fence = if let Some(fence) = &self.expected_session_fence {
fence.clone()
} else {
let workspace_id = self.client.workspace_id().ok_or_else(|| {
WorkdirError::Unavailable("Workspace identity is unavailable".to_string())
})?;
let response = self
.client
.execute(WorkspaceRequest {
method: WorkspaceRequestMethod::Get,
path: format!(
"/api/w/{}/workers/self/workdir-session/fence",
encode_path_segment(workspace_id)
),
body: None,
})
.map_err(|error| {
WorkdirError::Unavailable(format!(
"failed to capture Workdir attachment fence: {error}"
))
})?;
let fence: WorkspaceWorkdirSessionFence = serde_json::from_str(&response.body)
.map_err(|error| {
WorkdirError::Unavailable(format!(
"invalid Workdir attachment fence response: {error}"
))
})?;
fence.value
};
let mut delegations = self.delegations.clone();
delegations.push(request.clone());
let candidate = Arc::new(Self {
client: self.client.clone(),
workdir: self.workdir.clone(),
expected_session_fence: Some(expected_session_fence),
delegations,
});
candidate
.stat(StatRequest {
path: workdir::WorkdirPath::new("").expect("empty Workdir path is valid"),
})
.await?;
Ok(candidate)
}
async fn stat(&self, request: StatRequest) -> Result<StatResult, WorkdirError> {
match self.operate(WorkdirSessionOperation::Stat(request))? {
WorkdirSessionOperationResult::Stat(result) => Ok(result),
@@ -387,7 +376,21 @@ impl WorkdirSession for WorkspaceAttachedWorkdirSession {
impl WorkspaceHttpWorkdirBackend {
fn new(client: Arc<dyn WorkspaceClient>) -> Self {
Self { client }
Self {
client,
before_workdir_release: None,
after_workdir_attach: None,
}
}
fn with_child_lifecycle(
mut self,
before_workdir_release: Option<BeforeWorkdirRelease>,
after_workdir_attach: Option<AfterWorkdirAttach>,
) -> Self {
self.before_workdir_release = before_workdir_release;
self.after_workdir_attach = after_workdir_attach;
self
}
fn workspace_id(&self) -> Result<&str, ToolError> {
@@ -565,11 +568,26 @@ impl Tool for WorkspaceHttpWorkdirTool {
parse_input::<WorkdirCreateInput>(input_json)?,
ctx.call_id.to_string(),
),
WorkdirOperation::Attach => self
.backend
.attach(parse_input::<WorkdirAttachInput>(input_json)?),
WorkdirOperation::Attach => {
let result = self
.backend
.attach(parse_input::<WorkdirAttachInput>(input_json)?);
if result.is_ok()
&& let Some(after_attach) = &self.backend.after_workdir_attach
{
after_attach();
}
result
}
WorkdirOperation::Detach => {
let _input = parse_input::<WorkdirDetachInput>(input_json)?;
if let Some(before_release) = &self.backend.before_workdir_release {
before_release().await.map_err(|error| {
ToolError::ExecutionFailed(format!(
"stop Internal SubWorkers before Workdir detach: {error}"
))
})?;
}
self.backend.detach()
}
WorkdirOperation::Delete => self
@@ -765,6 +783,7 @@ struct WorkdirDeleteInput {
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};
use super::*;
use crate::feature::{FeatureModule, FeatureRegistryBuilder};
@@ -1155,6 +1174,7 @@ mod tests {
command: "true".to_string(),
timeout_secs: 120,
output_limit: 1024,
cwd: None,
spill_dir: Some("/worker-local/bash-output".into()),
tool_call_id: Some("call-1".to_string()),
})
@@ -1178,83 +1198,6 @@ mod tests {
);
}
#[tokio::test]
async fn delegated_attached_session_carries_captured_fence_on_operations() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![
response(json!({"value": "attachment-fence"})),
response(json!({
"operation": "stat",
"result": {"path": "", "kind": "directory", "size": 0}
})),
response(json!({
"operation": "stat",
"result": {"path": "visible.txt", "kind": "file", "size": 8}
})),
]));
let parent = workdir::delegation_capable_session(WorkspaceAttachedWorkdirSession::handle(
client.clone(),
));
let delegation = parent
.delegate(workdir::WorkdirDelegationRequest {
rules: vec![workdir::WorkdirDelegationRule {
target: workdir::WorkdirPath::new("").unwrap(),
permission: workdir::WorkdirDelegationPermission::Read,
recursive: false,
}],
cwd: workdir::WorkdirPath::new("").unwrap(),
})
.await
.unwrap();
delegation
.scoped_session
.stat(StatRequest {
path: workdir::WorkdirPath::new("visible.txt").unwrap(),
})
.await
.unwrap();
let requests = client.requests();
assert_eq!(requests.len(), 3);
assert_eq!(
requests[0].path,
"/api/w/workspace%2Ftest/workers/self/workdir-session/fence"
);
let body: serde_json::Value =
serde_json::from_str(requests[2].body.as_deref().unwrap()).unwrap();
assert_eq!(body["expected_session_fence"], "attachment-fence");
assert_eq!(body["operation"]["operation"], "stat");
assert_eq!(body["delegations"][0]["rules"][0]["target"], "");
}
#[tokio::test]
async fn attached_provider_rejection_happens_before_delegation_is_returned() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![
response(json!({"value": "attachment-fence"})),
response(json!({"error": "provider rejected delegated write target"})),
]));
let parent = workdir::delegation_capable_session(WorkspaceAttachedWorkdirSession::handle(
client.clone(),
));
let result = parent
.delegate(workdir::WorkdirDelegationRequest {
rules: vec![workdir::WorkdirDelegationRule {
target: workdir::WorkdirPath::new("linked-target").unwrap(),
permission: workdir::WorkdirDelegationPermission::Write,
recursive: true,
}],
cwd: workdir::WorkdirPath::new("linked-target").unwrap(),
})
.await;
assert!(result.is_err(), "provider rejection must fail before lease");
let requests = client.requests();
assert_eq!(requests.len(), 2);
let validation: serde_json::Value =
serde_json::from_str(requests[1].body.as_deref().unwrap()).unwrap();
assert_eq!(validation["operation"]["operation"], "stat");
assert_eq!(validation["delegations"].as_array().unwrap().len(), 1);
}
#[tokio::test]
async fn attached_session_preserves_typed_provider_validation_error() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![error_response(
@@ -1298,73 +1241,52 @@ mod tests {
}
#[tokio::test]
async fn nested_attached_session_preserves_full_delegation_chain() {
async fn scoped_broker_operations_carry_no_child_context() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![
response(json!({"value": "attachment-fence"})),
response(json!({
"operation": "stat",
"result": {"path": "", "kind": "directory", "size": 0}
"result": {"path": "visible.txt", "kind": "file", "size": 8}
})),
response(json!({
"operation": "stat",
"result": {"path": "nested", "kind": "directory", "size": 0}
})),
response(json!({
"operation": "stat",
"result": {"path": "nested/file", "kind": "file", "size": 1}
"result": {"path": "visible.txt", "kind": "file", "size": 8}
})),
]));
let parent = workdir::delegation_capable_session(WorkspaceAttachedWorkdirSession::handle(
let broker = workdir::WorkdirToolBroker::new(WorkspaceAttachedWorkdirSession::handle(
client.clone(),
));
let outer = parent
.delegate(workdir::WorkdirDelegationRequest {
rules: vec![workdir::WorkdirDelegationRule {
let scoped = broker
.scope(workdir::WorkdirToolScope {
rules: vec![workdir::WorkdirToolScopeRule {
target: workdir::WorkdirPath::new("").unwrap(),
permission: workdir::WorkdirDelegationPermission::Read,
permission: workdir::WorkdirToolScopePermission::Read,
recursive: true,
}],
cwd: workdir::WorkdirPath::new("").unwrap(),
command: false,
})
.await
.unwrap();
let nested = outer
.scoped_session
.delegate(workdir::WorkdirDelegationRequest {
rules: vec![workdir::WorkdirDelegationRule {
target: workdir::WorkdirPath::new("nested").unwrap(),
permission: workdir::WorkdirDelegationPermission::Read,
recursive: true,
}],
cwd: workdir::WorkdirPath::new("nested").unwrap(),
})
.await
.unwrap();
nested
.scoped_session
scoped
.stat(StatRequest {
path: workdir::WorkdirPath::new("file").unwrap(),
path: workdir::WorkdirPath::new("visible.txt").unwrap(),
})
.await
.unwrap();
let requests = client.requests();
assert_eq!(requests.len(), 4);
let outer_validation: serde_json::Value =
serde_json::from_str(requests[1].body.as_deref().unwrap()).unwrap();
let nested_validation: serde_json::Value =
serde_json::from_str(requests[2].body.as_deref().unwrap()).unwrap();
assert_eq!(outer_validation["delegations"].as_array().unwrap().len(), 1);
assert_eq!(
nested_validation["delegations"].as_array().unwrap().len(),
2
);
let body: serde_json::Value =
serde_json::from_str(requests[3].body.as_deref().unwrap()).unwrap();
assert_eq!(body["delegations"].as_array().unwrap().len(), 2);
assert_eq!(body["delegations"][0]["rules"][0]["target"], "");
assert_eq!(body["delegations"][1]["rules"][0]["target"], "nested");
assert_eq!(body["operation"]["request"]["path"], "file");
assert_eq!(requests.len(), 2);
for request in requests {
assert_eq!(
request.path,
"/api/w/workspace%2Ftest/workers/self/workdir-session/operations"
);
let body: serde_json::Value =
serde_json::from_str(request.body.as_deref().unwrap()).unwrap();
assert!(body.get("delegations").is_none());
assert!(body.get("child").is_none());
assert!(body.get("expected_session_fence").is_none());
}
}
#[test]
@@ -1416,4 +1338,86 @@ mod tests {
assert!(client.requests().is_empty());
assert!(parse_input::<WorkdirListInput>(r#"{"path":"/tmp"}"#).is_err());
}
#[tokio::test]
async fn detach_stops_internal_subworkers_before_backend_release() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![response(json!({
"workspace_id": "workspace/test",
"workdir_id": "wd-attached",
"attached": false
}))]));
let cleanup_calls = Arc::new(AtomicUsize::new(0));
let cleanup_calls_for_guard = cleanup_calls.clone();
let before_release: BeforeWorkdirRelease = Arc::new(move || {
let cleanup_calls = cleanup_calls_for_guard.clone();
Box::pin(async move {
cleanup_calls.fetch_add(1, Ordering::SeqCst);
Ok(())
})
});
let tool = WorkspaceHttpWorkdirTool {
backend: WorkspaceHttpWorkdirBackend::new(client.clone())
.with_child_lifecycle(Some(before_release), None),
operation: WorkdirOperation::Detach,
};
tool.execute("{}", ToolExecutionContext::default())
.await
.unwrap();
assert_eq!(cleanup_calls.load(Ordering::SeqCst), 1);
assert_eq!(client.requests().len(), 1);
assert_eq!(
client.requests()[0].path,
"/api/w/workspace%2Ftest/workers/self/workdir-attachment"
);
}
#[tokio::test]
async fn detach_does_not_release_backend_when_child_cleanup_fails() {
let client = Arc::new(RecordingWorkspaceClient::new(Vec::new()));
let before_release: BeforeWorkdirRelease =
Arc::new(|| Box::pin(async { Err(std::io::Error::other("child cleanup failed")) }));
let tool = WorkspaceHttpWorkdirTool {
backend: WorkspaceHttpWorkdirBackend::new(client.clone())
.with_child_lifecycle(Some(before_release), None),
operation: WorkdirOperation::Detach,
};
let error = tool
.execute("{}", ToolExecutionContext::default())
.await
.unwrap_err();
assert!(error.to_string().contains("stop Internal SubWorkers"));
assert!(client.requests().is_empty());
}
#[tokio::test]
async fn successful_attach_reopens_internal_subworker_admission() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![response(json!({
"workspace_id": "workspace/test",
"workdir_id": "wd-attached",
"attached": true
}))]));
let reopen_calls = Arc::new(AtomicUsize::new(0));
let reopen_calls_for_hook = reopen_calls.clone();
let after_attach: AfterWorkdirAttach = Arc::new(move || {
reopen_calls_for_hook.fetch_add(1, Ordering::SeqCst);
});
let tool = WorkspaceHttpWorkdirTool {
backend: WorkspaceHttpWorkdirBackend::new(client)
.with_child_lifecycle(None, Some(after_attach)),
operation: WorkdirOperation::Attach,
};
tool.execute(
r#"{"workdir_id":"wd-attached"}"#,
ToolExecutionContext::default(),
)
.await
.unwrap();
assert_eq!(reopen_calls.load(Ordering::SeqCst), 1);
}
}
crates/worker/src/internal_worker.rs
+6 -2
View File
@@ -744,7 +744,7 @@ pub(crate) fn prepare_internal_worker_from_spec(
}
Box::pin(prepare_internal_worker_session(
worker, store, visibility, None, None,
worker, store, visibility, None, None, None,
))
.await
})
@@ -781,13 +781,16 @@ pub(crate) async fn prepare_internal_worker_session(
visibility: InternalWorkerVisibility,
child_registry: Option<Arc<SpawnedWorkerRegistry>>,
on_turn_end: Option<Arc<dyn Fn(InternalWorkerSessionStatus) + Send + Sync>>,
command_event_broker: Option<workdir::WorkdirToolBroker>,
) -> Result<InternalWorkerSessionHandle, InternalWorkerSessionError> {
let (event_tx, _event_rx) = broadcast::channel(256);
let sink = worker.sink();
spawn_internal_log_event_bridge(sink.clone(), event_tx.clone());
let alerter = Alerter::new(event_tx.clone());
let in_flight = InFlightEvents::new(event_tx.clone());
if let Some(session) = worker.workdir_session() {
if let Some(broker) = command_event_broker.as_ref() {
wire_workdir_command_events(&broker.tool_session(), &in_flight);
} else if let Some(session) = worker.workdir_session() {
wire_workdir_command_events(session, &in_flight);
}
let actor_in_flight = in_flight.clone();
@@ -918,6 +921,7 @@ pub(crate) async fn spawn_prepared_internal_worker_session(
InternalWorkerVisibility::ServicePrivate,
None,
on_turn_end,
None,
)
.await?;
handle.send(input).await?;
crates/worker/src/spawn/registry.rs
+294 -27
View File
@@ -12,7 +12,7 @@ use std::collections::{BTreeMap, HashSet};
use std::io;
use std::sync::{
Arc, Mutex,
atomic::{AtomicBool, AtomicU64, Ordering},
atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
};
use std::time::Instant;
@@ -23,9 +23,9 @@ use protocol::{Event, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnaps
use session_store::{
LoggedItem, WorkerMetadataStore, WorkerReclaimedChild, WorkerSpawnedChild, WorkerStoreError,
};
use tokio::sync::broadcast;
use tokio::sync::{Notify, broadcast};
use tracing::warn;
use workdir::WorkdirDelegation;
use workdir::WorkdirScopeLease;
use crate::internal_worker::{InternalWorkerSessionHandle, InternalWorkerVisibility};
use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord};
@@ -68,10 +68,11 @@ pub(crate) struct SubWorkerStopSummary {
pub(crate) struct InternalSpawnedWorkerRecord {
pub worker_name: String,
pub scope_delegated: Vec<ScopeRule>,
pub workdir_delegation: Arc<WorkdirDelegation>,
pub workdir_tool_scope: Arc<WorkdirScopeLease>,
#[cfg(test)]
pub installed_tools: Arc<[String]>,
pub session: InternalWorkerSessionHandle,
pub child_registry: Arc<SpawnedWorkerRegistry>,
change_tracker: Option<tools::Tracker>,
started_at: Instant,
stop_lock: Arc<tokio::sync::Mutex<()>>,
@@ -86,18 +87,20 @@ impl InternalSpawnedWorkerRecord {
pub(crate) fn new(
worker_name: String,
scope_delegated: Vec<ScopeRule>,
workdir_delegation: WorkdirDelegation,
workdir_tool_scope: WorkdirScopeLease,
#[cfg(test)] installed_tools: Vec<String>,
session: InternalWorkerSessionHandle,
child_registry: Arc<SpawnedWorkerRegistry>,
change_tracker: Option<tools::Tracker>,
) -> Self {
Self {
worker_name,
scope_delegated,
workdir_delegation: Arc::new(workdir_delegation),
workdir_tool_scope: Arc::new(workdir_tool_scope),
#[cfg(test)]
installed_tools: installed_tools.into(),
session,
child_registry,
change_tracker,
started_at: Instant::now(),
stop_lock: Arc::new(tokio::sync::Mutex::new(())),
@@ -235,18 +238,56 @@ pub(crate) struct InternalSpawnReservation {
}
impl InternalSpawnReservation {
pub(crate) fn commit(mut self, record: InternalSpawnedWorkerRecord) -> io::Result<()> {
if record.worker_name != self.worker_name {
return Err(io::Error::new(
pub(crate) async fn commit(mut self, record: InternalSpawnedWorkerRecord) -> io::Result<()> {
let rejection = if record.worker_name != self.worker_name {
Some(io::Error::new(
io::ErrorKind::InvalidInput,
"internal SubWorker reservation name does not match record name",
));
))
} else {
match self.registry.internal_records.lock() {
Ok(mut records) => {
if self.registry.internal_shutting_down.load(Ordering::Acquire) {
Some(io::Error::new(
io::ErrorKind::Interrupted,
"internal SubWorker registry is shutting down",
))
} else {
records.push(record.clone());
None
}
}
Err(_) => Some(io::Error::other(
"internal spawned-worker registry lock poisoned",
)),
}
};
if let Some(error) = rejection {
let mut cleanup_failures = Vec::new();
if let Err(cleanup) = record.session.stop().await {
cleanup_failures.push(format!("stop rejected Internal SubWorker: {cleanup}"));
}
if let Err(cleanup) = Box::pin(record.child_registry.shutdown_internal()).await {
cleanup_failures.push(format!(
"stop rejected Internal SubWorker descendants: {cleanup}"
));
}
if let Err(cleanup) = record.workdir_tool_scope.close().await {
cleanup_failures.push(format!(
"close rejected Internal SubWorker Workdir tools: {cleanup}"
));
}
if cleanup_failures.is_empty() {
return Err(error);
}
self.registry
.internal_spawn_cleanup_failed
.store(true, Ordering::Release);
return Err(io::Error::other(format!(
"{error}; {}",
cleanup_failures.join("; ")
)));
}
self.registry
.internal_records
.lock()
.map_err(|_| io::Error::other("internal spawned-worker registry lock poisoned"))?
.push(record.clone());
self.registry.start_protocol_forwarding(record);
self.committed = true;
Ok(())
@@ -260,6 +301,10 @@ impl Drop for InternalSpawnReservation {
names.remove(&self.worker_name);
}
}
self.registry
.pending_internal_spawns
.fetch_sub(1, Ordering::AcqRel);
self.registry.pending_internal_notify.notify_waiters();
}
}
@@ -267,6 +312,10 @@ pub struct SpawnedWorkerRegistry {
internal_records: std::sync::Mutex<Vec<InternalSpawnedWorkerRecord>>,
service_records: std::sync::Mutex<Vec<InternalServiceWorkerRecord>>,
internal_names: std::sync::Mutex<HashSet<String>>,
internal_shutting_down: AtomicBool,
pending_internal_spawns: AtomicUsize,
pending_internal_notify: Notify,
internal_spawn_cleanup_failed: AtomicBool,
parent_scope: Option<SharedScope>,
parent_protocol: Mutex<Option<(broadcast::Sender<Event>, String)>>,
}
@@ -283,6 +332,10 @@ impl SpawnedWorkerRegistry {
internal_records: std::sync::Mutex::new(Vec::new()),
service_records: std::sync::Mutex::new(Vec::new()),
internal_names: std::sync::Mutex::new(HashSet::new()),
internal_shutting_down: AtomicBool::new(false),
pending_internal_spawns: AtomicUsize::new(0),
pending_internal_notify: Notify::new(),
internal_spawn_cleanup_failed: AtomicBool::new(false),
parent_scope: None,
parent_protocol: Mutex::new(None),
})
@@ -294,6 +347,10 @@ impl SpawnedWorkerRegistry {
internal_records: std::sync::Mutex::new(Vec::new()),
service_records: std::sync::Mutex::new(Vec::new()),
internal_names: std::sync::Mutex::new(HashSet::new()),
internal_shutting_down: AtomicBool::new(false),
pending_internal_spawns: AtomicUsize::new(0),
pending_internal_notify: Notify::new(),
internal_spawn_cleanup_failed: AtomicBool::new(false),
parent_scope: None,
parent_protocol: Mutex::new(None),
})
@@ -304,6 +361,10 @@ impl SpawnedWorkerRegistry {
internal_records: std::sync::Mutex::new(Vec::new()),
service_records: std::sync::Mutex::new(Vec::new()),
internal_names: std::sync::Mutex::new(HashSet::new()),
internal_shutting_down: AtomicBool::new(false),
pending_internal_spawns: AtomicUsize::new(0),
pending_internal_notify: Notify::new(),
internal_spawn_cleanup_failed: AtomicBool::new(false),
parent_scope: Some(parent_scope),
parent_protocol: Mutex::new(None),
})
@@ -383,6 +444,10 @@ impl SpawnedWorkerRegistry {
internal_records: std::sync::Mutex::new(Vec::new()),
service_records: std::sync::Mutex::new(Vec::new()),
internal_names: std::sync::Mutex::new(HashSet::new()),
internal_shutting_down: AtomicBool::new(false),
pending_internal_spawns: AtomicUsize::new(0),
pending_internal_notify: Notify::new(),
internal_spawn_cleanup_failed: AtomicBool::new(false),
parent_scope,
parent_protocol: Mutex::new(None),
}),
@@ -394,6 +459,16 @@ impl SpawnedWorkerRegistry {
self: &Arc<Self>,
worker_name: String,
) -> io::Result<InternalSpawnReservation> {
let records = self
.internal_records
.lock()
.map_err(|_| io::Error::other("internal Worker registry lock poisoned"))?;
if self.internal_shutting_down.load(Ordering::Acquire) {
return Err(io::Error::new(
io::ErrorKind::Interrupted,
"internal SubWorker registry is shutting down",
));
}
let mut names = self
.internal_names
.lock()
@@ -404,7 +479,9 @@ impl SpawnedWorkerRegistry {
format!("spawned worker `{worker_name}` is already registered"),
));
}
self.pending_internal_spawns.fetch_add(1, Ordering::AcqRel);
drop(names);
drop(records);
Ok(InternalSpawnReservation {
registry: Arc::clone(self),
worker_name,
@@ -679,18 +756,11 @@ impl SpawnedWorkerRegistry {
.unwrap_or_default()
}
pub(crate) fn reclaim_internal_scope(&self, worker_name: &str) -> io::Result<bool> {
let record = self.get_internal(worker_name).ok_or_else(|| {
io::Error::new(io::ErrorKind::NotFound, "internal SubWorker not found")
})?;
self.reclaim_record_scope(&record)
}
fn reclaim_record_scope(&self, record: &InternalSpawnedWorkerRecord) -> io::Result<bool> {
if !record.claim_scope_reclaim() {
return Ok(false);
}
record.workdir_delegation.release();
record.workdir_tool_scope.revoke();
let result = if let Some(parent_scope) = &self.parent_scope {
parent_scope
.update(|current| current.with_removed_deny_rules(delegated_write_rules(record)))
@@ -705,6 +775,58 @@ impl SpawnedWorkerRegistry {
result
}
pub(crate) async fn close_internal_scope(&self, name: &str) -> io::Result<bool> {
let Some(record) = self.get_internal(name) else {
return Ok(false);
};
Box::pin(record.child_registry.shutdown_internal()).await?;
record
.workdir_tool_scope
.close()
.await
.map_err(|error| io::Error::other(error.to_string()))?;
self.reclaim_record_scope(&record)
}
pub(crate) async fn shutdown_internal(&self) -> io::Result<()> {
let names = {
let records = self
.internal_records
.lock()
.map_err(|_| io::Error::other("internal Worker registry lock poisoned"))?;
self.internal_shutting_down.store(true, Ordering::Release);
records
.iter()
.map(|record| record.worker_name.clone())
.collect::<Vec<_>>()
};
loop {
let notified = self.pending_internal_notify.notified();
if self.pending_internal_spawns.load(Ordering::Acquire) == 0 {
break;
}
notified.await;
}
let mut first_error = None;
for name in names {
if let Err(error) = self.remove_internal(&name).await {
first_error.get_or_insert(error);
}
}
if first_error.is_none() && self.internal_spawn_cleanup_failed.load(Ordering::Acquire) {
first_error = Some(io::Error::other(
"an in-flight Internal SubWorker failed cleanup during shutdown",
));
}
first_error.map_or(Ok(()), Err)
}
pub(crate) fn reopen_internal(&self) {
self.internal_shutting_down.store(false, Ordering::Release);
self.internal_spawn_cleanup_failed
.store(false, Ordering::Release);
}
/// Stop one direct Internal SubWorker and discard its registry/scope state.
///
/// The child actor must acknowledge its stop before the registry is removed.
@@ -731,6 +853,12 @@ impl SpawnedWorkerRegistry {
.stop()
.await
.map_err(|error| io::Error::other(error.to_string()))?;
Box::pin(record.child_registry.shutdown_internal()).await?;
record
.workdir_tool_scope
.close()
.await
.map_err(|error| io::Error::other(error.to_string()))?;
let summary = record.stop_summary();
self.reclaim_record_scope(&record)?;
let removed =
@@ -966,7 +1094,7 @@ mod tests {
deny: Vec::new(),
})
.unwrap();
let source = workdir::delegation_capable_session(Arc::new(
let source = workdir::WorkdirToolBroker::new(Arc::new(
workdir::LocalWorkdirSession::materialized_bound(
workdir::Workdir::new("registry-test"),
root.clone(),
@@ -976,13 +1104,14 @@ mod tests {
),
));
let delegation = source
.delegate(workdir::WorkdirDelegationRequest {
rules: vec![workdir::WorkdirDelegationRule {
.scope(workdir::WorkdirToolScope {
rules: vec![workdir::WorkdirToolScopeRule {
target: workdir::WorkdirPath::new("").unwrap(),
permission: workdir::WorkdirDelegationPermission::Read,
permission: workdir::WorkdirToolScopePermission::Read,
recursive: true,
}],
cwd: workdir::WorkdirPath::new("").unwrap(),
command: false,
})
.await
.unwrap();
@@ -993,6 +1122,7 @@ mod tests {
delegation,
Vec::new(),
session,
registry(),
None,
),
sender,
@@ -1230,6 +1360,143 @@ mod tests {
}
}
#[tokio::test]
async fn parent_shutdown_stops_all_internal_workers_before_returning() {
let registry = registry();
for name in ["first", "second"] {
let (record, _events) = record(name, InternalWorkerVisibility::ParentClient).await;
record
.session
.force_status(InternalWorkerSessionStatus::Running);
install_record(&registry, record);
}
registry.shutdown_internal().await.unwrap();
assert!(registry.list_internal().is_empty());
assert!(registry.get_internal("first").is_none());
assert!(registry.get_internal("second").is_none());
}
#[tokio::test]
async fn shutdown_rejects_new_reservations_until_reopened() {
let registry = registry();
registry.shutdown_internal().await.unwrap();
assert!(registry.reserve_internal_name("late-child".into()).is_err());
registry.reopen_internal();
let reservation = registry.reserve_internal_name("late-child".into()).unwrap();
drop(reservation);
}
#[tokio::test]
async fn concurrent_commit_and_shutdown_leave_no_live_internal_worker() {
let registry = registry();
let reservation = registry
.reserve_internal_name("racing-child".into())
.unwrap();
let (record, _events) =
record("racing-child", InternalWorkerVisibility::ParentClient).await;
let scope = record.workdir_tool_scope.clone();
let barrier = Arc::new(tokio::sync::Barrier::new(2));
let commit_barrier = barrier.clone();
let commit = tokio::spawn(async move {
commit_barrier.wait().await;
reservation.commit(record).await
});
let shutdown_registry = registry.clone();
let shutdown = tokio::spawn(async move {
barrier.wait().await;
shutdown_registry.shutdown_internal().await
});
let commit = commit.await.unwrap();
shutdown.await.unwrap().unwrap();
if let Err(error) = commit {
assert_eq!(error.kind(), io::ErrorKind::Interrupted);
}
assert!(registry.list_internal().is_empty());
assert!(!scope.is_active());
}
#[tokio::test]
async fn shutdown_fences_a_reservation_that_has_not_committed() {
let registry = registry();
let reservation = registry
.reserve_internal_name("racing-child".into())
.unwrap();
let (record, _events) =
record("racing-child", InternalWorkerVisibility::ParentClient).await;
let mut shutdown = {
let registry = registry.clone();
tokio::spawn(async move { registry.shutdown_internal().await })
};
while !registry.internal_shutting_down.load(Ordering::Acquire) {
tokio::task::yield_now().await;
}
assert!(
tokio::time::timeout(std::time::Duration::from_millis(50), &mut shutdown)
.await
.is_err(),
"shutdown must wait for the pending spawn to roll back"
);
let error = reservation.commit(record).await.unwrap_err();
assert_eq!(error.kind(), io::ErrorKind::Interrupted);
shutdown.await.unwrap().unwrap();
}
#[tokio::test]
async fn rejected_spawn_cleanup_failure_keeps_shutdown_failed_closed() {
let registry = registry();
let reservation = registry
.reserve_internal_name("cleanup-failure".into())
.unwrap();
let (record, _events) =
record("cleanup-failure", InternalWorkerVisibility::ParentClient).await;
record.session.force_stop_failure();
let shutdown = {
let registry = registry.clone();
tokio::spawn(async move { registry.shutdown_internal().await })
};
while !registry.internal_shutting_down.load(Ordering::Acquire) {
tokio::task::yield_now().await;
}
let error = reservation.commit(record).await.unwrap_err();
assert!(
error
.to_string()
.contains("stop rejected Internal SubWorker")
);
let shutdown_error = shutdown.await.unwrap().unwrap_err();
assert!(
shutdown_error
.to_string()
.contains("failed cleanup during shutdown")
);
assert!(registry.internal_shutting_down.load(Ordering::Acquire));
}
#[tokio::test]
async fn shutdown_recursively_stops_grandchildren_before_parent_scope_release() {
let registry = registry();
let (child, _child_events) = record("child", InternalWorkerVisibility::ParentClient).await;
let child_registry = child.child_registry.clone();
let (grandchild, _grandchild_events) =
record("grandchild", InternalWorkerVisibility::ParentClient).await;
let grandchild_scope = grandchild.workdir_tool_scope.clone();
install_record(&child_registry, grandchild);
install_record(&registry, child);
registry.shutdown_internal().await.unwrap();
assert!(registry.list_internal().is_empty());
assert!(child_registry.list_internal().is_empty());
assert!(!grandchild_scope.is_active());
}
#[tokio::test]
async fn running_worker_is_stopped_before_removal() {
let registry = registry();
crates/worker/src/spawn/tool.rs
+85 -176
View File
@@ -22,8 +22,7 @@ use manifest::{
use serde::Deserialize;
use tokio::sync::mpsc;
use workdir::{
WorkdirDelegationPermission, WorkdirDelegationRequest, WorkdirDelegationRule, WorkdirPath,
WorkdirSessionHandle,
WorkdirToolBroker, WorkdirToolScope, WorkdirToolScopePermission, WorkdirToolScopeRule,
};
use crate::PromptCatalogSource;
@@ -64,6 +63,9 @@ struct SubWorkerSpawnInput {
/// spawner's explicit delegation authority; direct tool scope alone is not
/// sufficient. Omit `recursive` for normal workspace/worktree delegation; it defaults to true.
scope: Vec<ScopeRuleInput>,
/// Explicitly grant command execution through the parent-owned Workdir tool broker.
#[serde(default)]
command: bool,
/// Binds an actual read-only builtin Reviewer child to the current Merge Request candidate.
/// Review capability material is generated by the trusted spawn layer.
#[serde(default)]
@@ -284,8 +286,8 @@ pub struct SubWorkerSpawnTool {
workspace_root: PathBuf,
/// Directory the spawned SubWorker's tools should use when the LLM did not
/// override it. Defaults to the spawner's cwd.
/// Active provider-backed Workdir session from which child leases are captured.
source_workdir_session: Option<WorkdirSessionHandle>,
/// Parent-owned broker for scoped Workdir tool execution.
workdir_tool_broker: Option<WorkdirToolBroker>,
/// Parent-owned in-memory registry shared by the five SubWorker tools.
registry: Arc<SpawnedWorkerRegistry>,
/// Spawner's resolved Manifest. `profile = "inherit"` derives the
@@ -312,7 +314,7 @@ impl SubWorkerSpawnTool {
runtime_base: PathBuf,
bash_output_dir: PathBuf,
workspace_root: PathBuf,
source_workdir_session: Option<WorkdirSessionHandle>,
workdir_tool_broker: Option<WorkdirToolBroker>,
registry: Arc<SpawnedWorkerRegistry>,
spawner_manifest: WorkerManifest,
prompt_loader: PromptCatalogSource,
@@ -325,7 +327,7 @@ impl SubWorkerSpawnTool {
runtime_base,
bash_output_dir,
workspace_root,
source_workdir_session,
workdir_tool_broker,
registry,
spawner_manifest,
prompt_loader,
@@ -358,6 +360,11 @@ fn validate_reviewer_handoff(input: &SubWorkerSpawnInput) -> Result<(), ToolErro
"Merge Request Reviewer SubWorkers must include writable delegated scope".to_string(),
));
}
if !input.command {
return Err(ToolError::InvalidArgument(
"Merge Request Reviewer SubWorkers require an explicit command grant".to_string(),
));
}
Ok(())
}
@@ -387,7 +394,7 @@ impl Tool for SubWorkerSpawnTool {
.reserve_internal_name(input.name.clone())
.map_err(|error| ToolError::InvalidArgument(error.to_string()))?;
let mut workdir_rules = parse_workdir_scope(&input.scope)?;
let workdir_rules = parse_workdir_scope(&input.scope)?;
let child_bash_output_dir = self.bash_output_dir.join("sub-workers").join(&input.name);
tokio::fs::create_dir_all(&child_bash_output_dir)
.await
@@ -397,28 +404,15 @@ impl Tool for SubWorkerSpawnTool {
child_bash_output_dir.display()
))
})?;
let source_workdir_session =
require_active_workdir_session(self.source_workdir_session.as_ref())?;
let transports_delegation_context = source_workdir_session.transports_delegation_context();
// Provider-transported sessions resolve every delegation rule in the
// receiving Workdir namespace. The Bash spill directory instead belongs
// to this Worker host, so forwarding it would widen the request with a
// foreign absolute path and fail the provider's existing scope check.
if !transports_delegation_context {
workdir_rules.push(WorkdirDelegationRule {
target: WorkdirPath::new_scoped(child_bash_output_dir.to_string_lossy())
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?,
permission: WorkdirDelegationPermission::Read,
recursive: true,
});
}
let delegation_request = workdir_delegation_request(input.cwd.as_deref(), workdir_rules)?;
let workdir_delegation = source_workdir_session
.delegate(delegation_request)
let workdir_tool_broker = require_workdir_tool_broker(self.workdir_tool_broker.as_ref())?;
let tool_scope = workdir_tool_scope(input.cwd.as_deref(), workdir_rules, input.command)?;
let workdir_scope = workdir_tool_broker
.scope(tool_scope)
.await
.map_err(|error| {
ToolError::InvalidArgument(format!("delegate Workdir session: {error}"))
ToolError::InvalidArgument(format!("scope parent-owned Workdir tools: {error}"))
})?;
let child_workdir_tool_broker = workdir_scope.broker();
let spawn_selector =
parse_spawn_profile_selector(input.profile.as_deref()).map_err(|msg| {
@@ -507,7 +501,6 @@ impl Tool for SubWorkerSpawnTool {
)
.await
.map_err(|error| ToolError::ExecutionFailed(format!("build Internal Worker: {error}")))?;
child.bind_workdir_session(Some(workdir_delegation.scoped_session.clone()));
child
.add_scope_rules([ScopeRule {
target: child_bash_output_dir.clone(),
@@ -527,6 +520,7 @@ impl Tool for SubWorkerSpawnTool {
self.runtime_base.clone(),
child_registry.clone(),
None,
Some(child_workdir_tool_broker.clone()),
)
.await
.map_err(|error| {
@@ -555,13 +549,16 @@ impl Tool for SubWorkerSpawnTool {
InternalWorkerSessionStatus::Failed | InternalWorkerSessionStatus::Stopped
) {
if let Some(registry) = registry.upgrade() {
if let Err(error) = registry.reclaim_internal_scope(&child_name) {
tracing::warn!(
child_name,
%error,
"failed to reclaim delegated scope after Internal SubWorker failure"
);
}
let child_name = child_name.clone();
tokio::spawn(async move {
if let Err(error) = registry.close_internal_scope(&child_name).await {
tracing::warn!(
child_name,
%error,
"failed to close parent-owned Workdir tools after Internal SubWorker failure"
);
}
});
}
}
let message = format!(
@@ -569,6 +566,7 @@ impl Tool for SubWorkerSpawnTool {
);
parent_notifications.notify(child_name.clone(), message, true);
})),
Some(child_workdir_tool_broker.clone()),
)
.await;
let session = session_result.map_err(|error| {
@@ -619,15 +617,19 @@ impl Tool for SubWorkerSpawnTool {
),
body.to_string(),
);
let response = self
.workspace_context
.client()
.execute(request)
.map_err(|error| {
ToolError::ExecutionFailed(format!("register review capability: {error}"))
})?;
let response = match self.workspace_context.client().execute(request) {
Ok(response) => response,
Err(error) => {
let _ = session.stop().await;
let _ = workdir_scope.close().await;
return Err(ToolError::ExecutionFailed(format!(
"register review capability: {error}"
)));
}
};
if !response.is_success() {
let _ = session.stop().await;
let _ = workdir_scope.close().await;
return Err(ToolError::ExecutionFailed(format!(
"register review capability failed with status {}: {}",
response.status, response.body
@@ -638,14 +640,14 @@ impl Tool for SubWorkerSpawnTool {
let record = crate::spawn::registry::InternalSpawnedWorkerRecord::new(
input.name.clone(),
scope_allow,
workdir_delegation,
workdir_scope,
#[cfg(test)]
installed_tools,
session.clone(),
child_registry,
child_change_tracker,
);
if let Err(error) = name_reservation.commit(record) {
let _ = session.stop().await;
if let Err(error) = name_reservation.commit(record).await {
return Err(ToolError::ExecutionFailed(format!(
"register Internal Worker session: {error}"
)));
@@ -691,18 +693,18 @@ fn logical_workdir_path(value: &str, field: &str) -> Result<FsPath, ToolError> {
})
}
fn parse_workdir_scope(rules: &[ScopeRuleInput]) -> Result<Vec<WorkdirDelegationRule>, ToolError> {
fn parse_workdir_scope(rules: &[ScopeRuleInput]) -> Result<Vec<WorkdirToolScopeRule>, ToolError> {
if rules.is_empty() {
return Err(ToolError::InvalidArgument("scope must not be empty".into()));
}
rules
.iter()
.map(|rule| {
Ok(WorkdirDelegationRule {
Ok(WorkdirToolScopeRule {
target: logical_workdir_path(&rule.target, "scope.target")?,
permission: match rule.permission {
PermissionInput::Read => WorkdirDelegationPermission::Read,
PermissionInput::Write => WorkdirDelegationPermission::Write,
PermissionInput::Read => WorkdirToolScopePermission::Read,
PermissionInput::Write => WorkdirToolScopePermission::Write,
},
recursive: rule.recursive,
})
@@ -710,22 +712,24 @@ fn parse_workdir_scope(rules: &[ScopeRuleInput]) -> Result<Vec<WorkdirDelegation
.collect()
}
fn workdir_delegation_request(
fn workdir_tool_scope(
cwd: Option<&str>,
rules: Vec<WorkdirDelegationRule>,
) -> Result<WorkdirDelegationRequest, ToolError> {
Ok(WorkdirDelegationRequest {
rules: Vec<WorkdirToolScopeRule>,
command: bool,
) -> Result<WorkdirToolScope, ToolError> {
Ok(WorkdirToolScope {
rules,
cwd: logical_workdir_path(cwd.unwrap_or("."), "cwd")?,
command,
})
}
fn require_active_workdir_session(
session: Option<&WorkdirSessionHandle>,
) -> Result<&WorkdirSessionHandle, ToolError> {
session.ok_or_else(|| {
fn require_workdir_tool_broker(
broker: Option<&WorkdirToolBroker>,
) -> Result<&WorkdirToolBroker, ToolError> {
broker.ok_or_else(|| {
ToolError::InvalidArgument(
"SubWorkerSpawn requires an active Workdir session; attach a Workdir before delegating filesystem access"
"SubWorkerSpawn requires parent-owned Workdir tools; attach a Workdir before granting filesystem access"
.to_string(),
)
})
@@ -963,7 +967,7 @@ pub(crate) fn sub_worker_spawn_tool(
runtime_base: PathBuf,
bash_output_dir: PathBuf,
workspace_root: PathBuf,
source_workdir_session: Option<WorkdirSessionHandle>,
workdir_tool_broker: Option<WorkdirToolBroker>,
registry: Arc<SpawnedWorkerRegistry>,
spawner_manifest: WorkerManifest,
prompts: Arc<ArcSwap<PromptCatalog>>,
@@ -975,7 +979,7 @@ pub(crate) fn sub_worker_spawn_tool(
runtime_base,
bash_output_dir,
workspace_root,
source_workdir_session,
workdir_tool_broker,
registry,
spawner_manifest,
prompts,
@@ -989,7 +993,7 @@ fn sub_worker_spawn_tool_impl(
runtime_base: PathBuf,
bash_output_dir: PathBuf,
workspace_root: PathBuf,
source_workdir_session: Option<WorkdirSessionHandle>,
workdir_tool_broker: Option<WorkdirToolBroker>,
registry: Arc<SpawnedWorkerRegistry>,
spawner_manifest: WorkerManifest,
prompts: Arc<ArcSwap<PromptCatalog>>,
@@ -1021,7 +1025,7 @@ fn sub_worker_spawn_tool_impl(
runtime_base.clone(),
bash_output_dir.clone(),
workspace_root.clone(),
source_workdir_session.clone(),
workdir_tool_broker.clone(),
registry.clone(),
spawner_manifest.clone(),
prompts.load_full().source(),
@@ -1054,12 +1058,12 @@ mod tests {
};
#[test]
fn missing_active_workdir_session_fails_deterministically() {
let error = require_active_workdir_session(None).unwrap_err();
fn missing_parent_workdir_tool_broker_fails_deterministically() {
let error = require_workdir_tool_broker(None).unwrap_err();
assert!(matches!(
error,
ToolError::InvalidArgument(message)
if message.contains("requires an active Workdir session")
if message.contains("requires parent-owned Workdir tools")
));
}
@@ -1096,6 +1100,7 @@ mod tests {
let valid: SubWorkerSpawnInput = serde_json::from_value(serde_json::json!({
"name":"reviewer","task":"review","profile":"builtin:reviewer",
"scope":[{"target":"work","permission":"write"}],
"command":true,
"review":{"ticket_id":"T1"}
}))
.unwrap();
@@ -1219,7 +1224,7 @@ enabled = false
let fail_requests = Arc::new(AtomicBool::new(false));
let prompt_loader = PromptCatalogSource::builtins_only();
let (parent_method_tx, mut parent_method_rx) = mpsc::channel(8);
let source_workdir_session = workdir::delegation_capable_session(Arc::new(
let workdir_tool_broker = workdir::WorkdirToolBroker::new(Arc::new(
workdir::LocalWorkdirSession::materialized_bound(
workdir::Workdir::new("test-workdir"),
workspace_root.clone(),
@@ -1238,7 +1243,7 @@ enabled = false
runtime.path().to_path_buf(),
bash_output_dir.clone(),
workspace_root.clone(),
Some(source_workdir_session),
Some(workdir_tool_broker),
registry.clone(),
manifest.clone(),
prompt_loader,
@@ -1261,7 +1266,8 @@ enabled = false
"target": ".",
"permission": "write",
"recursive": true
}]
}],
"command": true
});
assert!(spawner_scope.snapshot().is_writable(&workspace_root));
@@ -1296,15 +1302,6 @@ enabled = false
let record = registry
.get_internal("reviewer-child")
.expect("Internal reviewer registry record");
let child_bash_output_dir = bash_output_dir.join("sub-workers").join("reviewer-child");
record
.workdir_delegation
.scoped_session
.stat(workdir::StatRequest {
path: WorkdirPath::new_scoped(child_bash_output_dir.to_string_lossy()).unwrap(),
})
.await
.expect("local child retains read scope for its Bash output directory");
for required in ["Read", "Write", "Edit", "Glob", "Grep", "Bash"] {
assert!(
record.installed_tools.iter().any(|name| name == required),
@@ -1423,7 +1420,7 @@ enabled = false
"Stopped terminal child must release its delegated Workdir session"
);
assert!(
!record.workdir_delegation.is_active(),
!record.workdir_tool_scope.is_active(),
"stopped child must revoke cloned scoped sessions"
);
assert!(registry.get_internal("reviewer-child").is_some());
@@ -1478,7 +1475,7 @@ enabled = false
Arc::new(AvailableWorkspaceClient),
);
let remote_client = Arc::new(StrictRemoteWorkdirWorkspaceClient::default());
let source_workdir_session = workdir::delegation_capable_session(
let workdir_tool_broker = workdir::WorkdirToolBroker::new(
WorkspaceAttachedWorkdirSession::handle(remote_client.clone()),
);
let calls = Arc::new(AtomicUsize::new(0));
@@ -1493,7 +1490,7 @@ enabled = false
runtime.path().to_path_buf(),
bash_output_dir.clone(),
workspace_root.clone(),
Some(source_workdir_session),
Some(workdir_tool_broker),
registry.clone(),
manifest,
PromptCatalogSource::builtins_only(),
@@ -1533,51 +1530,12 @@ enabled = false
record.session.wait_until_idle().await,
crate::internal_worker::InternalWorkerSessionStatus::Idle
);
assert!(record.installed_tools.iter().any(|tool| tool == "Write"));
assert!(!record.installed_tools.iter().any(|tool| tool == "Bash"));
assert_eq!(calls.load(Ordering::SeqCst), 1);
assert_eq!(
remote_client
.foreign_scope_rejections
.load(Ordering::SeqCst),
0
);
let child_bash_output_dir = bash_output_dir.join("sub-workers").join("remote-child");
assert!(child_bash_output_dir.is_dir());
for required in ["Read", "Write", "Edit", "Glob", "Grep", "Bash"] {
assert!(
record.installed_tools.iter().any(|name| name == required),
"remote write-scoped child is missing {required}: {:?}",
record.installed_tools
);
}
let remote_requests = remote_client.requests();
let operate_requests = remote_requests
.iter()
.filter(|request| request.body.is_some())
.collect::<Vec<_>>();
assert_eq!(
operate_requests.len(),
1,
"remote requests: {remote_requests:?}"
);
let operation_body: serde_json::Value = serde_json::from_str(
operate_requests[0]
.body
.as_deref()
.expect("remote operation body"),
)
.unwrap();
let rules = operation_body["delegations"][0]["rules"]
.as_array()
.expect("delegation rules");
assert_eq!(rules.len(), 1, "remote operation body: {operation_body}");
assert_eq!(rules[0]["target"], "");
assert!(
!operation_body.to_string().contains(
child_bash_output_dir
.to_str()
.expect("UTF-8 test output directory")
)
remote_client.requests().is_empty(),
"spawning a child must not open or delegate a provider Workdir session"
);
}
@@ -1589,6 +1547,7 @@ enabled = false
.and_then(serde_json::Value::as_object)
.expect("schema properties");
assert!(properties.contains_key("cwd"), "schema: {schema}");
assert!(properties.contains_key("command"), "schema: {schema}");
let required = schema
.get("required")
.and_then(serde_json::Value::as_array)
@@ -1718,7 +1677,6 @@ enabled = false
#[derive(Debug, Default)]
struct StrictRemoteWorkdirWorkspaceClient {
requests: Mutex<Vec<WorkspaceRequest>>,
foreign_scope_rejections: AtomicUsize,
}
impl StrictRemoteWorkdirWorkspaceClient {
@@ -1750,59 +1708,10 @@ enabled = false
self.requests
.lock()
.expect("remote Workdir request lock")
.push(request.clone());
if request.path.ends_with("/fence") {
return Ok(WorkspaceResponse {
status: 200,
body: serde_json::json!({ "value": "remote-fence-1" }).to_string(),
});
}
let body: serde_json::Value = serde_json::from_str(
request
.body
.as_deref()
.ok_or_else(|| WorkspaceClientError::Request("missing request body".into()))?,
)
.map_err(|error| WorkspaceClientError::Request(error.to_string()))?;
let has_foreign_scope = body
.get("delegations")
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
.flat_map(|delegation| {
delegation
.get("rules")
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
})
.filter_map(|rule| rule.get("target").and_then(serde_json::Value::as_str))
.any(|target| Path::new(target).is_absolute());
if has_foreign_scope {
self.foreign_scope_rejections.fetch_add(1, Ordering::SeqCst);
return Ok(WorkspaceResponse {
status: 403,
body: serde_json::json!({
"code": "out_of_scope",
"message": "Worker-host path is outside the remote Workdir namespace"
})
.to_string(),
});
}
Ok(WorkspaceResponse {
status: 200,
body: serde_json::json!({
"operation": "stat",
"result": {
"path": "",
"kind": "directory",
"size": 0
}
})
.to_string(),
})
.push(request);
Err(WorkspaceClientError::Request(
"SubWorker spawn must not call the remote Workdir provider".into(),
))
}
}
crates/worker/tests/controller_test.rs
+4
View File
@@ -346,6 +346,7 @@ async fn shutdown_closes_bound_workdir_session() {
command: "sleep 30".to_owned(),
timeout_secs: 60,
output_limit: 1024,
cwd: None,
spill_dir: None,
tool_call_id: None,
})
@@ -395,6 +396,7 @@ async fn controller_projects_workdir_command_events_and_snapshot_state() {
command: "printf ready; sleep 0.3; printf done".to_owned(),
timeout_secs: 5,
output_limit: 1024,
cwd: None,
spill_dir: None,
tool_call_id: Some("tool-command-1".into()),
})
@@ -508,6 +510,7 @@ async fn controller_refreshes_command_snapshot_after_high_output_provider_lag()
.to_owned(),
timeout_secs: 10,
output_limit: 1024,
cwd: None,
spill_dir: None,
tool_call_id: Some("tool-high-output".into()),
})
@@ -589,6 +592,7 @@ async fn controller_startup_failure_closes_bound_workdir_session() {
command: "printf unreachable".to_owned(),
timeout_secs: 5,
output_limit: 1024,
cwd: None,
spill_dir: None,
tool_call_id: None,
})