From 68f00bc948af0d09c114b30324f269594a182926 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 02:17:44 +0900 Subject: [PATCH 1/6] refactor: broker SubWorker Workdir tools through parent --- crates/tools/src/bash.rs | 1 + crates/workdir/src/http.rs | 44 +- crates/workdir/src/lib.rs | 44 +- crates/workdir/src/local.rs | 92 +- crates/workdir/src/operation.rs | 4 + .../workdir/src/{delegation.rs => scope.rs} | 865 +++++++++++------- crates/workdir/src/workspace.rs | 10 - crates/worker-runtime/src/http_server.rs | 92 +- crates/worker/src/controller.rs | 24 +- .../src/feature/builtin/manage_workdir.rs | 221 +---- crates/worker/src/internal_worker.rs | 8 +- crates/worker/src/spawn/registry.rs | 19 +- crates/worker/src/spawn/tool.rs | 222 ++--- crates/worker/tests/controller_test.rs | 4 + crates/workspace-server/src/server.rs | 176 +--- resources/flows/coder-review.dcdl | 2 +- .../sub_worker_spawn_tool_description.md | 4 +- 17 files changed, 715 insertions(+), 1117 deletions(-) rename crates/workdir/src/{delegation.rs => scope.rs} (58%) diff --git a/crates/tools/src/bash.rs b/crates/tools/src/bash.rs index fb68fcbc..433dec0c 100644 --- a/crates/tools/src/bash.rs +++ b/crates/tools/src/bash.rs @@ -118,6 +118,7 @@ impl Tool for BashTool { command: params.command, timeout_secs, output_limit: INLINE_BYTE_BUDGET, + cwd: None, spill_dir: Some(self.output_dir.clone()), tool_call_id: Some(call_id.clone()), }) diff --git a/crates/workdir/src/http.rs b/crates/workdir/src/http.rs index 72733e8b..f89b608b 100644 --- a/crates/workdir/src/http.rs +++ b/crates/workdir/src/http.rs @@ -68,12 +68,10 @@ pub enum WorkdirSessionOperation { CommandCancel(CommandHandle), } -/// Wire envelope for an operation and its optional provider-enforced child scope. +/// Wire envelope for one provider operation. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct WorkdirSessionOperationRequest { - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub delegations: Vec, pub operation: WorkdirSessionOperation, } @@ -289,7 +287,7 @@ mod client { use reqwest::{Client, StatusCode, Url}; use super::*; - use crate::{Workdir, WorkdirSession, WorkdirSessionHandle}; + use crate::{Workdir, WorkdirSession}; /// Provides a fresh bearer token for each Runtime request. Backend /// implementations can mint short-lived capability tokens without making a @@ -324,7 +322,6 @@ mod client { workdir: Workdir, session_id: WorkdirSessionId, capabilities: WorkdirSessionCapabilities, - delegations: Vec, closed: AtomicBool, } @@ -377,7 +374,6 @@ mod client { workdir: Workdir::new(opened.workdir_id.as_str()), session_id: opened.session_id, capabilities: opened.capabilities, - delegations: Vec::new(), closed: AtomicBool::new(false), }) } @@ -404,10 +400,7 @@ mod client { "operations", ], )?; - let operation = WorkdirSessionOperationRequest { - delegations: self.delegations.clone(), - operation, - }; + let operation = WorkdirSessionOperationRequest { operation }; let response = self .client .post(url) @@ -436,37 +429,6 @@ mod client { self.capabilities } - fn transports_delegation_context(&self) -> bool { - true - } - - async fn capture_delegation_source( - &self, - request: &crate::WorkdirDelegationRequest, - ) -> Result { - if self.closed.load(Ordering::Acquire) { - return Err(WorkdirError::SessionClosed); - } - let mut delegations = self.delegations.clone(); - delegations.push(request.clone()); - let candidate = Arc::new(Self { - client: self.client.clone(), - base_url: self.base_url.clone(), - authorization: self.authorization.clone(), - workdir: self.workdir.clone(), - session_id: self.session_id.clone(), - capabilities: self.capabilities, - delegations, - closed: AtomicBool::new(false), - }); - candidate - .stat(StatRequest { - path: fs_operation::FsPath::new("").expect("empty Workdir path is valid"), - }) - .await?; - Ok(candidate) - } - async fn stat(&self, request: StatRequest) -> Result { match self.operate(WorkdirSessionOperation::Stat(request)).await? { WorkdirSessionOperationResult::Stat(result) => Ok(result), diff --git a/crates/workdir/src/lib.rs b/crates/workdir/src/lib.rs index 5d58da26..4cfa91c1 100644 --- a/crates/workdir/src/lib.rs +++ b/crates/workdir/src/lib.rs @@ -5,10 +5,10 @@ //! bound to one Worker. Tools consume sessions; they do not own Workdir //! materialization or cleanup. -mod delegation; pub mod http; mod local; mod operation; +mod scope; pub mod workspace; use std::path::{Path, PathBuf}; @@ -18,11 +18,6 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; use tokio::sync::broadcast; -pub use delegation::{ - AppliedWorkdirDelegation, ReadOnlyWorkdirSession, WorkdirDelegation, - WorkdirDelegationPermission, WorkdirDelegationRequest, WorkdirDelegationRule, - apply_delegation_chain, delegation_capable_session, -}; pub use fs_operation::{ ContentHash, EditRequest, EditResult, EntryKind, FsPath as WorkdirPath, GlobRequest, GlobResult, GrepOutputMode, GrepRequest, GrepResult, ListEntry, ListRequest, ListResult, @@ -32,6 +27,10 @@ pub use local::{ LocalWorkdirSession, SymlinkInfo, WorkdirSessionResource, direct_symlink, first_symlink, }; pub use operation::*; +pub use scope::{ + ReadOnlyWorkdirSession, WorkdirScopeLease, WorkdirToolBroker, WorkdirToolScope, + WorkdirToolScopePermission, WorkdirToolScopeRule, +}; /// Persistent, opaque identity of one materialized Workdir. #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -148,39 +147,6 @@ pub trait WorkdirSession: std::fmt::Debug + Send + Sync { fn workdir(&self) -> &Workdir; fn capabilities(&self) -> WorkdirSessionCapabilities; - fn is_delegation_capable(&self) -> bool { - false - } - - /// Whether this session transports the delegation chain to another - /// provider boundary that will apply logical cwd/path resolution there. - fn transports_delegation_context(&self) -> bool { - false - } - - /// Capture a provider-specific source for a delegated child session. - /// Remote providers use this boundary to pin attachment identity without - /// exposing transport handles or host paths. - async fn capture_delegation_source( - &self, - _request: &WorkdirDelegationRequest, - ) -> Result { - Err(WorkdirError::Denied( - "workdir provider does not support delegated sessions".into(), - )) - } - - /// Attenuate this session into a revocable child lease. Only sessions - /// created with [`delegation_capable_session`] implement this operation. - async fn delegate( - &self, - _request: WorkdirDelegationRequest, - ) -> Result { - Err(WorkdirError::Denied( - "workdir session is not delegation-capable".into(), - )) - } - async fn stat(&self, request: StatRequest) -> Result; async fn read(&self, request: ReadRequest) -> Result; async fn write(&self, request: WriteRequest) -> Result; diff --git a/crates/workdir/src/local.rs b/crates/workdir/src/local.rs index 8f9556dd..94efeceb 100644 --- a/crates/workdir/src/local.rs +++ b/crates/workdir/src/local.rs @@ -18,7 +18,7 @@ use std::sync::{Arc, Mutex as StdMutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use async_trait::async_trait; -use manifest::{Permission, Scope, ScopeConfig, ScopeRule, SharedScope}; +use manifest::{Scope, SharedScope}; use sha2::{Digest, Sha256}; use tokio::process::Command; use tokio::sync::{Mutex, broadcast, watch}; @@ -28,10 +28,8 @@ use crate::{ CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandSnapshot, CommandStatus, CommandStream, CommandStreamSlice, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult, ReadRequest, - ReadResult, StatRequest, StatResult, Workdir, WorkdirDelegationPermission, - WorkdirDelegationRequest, WorkdirError, WorkdirPath, WorkdirSession, - WorkdirSessionCapabilities, WorkdirSessionCapability, WorkdirSessionHandle, WriteRequest, - WriteResult, + ReadResult, StatRequest, StatResult, Workdir, WorkdirError, WorkdirPath, WorkdirSession, + WorkdirSessionCapabilities, WorkdirSessionCapability, WriteRequest, WriteResult, }; #[cfg(test)] use crate::{EntryKind, WriteOutcome}; @@ -558,69 +556,6 @@ impl WorkdirSession for LocalWorkdirSession { self.inner.capabilities } - async fn capture_delegation_source( - &self, - request: &WorkdirDelegationRequest, - ) -> Result { - let host_rules = request - .rules - .iter() - .map(|rule| ScopeRule { - target: self.inner.root.join(rule.target.as_str()), - permission: match rule.permission { - WorkdirDelegationPermission::Read => Permission::Read, - WorkdirDelegationPermission::Write => Permission::Write, - }, - recursive: rule.recursive, - }) - .collect::>(); - for (logical, host) in request.rules.iter().zip(&host_rules) { - if logical.permission == WorkdirDelegationPermission::Write { - let resolved = Scope::resolved_target(host) - .map_err(|error| WorkdirError::Denied(error.to_string()))?; - if resolved != host.target { - return Err(WorkdirError::Denied(format!( - "write delegation target `{}` traverses a symlink", - logical.target - ))); - } - } - } - let parent_scope = self.inner.scope.snapshot(); - for rule in &host_rules { - if !parent_scope - .allows_rule(rule) - .map_err(|error| WorkdirError::Denied(error.to_string()))? - { - return Err(WorkdirError::Denied(format!( - "delegated provider scope `{}` exceeds the parent session", - rule.target.display() - ))); - } - } - let child_scope = Scope::from_config(&ScopeConfig { - allow: host_rules, - deny: Vec::new(), - }) - .map_err(|error| WorkdirError::Denied(error.to_string()))?; - let child_cwd = self.inner.root.join(request.cwd.as_str()); - if !child_scope.is_readable(&child_cwd) - || !std::fs::metadata(&child_cwd).is_ok_and(|metadata| metadata.is_dir()) - { - return Err(WorkdirError::Denied(format!( - "delegated cwd `{}` is not a readable Workdir directory", - request.cwd - ))); - } - Ok(Arc::new(LocalWorkdirSession::materialized_bound( - self.inner.workdir.clone(), - self.inner.root.clone(), - self.inner.root.clone(), - SharedScope::new(child_scope), - self.inner.capabilities, - ))) - } - async fn stat(&self, request: StatRequest) -> Result { self.ensure_capability(WorkdirSessionCapability::Read)?; let logical = request.path.clone(); @@ -694,9 +629,20 @@ impl WorkdirSession for LocalWorkdirSession { { return Err(WorkdirError::OutOfScope(spill_dir.to_path_buf())); } + let cwd = if let Some(logical_cwd) = request.cwd.as_ref() { + let cwd = self.resolve(logical_cwd); + let scope = self.inner.scope.snapshot(); + if !scope.is_readable(&cwd) + || !std::fs::metadata(&cwd).is_ok_and(|metadata| metadata.is_dir()) + { + return Err(WorkdirError::OutOfScope(cwd)); + } + cwd + } else { + self.inner.cwd.clone() + }; let id = self.inner.next_command_id.fetch_add(1, Ordering::Relaxed); let handle = CommandHandle(format!("command-{id}")); - let cwd = self.inner.cwd.clone(); let (completion_tx, completion) = watch::channel(false); let command_id = handle.0.clone(); let telemetry = self.inner.command_telemetry.clone(); @@ -1516,6 +1462,7 @@ mod tests { command: "sleep 30".to_owned(), timeout_secs: 60, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: None, }, @@ -2043,6 +1990,7 @@ mod tests { command: "pwd && printf provider-command".into(), timeout_secs: 5, output_limit: 4096, + cwd: None, spill_dir: None, tool_call_id: None, }, @@ -2141,6 +2089,7 @@ mod tests { command: "printf hidden".into(), timeout_secs: 5, output_limit: 1, + cwd: None, spill_dir: Some(spill.path().to_path_buf()), tool_call_id: None, }, @@ -2178,6 +2127,7 @@ mod tests { command: "i=0; while [ $i -lt 200 ]; do printf 'line-%03d\\n' \"$i\"; i=$((i+1)); done; printf 'FINAL-NEEDLE\\n'".into(), timeout_secs: 5, output_limit: 64, + cwd: None, spill_dir: Some(spill.path().to_path_buf()), tool_call_id: None, }, @@ -2224,6 +2174,7 @@ mod tests { command: "printf 'aéz'".into(), timeout_secs: 5, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: None, }, @@ -2449,6 +2400,7 @@ mod tests { command: "printf ready; printf warning >&2; sleep 0.2; printf done".into(), timeout_secs: 5, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: Some("tool-7".into()), }, @@ -2553,6 +2505,7 @@ mod tests { command: "sleep 30".into(), timeout_secs: 1, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: None, }, @@ -2623,6 +2576,7 @@ mod tests { command: "sleep 30".into(), timeout_secs: 60, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: None, }, diff --git a/crates/workdir/src/operation.rs b/crates/workdir/src/operation.rs index 67858685..5af8b54a 100644 --- a/crates/workdir/src/operation.rs +++ b/crates/workdir/src/operation.rs @@ -11,6 +11,10 @@ pub struct CommandRequest { pub command: String, pub timeout_secs: u64, pub output_limit: usize, + /// Workdir-relative command directory. Providers validate it against the + /// active session before process start. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cwd: Option, /// Provider-local directory where complete output is retained when the /// inline result exceeds `output_limit`. pub spill_dir: Option, diff --git a/crates/workdir/src/delegation.rs b/crates/workdir/src/scope.rs similarity index 58% rename from crates/workdir/src/delegation.rs rename to crates/workdir/src/scope.rs index 17598dab..7b024696 100644 --- a/crates/workdir/src/delegation.rs +++ b/crates/workdir/src/scope.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::path::Path; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex, Weak}; @@ -18,94 +18,153 @@ use crate::{ #[derive(Clone, Copy, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] #[serde(rename_all = "snake_case")] -pub enum WorkdirDelegationPermission { +pub enum WorkdirToolScopePermission { Read, Write, } #[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] #[serde(deny_unknown_fields)] -pub struct WorkdirDelegationRule { +pub struct WorkdirToolScopeRule { pub target: FsPath, - pub permission: WorkdirDelegationPermission, + pub permission: WorkdirToolScopePermission, pub recursive: bool, } #[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] #[serde(deny_unknown_fields)] -pub struct WorkdirDelegationRequest { - pub rules: Vec, +pub struct WorkdirToolScope { + pub rules: Vec, pub cwd: FsPath, + pub command: bool, } -pub struct WorkdirDelegation { - pub scoped_session: WorkdirSessionHandle, +#[derive(Clone)] +pub struct WorkdirToolBroker { + authority: Arc, + session: WorkdirSessionHandle, + event_forwarder: Option>>>>, +} + +impl std::fmt::Debug for WorkdirToolBroker { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkdirToolBroker") + .field("workdir", self.session.workdir()) + .field("capabilities", &self.session.capabilities()) + .finish_non_exhaustive() + } +} + +impl WorkdirToolBroker { + /// Own the parent Worker's active session and mediate every scoped child operation. + pub fn new(source: WorkdirSessionHandle) -> Self { + let capabilities = source.capabilities(); + let (command_events, _) = broadcast::channel(64); + let authority = Arc::new(ScopedWorkdirSession { + source, + cwd: FsPath::new("").expect("empty Workdir path is valid"), + scope: None, + capabilities, + validity: SessionValidity::root(), + child_write_leases: Mutex::new(HashMap::new()), + next_lease_id: AtomicU64::new(1), + owned_commands: Arc::new(Mutex::new(HashSet::new())), + command_events, + closes_source: true, + }); + Self { + session: authority.clone(), + authority, + event_forwarder: None, + } + } + + /// Session used only by tools registered by the owning Worker. + pub fn tool_session(&self) -> WorkdirSessionHandle { + self.session.clone() + } + + /// Create a revocable, attenuated tool route without delegating a provider session. + pub async fn scope( + &self, + request: WorkdirToolScope, + ) -> Result { + self.authority.scope(request).await + } +} + +impl std::ops::Deref for WorkdirToolBroker { + type Target = WorkdirSessionHandle; + + fn deref(&self) -> &Self::Target { + &self.session + } +} + +pub struct WorkdirScopeLease { + broker: WorkdirToolBroker, pub capabilities: WorkdirSessionCapabilities, validity: Arc, } -impl std::fmt::Debug for WorkdirDelegation { +impl std::fmt::Debug for WorkdirScopeLease { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("WorkdirDelegation") - .field("workdir", &self.scoped_session.workdir()) + f.debug_struct("WorkdirScopeLease") + .field("workdir", self.broker.session.workdir()) .field("capabilities", &self.capabilities) .field("active", &self.is_active()) .finish() } } -impl WorkdirDelegation { +impl WorkdirScopeLease { + pub fn broker(&self) -> WorkdirToolBroker { + self.broker.clone() + } + + pub fn tool_session(&self) -> WorkdirSessionHandle { + self.broker.tool_session() + } + + pub async fn scope( + &self, + request: WorkdirToolScope, + ) -> Result { + self.broker.scope(request).await + } + pub fn is_active(&self) -> bool { self.validity.is_active() } pub fn release(&self) { self.validity.active.store(false, Ordering::Release); + if let Some(forwarder) = &self.broker.event_forwarder + && let Some(handle) = forwarder + .lock() + .expect("scoped command forwarder mutex poisoned") + .take() + { + handle.abort(); + } } } -impl Drop for WorkdirDelegation { +impl std::ops::Deref for WorkdirScopeLease { + type Target = WorkdirSessionHandle; + + fn deref(&self) -> &Self::Target { + &self.broker.session + } +} + +impl Drop for WorkdirScopeLease { fn drop(&mut self) { self.release(); } } -pub struct AppliedWorkdirDelegation { - pub scoped_session: WorkdirSessionHandle, - _leases: Vec, -} - -impl std::fmt::Debug for AppliedWorkdirDelegation { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("AppliedWorkdirDelegation") - .field("workdir", self.scoped_session.workdir()) - .field("lease_count", &self._leases.len()) - .finish() - } -} - -pub async fn apply_delegation_chain( - source: WorkdirSessionHandle, - requests: impl IntoIterator, -) -> Result { - let mut current = source; - let mut leases = Vec::new(); - for request in requests { - let authority = if current.is_delegation_capable() { - current.clone() - } else { - delegation_capable_session(current.clone()) - }; - let lease = authority.delegate(request).await?; - current = lease.scoped_session.clone(); - leases.push(lease); - } - Ok(AppliedWorkdirDelegation { - scoped_session: current, - _leases: leases, - }) -} - #[derive(Debug)] struct SessionValidity { active: AtomicBool, @@ -136,23 +195,25 @@ impl SessionValidity { #[derive(Clone, Debug)] struct ActiveWriteLease { validity: Weak, - rules: Vec, + rules: Vec, } -struct DelegatingWorkdirSession { +struct ScopedWorkdirSession { source: WorkdirSessionHandle, cwd: FsPath, - scope: Option>, + scope: Option>, capabilities: WorkdirSessionCapabilities, validity: Arc, child_write_leases: Mutex>, next_lease_id: AtomicU64, + owned_commands: Arc>>, + command_events: broadcast::Sender, closes_source: bool, } -impl std::fmt::Debug for DelegatingWorkdirSession { +impl std::fmt::Debug for ScopedWorkdirSession { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("DelegatingWorkdirSession") + f.debug_struct("ScopedWorkdirSession") .field("workdir", &self.source.workdir()) .field("scope", &self.scope) .field("capabilities", &self.capabilities) @@ -161,22 +222,7 @@ impl std::fmt::Debug for DelegatingWorkdirSession { } } -/// Wrap a provider session with logical-path delegation and parent write gates. -pub fn delegation_capable_session(source: WorkdirSessionHandle) -> WorkdirSessionHandle { - let capabilities = source.capabilities(); - Arc::new(DelegatingWorkdirSession { - source, - cwd: FsPath::new("").expect("empty Workdir path is valid"), - scope: None, - capabilities, - validity: SessionValidity::root(), - child_write_leases: Mutex::new(HashMap::new()), - next_lease_id: AtomicU64::new(1), - closes_source: true, - }) -} - -impl DelegatingWorkdirSession { +impl ScopedWorkdirSession { fn ensure_active(&self) -> Result<(), WorkdirError> { if self.validity.is_active() { Ok(()) @@ -195,7 +241,7 @@ impl DelegatingWorkdirSession { Ok(()) } else { Err(WorkdirError::Denied(format!( - "delegated workdir session does not permit {operation}" + "scoped Workdir tools do not permit {operation}" ))) } } @@ -203,7 +249,7 @@ impl DelegatingWorkdirSession { fn ensure_path( &self, path: &FsPath, - permission: WorkdirDelegationPermission, + permission: WorkdirToolScopePermission, ) -> Result<(), WorkdirError> { self.ensure_active()?; if let Some(scope) = &self.scope { @@ -212,11 +258,11 @@ impl DelegatingWorkdirSession { .any(|rule| rule_allows_path(rule, path, permission)) { return Err(WorkdirError::Denied(format!( - "logical workdir path `{path}` is outside the delegated {permission:?} scope" + "logical workdir path `{path}` is outside the scoped {permission:?} scope" ))); } } - if permission == WorkdirDelegationPermission::Write { + if permission == WorkdirToolScopePermission::Write { self.ensure_parent_write_available(path)?; } Ok(()) @@ -239,7 +285,7 @@ impl DelegatingWorkdirSession { capability: WorkdirSessionCapability, ) -> Result<(), WorkdirError> { self.ensure_capability(capability, "read operations")?; - self.ensure_path(path, WorkdirDelegationPermission::Read) + self.ensure_path(path, WorkdirToolScopePermission::Read) } fn ensure_write( @@ -248,58 +294,133 @@ impl DelegatingWorkdirSession { capability: WorkdirSessionCapability, ) -> Result<(), WorkdirError> { self.ensure_capability(capability, "write operations")?; - self.ensure_path(path, WorkdirDelegationPermission::Write) + self.ensure_path(path, WorkdirToolScopePermission::Write) } fn ensure_command(&self) -> Result<(), WorkdirError> { self.ensure_capability(WorkdirSessionCapability::Command, "command execution") } + fn ensure_owned_command(&self, handle: &CommandHandle) -> Result<(), WorkdirError> { + self.ensure_command()?; + if self.scope.is_none() + || self + .owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .contains(&handle.0) + { + Ok(()) + } else { + Err(WorkdirError::UnknownCommand(handle.0.clone())) + } + } + fn ensure_parent_write_available(&self, path: &FsPath) -> Result<(), WorkdirError> { let mut leases = self .child_write_leases .lock() - .expect("workdir delegation lease mutex poisoned"); + .expect("Workdir tool scope lease mutex poisoned"); leases.retain(|_, lease| lease.validity.upgrade().is_some_and(|v| v.is_active())); if leases.values().any(|lease| { lease.rules.iter().any(|rule| { - rule.permission == WorkdirDelegationPermission::Write - && rule_allows_path(rule, path, WorkdirDelegationPermission::Write) + rule.permission == WorkdirToolScopePermission::Write + && rule_allows_path(rule, path, WorkdirToolScopePermission::Write) }) }) { Err(WorkdirError::Denied(format!( - "logical workdir path `{path}` is leased to a child session" + "logical workdir path `{path}` is leased to child Workdir tools" ))) } else { Ok(()) } } - fn validate_delegation_rules( + async fn ensure_source_path_has_no_symlink(&self, path: &FsPath) -> Result<(), WorkdirError> { + let mut current = String::new(); + for component in Path::new(path.as_str()).components() { + let component = component.as_os_str().to_string_lossy(); + if component.is_empty() || component == "." { + continue; + } + if !current.is_empty() { + current.push('/'); + } + current.push_str(&component); + let current = FsPath::new(¤t).map_err(|error| { + WorkdirError::Denied(format!("invalid scoped Workdir path: {error}")) + })?; + match self.source.stat(StatRequest { path: current }).await { + Ok(result) if result.kind == fs_operation::EntryKind::Symlink => { + return Err(WorkdirError::Denied(format!( + "scoped Workdir path `{path}` traverses a symlink" + ))); + } + Ok(_) => {} + Err(WorkdirError::NotFound(_)) => break, + Err(error) => return Err(error), + } + } + Ok(()) + } + + async fn ensure_scope_targets_do_not_traverse_symlinks( &self, - rules: &[WorkdirDelegationRule], + rules: &[WorkdirToolScopeRule], + ) -> Result<(), WorkdirError> { + for rule in rules { + self.ensure_source_path_has_no_symlink(&rule.target).await?; + } + Ok(()) + } + + async fn resolve_operation_path(&self, path: &FsPath) -> Result { + self.ensure_active()?; + let resolved = self.resolve_path(path)?; + if self.scope.is_some() { + self.ensure_source_path_has_no_symlink(&resolved).await?; + } + Ok(resolved) + } + + fn validate_scope( + &self, + rules: &[WorkdirToolScopeRule], + command: bool, ) -> Result { self.ensure_active()?; if rules.is_empty() { return Err(WorkdirError::Denied( - "workdir delegation requires at least one logical scope rule".into(), + "workdir tool scope requires at least one logical scope rule".into(), )); } let writable = rules .iter() - .any(|rule| rule.permission == WorkdirDelegationPermission::Write); + .any(|rule| rule.permission == WorkdirToolScopePermission::Write); if !self.capabilities.supports(WorkdirSessionCapability::Read) || (writable && (!self.capabilities.supports(WorkdirSessionCapability::Write) - || !self.capabilities.supports(WorkdirSessionCapability::Edit) - || !self - .capabilities - .supports(WorkdirSessionCapability::Command))) + || !self.capabilities.supports(WorkdirSessionCapability::Edit))) { return Err(WorkdirError::Denied( - "parent workdir session cannot delegate the requested capabilities".into(), + "parent Workdir session cannot scope the requested capabilities".into(), )); } + if command { + if !writable { + return Err(WorkdirError::Denied( + "command execution requires a writable scoped path".into(), + )); + } + if !self + .capabilities + .supports(WorkdirSessionCapability::Command) + { + return Err(WorkdirError::Denied( + "parent Workdir session does not support Command".into(), + )); + } + } for requested in rules { if let Some(scope) = &self.scope { if !scope @@ -307,7 +428,7 @@ impl DelegatingWorkdirSession { .any(|parent| rule_contains_rule(parent, requested)) { return Err(WorkdirError::Denied(format!( - "logical workdir scope `{}` exceeds the parent delegation", + "logical workdir scope `{}` exceeds the parent tool scope", requested.target ))); } @@ -325,69 +446,40 @@ impl DelegatingWorkdirSession { if writable { delegated.push(WorkdirSessionCapability::Write); delegated.push(WorkdirSessionCapability::Edit); + } + if command { delegated.push(WorkdirSessionCapability::Command); } Ok(WorkdirSessionCapabilities::from_capabilities(delegated)) } -} -#[async_trait] -impl WorkdirSession for DelegatingWorkdirSession { - fn workdir(&self) -> &Workdir { - self.source.workdir() - } - - fn capabilities(&self) -> WorkdirSessionCapabilities { - self.capabilities - } - - fn is_delegation_capable(&self) -> bool { - true - } - - fn transports_delegation_context(&self) -> bool { - self.source.transports_delegation_context() - } - - async fn capture_delegation_source( - &self, - request: &WorkdirDelegationRequest, - ) -> Result { - self.ensure_active()?; - if self.scope.is_some() { - return Err(WorkdirError::Denied( - "scoped Workdir sessions cannot expose their provider source".into(), - )); - } - self.source.capture_delegation_source(request).await - } - - async fn delegate( - &self, - request: WorkdirDelegationRequest, - ) -> Result { - let capabilities = self.validate_delegation_rules(&request.rules)?; + async fn scope( + self: &Arc, + request: WorkdirToolScope, + ) -> Result { + let capabilities = self.validate_scope(&request.rules, request.command)?; if !request .rules .iter() - .any(|rule| rule_allows_path(rule, &request.cwd, WorkdirDelegationPermission::Read)) + .any(|rule| rule_allows_path(rule, &request.cwd, WorkdirToolScopePermission::Read)) { return Err(WorkdirError::Denied(format!( - "delegated cwd `{}` is outside the delegated readable scope", + "scoped tool cwd `{}` is outside the readable scope", request.cwd ))); } - let source = self.source.capture_delegation_source(&request).await?; + self.ensure_scope_targets_do_not_traverse_symlinks(&request.rules) + .await?; let validity = SessionValidity::child(self.validity.clone()); let id = self.next_lease_id.fetch_add(1, Ordering::Relaxed); if request .rules .iter() - .any(|rule| rule.permission == WorkdirDelegationPermission::Write) + .any(|rule| rule.permission == WorkdirToolScopePermission::Write) { self.child_write_leases .lock() - .expect("workdir delegation lease mutex poisoned") + .expect("Workdir tool scope lease mutex poisoned") .insert( id, ActiveWriteLease { @@ -396,99 +488,125 @@ impl WorkdirSession for DelegatingWorkdirSession { }, ); } - let child: WorkdirSessionHandle = Arc::new(DelegatingWorkdirSession { - source, + let owned_commands = Arc::new(Mutex::new(HashSet::new())); + let (command_events, _) = broadcast::channel(64); + let event_forwarder = forward_owned_command_events( + self.source.subscribe_command_events(), + owned_commands.clone(), + command_events.clone(), + ) + .map(|handle| Arc::new(Mutex::new(Some(handle)))); + let child = Arc::new(ScopedWorkdirSession { + source: self.source.clone(), cwd: request.cwd, scope: Some(request.rules), capabilities, validity: validity.clone(), child_write_leases: Mutex::new(HashMap::new()), next_lease_id: AtomicU64::new(1), + owned_commands, + command_events, closes_source: false, }); - let scoped_session: WorkdirSessionHandle = - if capabilities == WorkdirSessionCapabilities::READ_ONLY { - Arc::new(ReadOnlyWorkdirSession::new(child)) - } else { - child - }; - Ok(WorkdirDelegation { - scoped_session, + let broker = WorkdirToolBroker { + session: child.clone(), + authority: child, + event_forwarder, + }; + Ok(WorkdirScopeLease { + broker, capabilities, validity, }) } +} + +#[async_trait] +impl WorkdirSession for ScopedWorkdirSession { + fn workdir(&self) -> &Workdir { + self.source.workdir() + } + + fn capabilities(&self) -> WorkdirSessionCapabilities { + self.capabilities + } async fn stat(&self, mut request: StatRequest) -> Result { - let path = self.resolve_path(&request.path)?; + let path = self.resolve_operation_path(&request.path).await?; self.ensure_read(&path, WorkdirSessionCapability::Read)?; - if !self.source.transports_delegation_context() { - request.path = path; - } + request.path = path; self.source.stat(request).await } async fn read(&self, mut request: ReadRequest) -> Result { - let path = self.resolve_path(&request.path)?; + let path = self.resolve_operation_path(&request.path).await?; self.ensure_read(&path, WorkdirSessionCapability::Read)?; - if !self.source.transports_delegation_context() { - request.path = path; - } + request.path = path; self.source.read(request).await } async fn write(&self, mut request: WriteRequest) -> Result { - let path = self.resolve_path(&request.path)?; + let path = self.resolve_operation_path(&request.path).await?; self.ensure_write(&path, WorkdirSessionCapability::Write)?; - if !self.source.transports_delegation_context() { - request.path = path; - } + request.path = path; self.source.write(request).await } async fn edit(&self, mut request: EditRequest) -> Result { - let path = self.resolve_path(&request.path)?; + let path = self.resolve_operation_path(&request.path).await?; self.ensure_write(&path, WorkdirSessionCapability::Edit)?; - if !self.source.transports_delegation_context() { - request.path = path; - } + request.path = path; self.source.edit(request).await } async fn list(&self, mut request: ListRequest) -> Result { - let path = self.resolve_path(&request.path)?; + let path = self.resolve_operation_path(&request.path).await?; self.ensure_read(&path, WorkdirSessionCapability::Read)?; - if !self.source.transports_delegation_context() { - request.path = path; - } + request.path = path; self.source.list(request).await } async fn glob(&self, mut request: GlobRequest) -> Result { - let path = self.resolve_path(&request.path)?; + let path = self.resolve_operation_path(&request.path).await?; self.ensure_read(&path, WorkdirSessionCapability::Glob)?; - if !self.source.transports_delegation_context() { - request.path = path; - } + request.path = path; self.source.glob(request).await } async fn grep(&self, mut request: GrepRequest) -> Result { - let path = self.resolve_path(&request.path)?; + let path = self.resolve_operation_path(&request.path).await?; self.ensure_read(&path, WorkdirSessionCapability::Grep)?; - if !self.source.transports_delegation_context() { - request.path = path; - } + request.path = path; self.source.grep(request).await } - async fn start_command(&self, request: CommandRequest) -> Result { + async fn start_command( + &self, + mut request: CommandRequest, + ) -> Result { self.ensure_command()?; - self.source.start_command(request).await + let tool_call_id = request.tool_call_id.clone(); + if self.scope.is_some() { + request.cwd = Some(match request.cwd.as_ref() { + Some(cwd) => self.resolve_path(cwd)?, + None => self.cwd.clone(), + }); + } + let handle = self.source.start_command(request).await?; + self.owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .insert(handle.0.clone()); + let _ = self.command_events.send(CommandEvent::Started { + command_id: handle.0.clone(), + tool_call_id, + observed_at_ms: unix_timestamp_ms(), + }); + Ok(handle) } async fn command_status(&self, handle: CommandHandle) -> Result { - self.ensure_command()?; + self.ensure_owned_command(&handle)?; self.source.command_status(handle).await } @@ -496,29 +614,56 @@ impl WorkdirSession for DelegatingWorkdirSession { &self, request: CommandOutputRequest, ) -> Result { - self.ensure_command()?; - self.source.command_output(request).await + self.ensure_owned_command(&request.handle)?; + let command_id = request.handle.0.clone(); + let output = self.source.command_output(request).await?; + if !matches!(output.status, CommandStatus::Running) { + self.owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .remove(&command_id); + } + Ok(output) } async fn cancel_command(&self, handle: CommandHandle) -> Result<(), WorkdirError> { - self.ensure_command()?; + self.ensure_owned_command(&handle)?; self.source.cancel_command(handle).await } fn subscribe_command_events(&self) -> Option> { - self.ensure_capability(WorkdirSessionCapability::Command, "command observation") - .ok()?; - self.source.subscribe_command_events() + if !self + .capabilities + .supports(WorkdirSessionCapability::Command) + { + return None; + } + if self.scope.is_none() { + self.source.subscribe_command_events() + } else { + Some(self.command_events.subscribe()) + } } fn command_snapshot(&self) -> Vec { - if self - .ensure_capability(WorkdirSessionCapability::Command, "command observation") - .is_err() + if !self + .capabilities + .supports(WorkdirSessionCapability::Command) { return Vec::new(); } - self.source.command_snapshot() + if self.scope.is_none() { + return self.source.command_snapshot(); + } + let owned = self + .owned_commands + .lock() + .expect("scoped command set mutex poisoned"); + self.source + .command_snapshot() + .into_iter() + .filter(|snapshot| owned.contains(&snapshot.command_id)) + .collect() } async fn close(&self) -> Result<(), WorkdirError> { @@ -531,7 +676,7 @@ impl WorkdirSession for DelegatingWorkdirSession { } } -/// A fail-closed read-only view over an already scoped delegated session. +/// A fail-closed read-only view over an already scoped scoped tool route. #[derive(Debug)] pub struct ReadOnlyWorkdirSession { inner: WorkdirSessionHandle, @@ -553,30 +698,6 @@ impl WorkdirSession for ReadOnlyWorkdirSession { WorkdirSessionCapabilities::READ_ONLY } - fn is_delegation_capable(&self) -> bool { - true - } - - fn transports_delegation_context(&self) -> bool { - self.inner.transports_delegation_context() - } - - async fn delegate( - &self, - request: WorkdirDelegationRequest, - ) -> Result { - if request - .rules - .iter() - .any(|rule| rule.permission == WorkdirDelegationPermission::Write) - { - return Err(WorkdirError::Denied( - "read-only workdir session cannot delegate write access".into(), - )); - } - self.inner.delegate(request).await - } - async fn stat(&self, request: StatRequest) -> Result { self.inner.stat(request).await } @@ -629,20 +750,57 @@ impl WorkdirSession for ReadOnlyWorkdirSession { } } +fn forward_owned_command_events( + receiver: Option>, + owned_commands: Arc>>, + sender: broadcast::Sender, +) -> Option> { + let mut receiver = receiver?; + Some(tokio::spawn(async move { + loop { + let event = match receiver.recv().await { + Ok(event) => event, + Err(broadcast::error::RecvError::Lagged(_)) => continue, + Err(broadcast::error::RecvError::Closed) => break, + }; + let command_id = match &event { + CommandEvent::Started { .. } => continue, + CommandEvent::Output { command_id, .. } + | CommandEvent::Terminal { command_id, .. } => command_id.clone(), + }; + let owned = owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .contains(&command_id); + if owned { + let _ = sender.send(event); + } + } + })) +} + +fn unix_timestamp_ms() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .min(u128::from(u64::MAX)) as u64 +} + fn rule_allows_path( - rule: &WorkdirDelegationRule, + rule: &WorkdirToolScopeRule, path: &FsPath, - required: WorkdirDelegationPermission, + required: WorkdirToolScopePermission, ) -> bool { - if required == WorkdirDelegationPermission::Write - && rule.permission != WorkdirDelegationPermission::Write + if required == WorkdirToolScopePermission::Write + && rule.permission != WorkdirToolScopePermission::Write { return false; } path_in_rule(rule, path) } -fn path_in_rule(rule: &WorkdirDelegationRule, path: &FsPath) -> bool { +fn path_in_rule(rule: &WorkdirToolScopeRule, path: &FsPath) -> bool { let target = Path::new(rule.target.as_str()); let path = Path::new(path.as_str()); if path == target { @@ -655,9 +813,9 @@ fn path_in_rule(rule: &WorkdirDelegationRule, path: &FsPath) -> bool { rule.recursive || depth <= 1 } -fn rule_contains_rule(parent: &WorkdirDelegationRule, child: &WorkdirDelegationRule) -> bool { - if child.permission == WorkdirDelegationPermission::Write - && parent.permission != WorkdirDelegationPermission::Write +fn rule_contains_rule(parent: &WorkdirToolScopeRule, child: &WorkdirToolScopeRule) -> bool { + if child.permission == WorkdirToolScopePermission::Write + && parent.permission != WorkdirToolScopePermission::Write { return false; } @@ -684,7 +842,7 @@ mod tests { FsPath::new(path).unwrap() } - fn session(root: &Path) -> WorkdirSessionHandle { + fn session(root: &Path) -> WorkdirToolBroker { let scope = SharedScope::new( Scope::from_config(&ScopeConfig { allow: vec![ScopeRule { @@ -696,7 +854,7 @@ mod tests { }) .unwrap(), ); - delegation_capable_session(Arc::new(LocalWorkdirSession::materialized_bound( + WorkdirToolBroker::new(Arc::new(LocalWorkdirSession::materialized_bound( Workdir::new("delegation-test"), root.to_path_buf(), root.to_path_buf(), @@ -705,14 +863,15 @@ mod tests { ))) } - fn request(path: &str, permission: WorkdirDelegationPermission) -> WorkdirDelegationRequest { - WorkdirDelegationRequest { - rules: vec![WorkdirDelegationRule { + fn request(path: &str, permission: WorkdirToolScopePermission) -> WorkdirToolScope { + WorkdirToolScope { + rules: vec![WorkdirToolScopeRule { target: fs_path(path), permission, recursive: true, }], cwd: fs_path(path), + command: permission == WorkdirToolScopePermission::Write, } } @@ -743,6 +902,7 @@ mod tests { command: command.into(), timeout_secs: 5, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: Some(tool_call_id.into()), }) @@ -760,7 +920,7 @@ mod tests { } #[tokio::test] - async fn delegation_capable_session_forwards_command_telemetry() { + async fn workdir_tool_broker_session_forwards_command_telemetry() { let root = TempDir::new().unwrap(); let parent = session(root.path()); let mut events = parent @@ -771,6 +931,7 @@ mod tests { command: "printf ready; sleep 0.2; printf done".into(), timeout_secs: 5, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: Some("tool-delegated".into()), }) @@ -807,11 +968,109 @@ mod tests { assert!(parent.command_snapshot().is_empty()); } + #[tokio::test] + async fn write_scope_without_command_grant_has_no_command_capability() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("work")).unwrap(); + let parent = session(root.path()); + let child = parent + .scope(WorkdirToolScope { + rules: vec![WorkdirToolScopeRule { + target: fs_path("work"), + permission: WorkdirToolScopePermission::Write, + recursive: true, + }], + cwd: fs_path("work"), + command: false, + }) + .await + .unwrap(); + + assert!(child.capabilities.supports(WorkdirSessionCapability::Write)); + assert!( + !child + .capabilities + .supports(WorkdirSessionCapability::Command) + ); + let error = child + .start_command(CommandRequest { + command: "pwd".into(), + timeout_secs: 5, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: None, + }) + .await + .unwrap_err(); + assert!(matches!(error, WorkdirError::Denied(_))); + } + + #[tokio::test] + async fn scoped_commands_use_child_cwd_and_do_not_leak_between_siblings() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("one")).unwrap(); + fs::create_dir_all(root.path().join("two")).unwrap(); + let parent = session(root.path()); + let first = parent + .scope(request("one", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + let second = parent + .scope(request("two", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + let mut first_events = first.subscribe_command_events().unwrap(); + let mut second_events = second.subscribe_command_events().unwrap(); + + let handle = first + .start_command(CommandRequest { + command: "pwd; sleep 0.2".into(), + timeout_secs: 5, + output_limit: 4096, + cwd: None, + spill_dir: None, + tool_call_id: Some("first-command".into()), + }) + .await + .unwrap(); + assert!(matches!( + first_events.recv().await.unwrap(), + CommandEvent::Started { .. } + )); + assert!(matches!( + tokio::time::timeout(std::time::Duration::from_millis(50), second_events.recv()).await, + Err(_) + )); + assert!(matches!( + second.command_status(handle.clone()).await, + Err(WorkdirError::UnknownCommand(_)) + )); + + let output = first + .command_output(CommandOutputRequest { + handle, + cursor: 0, + limit: 4096, + wait: true, + }) + .await + .unwrap(); + let expected = root.path().join("one").to_string_lossy().into_owned(); + assert!( + output + .content + .lines() + .next() + .is_some_and(|line| line == expected) + ); + } + #[test] fn non_recursive_rule_covers_target_and_direct_children_only() { - let rule = WorkdirDelegationRule { + let rule = WorkdirToolScopeRule { target: fs_path("docs"), - permission: WorkdirDelegationPermission::Read, + permission: WorkdirToolScopePermission::Read, recursive: false, }; assert!(path_in_rule(&rule, &fs_path("docs"))); @@ -829,21 +1088,16 @@ mod tests { let parent = session(root.path()); let child = parent - .delegate(request("docs", WorkdirDelegationPermission::Read)) + .scope(request("docs", WorkdirToolScopePermission::Read)) .await .unwrap(); assert_eq!(child.capabilities, WorkdirSessionCapabilities::READ_ONLY); assert_eq!( - child - .scoped_session - .read(read("readme.md")) - .await - .unwrap() - .bytes, + child.read(read("readme.md")).await.unwrap().bytes, b"visible" ); assert!(matches!( - child.scoped_session.write(write("new.md", "no")).await, + child.write(write("new.md", "no")).await, Err(WorkdirError::Denied(_)) )); assert!( @@ -851,15 +1105,15 @@ mod tests { .capabilities .supports(WorkdirSessionCapability::Command) ); - assert!(child.scoped_session.subscribe_command_events().is_none()); - assert!(child.scoped_session.command_snapshot().is_empty()); + assert!(child.subscribe_command_events().is_none()); + assert!(child.command_snapshot().is_empty()); assert!(matches!( child - .scoped_session .start_command(CommandRequest { command: "printf denied".into(), timeout_secs: 5, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: Some("read-only-command".into()), }) @@ -880,11 +1134,11 @@ mod tests { symlink("../secret/key", root.path().join("granted/link")).unwrap(); let parent = session(root.path()); let child = parent - .delegate(request("granted", WorkdirDelegationPermission::Read)) + .scope(request("granted", WorkdirToolScopePermission::Read)) .await .unwrap(); - let result = child.scoped_session.read(read("link")).await; + let result = child.read(read("link")).await; assert!( result.is_err(), "symlink read escaped provider scope: {result:?}" @@ -902,14 +1156,11 @@ mod tests { symlink("../secret", root.path().join("granted/outside")).unwrap(); let parent = session(root.path()); let child = parent - .delegate(request("granted", WorkdirDelegationPermission::Write)) + .scope(request("granted", WorkdirToolScopePermission::Write)) .await .unwrap(); - let result = child - .scoped_session - .write(write("outside/new", "forbidden")) - .await; + let result = child.write(write("outside/new", "forbidden")).await; assert!( result.is_err(), "symlink write escaped provider scope: {result:?}" @@ -930,9 +1181,9 @@ mod tests { assert!(matches!( parent - .delegate(request( + .scope(request( "granted/outside", - WorkdirDelegationPermission::Write + WorkdirToolScopePermission::Write )) .await, Err(WorkdirError::Denied(_)) @@ -950,7 +1201,7 @@ mod tests { fs::create_dir_all(root.path().join("other")).unwrap(); let parent = session(root.path()); let child = parent - .delegate(request("leased", WorkdirDelegationPermission::Write)) + .scope(request("leased", WorkdirToolScopePermission::Write)) .await .unwrap(); assert!( @@ -958,12 +1209,8 @@ mod tests { .capabilities .supports(WorkdirSessionCapability::Command) ); - let child_output = run_command( - &child.scoped_session, - "printf child-command", - "delegated-child-command", - ) - .await; + let child_output = + run_command(&child, "printf child-command", "delegated-child-command").await; assert_eq!(child_output.content, "child-command"); let parent_output = run_command( &parent, @@ -983,19 +1230,15 @@ mod tests { Err(WorkdirError::Denied(_)) )); parent.write(write("other/file", "parent")).await.unwrap(); - child - .scoped_session - .write(write("file", "child")) - .await - .unwrap(); + child.write(write("file", "child")).await.unwrap(); child.release(); assert!(matches!( child - .scoped_session .start_command(CommandRequest { command: "printf revoked".into(), timeout_secs: 5, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: Some("revoked-child-command".into()), }) @@ -1007,7 +1250,7 @@ mod tests { .await .unwrap(); assert!(matches!( - child.scoped_session.read(read("file")).await, + child.read(read("file")).await, Err(WorkdirError::SessionClosed) )); } @@ -1021,34 +1264,31 @@ mod tests { fs::write(root.path().join("docs/peer/b"), "b").unwrap(); let root_session = session(root.path()); let child = root_session - .delegate(request("docs", WorkdirDelegationPermission::Read)) + .scope(request("docs", WorkdirToolScopePermission::Read)) .await .unwrap(); let nested = child - .scoped_session - .delegate(request("docs/sub", WorkdirDelegationPermission::Read)) + .scope(request("docs/sub", WorkdirToolScopePermission::Read)) .await .unwrap(); - nested.scoped_session.read(read("a")).await.unwrap(); + nested.read(read("a")).await.unwrap(); assert!( child - .scoped_session - .delegate(request("other", WorkdirDelegationPermission::Read)) + .scope(request("other", WorkdirToolScopePermission::Read)) .await .is_err() ); assert!( child - .scoped_session - .delegate(request("docs/sub", WorkdirDelegationPermission::Write)) + .scope(request("docs/sub", WorkdirToolScopePermission::Write)) .await .is_err() ); child.release(); assert!(matches!( - nested.scoped_session.read(read("a")).await, + nested.read(read("a")).await, Err(WorkdirError::SessionClosed) )); } @@ -1059,22 +1299,21 @@ mod tests { fs::create_dir_all(root.path().join("docs/sub")).unwrap(); let root_session = session(root.path()); let child = root_session - .delegate(request("docs", WorkdirDelegationPermission::Write)) + .scope(request("docs", WorkdirToolScopePermission::Write)) .await .unwrap(); let nested = child - .scoped_session - .delegate(request("docs/sub", WorkdirDelegationPermission::Write)) + .scope(request("docs/sub", WorkdirToolScopePermission::Write)) .await .unwrap(); for (session, label) in [ - (&root_session, "root"), - (&child.scoped_session, "child"), - (&nested.scoped_session, "nested"), + (root_session.tool_session(), "root"), + (child.tool_session(), "child"), + (nested.tool_session(), "nested"), ] { let output = run_command( - session, + &session, format!("printf {label}"), format!("{label}-command-during-nested-write"), ) @@ -1088,83 +1327,23 @@ mod tests { Err(WorkdirError::Denied(_)) )); assert!(matches!( - child - .scoped_session - .write(write("sub/child", "blocked")) - .await, + child.write(write("sub/child", "blocked")).await, Err(WorkdirError::Denied(_)) )); - nested - .scoped_session - .write(write("nested", "allowed")) - .await - .unwrap(); + nested.write(write("nested", "allowed")).await.unwrap(); nested.release(); child.release(); } #[tokio::test] - async fn reapplied_write_delegation_chain_forwards_command_lifecycle() { - let root = TempDir::new().unwrap(); - fs::create_dir_all(root.path().join("delegated")).unwrap(); - let applied = apply_delegation_chain( - session(root.path()), - [request("delegated", WorkdirDelegationPermission::Write)], - ) - .await - .unwrap(); - - let output = run_command( - &applied.scoped_session, - "printf reapplied", - "reapplied-command", - ) - .await; - assert_eq!(output.status, CommandStatus::Completed); - assert_eq!(output.content, "reapplied"); - } - - #[tokio::test] - async fn applied_chain_cannot_replace_outer_provider_attenuation() { - let root = TempDir::new().unwrap(); - fs::create_dir_all(root.path().join("outer")).unwrap(); - fs::create_dir_all(root.path().join("outside")).unwrap(); - let result = apply_delegation_chain( - Arc::new(LocalWorkdirSession::materialized_bound( - Workdir::new("delegation-chain-test"), - root.path().to_path_buf(), - root.path().to_path_buf(), - SharedScope::new( - Scope::from_config(&ScopeConfig { - allow: vec![ScopeRule { - target: root.path().to_path_buf(), - permission: Permission::Write, - recursive: true, - }], - deny: Vec::new(), - }) - .unwrap(), - ), - WorkdirSessionCapabilities::ALL, - )), - [ - request("outer", WorkdirDelegationPermission::Read), - request("outside", WorkdirDelegationPermission::Read), - ], - ) - .await; - assert!(matches!(result, Err(WorkdirError::Denied(_)))); - } - - #[tokio::test] - async fn closing_parent_invalidates_delegated_sessions() { + async fn closing_parent_invalidates_scoped_tools() { let root = TempDir::new().unwrap(); fs::create_dir_all(root.path().join("docs")).unwrap(); fs::write(root.path().join("docs/a"), "a").unwrap(); let parent = session(root.path()); let child = parent - .delegate(request("docs", WorkdirDelegationPermission::Read)) + .scope(request("docs", WorkdirToolScopePermission::Read)) .await .unwrap(); @@ -1175,15 +1354,17 @@ mod tests { command: "printf closed".into(), timeout_secs: 5, output_limit: 1024, + cwd: None, spill_dir: None, tool_call_id: Some("closed-parent-command".into()), }) .await, Err(WorkdirError::SessionClosed) )); - assert!(matches!( - child.scoped_session.read(read("a")).await, - Err(WorkdirError::SessionClosed) - )); + let child_result = child.read(read("a")).await; + assert!( + matches!(child_result, Err(WorkdirError::SessionClosed)), + "child result after parent close: {child_result:?}" + ); } } diff --git a/crates/workdir/src/workspace.rs b/crates/workdir/src/workspace.rs index 3872bc9d..fc696f99 100644 --- a/crates/workdir/src/workspace.rs +++ b/crates/workdir/src/workspace.rs @@ -104,15 +104,5 @@ mod tests { #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct WorkspaceWorkdirSessionOperationRequest { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub expected_session_fence: Option, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub delegations: Vec, pub operation: crate::http::WorkdirSessionOperation, } - -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct WorkspaceWorkdirSessionFence { - pub value: String, -} diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index 2217e973..51e7b7a1 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -718,8 +718,7 @@ async fn run_workdir_session_operation( .ok_or_else(RuntimeHttpWorkdirError::not_found)?; record.session.clone() }; - let applied = workdir::apply_delegation_chain(source, request.delegations).await?; - let session = applied.scoped_session.as_ref(); + let session = source.as_ref(); let operation = request.operation; let result = match operation { @@ -2080,8 +2079,8 @@ mod tests { use manifest::{Scope, SharedScope}; use tower::ServiceExt; use workdir::{ - GrepOutputMode, GrepRequest, LocalWorkdirSession, ReadRequest, StatRequest, Workdir, - WorkdirPath, WorkdirSessionCapabilities, + GrepOutputMode, GrepRequest, LocalWorkdirSession, StatRequest, Workdir, WorkdirPath, + WorkdirSessionCapabilities, }; #[test] @@ -2502,16 +2501,6 @@ mod tests { async fn workdir_session_operations_enforce_owner_and_close_terminally() { let temp = tempfile::tempdir().expect("tempdir"); std::fs::write(temp.path().join("hello.txt"), "hello").expect("write fixture"); - #[cfg(unix)] - { - use std::os::unix::fs::symlink; - std::fs::create_dir(temp.path().join("granted")).expect("granted directory"); - std::fs::write(temp.path().join("granted/visible"), "visible") - .expect("visible fixture"); - std::fs::create_dir(temp.path().join("secret")).expect("secret directory"); - std::fs::write(temp.path().join("secret/key"), "hidden").expect("secret fixture"); - symlink("../secret/key", temp.path().join("granted/link")).expect("symlink fixture"); - } let scope = SharedScope::new(Scope::writable(temp.path()).expect("scope")); let session: WorkdirSessionHandle = Arc::new(LocalWorkdirSession::materialized_bound( Workdir::new("wd-1"), @@ -2545,7 +2534,6 @@ mod tests { expires_at: u64::MAX, }; let operation = WorkdirSessionOperationRequest { - delegations: Vec::new(), operation: WorkdirSessionOperation::Stat(StatRequest { path: WorkdirPath::new("hello.txt").expect("logical path"), }), @@ -2562,7 +2550,6 @@ mod tests { assert!(matches!(result, WorkdirSessionOperationResult::Stat(_))); let grep = WorkdirSessionOperationRequest { - delegations: Vec::new(), operation: WorkdirSessionOperation::Grep(GrepRequest { pattern: "hello".into(), path: WorkdirPath::new("hello.txt").unwrap(), @@ -2585,78 +2572,7 @@ mod tests { ) .await .expect("grep direct file through provider operation"); - match result { - WorkdirSessionOperationResult::Grep(result) => { - assert_eq!(result.match_count, 1); - assert_eq!(result.matched_files, 1); - assert!(result.output.starts_with("hello.txt\n")); - assert!(result.output.contains("> 1 │ hello")); - } - other => panic!("unexpected workdir grep result: {other:?}"), - } - - #[cfg(unix)] - { - let delegated_visible = WorkdirSessionOperationRequest { - delegations: vec![workdir::WorkdirDelegationRequest { - rules: vec![workdir::WorkdirDelegationRule { - target: WorkdirPath::new("granted").unwrap(), - permission: workdir::WorkdirDelegationPermission::Read, - recursive: true, - }], - cwd: WorkdirPath::new("granted").unwrap(), - }], - operation: WorkdirSessionOperation::Read(ReadRequest { - path: WorkdirPath::new("visible").unwrap(), - offset: 0, - limit: 20, - max_bytes: 1024, - }), - }; - let visible = run_workdir_session_operation( - State(state.clone()), - Path("session-1".to_string()), - Some(Extension(auth.clone())), - Ok(Json(delegated_visible)), - ) - .await - .expect("non-root delegated cwd should resolve once") - .0; - assert!(matches!( - visible, - WorkdirSessionOperationResult::Read(result) if result.bytes == b"visible" - )); - - let delegated_read = WorkdirSessionOperationRequest { - delegations: vec![workdir::WorkdirDelegationRequest { - rules: vec![workdir::WorkdirDelegationRule { - target: WorkdirPath::new("granted").unwrap(), - permission: workdir::WorkdirDelegationPermission::Read, - recursive: true, - }], - cwd: WorkdirPath::new("granted").unwrap(), - }], - operation: WorkdirSessionOperation::Read(ReadRequest { - path: WorkdirPath::new("link").unwrap(), - offset: 0, - limit: 20, - max_bytes: 1024, - }), - }; - let error = run_workdir_session_operation( - State(state.clone()), - Path("session-1".to_string()), - Some(Extension(auth.clone())), - Ok(Json(delegated_read)), - ) - .await - .expect_err("provider must reject delegated symlink escape"); - assert_ne!(error.status, StatusCode::OK); - assert_eq!( - std::fs::read_to_string(temp.path().join("secret/key")).unwrap(), - "hidden" - ); - } + assert!(matches!(result, WorkdirSessionOperationResult::Grep(_))); let wrong_owner = RuntimeAuthContext { workspace_id: "workspace-b".to_string(), diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 70584074..214b1d55 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -514,6 +514,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() { @@ -911,6 +912,7 @@ pub(crate) async fn register_worker_tools( runtime_base: PathBuf, spawned_registry: Arc, parent_method_tx: Option>, + inherited_workdir_tool_broker: Option, ) -> std::io::Result> where C: LlmClient + Clone + 'static, @@ -919,21 +921,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(); @@ -1157,7 +1164,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(); @@ -1199,7 +1205,7 @@ where runtime_base.clone(), bash_output_dir.clone(), spawner_workspace_root, - source_workdir_session, + workdir_tool_broker, spawned_registry.clone(), spawner_manifest, prompts, diff --git a/crates/worker/src/feature/builtin/manage_workdir.rs b/crates/worker/src/feature/builtin/manage_workdir.rs index f00093eb..ab252835 100644 --- a/crates/worker/src/feature/builtin/manage_workdir.rs +++ b/crates/worker/src/feature/builtin/manage_workdir.rs @@ -12,7 +12,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, @@ -156,8 +156,6 @@ struct WorkspaceHttpWorkdirBackend { pub struct WorkspaceAttachedWorkdirSession { client: Arc, workdir: Workdir, - expected_session_fence: Option, - delegations: Vec, } impl WorkspaceAttachedWorkdirSession { @@ -165,8 +163,6 @@ impl WorkspaceAttachedWorkdirSession { Arc::new(Self { client, workdir: Workdir::new("workspace-attachment"), - expected_session_fence: None, - delegations: Vec::new(), }) } @@ -183,16 +179,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 +234,6 @@ impl WorkdirSession for WorkspaceAttachedWorkdirSession { WorkdirSessionCapabilities::ALL } - fn transports_delegation_context(&self) -> bool { - true - } - - async fn capture_delegation_source( - &self, - request: &workdir::WorkdirDelegationRequest, - ) -> Result { - 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 { match self.operate(WorkdirSessionOperation::Stat(request))? { WorkdirSessionOperationResult::Stat(result) => Ok(result), @@ -1155,6 +1095,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 +1119,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 +1162,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] diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index ec7ef586..52427800 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -709,7 +709,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 }) @@ -746,13 +746,16 @@ pub(crate) async fn prepare_internal_worker_session( visibility: InternalWorkerVisibility, child_registry: Option>, on_turn_end: Option>, + command_event_broker: Option, ) -> Result { 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(); @@ -887,6 +890,7 @@ pub(crate) async fn spawn_prepared_internal_worker_session( InternalWorkerVisibility::ServicePrivate, None, on_turn_end, + None, ) .await?; handle.send(input).await?; diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 62fec8d4..a9d670b3 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -25,7 +25,7 @@ use session_store::{ }; use tokio::sync::broadcast; use tracing::warn; -use workdir::WorkdirDelegation; +use workdir::WorkdirScopeLease; use crate::internal_worker::{InternalWorkerSessionHandle, InternalWorkerVisibility}; use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord}; @@ -68,7 +68,7 @@ pub(crate) struct SubWorkerStopSummary { pub(crate) struct InternalSpawnedWorkerRecord { pub worker_name: String, pub scope_delegated: Vec, - pub workdir_delegation: Arc, + pub workdir_tool_scope: Arc, #[cfg(test)] pub installed_tools: Arc<[String]>, pub session: InternalWorkerSessionHandle, @@ -86,7 +86,7 @@ impl InternalSpawnedWorkerRecord { pub(crate) fn new( worker_name: String, scope_delegated: Vec, - workdir_delegation: WorkdirDelegation, + workdir_tool_scope: WorkdirScopeLease, #[cfg(test)] installed_tools: Vec, session: InternalWorkerSessionHandle, change_tracker: Option, @@ -94,7 +94,7 @@ impl InternalSpawnedWorkerRecord { 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, @@ -690,7 +690,7 @@ impl SpawnedWorkerRegistry { if !record.claim_scope_reclaim() { return Ok(false); } - record.workdir_delegation.release(); + record.workdir_tool_scope.release(); let result = if let Some(parent_scope) = &self.parent_scope { parent_scope .update(|current| current.with_removed_deny_rules(delegated_write_rules(record))) @@ -966,7 +966,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 +976,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(); diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index e6b8f872..3091642e 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -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, + /// 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)] @@ -267,8 +269,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, + /// Parent-owned broker for scoped Workdir tool execution. + workdir_tool_broker: Option, /// Parent-owned in-memory registry shared by the five SubWorker tools. registry: Arc, /// Spawner's resolved Manifest. `profile = "inherit"` derives the @@ -295,7 +297,7 @@ impl SubWorkerSpawnTool { runtime_base: PathBuf, bash_output_dir: PathBuf, workspace_root: PathBuf, - source_workdir_session: Option, + workdir_tool_broker: Option, registry: Arc, spawner_manifest: WorkerManifest, prompt_loader: PromptCatalogSource, @@ -308,7 +310,7 @@ impl SubWorkerSpawnTool { runtime_base, bash_output_dir, workspace_root, - source_workdir_session, + workdir_tool_broker, registry, spawner_manifest, prompt_loader, @@ -341,6 +343,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(()) } @@ -370,7 +377,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 @@ -380,28 +387,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| { @@ -490,7 +484,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(), @@ -510,6 +503,7 @@ impl Tool for SubWorkerSpawnTool { self.runtime_base.clone(), child_registry.clone(), None, + Some(child_workdir_tool_broker.clone()), ) .await .map_err(|error| { @@ -552,6 +546,7 @@ impl Tool for SubWorkerSpawnTool { ); parent_notifications.notify(message, true); })), + Some(child_workdir_tool_broker.clone()), ) .await; let session = session_result.map_err(|error| { @@ -621,7 +616,7 @@ 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(), @@ -674,18 +669,18 @@ fn logical_workdir_path(value: &str, field: &str) -> Result { }) } -fn parse_workdir_scope(rules: &[ScopeRuleInput]) -> Result, ToolError> { +fn parse_workdir_scope(rules: &[ScopeRuleInput]) -> Result, 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, }) @@ -693,22 +688,24 @@ fn parse_workdir_scope(rules: &[ScopeRuleInput]) -> Result, - rules: Vec, -) -> Result { - Ok(WorkdirDelegationRequest { + rules: Vec, + command: bool, +) -> Result { + 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(), ) }) @@ -946,7 +943,7 @@ pub(crate) fn sub_worker_spawn_tool( runtime_base: PathBuf, bash_output_dir: PathBuf, workspace_root: PathBuf, - source_workdir_session: Option, + workdir_tool_broker: Option, registry: Arc, spawner_manifest: WorkerManifest, prompts: Arc>, @@ -958,7 +955,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, @@ -972,7 +969,7 @@ fn sub_worker_spawn_tool_impl( runtime_base: PathBuf, bash_output_dir: PathBuf, workspace_root: PathBuf, - source_workdir_session: Option, + workdir_tool_broker: Option, registry: Arc, spawner_manifest: WorkerManifest, prompts: Arc>, @@ -1004,7 +1001,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(), @@ -1037,12 +1034,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") )); } @@ -1079,6 +1076,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(); @@ -1173,7 +1171,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(), @@ -1189,7 +1187,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, @@ -1212,7 +1210,8 @@ enabled = false "target": ".", "permission": "write", "recursive": true - }] + }], + "command": true }); assert!(spawner_scope.snapshot().is_writable(&workspace_root)); @@ -1247,15 +1246,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), @@ -1371,7 +1361,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()); @@ -1426,7 +1416,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)); @@ -1438,7 +1428,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(), @@ -1478,51 +1468,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::>(); - 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" ); } @@ -1534,6 +1485,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) @@ -1663,7 +1615,6 @@ enabled = false #[derive(Debug, Default)] struct StrictRemoteWorkdirWorkspaceClient { requests: Mutex>, - foreign_scope_rejections: AtomicUsize, } impl StrictRemoteWorkdirWorkspaceClient { @@ -1695,59 +1646,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(), + )) } } diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index 143c1da9..42bba8f6 100644 --- a/crates/worker/tests/controller_test.rs +++ b/crates/worker/tests/controller_test.rs @@ -332,6 +332,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, }) @@ -376,6 +377,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()), }) @@ -484,6 +486,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()), }) @@ -560,6 +563,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, }) diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 3599f26d..aeee139e 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -48,8 +48,7 @@ use workdir::http::{ }; use workdir::workspace::{ MaterializerKind, WorkingDirectoryCleanupTarget, WorkingDirectoryOccupancy, - WorkingDirectoryStatusKind, WorkingDirectorySummary, WorkspaceWorkdirSessionFence, - WorkspaceWorkdirSessionOperationRequest, + WorkingDirectoryStatusKind, WorkingDirectorySummary, WorkspaceWorkdirSessionOperationRequest, }; use workdir::{CommandHandle, WorkdirSessionHandle}; use worker::feature::builtin::{WorkerObservationSubject, WorkerObservationSubjectRef}; @@ -355,7 +354,6 @@ static EMBEDDED_RUNTIME_REQUEST_IDENTITY: std::sync::LazyLock< struct WorkdirCommandSession { source: WorkdirSessionHandle, provider_handle: CommandHandle, - delegations: Vec, } enum RegisteredWorkdirSession { @@ -398,7 +396,6 @@ impl WorkdirSessionRegistry { worker: RuntimeWorkerRef, source: WorkdirSessionHandle, provider_handle: CommandHandle, - delegations: Vec, ) -> CommandHandle { let external_handle = loop { let candidate = CommandHandle(Uuid::now_v7().to_string()); @@ -414,7 +411,6 @@ impl WorkdirSessionRegistry { WorkdirCommandSession { source, provider_handle, - delegations, }, ); external_handle @@ -2586,10 +2582,6 @@ fn build_inner_router(api: WorkspaceApi) -> Router { post(scoped_attach_current_worker_workdir) .delete(scoped_detach_current_worker_workdir), ) - .route( - "/api/w/{workspace_id}/workers/self/workdir-session/fence", - get(scoped_current_worker_workdir_session_fence), - ) .route( "/api/w/{workspace_id}/workers/self/workdir-session/operations", post(scoped_execute_current_worker_workdir_operation), @@ -7341,46 +7333,11 @@ async fn scoped_detach_current_worker_workdir( })) } -async fn scoped_current_worker_workdir_session_fence( - State(api): State, - AxumPath(path): AxumPath, - headers: HeaderMap, -) -> ApiResult> { - validate_workspace_scope(&api, &path.workspace_id)?; - let worker = current_worker_identity(&api, &path.workspace_id, &headers)?; - let session_lock = current_worker_session_lock(&api, &worker); - let _session_guard = session_lock.lock().await; - let link = current_worker_active_attachment(&api, &worker)?; - Ok(Json(WorkspaceWorkdirSessionFence { - value: current_worker_workdir_session_fence(&link), - })) -} - -fn current_worker_workdir_session_fence(link: &WorkerWorkdirLinkRecord) -> String { - format!("v1:{}\0{}", link.workdir_id, link.linked_at) -} - -fn validate_current_worker_workdir_session_fence( - link: &WorkerWorkdirLinkRecord, - expected: Option<&str>, -) -> Result<()> { - if expected.is_some_and(|expected| expected != current_worker_workdir_session_fence(link)) { - Err(Error::WorkdirAttachmentConflict( - "delegated Workdir session attachment changed".to_string(), - )) - } else { - Ok(()) - } -} - fn validated_current_worker_attachment( api: &WorkspaceApi, worker: &RuntimeWorkerRef, - expected_session_fence: Option<&str>, ) -> ApiResult { - let link = current_worker_active_attachment(api, worker)?; - validate_current_worker_workdir_session_fence(&link, expected_session_fence)?; - Ok(link) + current_worker_active_attachment(api, worker) } #[derive(Debug)] @@ -7435,23 +7392,13 @@ async fn scoped_execute_current_worker_workdir_operation( ) -> std::result::Result, WorkdirOperationApiError> { validate_workspace_scope(&api, &path.workspace_id)?; let worker = current_worker_identity(&api, &path.workspace_id, &headers)?; - let expected_session_fence = request.expected_session_fence; - let delegations = request.delegations; let result = match request.operation { WorkdirSessionOperation::CommandStart(command) => { let session_lock = current_worker_session_lock(&api, &worker); let _session_guard = session_lock.lock().await; - let link = validated_current_worker_attachment( - &api, - &worker, - expected_session_fence.as_deref(), - )?; + let link = validated_current_worker_attachment(&api, &worker)?; let source = open_current_worker_workdir_session_locked(&api, &worker, &link).await?; - let applied = - apply_current_worker_delegations(&worker, source.clone(), delegations.clone()) - .await?; - let provider_handle = applied - .scoped_session + let provider_handle = source .start_command(command) .await .map_err(|error| current_worker_workdir_operation_error(&worker, error))?; @@ -7470,58 +7417,32 @@ async fn scoped_execute_current_worker_workdir_operation( .workdir_sessions .lock() .expect("Workdir session registry lock poisoned") - .register_command( - worker.clone(), - registered_source, - provider_handle, - delegations, - ); + .register_command(worker.clone(), registered_source, provider_handle); WorkdirSessionOperationResult::CommandStart(external_handle) } WorkdirSessionOperation::CommandStatus(external_handle) => { - let (session, provider_handle) = current_worker_command_session( - &api, - &worker, - &external_handle, - &delegations, - expected_session_fence.as_deref(), - ) - .await?; + let (session, provider_handle) = + current_worker_command_session(&api, &worker, &external_handle)?; session - .scoped_session .command_status(provider_handle) .await .map(WorkdirSessionOperationResult::CommandStatus) .map_err(|error| current_worker_workdir_operation_error(&worker, error))? } WorkdirSessionOperation::CommandOutput(mut output) => { - let (session, provider_handle) = current_worker_command_session( - &api, - &worker, - &output.handle, - &delegations, - expected_session_fence.as_deref(), - ) - .await?; + let (session, provider_handle) = + current_worker_command_session(&api, &worker, &output.handle)?; output.handle = provider_handle; session - .scoped_session .command_output(output) .await .map(WorkdirSessionOperationResult::CommandOutput) .map_err(|error| current_worker_workdir_operation_error(&worker, error))? } WorkdirSessionOperation::CommandCancel(external_handle) => { - let (session, provider_handle) = current_worker_command_session( - &api, - &worker, - &external_handle, - &delegations, - expected_session_fence.as_deref(), - ) - .await?; + let (session, provider_handle) = + current_worker_command_session(&api, &worker, &external_handle)?; session - .scoped_session .cancel_command(provider_handle) .await .map(|()| WorkdirSessionOperationResult::CommandCancel) @@ -7536,14 +7457,9 @@ async fn scoped_execute_current_worker_workdir_operation( | WorkdirSessionOperation::Grep(_)) => { let session_lock = current_worker_session_lock(&api, &worker); let _session_guard = session_lock.lock().await; - let link = validated_current_worker_attachment( - &api, - &worker, - expected_session_fence.as_deref(), - )?; + let link = validated_current_worker_attachment(&api, &worker)?; let source = open_current_worker_workdir_session_locked(&api, &worker, &link).await?; - let applied = apply_current_worker_delegations(&worker, source, delegations).await?; - execute_workdir_session_operation(&applied.scoped_session, operation) + execute_workdir_session_operation(&source, operation) .await .map_err(|error| current_worker_workdir_operation_error(&worker, error))? } @@ -7551,29 +7467,12 @@ async fn scoped_execute_current_worker_workdir_operation( Ok(Json(result)) } -async fn apply_current_worker_delegations( - worker: &RuntimeWorkerRef, - source: WorkdirSessionHandle, - delegations: Vec, -) -> Result { - workdir::apply_delegation_chain(source, delegations) - .await - .map_err(|error| Error::RuntimeOperationFailed { - runtime_id: worker.runtime_id.clone(), - code: "workdir_session_delegation_failed".to_string(), - message: error.to_string(), - }) -} - -async fn current_worker_command_session( +fn current_worker_command_session( api: &WorkspaceApi, worker: &RuntimeWorkerRef, external_handle: &CommandHandle, - delegations: &[workdir::WorkdirDelegationRequest], - expected_session_fence: Option<&str>, -) -> std::result::Result<(workdir::AppliedWorkdirDelegation, CommandHandle), WorkdirOperationApiError> -{ - let _link = validated_current_worker_attachment(api, worker, expected_session_fence)?; +) -> std::result::Result<(WorkdirSessionHandle, CommandHandle), WorkdirOperationApiError> { + let _link = validated_current_worker_attachment(api, worker)?; let command = api .workdir_sessions .lock() @@ -7585,15 +7484,7 @@ async fn current_worker_command_session( workdir::WorkdirError::UnknownCommand(external_handle.0.clone()), )) })?; - if command.delegations != delegations { - return Err(Error::WorkdirAttachmentConflict( - "command lifecycle delegation differs from CommandStart".to_string(), - ) - .into()); - } - let session = - apply_current_worker_delegations(worker, command.source, command.delegations).await?; - Ok((session, command.provider_handle)) + Ok((command.source, command.provider_handle)) } fn current_worker_workdir_operation_error( @@ -16775,6 +16666,7 @@ mod tests { command: "printf ready; sleep 30".to_string(), timeout_secs: 60, output_limit: 4096, + cwd: None, spill_dir: None, tool_call_id: Some("tool-call-command-session".to_string()), }) @@ -16784,12 +16676,8 @@ mod tests { let mut registry = WorkdirSessionRegistry::default(); registry.insert_attachment(worker.clone(), source.clone()); let registered_source = registry.remove_attachment(&worker).unwrap(); - let external_handle = registry.register_command( - worker.clone(), - registered_source, - provider_handle.clone(), - Vec::new(), - ); + let external_handle = + registry.register_command(worker.clone(), registered_source, provider_handle.clone()); assert_ne!(external_handle, provider_handle); let refreshed: WorkdirSessionHandle = Arc::new(workdir::LocalWorkdirSession::new( @@ -23503,30 +23391,6 @@ mod tests { assert_eq!(response.status(), StatusCode::BAD_REQUEST); } - #[test] - fn delegated_workdir_session_fence_rejects_reattached_link() { - let first = WorkerWorkdirLinkRecord { - workspace_id: "workspace-a".to_string(), - worker: workdir::workspace::RuntimeWorkerRef::new("runtime-a", "worker-a"), - workdir_id: "workdir-a".to_string(), - role: "primary".to_string(), - linked_at: "2026-01-01T00:00:00Z".to_string(), - unlinked_at: None, - }; - let expected = current_worker_workdir_session_fence(&first); - assert!(validate_current_worker_workdir_session_fence(&first, None).is_ok()); - assert!(validate_current_worker_workdir_session_fence(&first, Some(&expected)).is_ok()); - - let reattached = WorkerWorkdirLinkRecord { - linked_at: "2026-01-01T00:00:01Z".to_string(), - ..first - }; - assert!(matches!( - validate_current_worker_workdir_session_fence(&reattached, Some(&expected)), - Err(Error::WorkdirAttachmentConflict(_)) - )); - } - #[tokio::test] async fn backend_workdir_session_proxy_executes_typed_operations() { use manifest::Scope; diff --git a/resources/flows/coder-review.dcdl b/resources/flows/coder-review.dcdl index b42dbd76..4de10077 100644 --- a/resources/flows/coder-review.dcdl +++ b/resources/flows/coder-review.dcdl @@ -15,7 +15,7 @@ }; review = { - instructions = "Use the current Ticket Merge Request as review authority. Call `ShowMergeRequest` and confirm its source selector resolves to exact committed implementation HEAD, then spawn one actual direct-child SubWorker with profile builtin:reviewer, write scope for Workdir inspection and command validation, and only the Ticket id in the structured review handoff. The trusted spawn layer records `ReviewRequested` with the exact source ref and injects review capability; do not place commit/ref identity, capability material, or a prewritten verdict in model input. The child must commit `ReviewMergeRequest`; prose output and Worker observation are not approval authority. After the structured result for the exact current source ref exists, request a Flow transition."; + instructions = "Use the current Ticket Merge Request as review authority. Call `ShowMergeRequest` and confirm its source selector resolves to exact committed implementation HEAD, then spawn one actual direct-child SubWorker with profile builtin:reviewer, write scope plus an explicit command grant for Workdir inspection and command validation, and only the Ticket id in the structured review handoff. The trusted spawn layer records `ReviewRequested` with the exact source ref and injects review capability; do not place commit/ref identity, capability material, or a prewritten verdict in model input. The child must commit `ReviewMergeRequest`; prose output and Worker observation are not approval authority. After the structured result for the exact current source ref exists, request a Flow transition."; transitions = { approved = { target = "complete"; diff --git a/resources/prompts/internal/sub_worker_spawn_tool_description.md b/resources/prompts/internal/sub_worker_spawn_tool_description.md index a652f551..d6d50df1 100644 --- a/resources/prompts/internal/sub_worker_spawn_tool_description.md +++ b/resources/prompts/internal/sub_worker_spawn_tool_description.md @@ -1,8 +1,8 @@ Spawn a parent-owned Internal SubWorker session to split context for a delegated task. The parent Worker's write scope is reduced by the scope passed here; the Internal SubWorker starts running `task` immediately without creating a Runtime Worker record, OS process, PID, or Unix socket. It remains available for follow-up turns until explicitly stopped or its parent exits. -Optional `cwd`: when provided, the spawned SubWorker's tool default working directory only. It must be an absolute existing directory covered by the child's delegated readable scope, and it does not change workspace/Profile/memory/Ticket roots or grant authority. `name` must be unique among this Worker's direct children. +Optional `cwd`: when provided, the spawned SubWorker's tool default working directory only. It must be a Workdir-relative existing directory covered by the child's readable scope, and it does not change workspace/Profile/memory/Ticket roots or grant authority. `name` must be unique among this Worker's direct children. -Profile selection: `profile` may be omitted or set to `default` to use the effective child default profile, set to `inherit` to derive reusable child configuration from this Worker, or set to one of the registry selectors below. Raw/path profile selectors are not accepted by SubWorkerSpawn. `scope` is always the only delegated filesystem capability; profile scope is replaced by the explicit SubWorkerSpawn scope. +Profile selection: `profile` may be omitted or set to `default` to use the effective child default profile, set to `inherit` to derive reusable child configuration from this Worker, or set to one of the registry selectors below. Raw/path profile selectors are not accepted by SubWorkerSpawn. `scope` is the child's only filesystem capability and replaces profile scope. `command` is a separate explicit grant, defaults to false, and is accepted only with a writable scope; writable scope alone does not grant command execution. Default profile: {{ default_profile }} Special selector: inherit — derive reusable model/worker/tool policy from the spawner while replacing worker.name and scope. From e7079e223feac3e80a94c5d756a93b2f2824002b Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 02:37:13 +0900 Subject: [PATCH 2/6] fix: close scoped SubWorker command authority --- crates/workdir/src/scope.rs | 205 ++++++++++++++++++++++++++-- crates/worker/src/spawn/registry.rs | 7 +- 2 files changed, 196 insertions(+), 16 deletions(-) diff --git a/crates/workdir/src/scope.rs b/crates/workdir/src/scope.rs index 7b024696..cf6db8cc 100644 --- a/crates/workdir/src/scope.rs +++ b/crates/workdir/src/scope.rs @@ -106,6 +106,7 @@ pub struct WorkdirScopeLease { broker: WorkdirToolBroker, pub capabilities: WorkdirSessionCapabilities, validity: Arc, + cleanup_pending: Arc, } impl std::fmt::Debug for WorkdirScopeLease { @@ -134,12 +135,74 @@ impl WorkdirScopeLease { self.broker.scope(request).await } + pub async fn close(&self) -> Result<(), WorkdirError> { + self.validity.active.store(false, Ordering::Release); + let command_ids = self + .broker + .authority + .owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .iter() + .cloned() + .collect::>(); + let mut first_error = None; + for command_id in command_ids { + let handle = CommandHandle(command_id.clone()); + let cancel = self + .broker + .authority + .source + .cancel_command(handle.clone()) + .await; + let terminal = self + .broker + .authority + .source + .command_output(CommandOutputRequest { + handle, + cursor: 0, + limit: 1, + wait: true, + }) + .await; + match (cancel, terminal) { + (_, Ok(_)) + | (Ok(()), Err(WorkdirError::UnknownCommand(_))) + | (Err(WorkdirError::UnknownCommand(_)), Err(WorkdirError::UnknownCommand(_))) => { + self.broker + .authority + .owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .remove(&command_id); + } + (Err(error), _) | (_, Err(error)) => { + first_error.get_or_insert(error); + } + } + } + if let Some(error) = first_error { + return Err(error); + } + tokio::task::yield_now().await; + self.finish_release(); + Ok(()) + } + pub fn is_active(&self) -> bool { self.validity.is_active() } - pub fn release(&self) { + /// Revoke a scope whose owner has already terminalized every tool call. + /// Use [`Self::close`] when commands may still be live. + pub fn revoke(&self) { + self.finish_release(); + } + + fn finish_release(&self) { self.validity.active.store(false, Ordering::Release); + self.cleanup_pending.store(false, Ordering::Release); if let Some(forwarder) = &self.broker.event_forwarder && let Some(handle) = forwarder .lock() @@ -161,7 +224,7 @@ impl std::ops::Deref for WorkdirScopeLease { impl Drop for WorkdirScopeLease { fn drop(&mut self) { - self.release(); + self.finish_release(); } } @@ -195,6 +258,7 @@ impl SessionValidity { #[derive(Clone, Debug)] struct ActiveWriteLease { validity: Weak, + cleanup_pending: Weak, rules: Vec, } @@ -471,22 +535,52 @@ impl ScopedWorkdirSession { self.ensure_scope_targets_do_not_traverse_symlinks(&request.rules) .await?; let validity = SessionValidity::child(self.validity.clone()); + let cleanup_pending = Arc::new(AtomicBool::new(true)); let id = self.next_lease_id.fetch_add(1, Ordering::Relaxed); if request .rules .iter() .any(|rule| rule.permission == WorkdirToolScopePermission::Write) { - self.child_write_leases + let mut leases = self + .child_write_leases .lock() - .expect("Workdir tool scope lease mutex poisoned") - .insert( - id, - ActiveWriteLease { - validity: Arc::downgrade(&validity), - rules: request.rules.clone(), - }, - ); + .expect("Workdir tool scope lease mutex poisoned"); + leases.retain(|_, lease| { + lease + .validity + .upgrade() + .is_some_and(|validity| validity.is_active()) + || lease + .cleanup_pending + .upgrade() + .is_some_and(|pending| pending.load(Ordering::Acquire)) + }); + let requested_write_rules = request + .rules + .iter() + .filter(|rule| rule.permission == WorkdirToolScopePermission::Write); + for requested in requested_write_rules { + if leases.values().any(|lease| { + lease + .rules + .iter() + .any(|active| rules_overlap(active, requested)) + }) { + return Err(WorkdirError::Denied(format!( + "scoped write path `{}` overlaps an active child scope", + requested.target + ))); + } + } + leases.insert( + id, + ActiveWriteLease { + validity: Arc::downgrade(&validity), + cleanup_pending: Arc::downgrade(&cleanup_pending), + rules: request.rules.clone(), + }, + ); } let owned_commands = Arc::new(Mutex::new(HashSet::new())); let (command_events, _) = broadcast::channel(64); @@ -517,6 +611,7 @@ impl ScopedWorkdirSession { broker, capabilities, validity, + cleanup_pending, }) } } @@ -787,6 +882,13 @@ fn unix_timestamp_ms() -> u64 { .min(u128::from(u64::MAX)) as u64 } +fn rules_overlap(left: &WorkdirToolScopeRule, right: &WorkdirToolScopeRule) -> bool { + left.permission == WorkdirToolScopePermission::Write + && right.permission == WorkdirToolScopePermission::Write + && (rule_allows_path(left, &right.target, WorkdirToolScopePermission::Write) + || rule_allows_path(right, &left.target, WorkdirToolScopePermission::Write)) +} + fn rule_allows_path( rule: &WorkdirToolScopeRule, path: &FsPath, @@ -1231,7 +1333,7 @@ mod tests { )); parent.write(write("other/file", "parent")).await.unwrap(); child.write(write("file", "child")).await.unwrap(); - child.release(); + child.close().await.unwrap(); assert!(matches!( child .start_command(CommandRequest { @@ -1255,6 +1357,79 @@ mod tests { )); } + #[tokio::test] + async fn sibling_write_scopes_must_not_overlap() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("shared/one")).unwrap(); + fs::create_dir_all(root.path().join("other")).unwrap(); + let parent = session(root.path()); + let first = parent + .scope(request("shared", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + + assert!(matches!( + parent + .scope(request("shared/one", WorkdirToolScopePermission::Write)) + .await, + Err(WorkdirError::Denied(_)) + )); + let other = parent + .scope(request("other", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + other.close().await.unwrap(); + first.close().await.unwrap(); + } + + #[tokio::test] + async fn closing_scope_cancels_and_terminalizes_owned_commands() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("work")).unwrap(); + let parent = session(root.path()); + let child = parent + .scope(request("work", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + let mut events = child.subscribe_command_events().unwrap(); + let handle = child + .start_command(CommandRequest { + command: "sleep 30; printf leaked > marker".into(), + timeout_secs: 60, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: Some("owned-command".into()), + }) + .await + .unwrap(); + assert!(matches!( + events.recv().await.unwrap(), + CommandEvent::Started { .. } + )); + + child.close().await.unwrap(); + + assert!(matches!( + parent.command_status(handle).await, + Ok(CommandStatus::Cancelled | CommandStatus::Completed | CommandStatus::Failed) + | Err(WorkdirError::UnknownCommand(_)) + )); + assert!(!root.path().join("work/marker").exists()); + let terminal = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + if let CommandEvent::Terminal { .. } = events.recv().await.unwrap() { + break; + } + } + }) + .await; + assert!( + terminal.is_ok(), + "scope close must publish terminal command telemetry" + ); + } + #[tokio::test] async fn nested_delegation_is_attenuated_and_parent_revocation_cascades() { let root = TempDir::new().unwrap(); @@ -1286,7 +1461,7 @@ mod tests { .is_err() ); - child.release(); + child.close().await.unwrap(); assert!(matches!( nested.read(read("a")).await, Err(WorkdirError::SessionClosed) @@ -1332,8 +1507,8 @@ mod tests { )); nested.write(write("nested", "allowed")).await.unwrap(); - nested.release(); - child.release(); + nested.close().await.unwrap(); + child.close().await.unwrap(); } #[tokio::test] diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index a9d670b3..c270243f 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -690,7 +690,7 @@ impl SpawnedWorkerRegistry { if !record.claim_scope_reclaim() { return Ok(false); } - record.workdir_tool_scope.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))) @@ -731,6 +731,11 @@ impl SpawnedWorkerRegistry { .stop() .await .map_err(|error| io::Error::other(error.to_string()))?; + 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 = From 0f8d61188aad1842f14e5ebf94a3be9f48564225 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 03:12:55 +0900 Subject: [PATCH 3/6] fix: order SubWorker cleanup before Workdir release --- crates/workdir/src/scope.rs | 254 ++++++++++++++++-- crates/worker/src/controller.rs | 37 ++- .../src/feature/builtin/manage_workdir.rs | 124 ++++++++- crates/worker/src/spawn/registry.rs | 54 +++- crates/worker/src/spawn/tool.rs | 17 +- 5 files changed, 443 insertions(+), 43 deletions(-) diff --git a/crates/workdir/src/scope.rs b/crates/workdir/src/scope.rs index cf6db8cc..c48c1541 100644 --- a/crates/workdir/src/scope.rs +++ b/crates/workdir/src/scope.rs @@ -70,6 +70,9 @@ impl WorkdirToolBroker { child_write_leases: Mutex::new(HashMap::new()), next_lease_id: AtomicU64::new(1), owned_commands: Arc::new(Mutex::new(HashSet::new())), + pending_command_events: Arc::new(Mutex::new(HashMap::new())), + starting_tool_calls: Arc::new(Mutex::new(HashSet::new())), + forwarded_terminals: Arc::new(Mutex::new(HashSet::new())), command_events, closes_source: true, }); @@ -107,6 +110,7 @@ pub struct WorkdirScopeLease { pub capabilities: WorkdirSessionCapabilities, validity: Arc, cleanup_pending: Arc, + close_lock: Arc>, } impl std::fmt::Debug for WorkdirScopeLease { @@ -136,6 +140,10 @@ impl WorkdirScopeLease { } pub async fn close(&self) -> Result<(), WorkdirError> { + let _close_guard = self.close_lock.lock().await; + if !self.cleanup_pending.load(Ordering::Acquire) { + return Ok(()); + } self.validity.active.store(false, Ordering::Release); let command_ids = self .broker @@ -167,9 +175,28 @@ impl WorkdirScopeLease { }) .await; match (cancel, terminal) { - (_, Ok(_)) - | (Ok(()), Err(WorkdirError::UnknownCommand(_))) + (_, Ok(output)) => { + self.broker.authority.publish_terminal_if_missing( + &command_id, + output.status, + output.exit_code, + output.next_cursor.unwrap_or(output.content.len()) as u64, + ); + self.broker + .authority + .owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .remove(&command_id); + } + (Ok(()), Err(WorkdirError::UnknownCommand(_))) | (Err(WorkdirError::UnknownCommand(_)), Err(WorkdirError::UnknownCommand(_))) => { + self.broker.authority.publish_terminal_if_missing( + &command_id, + CommandStatus::Cancelled, + None, + 0, + ); self.broker .authority .owned_commands @@ -271,6 +298,9 @@ struct ScopedWorkdirSession { child_write_leases: Mutex>, next_lease_id: AtomicU64, owned_commands: Arc>>, + pending_command_events: Arc>>>, + starting_tool_calls: Arc>>, + forwarded_terminals: Arc>>, command_events: broadcast::Sender, closes_source: bool, } @@ -380,12 +410,42 @@ impl ScopedWorkdirSession { } } + fn publish_terminal_if_missing( + &self, + command_id: &str, + status: CommandStatus, + exit_code: Option, + offset: u64, + ) { + publish_owned_command_event( + &self.command_events, + &self.forwarded_terminals, + CommandEvent::Terminal { + command_id: command_id.to_string(), + status, + exit_code, + stdout_end_offset: offset, + stderr_end_offset: 0, + observed_at_ms: unix_timestamp_ms(), + }, + ); + } + fn ensure_parent_write_available(&self, path: &FsPath) -> Result<(), WorkdirError> { let mut leases = self .child_write_leases .lock() .expect("Workdir tool scope lease mutex poisoned"); - leases.retain(|_, lease| lease.validity.upgrade().is_some_and(|v| v.is_active())); + leases.retain(|_, lease| { + lease + .validity + .upgrade() + .is_some_and(|validity| validity.is_active()) + || lease + .cleanup_pending + .upgrade() + .is_some_and(|pending| pending.load(Ordering::Acquire)) + }); if leases.values().any(|lease| { lease.rules.iter().any(|rule| { rule.permission == WorkdirToolScopePermission::Write @@ -583,10 +643,16 @@ impl ScopedWorkdirSession { ); } let owned_commands = Arc::new(Mutex::new(HashSet::new())); + let pending_command_events = Arc::new(Mutex::new(HashMap::new())); + let starting_tool_calls = Arc::new(Mutex::new(HashSet::new())); + let forwarded_terminals = Arc::new(Mutex::new(HashSet::new())); let (command_events, _) = broadcast::channel(64); let event_forwarder = forward_owned_command_events( self.source.subscribe_command_events(), owned_commands.clone(), + pending_command_events.clone(), + starting_tool_calls.clone(), + forwarded_terminals.clone(), command_events.clone(), ) .map(|handle| Arc::new(Mutex::new(Some(handle)))); @@ -599,6 +665,9 @@ impl ScopedWorkdirSession { child_write_leases: Mutex::new(HashMap::new()), next_lease_id: AtomicU64::new(1), owned_commands, + pending_command_events, + starting_tool_calls, + forwarded_terminals, command_events, closes_source: false, }); @@ -612,6 +681,7 @@ impl ScopedWorkdirSession { capabilities, validity, cleanup_pending, + close_lock: Arc::new(tokio::sync::Mutex::new(())), }) } } @@ -687,16 +757,59 @@ impl WorkdirSession for ScopedWorkdirSession { None => self.cwd.clone(), }); } - let handle = self.source.start_command(request).await?; - self.owned_commands + if let Some(tool_call_id) = &tool_call_id { + self.starting_tool_calls + .lock() + .expect("starting tool call mutex poisoned") + .insert(tool_call_id.clone()); + } + let handle = match self.source.start_command(request).await { + Ok(handle) => handle, + Err(error) => { + if let Some(tool_call_id) = &tool_call_id { + self.starting_tool_calls + .lock() + .expect("starting tool call mutex poisoned") + .remove(tool_call_id); + } + return Err(error); + } + }; + let mut owned = self + .owned_commands .lock() - .expect("scoped command set mutex poisoned") - .insert(handle.0.clone()); - let _ = self.command_events.send(CommandEvent::Started { - command_id: handle.0.clone(), - tool_call_id, - observed_at_ms: unix_timestamp_ms(), - }); + .expect("scoped command set mutex poisoned"); + owned.insert(handle.0.clone()); + if let Some(tool_call_id) = &tool_call_id { + self.starting_tool_calls + .lock() + .expect("starting tool call mutex poisoned") + .remove(tool_call_id); + } + let pending = self + .pending_command_events + .lock() + .expect("pending scoped command event mutex poisoned") + .remove(&handle.0) + .unwrap_or_default(); + drop(owned); + if !pending + .iter() + .any(|event| matches!(event, CommandEvent::Started { .. })) + { + publish_owned_command_event( + &self.command_events, + &self.forwarded_terminals, + CommandEvent::Started { + command_id: handle.0.clone(), + tool_call_id, + observed_at_ms: unix_timestamp_ms(), + }, + ); + } + for event in pending { + publish_owned_command_event(&self.command_events, &self.forwarded_terminals, event); + } Ok(handle) } @@ -713,6 +826,12 @@ impl WorkdirSession for ScopedWorkdirSession { let command_id = request.handle.0.clone(); let output = self.source.command_output(request).await?; if !matches!(output.status, CommandStatus::Running) { + self.publish_terminal_if_missing( + &command_id, + output.status, + output.exit_code, + output.next_cursor.unwrap_or(output.content.len()) as u64, + ); self.owned_commands .lock() .expect("scoped command set mutex poisoned") @@ -848,6 +967,9 @@ impl WorkdirSession for ReadOnlyWorkdirSession { fn forward_owned_command_events( receiver: Option>, owned_commands: Arc>>, + pending_command_events: Arc>>>, + starting_tool_calls: Arc>>, + forwarded_terminals: Arc>>, sender: broadcast::Sender, ) -> Option> { let mut receiver = receiver?; @@ -858,22 +980,73 @@ fn forward_owned_command_events( Err(broadcast::error::RecvError::Lagged(_)) => continue, Err(broadcast::error::RecvError::Closed) => break, }; - let command_id = match &event { - CommandEvent::Started { .. } => continue, - CommandEvent::Output { command_id, .. } - | CommandEvent::Terminal { command_id, .. } => command_id.clone(), - }; - let owned = owned_commands + let command_id = command_event_id(&event).to_string(); + let mut owned = owned_commands .lock() - .expect("scoped command set mutex poisoned") - .contains(&command_id); - if owned { - let _ = sender.send(event); + .expect("scoped command set mutex poisoned"); + if !owned.contains(&command_id) { + let claimed = matches!( + &event, + CommandEvent::Started { + tool_call_id: Some(tool_call_id), + .. + } if starting_tool_calls + .lock() + .expect("starting tool call mutex poisoned") + .contains(tool_call_id) + ); + if !claimed { + continue; + } + owned.insert(command_id.clone()); + pending_command_events + .lock() + .expect("pending scoped command event mutex poisoned") + .entry(command_id) + .or_default() + .push(event); + continue; } + let mut pending = pending_command_events + .lock() + .expect("pending scoped command event mutex poisoned"); + if let Some(events) = pending.get_mut(&command_id) { + if events.len() < 64 { + events.push(event); + } + continue; + } + drop(pending); + drop(owned); + publish_owned_command_event(&sender, &forwarded_terminals, event); } })) } +fn command_event_id(event: &CommandEvent) -> &str { + match event { + CommandEvent::Started { command_id, .. } + | CommandEvent::Output { command_id, .. } + | CommandEvent::Terminal { command_id, .. } => command_id, + } +} + +fn publish_owned_command_event( + sender: &broadcast::Sender, + forwarded_terminals: &Mutex>, + event: CommandEvent, +) { + if let CommandEvent::Terminal { command_id, .. } = &event + && !forwarded_terminals + .lock() + .expect("forwarded terminal command mutex poisoned") + .insert(command_id.clone()) + { + return; + } + let _ = sender.send(event); +} + fn unix_timestamp_ms() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -1382,6 +1555,43 @@ mod tests { first.close().await.unwrap(); } + #[tokio::test] + async fn fast_command_keeps_started_output_terminal_event_order() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("work")).unwrap(); + let parent = session(root.path()); + let child = parent + .scope(request("work", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + let mut events = child.subscribe_command_events().unwrap(); + + let output = run_command(&child.tool_session(), "printf fast-output", "fast-command").await; + assert_eq!(output.content, "fast-output"); + + let mut kinds = Vec::new(); + let mut streamed = String::new(); + while kinds.last().is_none_or(|kind| *kind != "terminal") { + let event = tokio::time::timeout(std::time::Duration::from_secs(1), events.recv()) + .await + .expect("fast command event timeout") + .expect("fast command event channel"); + match event { + CommandEvent::Started { .. } => kinds.push("started"), + CommandEvent::Output { content, .. } => { + kinds.push("output"); + streamed.push_str(&content); + } + CommandEvent::Terminal { .. } => kinds.push("terminal"), + } + } + assert_eq!(kinds.first(), Some(&"started")); + assert_eq!(kinds.last(), Some(&"terminal")); + assert!(kinds.contains(&"output")); + assert!(streamed.contains("fast-output")); + child.close().await.unwrap(); + } + #[tokio::test] async fn closing_scope_cancels_and_terminalizes_owned_commands() { let root = TempDir::new().unwrap(); diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 214b1d55..a73a175a 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -1101,8 +1101,15 @@ where "manage Workdir tools require Backend Workspace API authority", )); } + let child_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_before_workdir_release( + workspace_client, + Arc::new(move || { + let child_registry = child_registry.clone(); + Box::pin(async move { child_registry.shutdown_internal().await }) + }), + ), ); } if feature_config.workspace_worker_discovery.enabled { @@ -1726,7 +1733,16 @@ async fn controller_loop( // 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"); @@ -2604,4 +2620,21 @@ mod tests { other => panic!("expected compact rejection error, got {other:?}"), } } + + #[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")); + } } diff --git a/crates/worker/src/feature/builtin/manage_workdir.rs b/crates/worker/src/feature/builtin/manage_workdir.rs index ab252835..24cbc0a9 100644 --- a/crates/worker/src/feature/builtin/manage_workdir.rs +++ b/crates/worker/src/feature/builtin/manage_workdir.rs @@ -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}; @@ -52,16 +54,43 @@ 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 Pin> + Send>> + 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, + before_workdir_release: Option, +} + +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) -> Self { - Self { client } + Self { + client, + before_workdir_release: None, + } + } + + pub(crate) fn with_before_workdir_release( + client: Arc, + before_workdir_release: BeforeWorkdirRelease, + ) -> Self { + Self { + client, + before_workdir_release: Some(before_workdir_release), + } } } @@ -81,7 +110,8 @@ 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_before_workdir_release(self.before_workdir_release.clone()); for (name, definition) in [ ( LIST_TOOL, @@ -142,9 +172,20 @@ impl FeatureModule for ManageWorkdirFeature { } } -#[derive(Clone, Debug)] +#[derive(Clone)] struct WorkspaceHttpWorkdirBackend { client: Arc, + before_workdir_release: Option, +} + +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. @@ -327,7 +368,18 @@ impl WorkdirSession for WorkspaceAttachedWorkdirSession { impl WorkspaceHttpWorkdirBackend { fn new(client: Arc) -> Self { - Self { client } + Self { + client, + before_workdir_release: None, + } + } + + fn with_before_workdir_release( + mut self, + before_workdir_release: Option, + ) -> Self { + self.before_workdir_release = before_workdir_release; + self } fn workspace_id(&self) -> Result<&str, ToolError> { @@ -510,6 +562,13 @@ impl Tool for WorkspaceHttpWorkdirTool { .attach(parse_input::(input_json)?), WorkdirOperation::Detach => { let _input = parse_input::(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 @@ -705,6 +764,7 @@ struct WorkdirDeleteInput { #[cfg(test)] mod tests { use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; use super::*; use crate::feature::{FeatureModule, FeatureRegistryBuilder}; @@ -1259,4 +1319,58 @@ mod tests { assert!(client.requests().is_empty()); assert!(parse_input::(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_before_workdir_release(Some(before_release)), + 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_before_workdir_release(Some(before_release)), + 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()); + } } diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index c270243f..3db26267 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -679,13 +679,6 @@ impl SpawnedWorkerRegistry { .unwrap_or_default() } - pub(crate) fn reclaim_internal_scope(&self, worker_name: &str) -> io::Result { - 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 { if !record.claim_scope_reclaim() { return Ok(false); @@ -705,6 +698,35 @@ impl SpawnedWorkerRegistry { result } + pub(crate) async fn close_internal_scope(&self, name: &str) -> io::Result { + let Some(record) = self.get_internal(name) else { + return Ok(false); + }; + 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 = self + .internal_records + .lock() + .expect("internal Worker registry lock poisoned") + .iter() + .map(|record| record.worker_name.clone()) + .collect::>(); + let mut first_error = None; + for name in names { + if let Err(error) = self.remove_internal(&name).await { + first_error.get_or_insert(error); + } + } + first_error.map_or(Ok(()), Err) + } + /// Stop one direct Internal SubWorker and discard its registry/scope state. /// /// The child actor must acknowledge its stop before the registry is removed. @@ -1236,6 +1258,24 @@ 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(®istry, 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 running_worker_is_stopped_before_removal() { let registry = registry(); diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index 3091642e..294983b0 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -532,13 +532,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!( From 7b1cf854f276ad766074d00db60cb5926e99df78 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 03:29:50 +0900 Subject: [PATCH 4/6] fix: serialize scoped command teardown --- crates/workdir/src/scope.rs | 201 ++++++++++++++++++++++++++++++++---- 1 file changed, 182 insertions(+), 19 deletions(-) diff --git a/crates/workdir/src/scope.rs b/crates/workdir/src/scope.rs index c48c1541..e730d9f0 100644 --- a/crates/workdir/src/scope.rs +++ b/crates/workdir/src/scope.rs @@ -10,9 +10,11 @@ use fs_operation::{ }; use tokio::sync::broadcast; +const MAX_SCOPED_COMMANDS: usize = 16; + use crate::{ CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, - CommandSnapshot, CommandStatus, Workdir, WorkdirError, WorkdirSession, + CommandSnapshot, CommandStatus, CommandStream, Workdir, WorkdirError, WorkdirSession, WorkdirSessionCapabilities, WorkdirSessionCapability, WorkdirSessionHandle, }; @@ -69,12 +71,15 @@ impl WorkdirToolBroker { validity: SessionValidity::root(), child_write_leases: Mutex::new(HashMap::new()), next_lease_id: AtomicU64::new(1), + close_lock: Arc::new(tokio::sync::Mutex::new(())), owned_commands: Arc::new(Mutex::new(HashSet::new())), pending_command_events: Arc::new(Mutex::new(HashMap::new())), starting_tool_calls: Arc::new(Mutex::new(HashSet::new())), forwarded_terminals: Arc::new(Mutex::new(HashSet::new())), command_events, closes_source: true, + #[cfg(test)] + command_start_gate: Mutex::new(None), }); Self { session: authority.clone(), @@ -181,6 +186,7 @@ impl WorkdirScopeLease { output.status, output.exit_code, output.next_cursor.unwrap_or(output.content.len()) as u64, + &output.content, ); self.broker .authority @@ -196,6 +202,7 @@ impl WorkdirScopeLease { CommandStatus::Cancelled, None, 0, + "", ); self.broker .authority @@ -289,6 +296,12 @@ struct ActiveWriteLease { rules: Vec, } +#[cfg(test)] +struct TestCommandStartGate { + entered: tokio::sync::Notify, + release: tokio::sync::Notify, +} + struct ScopedWorkdirSession { source: WorkdirSessionHandle, cwd: FsPath, @@ -297,12 +310,15 @@ struct ScopedWorkdirSession { validity: Arc, child_write_leases: Mutex>, next_lease_id: AtomicU64, + close_lock: Arc>, owned_commands: Arc>>, pending_command_events: Arc>>>, starting_tool_calls: Arc>>, forwarded_terminals: Arc>>, command_events: broadcast::Sender, closes_source: bool, + #[cfg(test)] + command_start_gate: Mutex>>, } impl std::fmt::Debug for ScopedWorkdirSession { @@ -416,19 +432,33 @@ impl ScopedWorkdirSession { status: CommandStatus, exit_code: Option, offset: u64, + fallback_output: &str, ) { - publish_owned_command_event( - &self.command_events, - &self.forwarded_terminals, - CommandEvent::Terminal { + let mut terminals = self + .forwarded_terminals + .lock() + .expect("forwarded terminal command mutex poisoned"); + if !terminals.insert(command_id.to_string()) { + return; + } + if !fallback_output.is_empty() { + let _ = self.command_events.send(CommandEvent::Output { command_id: command_id.to_string(), - status, - exit_code, - stdout_end_offset: offset, - stderr_end_offset: 0, + stream: CommandStream::Stdout, + start_offset: 0, + end_offset: fallback_output.len() as u64, + content: fallback_output.to_string(), observed_at_ms: unix_timestamp_ms(), - }, - ); + }); + } + let _ = self.command_events.send(CommandEvent::Terminal { + command_id: command_id.to_string(), + status, + exit_code, + stdout_end_offset: offset, + stderr_end_offset: 0, + observed_at_ms: unix_timestamp_ms(), + }); } fn ensure_parent_write_available(&self, path: &FsPath) -> Result<(), WorkdirError> { @@ -656,6 +686,7 @@ impl ScopedWorkdirSession { command_events.clone(), ) .map(|handle| Arc::new(Mutex::new(Some(handle)))); + let close_lock = Arc::new(tokio::sync::Mutex::new(())); let child = Arc::new(ScopedWorkdirSession { source: self.source.clone(), cwd: request.cwd, @@ -664,12 +695,15 @@ impl ScopedWorkdirSession { validity: validity.clone(), child_write_leases: Mutex::new(HashMap::new()), next_lease_id: AtomicU64::new(1), + close_lock: close_lock.clone(), owned_commands, pending_command_events, starting_tool_calls, forwarded_terminals, command_events, closes_source: false, + #[cfg(test)] + command_start_gate: Mutex::new(None), }); let broker = WorkdirToolBroker { session: child.clone(), @@ -681,7 +715,7 @@ impl ScopedWorkdirSession { capabilities, validity, cleanup_pending, - close_lock: Arc::new(tokio::sync::Mutex::new(())), + close_lock, }) } } @@ -749,7 +783,36 @@ impl WorkdirSession for ScopedWorkdirSession { &self, mut request: CommandRequest, ) -> Result { + let _admission_guard = self.close_lock.lock().await; + // Command is an explicit capability, not a typed path mutation. We + // intentionally keep an ancestor's Command capability available while + // a child holds a write scope; only typed Write/Edit operations use the + // best-effort overlapping-path guard below. self.ensure_command()?; + #[cfg(test)] + { + let gate = self + .command_start_gate + .lock() + .expect("command start gate mutex poisoned") + .clone(); + if let Some(gate) = gate { + gate.entered.notify_one(); + gate.release.notified().await; + } + } + if self.scope.is_some() + && self + .owned_commands + .lock() + .expect("scoped command set mutex poisoned") + .len() + >= MAX_SCOPED_COMMANDS + { + return Err(WorkdirError::Unavailable(format!( + "scoped command limit of {MAX_SCOPED_COMMANDS} is reached" + ))); + } let tool_call_id = request.tool_call_id.clone(); if self.scope.is_some() { request.cwd = Some(match request.cwd.as_ref() { @@ -831,6 +894,7 @@ impl WorkdirSession for ScopedWorkdirSession { output.status, output.exit_code, output.next_cursor.unwrap_or(output.content.len()) as u64, + &output.content, ); self.owned_commands .lock() @@ -1036,14 +1100,20 @@ fn publish_owned_command_event( forwarded_terminals: &Mutex>, event: CommandEvent, ) { - if let CommandEvent::Terminal { command_id, .. } = &event - && !forwarded_terminals - .lock() - .expect("forwarded terminal command mutex poisoned") - .insert(command_id.clone()) - { - return; + let command_id = command_event_id(&event); + let mut terminals = forwarded_terminals + .lock() + .expect("forwarded terminal command mutex poisoned"); + match &event { + CommandEvent::Terminal { .. } if !terminals.insert(command_id.to_string()) => return, + CommandEvent::Started { .. } | CommandEvent::Output { .. } + if terminals.contains(command_id) => + { + return; + } + _ => {} } + drop(terminals); let _ = sender.send(event); } @@ -1592,6 +1662,99 @@ mod tests { child.close().await.unwrap(); } + #[tokio::test] + async fn scoped_command_ceiling_rejects_the_seventeenth_live_command() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("work")).unwrap(); + let parent = session(root.path()); + let child = parent + .scope(request("work", WorkdirToolScopePermission::Write)) + .await + .unwrap(); + for index in 0..MAX_SCOPED_COMMANDS { + child + .start_command(CommandRequest { + command: "sleep 30".into(), + timeout_secs: 60, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: Some(format!("command-{index}")), + }) + .await + .unwrap(); + } + + let error = child + .start_command(CommandRequest { + command: "sleep 30".into(), + timeout_secs: 60, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: Some("command-over-limit".into()), + }) + .await + .unwrap_err(); + assert!(matches!(error, WorkdirError::Unavailable(message) if message.contains("limit"))); + child.close().await.unwrap(); + } + + #[tokio::test] + async fn close_serializes_with_inflight_command_admission() { + let root = TempDir::new().unwrap(); + fs::create_dir_all(root.path().join("work")).unwrap(); + let parent = session(root.path()); + let child = Arc::new( + parent + .scope(request("work", WorkdirToolScopePermission::Write)) + .await + .unwrap(), + ); + let gate = Arc::new(TestCommandStartGate { + entered: tokio::sync::Notify::new(), + release: tokio::sync::Notify::new(), + }); + *child.broker.authority.command_start_gate.lock().unwrap() = Some(gate.clone()); + let entered = gate.entered.notified(); + let command_child = child.clone(); + let command = tokio::spawn(async move { + command_child + .start_command(CommandRequest { + command: "sleep 30".into(), + timeout_secs: 60, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: Some("racing-command".into()), + }) + .await + }); + entered.await; + let close_child = child.clone(); + let mut close = tokio::spawn(async move { close_child.close().await }); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(50), &mut close) + .await + .is_err(), + "close must wait for command admission to commit or fail" + ); + + gate.release.notify_one(); + command.await.unwrap().unwrap(); + close.await.unwrap().unwrap(); + assert!(!child.is_active()); + assert!( + child + .broker + .authority + .owned_commands + .lock() + .unwrap() + .is_empty() + ); + } + #[tokio::test] async fn closing_scope_cancels_and_terminalizes_owned_commands() { let root = TempDir::new().unwrap(); From ca5fddf89b2f38e54ca8096ada6689404d058897 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 03:48:54 +0900 Subject: [PATCH 5/6] fix: fence recursive SubWorker shutdown --- crates/worker/src/controller.rs | 8 +- .../src/feature/builtin/manage_workdir.rs | 65 +++++++- crates/worker/src/spawn/registry.rs | 157 ++++++++++++++++-- crates/worker/src/spawn/tool.rs | 23 ++- 4 files changed, 217 insertions(+), 36 deletions(-) diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index a73a175a..a56c7c59 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -1101,14 +1101,16 @@ where "manage Workdir tools require Backend Workspace API authority", )); } - let child_registry = spawned_registry.clone(); + let shutdown_registry = spawned_registry.clone(); + let reopen_registry = spawned_registry.clone(); feature_registry.add_module( - crate::feature::builtin::manage_workdir::ManageWorkdirFeature::with_before_workdir_release( + crate::feature::builtin::manage_workdir::ManageWorkdirFeature::with_child_lifecycle( workspace_client, Arc::new(move || { - let child_registry = child_registry.clone(); + let child_registry = shutdown_registry.clone(); Box::pin(async move { child_registry.shutdown_internal().await }) }), + Arc::new(move || reopen_registry.reopen_internal()), ), ); } diff --git a/crates/worker/src/feature/builtin/manage_workdir.rs b/crates/worker/src/feature/builtin/manage_workdir.rs index 24cbc0a9..eb950bbc 100644 --- a/crates/worker/src/feature/builtin/manage_workdir.rs +++ b/crates/worker/src/feature/builtin/manage_workdir.rs @@ -56,6 +56,7 @@ const ATTACH_DESCRIPTION: &str = "Attach this Worker to one existing Workdir. Th 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 Pin> + Send>> + Send + Sync>; +pub(crate) type AfterWorkdirAttach = Arc; 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."; @@ -63,6 +64,7 @@ const DELETE_DESCRIPTION: &str = "Request removal of one persistent Workdir by i pub struct ManageWorkdirFeature { client: Arc, before_workdir_release: Option, + after_workdir_attach: Option, } impl std::fmt::Debug for ManageWorkdirFeature { @@ -80,16 +82,19 @@ impl ManageWorkdirFeature { Self { client, before_workdir_release: None, + after_workdir_attach: None, } } - pub(crate) fn with_before_workdir_release( + pub(crate) fn with_child_lifecycle( client: Arc, 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), } } } @@ -110,8 +115,10 @@ impl FeatureModule for ManageWorkdirFeature { } fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> { - let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone()) - .with_before_workdir_release(self.before_workdir_release.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, @@ -176,6 +183,7 @@ impl FeatureModule for ManageWorkdirFeature { struct WorkspaceHttpWorkdirBackend { client: Arc, before_workdir_release: Option, + after_workdir_attach: Option, } impl std::fmt::Debug for WorkspaceHttpWorkdirBackend { @@ -371,14 +379,17 @@ impl WorkspaceHttpWorkdirBackend { Self { client, before_workdir_release: None, + after_workdir_attach: None, } } - fn with_before_workdir_release( + fn with_child_lifecycle( mut self, before_workdir_release: Option, + after_workdir_attach: Option, ) -> Self { self.before_workdir_release = before_workdir_release; + self.after_workdir_attach = after_workdir_attach; self } @@ -557,9 +568,17 @@ impl Tool for WorkspaceHttpWorkdirTool { parse_input::(input_json)?, ctx.call_id.to_string(), ), - WorkdirOperation::Attach => self - .backend - .attach(parse_input::(input_json)?), + WorkdirOperation::Attach => { + let result = self + .backend + .attach(parse_input::(input_json)?); + if result.is_ok() + && let Some(after_attach) = &self.backend.after_workdir_attach + { + after_attach(); + } + result + } WorkdirOperation::Detach => { let _input = parse_input::(input_json)?; if let Some(before_release) = &self.backend.before_workdir_release { @@ -1338,7 +1357,7 @@ mod tests { }); let tool = WorkspaceHttpWorkdirTool { backend: WorkspaceHttpWorkdirBackend::new(client.clone()) - .with_before_workdir_release(Some(before_release)), + .with_child_lifecycle(Some(before_release), None), operation: WorkdirOperation::Detach, }; @@ -1361,7 +1380,7 @@ mod tests { Arc::new(|| Box::pin(async { Err(std::io::Error::other("child cleanup failed")) })); let tool = WorkspaceHttpWorkdirTool { backend: WorkspaceHttpWorkdirBackend::new(client.clone()) - .with_before_workdir_release(Some(before_release)), + .with_child_lifecycle(Some(before_release), None), operation: WorkdirOperation::Detach, }; @@ -1373,4 +1392,32 @@ mod tests { 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); + } } diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 3db26267..cba26327 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -72,6 +72,7 @@ pub(crate) struct InternalSpawnedWorkerRecord { #[cfg(test)] pub installed_tools: Arc<[String]>, pub session: InternalWorkerSessionHandle, + pub child_registry: Arc, change_tracker: Option, started_at: Instant, stop_lock: Arc>, @@ -89,6 +90,7 @@ impl InternalSpawnedWorkerRecord { workdir_tool_scope: WorkdirScopeLease, #[cfg(test)] installed_tools: Vec, session: InternalWorkerSessionHandle, + child_registry: Arc, change_tracker: Option, ) -> Self { Self { @@ -98,6 +100,7 @@ impl InternalSpawnedWorkerRecord { #[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,39 @@ pub(crate) struct InternalSpawnReservation { } impl InternalSpawnReservation { - pub(crate) fn commit(mut self, record: InternalSpawnedWorkerRecord) -> io::Result<()> { + pub(crate) fn commit( + mut self, + record: InternalSpawnedWorkerRecord, + ) -> Result<(), (io::Error, InternalSpawnedWorkerRecord)> { if record.worker_name != self.worker_name { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "internal SubWorker reservation name does not match record name", + return Err(( + io::Error::new( + io::ErrorKind::InvalidInput, + "internal SubWorker reservation name does not match record name", + ), + record, )); } - self.registry - .internal_records - .lock() - .map_err(|_| io::Error::other("internal spawned-worker registry lock poisoned"))? - .push(record.clone()); + let mut records = match self.registry.internal_records.lock() { + Ok(records) => records, + Err(_) => { + return Err(( + io::Error::other("internal spawned-worker registry lock poisoned"), + record, + )); + } + }; + if self.registry.internal_shutting_down.load(Ordering::Acquire) { + return Err(( + io::Error::new( + io::ErrorKind::Interrupted, + "internal SubWorker registry is shutting down", + ), + record, + )); + } + records.push(record.clone()); + drop(records); self.registry.start_protocol_forwarding(record); self.committed = true; Ok(()) @@ -267,6 +291,7 @@ pub struct SpawnedWorkerRegistry { internal_records: std::sync::Mutex>, service_records: std::sync::Mutex>, internal_names: std::sync::Mutex>, + internal_shutting_down: AtomicBool, parent_scope: Option, parent_protocol: Mutex, String)>>, } @@ -283,6 +308,7 @@ 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), parent_scope: None, parent_protocol: Mutex::new(None), }) @@ -294,6 +320,7 @@ 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), parent_scope: None, parent_protocol: Mutex::new(None), }) @@ -304,6 +331,7 @@ 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), parent_scope: Some(parent_scope), parent_protocol: Mutex::new(None), }) @@ -383,6 +411,7 @@ 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), parent_scope, parent_protocol: Mutex::new(None), }), @@ -394,6 +423,12 @@ impl SpawnedWorkerRegistry { self: &Arc, worker_name: String, ) -> io::Result { + 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() @@ -702,6 +737,7 @@ impl SpawnedWorkerRegistry { let Some(record) = self.get_internal(name) else { return Ok(false); }; + Box::pin(record.child_registry.shutdown_internal()).await?; record .workdir_tool_scope .close() @@ -711,13 +747,17 @@ impl SpawnedWorkerRegistry { } pub(crate) async fn shutdown_internal(&self) -> io::Result<()> { - let names = self - .internal_records - .lock() - .expect("internal Worker registry lock poisoned") - .iter() - .map(|record| record.worker_name.clone()) - .collect::>(); + 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::>() + }; let mut first_error = None; for name in names { if let Err(error) = self.remove_internal(&name).await { @@ -727,6 +767,10 @@ impl SpawnedWorkerRegistry { first_error.map_or(Ok(()), Err) } + pub(crate) fn reopen_internal(&self) { + self.internal_shutting_down.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. @@ -753,6 +797,7 @@ 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() @@ -1021,6 +1066,7 @@ mod tests { delegation, Vec::new(), session, + registry(), None, ), sender, @@ -1276,6 +1322,85 @@ mod tests { 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(std::sync::Barrier::new(2)); + let commit_barrier = barrier.clone(); + let commit = tokio::task::spawn_blocking(move || { + commit_barrier.wait(); + reservation.commit(record) + }); + let shutdown_registry = registry.clone(); + let shutdown = tokio::spawn(async move { + barrier.wait(); + shutdown_registry.shutdown_internal().await + }); + + let commit = commit.await.unwrap(); + shutdown.await.unwrap().unwrap(); + if let Err((_error, record)) = commit { + record.session.stop().await.unwrap(); + record.child_registry.shutdown_internal().await.unwrap(); + record.workdir_tool_scope.close().await.unwrap(); + } + + 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; + + registry.shutdown_internal().await.unwrap(); + let (error, record) = reservation.commit(record).unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::Interrupted); + record.session.stop().await.unwrap(); + record.child_registry.shutdown_internal().await.unwrap(); + record.workdir_tool_scope.close().await.unwrap(); + } + + #[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(®istry, 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(); diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index 294983b0..547d939d 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -600,15 +600,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 @@ -623,10 +627,13 @@ impl Tool for SubWorkerSpawnTool { #[cfg(test)] installed_tools, session.clone(), + child_registry, child_change_tracker, ); - if let Err(error) = name_reservation.commit(record) { + if let Err((error, record)) = name_reservation.commit(record) { let _ = session.stop().await; + let _ = record.child_registry.shutdown_internal().await; + let _ = record.workdir_tool_scope.close().await; return Err(ToolError::ExecutionFailed(format!( "register Internal Worker session: {error}" ))); From 052d60bd7d2f0f9cbe3da19ced380e3ce6af82d9 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 04:10:23 +0900 Subject: [PATCH 6/6] fix: fence late command events and spawn rollback --- crates/workdir/src/scope.rs | 34 ++++- crates/worker/src/spawn/registry.rs | 188 +++++++++++++++++++++------- crates/worker/src/spawn/tool.rs | 5 +- 3 files changed, 173 insertions(+), 54 deletions(-) diff --git a/crates/workdir/src/scope.rs b/crates/workdir/src/scope.rs index e730d9f0..131ae6cf 100644 --- a/crates/workdir/src/scope.rs +++ b/crates/workdir/src/scope.rs @@ -75,6 +75,7 @@ impl WorkdirToolBroker { owned_commands: Arc::new(Mutex::new(HashSet::new())), pending_command_events: Arc::new(Mutex::new(HashMap::new())), starting_tool_calls: Arc::new(Mutex::new(HashSet::new())), + forwarded_starts: Arc::new(Mutex::new(HashSet::new())), forwarded_terminals: Arc::new(Mutex::new(HashSet::new())), command_events, closes_source: true, @@ -314,6 +315,7 @@ struct ScopedWorkdirSession { owned_commands: Arc>>, pending_command_events: Arc>>>, starting_tool_calls: Arc>>, + forwarded_starts: Arc>>, forwarded_terminals: Arc>>, command_events: broadcast::Sender, closes_source: bool, @@ -675,6 +677,7 @@ impl ScopedWorkdirSession { let owned_commands = Arc::new(Mutex::new(HashSet::new())); let pending_command_events = Arc::new(Mutex::new(HashMap::new())); let starting_tool_calls = Arc::new(Mutex::new(HashSet::new())); + let forwarded_starts = Arc::new(Mutex::new(HashSet::new())); let forwarded_terminals = Arc::new(Mutex::new(HashSet::new())); let (command_events, _) = broadcast::channel(64); let event_forwarder = forward_owned_command_events( @@ -682,6 +685,7 @@ impl ScopedWorkdirSession { owned_commands.clone(), pending_command_events.clone(), starting_tool_calls.clone(), + forwarded_starts.clone(), forwarded_terminals.clone(), command_events.clone(), ) @@ -699,6 +703,7 @@ impl ScopedWorkdirSession { owned_commands, pending_command_events, starting_tool_calls, + forwarded_starts, forwarded_terminals, command_events, closes_source: false, @@ -862,6 +867,7 @@ impl WorkdirSession for ScopedWorkdirSession { { publish_owned_command_event( &self.command_events, + &self.forwarded_starts, &self.forwarded_terminals, CommandEvent::Started { command_id: handle.0.clone(), @@ -871,7 +877,12 @@ impl WorkdirSession for ScopedWorkdirSession { ); } for event in pending { - publish_owned_command_event(&self.command_events, &self.forwarded_terminals, event); + publish_owned_command_event( + &self.command_events, + &self.forwarded_starts, + &self.forwarded_terminals, + event, + ); } Ok(handle) } @@ -1033,6 +1044,7 @@ fn forward_owned_command_events( owned_commands: Arc>>, pending_command_events: Arc>>>, starting_tool_calls: Arc>>, + forwarded_starts: Arc>>, forwarded_terminals: Arc>>, sender: broadcast::Sender, ) -> Option> { @@ -1082,7 +1094,7 @@ fn forward_owned_command_events( } drop(pending); drop(owned); - publish_owned_command_event(&sender, &forwarded_terminals, event); + publish_owned_command_event(&sender, &forwarded_starts, &forwarded_terminals, event); } })) } @@ -1097,6 +1109,7 @@ fn command_event_id(event: &CommandEvent) -> &str { fn publish_owned_command_event( sender: &broadcast::Sender, + forwarded_starts: &Mutex>, forwarded_terminals: &Mutex>, event: CommandEvent, ) { @@ -1106,11 +1119,16 @@ fn publish_owned_command_event( .expect("forwarded terminal command mutex poisoned"); match &event { CommandEvent::Terminal { .. } if !terminals.insert(command_id.to_string()) => return, - CommandEvent::Started { .. } | CommandEvent::Output { .. } - if terminals.contains(command_id) => + CommandEvent::Started { .. } if terminals.contains(command_id) => return, + CommandEvent::Started { .. } + if !forwarded_starts + .lock() + .expect("forwarded command start mutex poisoned") + .insert(command_id.to_string()) => { return; } + CommandEvent::Output { .. } if terminals.contains(command_id) => return, _ => {} } drop(terminals); @@ -1657,8 +1675,16 @@ mod tests { } assert_eq!(kinds.first(), Some(&"started")); assert_eq!(kinds.last(), Some(&"terminal")); + assert_eq!(kinds.iter().filter(|kind| **kind == "started").count(), 1); assert!(kinds.contains(&"output")); assert!(streamed.contains("fast-output")); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), events.recv()) + .await + .is_err(), + "no provider event may follow the terminal event" + ); + assert!(child.command_snapshot().is_empty()); child.close().await.unwrap(); } diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index cba26327..fd946e54 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -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,7 +23,7 @@ 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::WorkdirScopeLease; @@ -238,39 +238,56 @@ pub(crate) struct InternalSpawnReservation { } impl InternalSpawnReservation { - pub(crate) fn commit( - mut self, - record: InternalSpawnedWorkerRecord, - ) -> Result<(), (io::Error, InternalSpawnedWorkerRecord)> { - if record.worker_name != self.worker_name { - return Err(( - io::Error::new( - io::ErrorKind::InvalidInput, - "internal SubWorker reservation name does not match record name", - ), - record, - )); - } - let mut records = match self.registry.internal_records.lock() { - Ok(records) => records, - Err(_) => { - return Err(( - io::Error::other("internal spawned-worker registry lock poisoned"), - record, - )); + 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 self.registry.internal_shutting_down.load(Ordering::Acquire) { - return Err(( - io::Error::new( - io::ErrorKind::Interrupted, - "internal SubWorker registry is shutting down", - ), - record, - )); + 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("; ") + ))); } - records.push(record.clone()); - drop(records); self.registry.start_protocol_forwarding(record); self.committed = true; Ok(()) @@ -284,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(); } } @@ -292,6 +313,9 @@ pub struct SpawnedWorkerRegistry { service_records: std::sync::Mutex>, internal_names: std::sync::Mutex>, internal_shutting_down: AtomicBool, + pending_internal_spawns: AtomicUsize, + pending_internal_notify: Notify, + internal_spawn_cleanup_failed: AtomicBool, parent_scope: Option, parent_protocol: Mutex, String)>>, } @@ -309,6 +333,9 @@ impl SpawnedWorkerRegistry { 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), }) @@ -321,6 +348,9 @@ impl SpawnedWorkerRegistry { 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), }) @@ -332,6 +362,9 @@ impl SpawnedWorkerRegistry { 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), }) @@ -412,6 +445,9 @@ impl SpawnedWorkerRegistry { 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), }), @@ -423,6 +459,10 @@ impl SpawnedWorkerRegistry { self: &Arc, worker_name: String, ) -> io::Result { + 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, @@ -439,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, @@ -758,17 +800,31 @@ impl SpawnedWorkerRegistry { .map(|record| record.worker_name.clone()) .collect::>() }; + 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. @@ -1342,24 +1398,22 @@ mod tests { let (record, _events) = record("racing-child", InternalWorkerVisibility::ParentClient).await; let scope = record.workdir_tool_scope.clone(); - let barrier = Arc::new(std::sync::Barrier::new(2)); + let barrier = Arc::new(tokio::sync::Barrier::new(2)); let commit_barrier = barrier.clone(); - let commit = tokio::task::spawn_blocking(move || { - commit_barrier.wait(); - reservation.commit(record) + 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(); + barrier.wait().await; shutdown_registry.shutdown_internal().await }); let commit = commit.await.unwrap(); shutdown.await.unwrap().unwrap(); - if let Err((_error, record)) = commit { - record.session.stop().await.unwrap(); - record.child_registry.shutdown_internal().await.unwrap(); - record.workdir_tool_scope.close().await.unwrap(); + if let Err(error) = commit { + assert_eq!(error.kind(), io::ErrorKind::Interrupted); } assert!(registry.list_internal().is_empty()); @@ -1375,12 +1429,54 @@ mod tests { let (record, _events) = record("racing-child", InternalWorkerVisibility::ParentClient).await; - registry.shutdown_internal().await.unwrap(); - let (error, record) = reservation.commit(record).unwrap_err(); + 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); - record.session.stop().await.unwrap(); - record.child_registry.shutdown_internal().await.unwrap(); - record.workdir_tool_scope.close().await.unwrap(); + 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] diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index 547d939d..6370653f 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -630,10 +630,7 @@ impl Tool for SubWorkerSpawnTool { child_registry, child_change_tracker, ); - if let Err((error, record)) = name_reservation.commit(record) { - let _ = session.stop().await; - let _ = record.child_registry.shutdown_internal().await; - let _ = record.workdir_tool_scope.close().await; + if let Err(error) = name_reservation.commit(record).await { return Err(ToolError::ExecutionFailed(format!( "register Internal Worker session: {error}" )));