diff --git a/crates/tools/src/bash.rs b/crates/tools/src/bash.rs index 34e979b56..a213a30eb 100644 --- a/crates/tools/src/bash.rs +++ b/crates/tools/src/bash.rs @@ -6,7 +6,9 @@ use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use async_trait::async_trait; use schemars::JsonSchema; use serde::Deserialize; -use workdir::{CommandHandle, CommandOutputRequest, CommandRequest, WorkdirSessionHandle}; +use workdir::{ + CommandHandle, CommandOutputRequest, CommandRequest, WorkdirSessionHandle, WorkdirSessionRouter, +}; const DEFAULT_TIMEOUT_SECS: u64 = 120; const MAX_TIMEOUT_SECS: u64 = 600; @@ -14,13 +16,16 @@ const INLINE_BYTE_BUDGET: usize = 12 * 1024; #[derive(Debug, Deserialize, JsonSchema)] struct BashParams { + /// Worker-local alias of the Workdir attachment to use. + #[serde(default)] + target_workdir: Option, command: String, #[serde(default)] timeout: Option, } pub(crate) struct BashTool { - session: WorkdirSessionHandle, + router: Arc, output_dir: PathBuf, state: Arc>, } @@ -29,6 +34,7 @@ pub(crate) struct BashTool { struct ActiveCommand { call_id: String, execution_nonce: u64, + session: WorkdirSessionHandle, handle: CommandHandle, } @@ -93,6 +99,12 @@ impl Tool for BashTool { ) -> Result { let params: BashParams = serde_json::from_str(input_json) .map_err(|error| ToolError::InvalidArgument(format!("invalid Bash input: {error}")))?; + let selected = crate::routing::resolve_session( + &self.router, + params.target_workdir.as_deref(), + workdir::WorkdirSessionCapability::Command, + )?; + let session = selected.session; let timeout_secs = params .timeout .unwrap_or(DEFAULT_TIMEOUT_SECS) @@ -106,14 +118,13 @@ impl Tool for BashTool { state.next_execution_nonce }; let mut guard = CommandGuard { - session: self.session.clone(), + session: session.clone(), state: self.state.clone(), execution_id: execution_id.clone(), execution_nonce, handle: None, }; - let handle = self - .session + let handle = session .start_command(CommandRequest { command: params.command, timeout_secs, @@ -131,6 +142,7 @@ impl Tool for BashTool { ActiveCommand { call_id: call_id.clone(), execution_nonce, + session: session.clone(), handle: handle.clone(), }, ); @@ -139,13 +151,12 @@ impl Tool for BashTool { }; guard.handle = Some(handle.clone()); if cancel_after_start { - self.session + session .cancel_command(handle.clone()) .await .map_err(crate::ToolsError::from)?; } - let output = self - .session + let output = session .command_output(CommandOutputRequest { handle, cursor: 0, @@ -222,12 +233,13 @@ impl Tool for BashTool { .active .values() .filter(|active| active.call_id == call_id) - .map(|active| active.handle.clone()) + .cloned() .collect::>() }; - for handle in handles { - self.session - .cancel_command(handle) + for active in handles { + active + .session + .cancel_command(active.handle) .await .map_err(crate::ToolsError::from)?; } @@ -242,14 +254,12 @@ impl Tool for BashTool { let handle = { let mut state = self.state.lock().unwrap(); state.cancellation_requested.insert(execution_id.clone()); - state - .active - .get(&execution_id) - .map(|active| active.handle.clone()) + state.active.get(&execution_id).cloned() }; - if let Some(handle) = handle { - self.session - .cancel_command(handle) + if let Some(active) = handle { + active + .session + .cancel_command(active.handle) .await .map_err(crate::ToolsError::from)?; } @@ -268,13 +278,20 @@ fn truncate_for_summary(command: &str) -> String { } pub fn bash_tool(session: WorkdirSessionHandle, output_dir: PathBuf) -> ToolDefinition { + routed_bash_tool(crate::routing::singleton_router(session), output_dir) +} + +pub(crate) fn routed_bash_tool( + router: Arc, + output_dir: PathBuf, +) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(BashParams); let meta = ToolMeta::new("Bash") - .description("Execute a shell command in the bound Workdir. Process start, bounded inline output, full-output spill, timeout and cancellation are owned by the WorkdirSession provider. This is not a sandbox.") + .description("Execute a shell command in one selected Workdir. Process start, bounded inline output, full-output spill, timeout and cancellation are owned by that WorkdirSession provider. This is not a sandbox.") .input_schema(serde_json::to_value(schema).expect("Bash schema serialization")); let tool: Arc = Arc::new(BashTool { - session: session.clone(), + router: router.clone(), output_dir: output_dir.clone(), state: Arc::new(Mutex::new(BashExecutionState::default())), }); diff --git a/crates/tools/src/edit.rs b/crates/tools/src/edit.rs index a303628fb..f8fd9c2e5 100644 --- a/crates/tools/src/edit.rs +++ b/crates/tools/src/edit.rs @@ -9,15 +9,18 @@ use serde::Deserialize; use crate::error::ToolsError; use crate::tracker::Tracker; -use workdir::{EditRequest, WorkdirPath, WorkdirSessionHandle}; +use workdir::{EditRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter}; -const DESCRIPTION: &str = "Replace a substring in an existing file. By default \ +const DESCRIPTION: &str = "Replace a substring in an existing file in the selected Workdir attachment. By default \ `old_string` must be unique in the file; set `replace_all: true` to replace \ -every occurrence. The file must have been read first (via the Read tool) in \ -this session. Paths are relative to the bound Workdir."; +every occurrence. The file must have been read first through the same attachment. \ +When exactly one Workdir is attached, target_workdir may be omitted."; #[derive(Debug, Deserialize, schemars::JsonSchema)] pub(crate) struct EditParams { + /// Worker-local alias of the Workdir attachment to use. + #[serde(default)] + pub target_workdir: Option, /// Logical path relative to the bound Workdir root. pub file_path: String, /// String to replace. Must be unique in the file unless `replace_all` is true. @@ -30,7 +33,7 @@ pub(crate) struct EditParams { } pub(crate) struct EditTool { - session: WorkdirSessionHandle, + router: Arc, tracker: Tracker, } @@ -44,6 +47,12 @@ impl Tool for EditTool { let params: EditParams = serde_json::from_str(input_json) .map_err(|e| ToolError::InvalidArgument(format!("invalid Edit input: {e}")))?; + let selected = crate::routing::resolve_session( + &self.router, + params.target_workdir.as_deref(), + workdir::WorkdirSessionCapability::Edit, + )?; + let tracker = self.tracker.scoped(selected.alias.as_str()); let path = WorkdirPath::new(¶ms.file_path).map_err(ToolsError::from)?; tracing::debug!(path = %path, replace_all = params.replace_all, "Edit"); @@ -59,9 +68,9 @@ impl Tool for EditTool { } let mutation_key = PathBuf::from(path.as_str()); - let _mutation_permit = self.tracker.acquire_mutation(&mutation_key, &ctx).await; - let expected_hash = self.tracker.expected_workdir_hash(&path)?; - let result = self + let _mutation_permit = tracker.acquire_mutation(&mutation_key, &ctx).await; + let expected_hash = tracker.expected_workdir_hash(&path)?; + let result = selected .session .edit(EditRequest { path: path.clone(), @@ -73,7 +82,7 @@ impl Tool for EditTool { .await .map_err(ToolsError::from)?; let replacements = result.replacements; - self.tracker.record_workdir_edit( + tracker.record_workdir_edit( &path, result.content_hash, replacements, @@ -122,8 +131,15 @@ fn make_preview(text: &str, needle: &str) -> String { out } -/// Factory for the `Edit` tool. +/// Factory for the `Edit` tool bound to one compatibility session. pub fn edit_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinition { + routed_edit_tool(crate::routing::singleton_router(session), tracker) +} + +pub(crate) fn routed_edit_tool( + router: Arc, + tracker: Tracker, +) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(EditParams); let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({})); @@ -131,7 +147,7 @@ pub fn edit_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinit .description(DESCRIPTION) .input_schema(schema_value); let tool: Arc = Arc::new(EditTool { - session: session.clone(), + router: router.clone(), tracker: tracker.clone(), }); (meta, tool) diff --git a/crates/tools/src/glob.rs b/crates/tools/src/glob.rs index f2e8616e7..0c2914dfd 100644 --- a/crates/tools/src/glob.rs +++ b/crates/tools/src/glob.rs @@ -4,7 +4,7 @@ use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use async_trait::async_trait; use schemars::JsonSchema; use serde::Deserialize; -use workdir::{GlobRequest, WorkdirPath, WorkdirSessionHandle}; +use workdir::{GlobRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter}; use crate::ToolsError; @@ -12,6 +12,9 @@ const RESULT_LIMIT: usize = 1000; #[derive(Debug, Deserialize, JsonSchema)] struct GlobParams { + /// Worker-local alias of the Workdir attachment to use. + #[serde(default)] + target_workdir: Option, /// Glob pattern, for example `**/*.rs` or `src/**/test_*.py`. pattern: String, /// Logical Workdir-relative directory. Defaults to the Workdir root. @@ -20,7 +23,7 @@ struct GlobParams { } struct GlobTool { - session: WorkdirSessionHandle, + router: Arc, } #[async_trait] @@ -32,13 +35,18 @@ impl Tool for GlobTool { ) -> Result { let params: GlobParams = serde_json::from_str(input_json) .map_err(|error| ToolError::InvalidArgument(format!("invalid Glob input: {error}")))?; + let selected = crate::routing::resolve_session( + &self.router, + params.target_workdir.as_deref(), + workdir::WorkdirSessionCapability::Glob, + )?; let path = match params.path { Some(path) => WorkdirPath::new(&path).map_err(ToolsError::from)?, None => WorkdirPath::root(), }; let pattern = params.pattern; tracing::debug!(%pattern, %path, "Glob"); - let result = self + let result = selected .session .glob(GlobRequest { pattern: pattern.clone(), @@ -75,13 +83,17 @@ impl Tool for GlobTool { } pub fn glob_tool(session: WorkdirSessionHandle) -> ToolDefinition { + routed_glob_tool(crate::routing::singleton_router(session)) +} + +pub(crate) fn routed_glob_tool(router: Arc) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(GlobParams); let meta = ToolMeta::new("Glob") - .description("Find files matching a glob pattern inside the bound Workdir. Results are sorted and capped at 1000 entries. Paths are Workdir-relative.") + .description("Find files matching a glob pattern inside the selected Workdir. Results are sorted and capped at 1000 entries. Paths are Workdir-relative.") .input_schema(serde_json::to_value(schema).expect("Glob schema serialization")); let tool: Arc = Arc::new(GlobTool { - session: session.clone(), + router: router.clone(), }); (meta, tool) }) diff --git a/crates/tools/src/grep.rs b/crates/tools/src/grep.rs index 4774e6ffa..190b8590e 100644 --- a/crates/tools/src/grep.rs +++ b/crates/tools/src/grep.rs @@ -4,7 +4,9 @@ use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use async_trait::async_trait; use schemars::JsonSchema; use serde::Deserialize; -use workdir::{GrepOutputMode, GrepRequest, WorkdirPath, WorkdirSessionHandle}; +use workdir::{ + GrepOutputMode, GrepRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter, +}; use crate::ToolsError; @@ -21,6 +23,9 @@ enum OutputMode { #[derive(Debug, Deserialize, JsonSchema)] struct GrepParams { + /// Worker-local alias of the Workdir attachment to use. + #[serde(default)] + target_workdir: Option, pattern: String, /// Workdir-relative path, or an absolute path covered by readable scope. Defaults to the Workdir root. #[serde(default)] @@ -48,7 +53,7 @@ struct GrepParams { } struct GrepTool { - session: WorkdirSessionHandle, + router: Arc, } #[async_trait] @@ -60,6 +65,11 @@ impl Tool for GrepTool { ) -> Result { let params: GrepParams = serde_json::from_str(input_json) .map_err(|error| ToolError::InvalidArgument(format!("invalid Grep input: {error}")))?; + let selected = crate::routing::resolve_session( + &self.router, + params.target_workdir.as_deref(), + workdir::WorkdirSessionCapability::Grep, + )?; let path = match params.path { Some(path) => WorkdirPath::new_scoped(&path).map_err(ToolsError::from)?, None => WorkdirPath::root(), @@ -74,7 +84,7 @@ impl Tool for GrepTool { .map(|context| (context, context)) .unwrap_or((params.before.unwrap_or(0), params.after.unwrap_or(0))); let head_limit = params.head_limit.unwrap_or(DEFAULT_HEAD_LIMIT); - let result = self + let result = selected .session .grep(GrepRequest { pattern: params.pattern, @@ -126,13 +136,17 @@ impl Tool for GrepTool { } pub fn grep_tool(session: WorkdirSessionHandle) -> ToolDefinition { + routed_grep_tool(crate::routing::singleton_router(session)) +} + +pub(crate) fn routed_grep_tool(router: Arc) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(GrepParams); let meta = ToolMeta::new("Grep") - .description("Search a Workdir file or directory with a regex. Content results group lines by file; `>` marks matching lines and unmarked lines are context. Directory traversal executes inside the WorkdirSession provider. Results are bounded and Workdir-relative.") + .description("Search a file or directory in the selected Workdir with a regex. Content results group lines by file; `>` marks matching lines and unmarked lines are context. Directory traversal executes inside the WorkdirSession provider. Results are bounded and Workdir-relative.") .input_schema(serde_json::to_value(schema).expect("Grep schema serialization")); let tool: Arc = Arc::new(GrepTool { - session: session.clone(), + router: router.clone(), }); (meta, tool) }) diff --git a/crates/tools/src/lib.rs b/crates/tools/src/lib.rs index 7b838dadd..0aaabe8d6 100644 --- a/crates/tools/src/lib.rs +++ b/crates/tools/src/lib.rs @@ -17,6 +17,7 @@ mod edit; mod glob; mod grep; mod read; +mod routing; mod view_image; mod web; mod write; @@ -28,11 +29,11 @@ pub use glob::glob_tool; pub use grep::grep_tool; pub use read::read_tool; pub use tracker::{ChangeStat, Tracker}; -pub use view_image::view_image_tool; +pub use view_image::{routed_view_image_tool, view_image_tool}; pub use web::{web_fetch_tool, web_search_tool}; pub use write::write_tool; -/// Build the local filesystem/command tool surface implemented by a WorkdirSession. +/// Build the filesystem/command tool surface for a single Workdir session. /// Profile/manifest policy may narrow this set further in the Engine. pub fn core_builtin_tools( session: workdir::WorkdirSessionHandle, @@ -42,24 +43,75 @@ pub fn core_builtin_tools( use workdir::WorkdirSessionCapability; let capabilities = session.capabilities(); + let router = routing::singleton_router(session); let mut tools = Vec::with_capacity(6); if capabilities.supports(WorkdirSessionCapability::Read) { - tools.push(read_tool(session.clone(), tracker.clone())); + tools.push(read::routed_read_tool(router.clone(), tracker.clone())); } if capabilities.supports(WorkdirSessionCapability::Write) { - tools.push(write_tool(session.clone(), tracker.clone())); + tools.push(write::routed_write_tool(router.clone(), tracker.clone())); } if capabilities.supports(WorkdirSessionCapability::Edit) { - tools.push(edit_tool(session.clone(), tracker)); + tools.push(edit::routed_edit_tool(router.clone(), tracker)); } if capabilities.supports(WorkdirSessionCapability::Glob) { - tools.push(glob_tool(session.clone())); + tools.push(glob::routed_glob_tool(router.clone())); } if capabilities.supports(WorkdirSessionCapability::Grep) { - tools.push(grep_tool(session.clone())); + tools.push(grep::routed_grep_tool(router.clone())); } if capabilities.supports(WorkdirSessionCapability::Command) { - tools.push(bash_tool(session, bash_output_dir)); + tools.push(bash::routed_bash_tool(router, bash_output_dir)); + } + tools +} + +/// Build the stable routed tool surface for a Worker's live attachment set. +/// +/// The definitions remain installed while the set is empty or changes at runtime; +/// each invocation resolves `target_workdir` against the router and then lets the +/// selected provider enforce its capabilities. +pub fn routed_builtin_tools( + router: std::sync::Arc, + tracker: Tracker, + bash_output_dir: std::path::PathBuf, +) -> Vec { + vec![ + read::routed_read_tool(router.clone(), tracker.clone()), + write::routed_write_tool(router.clone(), tracker.clone()), + edit::routed_edit_tool(router.clone(), tracker), + glob::routed_glob_tool(router.clone()), + grep::routed_grep_tool(router.clone()), + bash::routed_bash_tool(router, bash_output_dir), + ] +} + +pub fn routed_builtin_tools_for_current_capabilities( + router: std::sync::Arc, + tracker: Tracker, + bash_output_dir: std::path::PathBuf, +) -> Vec { + use workdir::WorkdirSessionCapability; + + let capabilities = router.capabilities(); + let mut tools = Vec::with_capacity(6); + if capabilities.supports(WorkdirSessionCapability::Read) { + tools.push(read::routed_read_tool(router.clone(), tracker.clone())); + } + if capabilities.supports(WorkdirSessionCapability::Write) { + tools.push(write::routed_write_tool(router.clone(), tracker.clone())); + } + if capabilities.supports(WorkdirSessionCapability::Edit) { + tools.push(edit::routed_edit_tool(router.clone(), tracker)); + } + if capabilities.supports(WorkdirSessionCapability::Glob) { + tools.push(glob::routed_glob_tool(router.clone())); + } + if capabilities.supports(WorkdirSessionCapability::Grep) { + tools.push(grep::routed_grep_tool(router.clone())); + } + if capabilities.supports(WorkdirSessionCapability::Command) { + tools.push(bash::routed_bash_tool(router, bash_output_dir)); } tools } @@ -75,6 +127,29 @@ pub fn read_only_builtin_tools( core_builtin_tools(session, Tracker::new(), std::path::PathBuf::new()) } +pub fn read_only_routed_builtin_tools( + source: std::sync::Arc, +) -> Vec { + let router = std::sync::Arc::new(workdir::WorkdirSessionRouter::new()); + for alias in source.aliases() { + let Ok(selected) = source.resolve(Some(alias.as_str())) else { + continue; + }; + router + .attach( + alias, + std::sync::Arc::new(workdir::ReadOnlyWorkdirSession::new(selected.session)), + ) + .expect("read-only projection preserves unique aliases"); + } + let tracker = Tracker::new(); + vec![ + read::routed_read_tool(router.clone(), tracker), + glob::routed_glob_tool(router.clone()), + grep::routed_grep_tool(router), + ] +} + pub fn web_builtin_tools( web_config: Option, ) -> Vec { diff --git a/crates/tools/src/read.rs b/crates/tools/src/read.rs index 668f4a909..2769d36be 100644 --- a/crates/tools/src/read.rs +++ b/crates/tools/src/read.rs @@ -8,18 +8,22 @@ use serde::Deserialize; use crate::error::ToolsError; use crate::tracker::Tracker; -use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle}; +use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter}; -const DESCRIPTION: &str = "Read a text file from the local filesystem. \ +const DESCRIPTION: &str = "Read a text file from a Workdir attachment selected by its Worker-local alias. \ Supports offset/limit for large files. Returns line-numbered output (1-based). \ Directories cannot be read. The file must be read before Write or Edit can \ -modify it. Paths are Workdir-relative unless an absolute path is explicitly readable."; +modify it. Paths are Workdir-relative unless an absolute path is explicitly readable. \ +When exactly one Workdir is attached, target_workdir may be omitted."; const DEFAULT_LIMIT: usize = 2000; const PROVIDER_BYTE_LIMIT: usize = 256 * 1024; #[derive(Debug, Deserialize, schemars::JsonSchema)] pub(crate) struct ReadParams { + /// Worker-local alias of the Workdir attachment to use. + #[serde(default)] + pub target_workdir: Option, /// Workdir-relative path, or an absolute path covered by readable scope. pub file_path: String, /// 0-based line offset from the start. Defaults to 0. @@ -31,7 +35,7 @@ pub(crate) struct ReadParams { } pub(crate) struct ReadTool { - session: WorkdirSessionHandle, + router: Arc, tracker: Tracker, } @@ -44,13 +48,19 @@ impl Tool for ReadTool { ) -> Result { let params: ReadParams = serde_json::from_str(input_json) .map_err(|e| ToolError::InvalidArgument(format!("invalid Read input: {e}")))?; + let selected = crate::routing::resolve_session( + &self.router, + params.target_workdir.as_deref(), + workdir::WorkdirSessionCapability::Read, + )?; + let tracker = self.tracker.scoped(selected.alias.as_str()); let offset = params.offset.unwrap_or(0); let limit = params.limit.unwrap_or(DEFAULT_LIMIT).max(1); let path = WorkdirPath::new_scoped(¶ms.file_path).map_err(ToolsError::from)?; tracing::debug!(path = %path, offset, limit, "Read"); - let result = self + let result = selected .session .read(ReadRequest { path: path.clone(), @@ -60,7 +70,7 @@ impl Tool for ReadTool { }) .await .map_err(ToolsError::from)?; - self.tracker.record_workdir_hash(&path, result.content_hash); + tracker.record_workdir_hash(&path, result.content_hash); let text = String::from_utf8_lossy(&result.bytes).into_owned(); let rendered = render_provider_read( @@ -144,8 +154,15 @@ fn render_numbered(text: &str, offset: usize, limit: usize) -> Rendered { } } -/// Factory for the `Read` tool. +/// Factory for the `Read` tool bound to one compatibility session. pub fn read_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinition { + routed_read_tool(crate::routing::singleton_router(session), tracker) +} + +pub(crate) fn routed_read_tool( + router: Arc, + tracker: Tracker, +) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(ReadParams); let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({})); @@ -153,7 +170,7 @@ pub fn read_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinit .description(DESCRIPTION) .input_schema(schema_value); let tool: Arc = Arc::new(ReadTool { - session: session.clone(), + router: router.clone(), tracker: tracker.clone(), }); (meta, tool) @@ -199,6 +216,7 @@ mod tests { // History recorded assert!( tracker + .scoped("workdir") .expected_workdir_hash(&WorkdirPath::new("a.txt").unwrap()) .is_ok() ); diff --git a/crates/tools/src/routing.rs b/crates/tools/src/routing.rs new file mode 100644 index 000000000..e586da4f8 --- /dev/null +++ b/crates/tools/src/routing.rs @@ -0,0 +1,32 @@ +use std::sync::Arc; + +use agen::tool::ToolError; +use workdir::{ + ResolvedWorkdirSession, WorkdirAttachmentAlias, WorkdirSessionCapability, WorkdirSessionHandle, + WorkdirSessionRouter, +}; + +pub(crate) fn singleton_router(session: WorkdirSessionHandle) -> Arc { + let router = Arc::new(WorkdirSessionRouter::new()); + router + .attach( + WorkdirAttachmentAlias::new("workdir").expect("static attachment alias is valid"), + session, + ) + .expect("fresh singleton router accepts its session"); + router +} + +pub(crate) fn resolve_session( + router: &WorkdirSessionRouter, + target_workdir: Option<&str>, + required: WorkdirSessionCapability, +) -> Result { + let selected = router + .resolve(target_workdir) + .map_err(|error| ToolError::InvalidArgument(error.to_string()))?; + if !selected.session.capabilities().supports(required) { + return Err(crate::ToolsError::from(workdir::WorkdirError::Unsupported(required)).into()); + } + Ok(selected) +} diff --git a/crates/tools/src/tracker.rs b/crates/tools/src/tracker.rs index a544aa67c..14be471c2 100644 --- a/crates/tools/src/tracker.rs +++ b/crates/tools/src/tracker.rs @@ -65,9 +65,15 @@ fn hash_bytes(bytes: &[u8]) -> ContentHash { hasher.finalize().into() } +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct TrackedPath { + namespace: Option>, + path: PathBuf, +} + #[derive(Debug, Clone, Default)] struct FileMutationCoordinator { - locks: Arc>>>>, + locks: Arc>>>>, } pub(crate) struct FileMutationPermit { @@ -75,8 +81,7 @@ pub(crate) struct FileMutationPermit { } impl FileMutationCoordinator { - async fn acquire(&self, path: &Path) -> FileMutationPermit { - let key = file_mutation_key(path); + async fn acquire(&self, key: TrackedPath) -> FileMutationPermit { let lock = { let mut locks = self.locks.lock().await; locks @@ -127,12 +132,12 @@ pub struct ChangeStat { #[derive(Debug, Default)] struct Inner { - /// Hash of each file's last observed contents, keyed by canonical path. - hashes: HashMap, + /// Hash of each file's last observed contents, keyed by attachment and canonical path. + hashes: HashMap, /// Line count paired with observations that included the file content. - line_counts: HashMap, + line_counts: HashMap, /// LRU list of touched files. Front = most recently touched. - recency: VecDeque, + recency: VecDeque, /// Successful Write/Edit mutations attributed to this session's tools. change_stat: ChangeStat, } @@ -146,6 +151,7 @@ struct Inner { pub struct Tracker { inner: Arc>, mutations: FileMutationCoordinator, + namespace: Option>, } impl Tracker { @@ -154,6 +160,24 @@ impl Tracker { Self::default() } + /// Return a view whose read-before-mutation state is isolated to one + /// Worker-local attachment alias while aggregate recency/change statistics + /// remain shared with the Worker tracker. + pub fn scoped(&self, namespace: impl Into) -> Self { + Self { + inner: self.inner.clone(), + mutations: self.mutations.clone(), + namespace: Some(Arc::from(namespace.into())), + } + } + + fn key(&self, path: PathBuf) -> TrackedPath { + TrackedPath { + namespace: self.namespace.clone(), + path, + } + } + /// Acquire the per-target-file mutation guard shared by `Write` and `Edit`. /// /// The guard is keyed by canonical target path where possible so equivalent @@ -170,7 +194,9 @@ impl Tracker { call_index = ctx.call_index, "acquire file mutation guard" ); - self.mutations.acquire(path).await + self.mutations + .acquire(self.key(file_mutation_key(path))) + .await } /// Record that `path` has been observed with the given content bytes. @@ -182,7 +208,7 @@ impl Tracker { /// Also bumps `path` to the front of the recency list. If the list /// grows past [`RECENCY_CAPACITY`], the oldest entry is evicted. pub fn record(&self, path: &Path, bytes: &[u8]) { - let key = canonicalize_or_owned(path); + let key = self.key(canonicalize_or_owned(path)); let hash = hash_bytes(bytes); let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); inner.hashes.insert(key.clone(), hash); @@ -198,7 +224,7 @@ impl Tracker { } pub fn record_workdir_content(&self, path: &workdir::WorkdirPath, content: &[u8]) { - let key = PathBuf::from(path.as_str()); + let key = self.key(PathBuf::from(path.as_str())); let hash = hash_bytes(content); let line_count = String::from_utf8_lossy(content).lines().count(); let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); @@ -212,16 +238,17 @@ impl Tracker { } pub fn observed_workdir_line_count(&self, path: &workdir::WorkdirPath) -> Option { + let key = self.key(PathBuf::from(path.as_str())); self.inner .lock() .unwrap_or_else(|e| e.into_inner()) .line_counts - .get(Path::new(path.as_str())) + .get(&key) .copied() } pub fn record_workdir_hash(&self, path: &workdir::WorkdirPath, hash: workdir::ContentHash) { - let key = PathBuf::from(path.as_str()); + let key = self.key(PathBuf::from(path.as_str())); let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); inner.hashes.insert(key.clone(), hash); inner.recency.retain(|candidate| candidate != &key); @@ -253,7 +280,7 @@ impl Tracker { ) { let added = added_lines_per_replacement.saturating_mul(replacements); let deleted = deleted_lines_per_replacement.saturating_mul(replacements); - let key = PathBuf::from(path.as_str()); + let key = self.key(PathBuf::from(path.as_str())); let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); inner.change_stat.added = inner.change_stat.added.saturating_add(added as u64); inner.change_stat.deleted = inner.change_stat.deleted.saturating_add(deleted as u64); @@ -279,14 +306,14 @@ impl Tracker { &self, path: &workdir::WorkdirPath, ) -> Result { - let key = PathBuf::from(path.as_str()); + let key = self.key(PathBuf::from(path.as_str())); self.inner .lock() .unwrap_or_else(|error| error.into_inner()) .hashes .get(&key) .copied() - .ok_or_else(|| ToolsError::NotRead(key)) + .ok_or_else(|| ToolsError::NotRead(key.path.clone())) } /// Verify that `path` was previously recorded and its current bytes @@ -296,7 +323,7 @@ impl Tracker { /// - If the current content hashes differ from the recorded value, /// returns [`ToolsError::ExternallyModified`]. pub fn verify(&self, path: &Path, current_bytes: &[u8]) -> Result<(), ToolsError> { - let key = canonicalize_or_owned(path); + let key = self.key(canonicalize_or_owned(path)); let guard = self.inner.lock().unwrap_or_else(|e| e.into_inner()); let recorded = guard .hashes @@ -316,13 +343,18 @@ impl Tracker { /// can pass them as default references to the compaction worker. pub fn recent_files(&self, n: usize) -> Vec { let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - inner.recency.iter().take(n).cloned().collect() + inner + .recency + .iter() + .take(n) + .map(|key| key.path.clone()) + .collect() } /// Returns true if `path` has a history entry. Test-only. #[cfg(test)] pub(crate) fn has(&self, path: &Path) -> bool { - let key = canonicalize_or_owned(path); + let key = self.key(canonicalize_or_owned(path)); self.inner .lock() .unwrap_or_else(|e| e.into_inner()) @@ -420,6 +452,22 @@ mod tests { } } + #[test] + fn attachment_namespaces_do_not_share_read_state() { + let tracker = Tracker::new(); + let left = tracker.scoped("left"); + let right = tracker.scoped("right"); + let path = workdir::WorkdirPath::new("same.txt").unwrap(); + + left.record_workdir_content(&path, b"left"); + + assert!(left.expected_workdir_hash(&path).is_ok()); + assert!(matches!( + right.expected_workdir_hash(&path), + Err(ToolsError::NotRead(_)) + )); + } + #[test] fn clone_shares_state() { let dir = TempDir::new().unwrap(); diff --git a/crates/tools/src/view_image.rs b/crates/tools/src/view_image.rs index 5c9adb520..1fc86ca44 100644 --- a/crates/tools/src/view_image.rs +++ b/crates/tools/src/view_image.rs @@ -7,24 +7,28 @@ use agen::tool::{ }; use async_trait::async_trait; use serde::Deserialize; -use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle}; +use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter}; use crate::error::ToolsError; /// Maximum image body accepted for one model request. pub const MAX_IMAGE_BYTES: usize = 10 * 1024 * 1024; -const DESCRIPTION: &str = "Attach an image from the bound Workdir to the next model request. \ -The path must be logical and Workdir-relative. Supported formats: PNG, JPEG, GIF, and WebP."; +const DESCRIPTION: &str = "Attach an image from the selected Workdir attachment to the next model request. \ +The path must be logical and Workdir-relative. Supported formats: PNG, JPEG, GIF, and WebP. \ +When exactly one Workdir is attached, target_workdir may be omitted."; #[derive(Debug, Deserialize, schemars::JsonSchema)] struct ViewImageParams { + /// Worker-local alias of the Workdir attachment to use. + #[serde(default)] + target_workdir: Option, /// Logical path relative to the bound Workdir root. path: String, } struct ViewImageTool { - session: WorkdirSessionHandle, + router: Arc, } #[async_trait] @@ -37,8 +41,13 @@ impl Tool for ViewImageTool { let input: ViewImageParams = serde_json::from_str(input_json).map_err(|error| { ToolError::InvalidArgument(format!("invalid ViewImage input: {error}")) })?; + let selected = crate::routing::resolve_session( + &self.router, + input.target_workdir.as_deref(), + workdir::WorkdirSessionCapability::Read, + )?; let path = WorkdirPath::new(&input.path).map_err(ToolsError::from)?; - let result = self + let result = selected .session .read(ReadRequest { path: path.clone(), @@ -75,6 +84,10 @@ impl Tool for ViewImageTool { } pub fn view_image_tool(session: WorkdirSessionHandle) -> ToolDefinition { + routed_view_image_tool(crate::routing::singleton_router(session)) +} + +pub fn routed_view_image_tool(router: Arc) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(ViewImageParams); let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({})); @@ -82,7 +95,7 @@ pub fn view_image_tool(session: WorkdirSessionHandle) -> ToolDefinition { .description(DESCRIPTION) .input_schema(schema_value); let tool: Arc = Arc::new(ViewImageTool { - session: session.clone(), + router: router.clone(), }); (meta, tool) }) diff --git a/crates/tools/src/write.rs b/crates/tools/src/write.rs index 4eb567556..5b8e33da6 100644 --- a/crates/tools/src/write.rs +++ b/crates/tools/src/write.rs @@ -9,15 +9,20 @@ use serde::Deserialize; use crate::error::ToolsError; use crate::tracker::Tracker; -use workdir::{StatRequest, WorkdirError, WorkdirPath, WorkdirSessionHandle, WriteRequest}; +use workdir::{ + StatRequest, WorkdirError, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter, + WriteRequest, +}; -const DESCRIPTION: &str = "Create a new file or overwrite an existing one with \ -the given content. Missing parent directories within scope are created \ -automatically. Existing files must have been read first (via the Read tool) \ -in this session. Paths are relative to the bound Workdir."; +const DESCRIPTION: &str = "Create a new file or overwrite an existing one in the selected Workdir attachment. \ +Missing parent directories within scope are created automatically. Existing files must have been read first \ +through the same attachment. When exactly one Workdir is attached, target_workdir may be omitted."; #[derive(Debug, Deserialize, schemars::JsonSchema)] pub(crate) struct WriteParams { + /// Worker-local alias of the Workdir attachment to use. + #[serde(default)] + pub target_workdir: Option, /// Logical path relative to the bound Workdir root. pub file_path: String, /// Full content to write. Overwrites any existing content. @@ -25,7 +30,7 @@ pub(crate) struct WriteParams { } pub(crate) struct WriteTool { - session: WorkdirSessionHandle, + router: Arc, tracker: Tracker, } @@ -39,19 +44,29 @@ impl Tool for WriteTool { let params: WriteParams = serde_json::from_str(input_json) .map_err(|e| ToolError::InvalidArgument(format!("invalid Write input: {e}")))?; + let selected = crate::routing::resolve_session( + &self.router, + params.target_workdir.as_deref(), + workdir::WorkdirSessionCapability::Write, + )?; + let tracker = self.tracker.scoped(selected.alias.as_str()); let path = WorkdirPath::new(¶ms.file_path).map_err(ToolsError::from)?; tracing::debug!(path = %path, bytes = params.content.len(), "Write"); let mutation_key = PathBuf::from(path.as_str()); - let _mutation_permit = self.tracker.acquire_mutation(&mutation_key, &ctx).await; - let expected_hash = match self.session.stat(StatRequest { path: path.clone() }).await { - Ok(_) => Some(self.tracker.expected_workdir_hash(&path)?), + let _mutation_permit = tracker.acquire_mutation(&mutation_key, &ctx).await; + let expected_hash = match selected + .session + .stat(StatRequest { path: path.clone() }) + .await + { + Ok(_) => Some(tracker.expected_workdir_hash(&path)?), Err(WorkdirError::NotFound(_)) => None, Err(error) => return Err(ToolsError::from(error).into()), }; - let old_line_count = self.tracker.observed_workdir_line_count(&path).unwrap_or(0); - let outcome = self + let old_line_count = tracker.observed_workdir_line_count(&path).unwrap_or(0); + let outcome = selected .session .write(WriteRequest { path: path.clone(), @@ -61,10 +76,8 @@ impl Tool for WriteTool { .await .map_err(ToolsError::from)?; - self.tracker - .record_change(params.content.lines().count(), old_line_count); - self.tracker - .record_workdir_content(&path, params.content.as_bytes()); + tracker.record_change(params.content.lines().count(), old_line_count); + tracker.record_workdir_content(&path, params.content.as_bytes()); let summary = format!( "{} {} ({} bytes)", @@ -84,8 +97,15 @@ impl Tool for WriteTool { } } -/// Factory for the `Write` tool. +/// Factory for the `Write` tool bound to one compatibility session. pub fn write_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinition { + routed_write_tool(crate::routing::singleton_router(session), tracker) +} + +pub(crate) fn routed_write_tool( + router: Arc, + tracker: Tracker, +) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(WriteParams); let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({})); @@ -93,7 +113,7 @@ pub fn write_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefini .description(DESCRIPTION) .input_schema(schema_value); let tool: Arc = Arc::new(WriteTool { - session: session.clone(), + router: router.clone(), tracker: tracker.clone(), }); (meta, tool) diff --git a/crates/tools/tests/integration.rs b/crates/tools/tests/integration.rs index 02b8c986c..f79b89b13 100644 --- a/crates/tools/tests/integration.rs +++ b/crates/tools/tests/integration.rs @@ -14,8 +14,12 @@ use agen::tool::{ use manifest::{Permission, Scope, ScopeConfig, ScopeRule}; use serde_json::json; use tempfile::TempDir; -use tools::{Tracker, core_builtin_tools, view_image_tool}; -use workdir::{LocalWorkdirSession, WorkdirSessionHandle}; +use tools::{ + Tracker, core_builtin_tools, routed_builtin_tools, routed_view_image_tool, view_image_tool, +}; +use workdir::{ + LocalWorkdirSession, WorkdirAttachmentAlias, WorkdirSessionHandle, WorkdirSessionRouter, +}; fn scope_with_spill(workspace: &Path, spill: &Path) -> Scope { let base = Scope::writable(workspace).unwrap(); @@ -66,6 +70,30 @@ fn setup() -> (TempDir, TempDir, Registry) { (dir, spill, reg) } +fn setup_routed() -> ( + TempDir, + TempDir, + TempDir, + Arc, + Registry, +) { + let left = TempDir::new().unwrap(); + let right = TempDir::new().unwrap(); + let spill = TempDir::new().unwrap(); + let router = Arc::new(WorkdirSessionRouter::new()); + for (alias, dir) in [("left", &left), ("right", &right)] { + let scope = scope_with_spill(dir.path(), spill.path()); + let session: WorkdirSessionHandle = + Arc::new(LocalWorkdirSession::new(scope, dir.path().to_path_buf())); + router + .attach(WorkdirAttachmentAlias::new(alias).unwrap(), session) + .unwrap(); + } + let definitions = + routed_builtin_tools(router.clone(), Tracker::new(), spill.path().to_path_buf()); + (left, right, spill, router, Registry::new(definitions)) +} + async fn call(tool: &Arc, input: serde_json::Value) -> agen::tool::ToolOutput { tool.execute(&input.to_string(), Default::default()) .await @@ -104,6 +132,213 @@ fn meta_has_description_and_schema() { } } +#[test] +fn routed_tool_schemas_share_optional_target_workdir() { + let (_left, _right, _spill, router, mut reg) = setup_routed(); + reg.entries.push(routed_view_image_tool(router)()); + for (meta, _) in ®.entries { + let target = &meta.input_schema["properties"]["target_workdir"]; + assert!(target.is_object(), "{} lacks target_workdir", meta.name); + assert_eq!(target["type"], json!(["string", "null"])); + } +} + +#[tokio::test] +async fn routed_tools_require_alias_only_when_attachment_set_is_ambiguous() { + let (_left, _right, _spill, router, reg) = setup_routed(); + let error = call_err(®.get("Read"), json!({ "file_path": "same.txt" })).await; + let message = error.to_string(); + assert!(message.contains("target_workdir_required"), "{message}"); + assert!(message.contains("left"), "{message}"); + assert!(message.contains("right"), "{message}"); + + let unknown = call_err( + ®.get("Glob"), + json!({ "target_workdir": "missing", "pattern": "*" }), + ) + .await + .to_string(); + assert!(unknown.contains("unknown_target_workdir"), "{unknown}"); + assert!(unknown.contains("available_aliases"), "{unknown}"); + + router.close_all().await.unwrap(); +} + +#[tokio::test] +async fn empty_routed_surface_returns_no_attachment_error() { + let router = Arc::new(WorkdirSessionRouter::new()); + let reg = Registry::new(routed_builtin_tools( + router, + Tracker::new(), + Path::new("unused").to_path_buf(), + )); + let message = call_err(®.get("Read"), json!({ "file_path": "file.txt" })) + .await + .to_string(); + assert!(message.contains("no_workdir_attached"), "{message}"); +} + +#[tokio::test] +async fn target_alias_routes_each_tool_and_isolates_read_before_edit() { + let (left, right, _spill, router, mut reg) = setup_routed(); + std::fs::write(left.path().join("same.txt"), "left value\n").unwrap(); + std::fs::write(right.path().join("same.txt"), "right value\n").unwrap(); + std::fs::write(left.path().join("only-left.log"), "LEFT-NEEDLE\n").unwrap(); + std::fs::write(right.path().join("only-right.log"), "RIGHT-NEEDLE\n").unwrap(); + + let read = reg.get("Read"); + let left_read = call( + &read, + json!({ "target_workdir": "left", "file_path": "same.txt" }), + ) + .await; + assert!(left_read.content.unwrap().contains("left value")); + + let edit = reg.get("Edit"); + let isolated = call_err( + &edit, + json!({ + "target_workdir": "right", + "file_path": "same.txt", + "old_string": "right", + "new_string": "edited" + }), + ) + .await; + assert!(isolated.to_string().contains("has not been read")); + call( + &read, + json!({ "target_workdir": "right", "file_path": "same.txt" }), + ) + .await; + call( + &edit, + json!({ + "target_workdir": "right", + "file_path": "same.txt", + "old_string": "right", + "new_string": "edited" + }), + ) + .await; + assert_eq!( + std::fs::read_to_string(right.path().join("same.txt")).unwrap(), + "edited value\n" + ); + assert_eq!( + std::fs::read_to_string(left.path().join("same.txt")).unwrap(), + "left value\n" + ); + + let glob = call( + ®.get("Glob"), + json!({ "target_workdir": "left", "pattern": "*.log" }), + ) + .await; + assert!(glob.content.unwrap().contains("only-left.log")); + let grep = call( + ®.get("Grep"), + json!({ + "target_workdir": "right", + "pattern": "RIGHT-NEEDLE", + "output_mode": "content" + }), + ) + .await; + assert!(grep.content.unwrap().contains("only-right.log")); + + let bash = call( + ®.get("Bash"), + json!({ "target_workdir": "right", "command": "pwd" }), + ) + .await; + assert_eq!( + std::fs::canonicalize(bash.content.unwrap().trim()).unwrap(), + std::fs::canonicalize(right.path()).unwrap() + ); + + let png = b"\x89PNG\r\n\x1a\nrouted"; + std::fs::write(left.path().join("image.png"), png).unwrap(); + reg.entries.push(routed_view_image_tool(router.clone())()); + let image = call( + ®.get("ViewImage"), + json!({ "target_workdir": "left", "path": "image.png" }), + ) + .await; + let agen::tool::Attachment::Image(image) = &image.attachments[0]; + assert_eq!(image.data(), png); + + router.close_all().await.unwrap(); +} + +#[tokio::test] +async fn selected_attachment_capability_is_checked_without_fallback() { + let left = TempDir::new().unwrap(); + let right = TempDir::new().unwrap(); + let spill = TempDir::new().unwrap(); + let router = Arc::new(WorkdirSessionRouter::new()); + router + .attach( + WorkdirAttachmentAlias::new("writable").unwrap(), + Arc::new(LocalWorkdirSession::new( + scope_with_spill(left.path(), spill.path()), + left.path().to_path_buf(), + )), + ) + .unwrap(); + router + .attach( + WorkdirAttachmentAlias::new("readonly").unwrap(), + Arc::new(LocalWorkdirSession::materialized_bound( + workdir::Workdir::new("readonly"), + right.path().to_path_buf(), + right.path().to_path_buf(), + manifest::SharedScope::new(scope_with_spill(right.path(), spill.path())), + workdir::WorkdirSessionCapabilities::READ_ONLY, + )), + ) + .unwrap(); + let reg = Registry::new(routed_builtin_tools( + router.clone(), + Tracker::new(), + spill.path().to_path_buf(), + )); + + let message = call_err( + ®.get("Write"), + json!({ + "target_workdir": "readonly", + "file_path": "created.txt", + "content": "must not be routed elsewhere" + }), + ) + .await + .to_string(); + assert!(message.contains("does not support Write"), "{message}"); + assert!(!left.path().join("created.txt").exists()); + assert!(!right.path().join("created.txt").exists()); + + router.close_all().await.unwrap(); +} + +#[tokio::test] +async fn routed_tools_observe_detach_after_registration() { + let (_left, _right, _spill, router, reg) = setup_routed(); + let alias = WorkdirAttachmentAlias::new("left").unwrap(); + router.detach(&alias).await.unwrap(); + + let message = call_err( + ®.get("Glob"), + json!({ "target_workdir": "left", "pattern": "*" }), + ) + .await + .to_string(); + assert!(message.contains("unknown_target_workdir"), "{message}"); + assert!(message.contains("right"), "{message}"); + + router.close_all().await.unwrap(); +} + #[tokio::test] async fn view_image_reads_scoped_bytes_into_durable_tool_detail() { let dir = TempDir::new().unwrap(); diff --git a/crates/workdir/src/lib.rs b/crates/workdir/src/lib.rs index 683f699a1..ba1070541 100644 --- a/crates/workdir/src/lib.rs +++ b/crates/workdir/src/lib.rs @@ -28,11 +28,14 @@ pub use local::{ LocalWorkdirSession, SymlinkInfo, WorkdirSessionResource, direct_symlink, first_symlink, }; pub use operation::*; -pub use router::{RoutedWorkdirSession, WorkdirAttachmentAlias, WorkdirSessionRouter}; +pub use router::{ + ResolvedWorkdirSession, RoutedWorkdirSession, WorkdirAttachmentAlias, WorkdirRouteError, + WorkdirRouteErrorCode, WorkdirSessionRouter, +}; pub use scope::{ ReadOnlyWorkdirSession, WorkdirScopeAuthorizationRequest, WorkdirScopeLease, - WorkdirScopeOverlapRequest, WorkdirToolBroker, WorkdirToolScope, WorkdirToolScopePermission, - WorkdirToolScopeRule, + WorkdirScopeLeaseSet, WorkdirScopeOverlapRequest, WorkdirToolBroker, WorkdirToolBrokerRouter, + WorkdirToolScope, WorkdirToolScopePermission, WorkdirToolScopeRule, }; /// Persistent, opaque identity of one materialized Workdir. @@ -124,6 +127,12 @@ impl WorkdirSessionCapabilities { bits: Self::READ | Self::GLOB | Self::GREP, }; + pub const fn union(self, other: Self) -> Self { + Self { + bits: self.bits | other.bits, + } + } + pub const fn supports(self, capability: WorkdirSessionCapability) -> bool { let bit = match capability { WorkdirSessionCapability::Read => Self::READ, diff --git a/crates/workdir/src/router.rs b/crates/workdir/src/router.rs index 8f5e3c3f1..d8e2c9a20 100644 --- a/crates/workdir/src/router.rs +++ b/crates/workdir/src/router.rs @@ -6,7 +6,7 @@ //! routing keys. use std::collections::{BTreeMap, HashSet}; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, RwLock}; use async_trait::async_trait; @@ -66,9 +66,56 @@ impl std::fmt::Display for WorkdirAttachmentAlias { } } +/// Stable machine-readable reason why a tool target could not be resolved. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum WorkdirRouteErrorCode { + NoWorkdirAttached, + TargetWorkdirRequired, + UnknownTargetWorkdir, +} + +/// Provider-neutral attachment routing failure returned before an operation starts. +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct WorkdirRouteError { + pub code: WorkdirRouteErrorCode, + #[serde(skip_serializing_if = "Option::is_none")] + pub target_workdir: Option, + pub available_aliases: Vec, +} + +impl std::fmt::Display for WorkdirRouteError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let encoded = serde_json::to_string(self).map_err(|_| std::fmt::Error)?; + formatter.write_str(&encoded) + } +} + +impl std::error::Error for WorkdirRouteError {} + +/// One attachment selected for a single tool invocation. +#[derive(Clone)] +pub struct ResolvedWorkdirSession { + pub alias: WorkdirAttachmentAlias, + pub generation: u64, + pub session: WorkdirSessionHandle, +} + +impl std::fmt::Debug for ResolvedWorkdirSession { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ResolvedWorkdirSession") + .field("alias", &self.alias) + .field("generation", &self.generation) + .field("workdir", self.session.workdir()) + .finish_non_exhaustive() + } +} + #[derive(Debug)] struct RoutedSessionState { session: WorkdirSessionHandle, + generation: u64, detached: AtomicBool, active_operations: AtomicUsize, active_commands: Mutex>, @@ -299,6 +346,7 @@ impl WorkdirSession for RoutedWorkdirSession { #[derive(Debug, Default)] pub struct WorkdirSessionRouter { sessions: RwLock>>, + next_generation: AtomicU64, } impl WorkdirSessionRouter { @@ -315,6 +363,16 @@ impl WorkdirSessionRouter { .collect() } + pub fn capabilities(&self) -> WorkdirSessionCapabilities { + self.sessions + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .values() + .fold(WorkdirSessionCapabilities::EMPTY, |capabilities, state| { + capabilities.union(state.session.capabilities()) + }) + } + pub fn len(&self) -> usize { self.sessions .read() @@ -326,6 +384,61 @@ impl WorkdirSessionRouter { self.len() == 0 } + /// Resolve one model-visible target without exposing provider identity. + /// + /// Omission is accepted only for the unambiguous single-attachment case. + /// Every error includes the current alias snapshot so callers can recover + /// without consulting Workdir ids, display names, or transport metadata. + pub fn resolve( + &self, + target_workdir: Option<&str>, + ) -> Result { + let sessions = self + .sessions + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let available_aliases = sessions + .keys() + .map(|alias| alias.as_str().to_owned()) + .collect::>(); + let selected = match target_workdir { + Some(target) => sessions + .iter() + .find(|(alias, _)| alias.as_str() == target) + .map(|(alias, state)| (alias.clone(), state.clone())) + .ok_or_else(|| WorkdirRouteError { + code: WorkdirRouteErrorCode::UnknownTargetWorkdir, + target_workdir: Some(target.to_owned()), + available_aliases: available_aliases.clone(), + })?, + None if sessions.is_empty() => { + return Err(WorkdirRouteError { + code: WorkdirRouteErrorCode::NoWorkdirAttached, + target_workdir: None, + available_aliases, + }); + } + None if sessions.len() > 1 => { + return Err(WorkdirRouteError { + code: WorkdirRouteErrorCode::TargetWorkdirRequired, + target_workdir: None, + available_aliases, + }); + } + None => { + let (alias, state) = sessions + .first_key_value() + .expect("single attachment exists after cardinality checks"); + (alias.clone(), state.clone()) + } + }; + Ok(ResolvedWorkdirSession { + alias: selected.0, + generation: selected.1.generation, + session: Arc::new(RoutedWorkdirSession { state: selected.1 }), + }) + } + pub fn attach( &self, alias: WorkdirAttachmentAlias, @@ -341,6 +454,7 @@ impl WorkdirSessionRouter { } let state = Arc::new(RoutedSessionState { session, + generation: self.next_generation.fetch_add(1, Ordering::AcqRel), detached: AtomicBool::new(false), active_operations: AtomicUsize::new(0), active_commands: Mutex::new(HashSet::new()), @@ -556,9 +670,40 @@ mod tests { ["left", "right"] ); assert!(router.only_session().is_none()); + let missing = router.resolve(None).unwrap_err(); + assert_eq!(missing.code, WorkdirRouteErrorCode::TargetWorkdirRequired); + assert_eq!(missing.available_aliases, ["left", "right"]); + assert_eq!( + router + .resolve(Some("right")) + .unwrap() + .session + .workdir() + .id() + .as_str(), + "wd-right" + ); + let unknown = router.resolve(Some("display-name")).unwrap_err(); + assert_eq!(unknown.code, WorkdirRouteErrorCode::UnknownTargetWorkdir); + assert_eq!(unknown.target_workdir.as_deref(), Some("display-name")); + assert_eq!(unknown.available_aliases, ["left", "right"]); router.close_all().await.unwrap(); } + #[test] + fn empty_router_returns_structured_no_attachment_error() { + let error = WorkdirSessionRouter::new().resolve(None).unwrap_err(); + assert_eq!(error.code, WorkdirRouteErrorCode::NoWorkdirAttached); + assert!(error.available_aliases.is_empty()); + assert_eq!( + serde_json::to_value(error).unwrap(), + serde_json::json!({ + "code": "no_workdir_attached", + "available_aliases": [] + }) + ); + } + #[tokio::test] async fn detach_allows_naturally_completed_command_without_final_poll() { let dir = TempDir::new().unwrap(); diff --git a/crates/workdir/src/scope.rs b/crates/workdir/src/scope.rs index 5e158d231..7eadad4e9 100644 --- a/crates/workdir/src/scope.rs +++ b/crates/workdir/src/scope.rs @@ -1,4 +1,4 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::path::Path; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex, Weak}; @@ -284,6 +284,204 @@ impl Drop for WorkdirScopeLease { } } +/// Alias-keyed parent authority used to attenuate one or more Workdir +/// attachments for an Internal SubWorker. +#[derive(Clone)] +pub struct WorkdirToolBrokerRouter { + sessions: Arc, + brokers: Arc>>, +} + +impl std::fmt::Debug for WorkdirToolBrokerRouter { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkdirToolBrokerRouter") + .field("aliases", &self.sessions.aliases()) + .finish_non_exhaustive() + } +} + +impl WorkdirToolBrokerRouter { + pub fn from_single( + alias: crate::WorkdirAttachmentAlias, + session: WorkdirSessionHandle, + ) -> Result { + let sessions = Arc::new(crate::WorkdirSessionRouter::new()); + sessions.attach(alias, session)?; + Ok(Self::new(sessions)) + } + + pub fn new(sessions: Arc) -> Self { + Self { + sessions, + brokers: Arc::new(Mutex::new(BTreeMap::new())), + } + } + + pub fn sessions(&self) -> Arc { + self.sessions.clone() + } + + pub fn resolve_alias( + &self, + target_workdir: Option<&str>, + ) -> Result { + self.sessions + .resolve(target_workdir) + .map(|selected| selected.alias) + } + + fn broker( + &self, + alias: &crate::WorkdirAttachmentAlias, + ) -> Result { + let selected = self + .sessions + .resolve(Some(alias.as_str())) + .map_err(|error| WorkdirError::InvalidArgument(error.to_string()))?; + let mut brokers = self.brokers.lock().map_err(|_| { + WorkdirError::Unavailable("Workdir tool broker router is poisoned".to_string()) + })?; + if let Some((generation, broker)) = brokers.get(alias) + && *generation == selected.generation + { + return Ok(broker.clone()); + } + let broker = WorkdirToolBroker::new(selected.session); + brokers.insert(alias.clone(), (selected.generation, broker.clone())); + Ok(broker) + } + + pub async fn scope( + &self, + targets: Vec<(crate::WorkdirAttachmentAlias, WorkdirToolScope)>, + ) -> Result { + if targets.is_empty() { + return Err(WorkdirError::InvalidArgument( + "at least one Workdir attachment scope is required".to_string(), + )); + } + let child_sessions = Arc::new(crate::WorkdirSessionRouter::new()); + let mut child_brokers = BTreeMap::new(); + let mut leases = Vec::with_capacity(targets.len()); + for (alias, request) in targets { + if leases.iter().any( + |(existing, _): &(crate::WorkdirAttachmentAlias, WorkdirScopeLease)| { + existing == &alias + }, + ) { + for (_, lease) in &leases { + let _ = lease.close().await; + } + return Err(WorkdirError::InvalidArgument(format!( + "duplicate Workdir attachment scope `{alias}`" + ))); + } + let lease = match self.broker(&alias)?.scope(request).await { + Ok(lease) => lease, + Err(error) => { + for (_, lease) in &leases { + let _ = lease.close().await; + } + return Err(error); + } + }; + if let Err(error) = child_sessions.attach(alias.clone(), lease.tool_session()) { + let _ = lease.close().await; + for (_, existing) in &leases { + let _ = existing.close().await; + } + return Err(error); + } + let child_generation = child_sessions + .resolve(Some(alias.as_str())) + .expect("newly attached child route resolves") + .generation; + child_brokers.insert(alias.clone(), (child_generation, lease.broker())); + leases.push((alias, lease)); + } + Ok(WorkdirScopeLeaseSet { + broker_router: WorkdirToolBrokerRouter { + sessions: child_sessions, + brokers: Arc::new(Mutex::new(child_brokers)), + }, + leases, + }) + } +} + +pub struct WorkdirScopeLeaseSet { + broker_router: WorkdirToolBrokerRouter, + leases: Vec<(crate::WorkdirAttachmentAlias, WorkdirScopeLease)>, +} + +impl std::fmt::Debug for WorkdirScopeLeaseSet { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkdirScopeLeaseSet") + .field( + "aliases", + &self + .leases + .iter() + .map(|(alias, _)| alias) + .collect::>(), + ) + .field("active", &self.is_active()) + .finish() + } +} + +impl WorkdirScopeLeaseSet { + pub fn from_single( + alias: crate::WorkdirAttachmentAlias, + lease: WorkdirScopeLease, + ) -> Result { + let sessions = Arc::new(crate::WorkdirSessionRouter::new()); + sessions.attach(alias.clone(), lease.tool_session())?; + let generation = sessions + .resolve(Some(alias.as_str())) + .expect("newly attached single route resolves") + .generation; + let brokers = BTreeMap::from([(alias.clone(), (generation, lease.broker()))]); + Ok(Self { + broker_router: WorkdirToolBrokerRouter { + sessions, + brokers: Arc::new(Mutex::new(brokers)), + }, + leases: vec![(alias, lease)], + }) + } + + pub fn broker_router(&self) -> WorkdirToolBrokerRouter { + self.broker_router.clone() + } + + pub fn session_router(&self) -> Arc { + self.broker_router.sessions() + } + + pub async fn close(&self) -> Result<(), WorkdirError> { + let mut first_error = None; + for (_, lease) in &self.leases { + if let Err(error) = lease.close().await { + first_error.get_or_insert(error); + } + } + first_error.map_or(Ok(()), Err) + } + + pub fn is_active(&self) -> bool { + self.leases.iter().all(|(_, lease)| lease.is_active()) + } + + pub fn revoke(&self) { + for (_, lease) in &self.leases { + lease.revoke(); + } + } +} + #[derive(Debug)] struct SessionValidity { active: AtomicBool, @@ -2303,6 +2501,69 @@ mod tests { child.close().await.unwrap(); } + #[tokio::test] + async fn broker_router_delegates_multiple_aliases_without_namespace_mixing() { + let left = TempDir::new().unwrap(); + let right = TempDir::new().unwrap(); + fs::write(left.path().join("same.txt"), "left").unwrap(); + fs::write(right.path().join("same.txt"), "right").unwrap(); + let sessions = Arc::new(crate::WorkdirSessionRouter::new()); + sessions + .attach( + crate::WorkdirAttachmentAlias::new("left").unwrap(), + Arc::new(crate::LocalWorkdirSession::new( + manifest::Scope::writable(left.path()).unwrap(), + left.path().to_path_buf(), + )), + ) + .unwrap(); + sessions + .attach( + crate::WorkdirAttachmentAlias::new("right").unwrap(), + Arc::new(crate::LocalWorkdirSession::new( + manifest::Scope::writable(right.path()).unwrap(), + right.path().to_path_buf(), + )), + ) + .unwrap(); + let brokers = WorkdirToolBrokerRouter::new(sessions); + let leases = brokers + .scope(vec![ + ( + crate::WorkdirAttachmentAlias::new("left").unwrap(), + request("", WorkdirToolScopePermission::Read), + ), + ( + crate::WorkdirAttachmentAlias::new("right").unwrap(), + request("", WorkdirToolScopePermission::Read), + ), + ]) + .await + .unwrap(); + + let child = leases.session_router(); + let left_content = child + .resolve(Some("left")) + .unwrap() + .session + .read(read("same.txt")) + .await + .unwrap(); + let right_content = child + .resolve(Some("right")) + .unwrap() + .session + .read(read("same.txt")) + .await + .unwrap(); + assert_eq!(left_content.bytes, b"left"); + assert_eq!(right_content.bytes, b"right"); + assert!(child.resolve(None).is_err()); + + leases.close().await.unwrap(); + assert!(!leases.is_active()); + } + #[tokio::test] async fn closing_parent_invalidates_scoped_tools() { let root = TempDir::new().unwrap(); diff --git a/crates/workdir/src/workspace.rs b/crates/workdir/src/workspace.rs index 5e2b72584..3bec80329 100644 --- a/crates/workdir/src/workspace.rs +++ b/crates/workdir/src/workspace.rs @@ -240,5 +240,7 @@ mod tests { #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct WorkspaceWorkdirSessionOperationRequest { + /// Worker-local attachment alias selected by the calling tool router. + pub target_workdir: String, pub operation: crate::http::WorkdirSessionOperation, } diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 57475f025..babae8e21 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -1143,7 +1143,7 @@ pub(crate) async fn register_worker_tools( runtime_base: PathBuf, spawned_registry: Arc, parent_method_tx: Option>, - inherited_workdir_tool_broker: Option, + inherited_workdir_tool_broker: Option, ) -> std::io::Result> where C: LlmClient + Clone + 'static, @@ -1152,23 +1152,13 @@ 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 { - let workspace_client = worker.workspace_client_handle(); - workdir_tool_broker = Some(workdir::WorkdirToolBroker::new( - crate::feature::builtin::manage_workdir::WorkspaceAttachedWorkdirSession::handle( - workspace_client, - ), - )); - } else if workdir_tool_broker.is_none() - && let Some(existing) = worker.workdir_session() - { - let broker = workdir::WorkdirToolBroker::new(existing); - workdir_tool_broker = Some(broker); - } - let worker_workdir = workdir_tool_broker - .as_ref() - .map(workdir::WorkdirToolBroker::tool_session); + let workdir_sessions = worker.workdir_sessions(); + let restrict_workdir_tools_to_current_capabilities = inherited_workdir_tool_broker.is_some(); + let workdir_tool_broker = Some( + inherited_workdir_tool_broker + .unwrap_or_else(|| workdir::WorkdirToolBrokerRouter::new(workdir_sessions.clone())), + ); + let workdir_for_view = workdir_sessions.only_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(); @@ -1203,28 +1193,29 @@ where paste_store, paste_session_id, )); - // Resolve the existing Worker–Workdir binding into the domain provider. - // Tools only consume the provider handle; they do not own its root, cwd, - // scope, or lifecycle. No-workdir Workers expose no local tools. - let (workdir_for_view, tracker) = if let Some(workdir) = worker_workdir { - let tracker = tools::Tracker::new(); + // Keep the alias-routed schemas installed for the Worker lifetime. The + // attachment set may be empty or mutate after registration, so target + // resolution and provider capability checks happen per invocation. + let tracker = tools::Tracker::new(); + let workdir_tools = if restrict_workdir_tools_to_current_capabilities { + tools::routed_builtin_tools_for_current_capabilities( + workdir_sessions.clone(), + tracker.clone(), + bash_output_dir.clone(), + ) + } else { + tools::routed_builtin_tools( + workdir_sessions.clone(), + tracker.clone(), + bash_output_dir.clone(), + ) + }; + worker.engine_mut().register_tools(workdir_tools); + if feature_config.image.enabled && model_supports_image_attachments(&spawner_manifest.model) { worker .engine_mut() - .register_tools(tools::core_builtin_tools( - workdir.clone(), - tracker.clone(), - bash_output_dir.clone(), - )); - if feature_config.image.enabled && model_supports_image_attachments(&spawner_manifest.model) - { - worker - .engine_mut() - .register_tool(tools::view_image_tool(workdir.clone())); - } - (Some(workdir), Some(tracker)) - } else { - (None, None) - }; + .register_tool(tools::routed_view_image_tool(workdir_sessions)); + } if feature_config.web.enabled { worker .engine_mut() @@ -1482,9 +1473,7 @@ where if let Some((resident_summary, system_prompt_override)) = memory_prompt_contribution { worker.install_system_prompt_contribution(resident_summary, system_prompt_override); } - if let Some(tracker) = tracker { - worker.attach_tracker(tracker); - } + worker.attach_tracker(tracker); Ok(workdir_for_view) } diff --git a/crates/worker/src/feature/builtin/manage_workdir.rs b/crates/worker/src/feature/builtin/manage_workdir.rs index fd4198fb4..2df835b34 100644 --- a/crates/worker/src/feature/builtin/manage_workdir.rs +++ b/crates/worker/src/feature/builtin/manage_workdir.rs @@ -213,14 +213,27 @@ impl std::fmt::Debug for WorkspaceHttpWorkdirBackend { #[derive(Debug)] pub struct WorkspaceAttachedWorkdirSession { client: Arc, + target_workdir: String, workdir: Workdir, } impl WorkspaceAttachedWorkdirSession { - pub fn handle(client: Arc) -> WorkdirSessionHandle { + pub fn handle( + client: Arc, + target_workdir: impl Into, + ) -> WorkdirSessionHandle { + Self::handle_for_workdir(client, target_workdir, "workspace-attachment") + } + + pub fn handle_for_workdir( + client: Arc, + target_workdir: impl Into, + workdir_id: impl Into, + ) -> WorkdirSessionHandle { Arc::new(Self { client, - workdir: Workdir::new("workspace-attachment"), + target_workdir: target_workdir.into(), + workdir: Workdir::new(workdir_id), }) } @@ -237,13 +250,15 @@ impl WorkspaceAttachedWorkdirSession { "/api/w/{}/workers/self/workdir-session/operations", encode_path_segment(workspace_id) ), - serde_json::to_string(&WorkspaceWorkdirSessionOperationRequest { operation }).map_err( - |error| { - WorkdirError::Transport(format!( - "failed to encode Workspace Workdir operation: {error}" - )) - }, - )?, + serde_json::to_string(&WorkspaceWorkdirSessionOperationRequest { + target_workdir: self.target_workdir.clone(), + operation, + }) + .map_err(|error| { + WorkdirError::Transport(format!( + "failed to encode Workspace Workdir operation: {error}" + )) + })?, ); let response = self .client @@ -507,7 +522,11 @@ impl WorkspaceHttpWorkdirBackend { let response = self.attach_response(alias, working_directory_id)?; if let Err(error) = self.session_router.attach( alias_key.clone(), - WorkspaceAttachedWorkdirSession::handle(self.client.clone()), + WorkspaceAttachedWorkdirSession::handle_for_workdir( + self.client.clone(), + alias, + working_directory_id, + ), ) { let _ = self.detach(alias); return Err(ToolError::ExecutionFailed(format!( @@ -1282,7 +1301,7 @@ mod tests { "operation": "stat", "result": {"path": "visible.txt", "kind": "file", "size": 8} }))])); - let session = WorkspaceAttachedWorkdirSession::handle(client.clone()); + let session = WorkspaceAttachedWorkdirSession::handle(client.clone(), "workdir"); let result = session .stat(StatRequest { path: workdir::WorkdirPath::new("visible.txt").unwrap(), @@ -1298,6 +1317,7 @@ mod tests { ); let body: serde_json::Value = serde_json::from_str(requests[0].body.as_deref().unwrap()).unwrap(); + assert_eq!(body["target_workdir"], "workdir"); assert_eq!(body["operation"]["operation"], "stat"); assert!(body.get("expected_session_fence").is_none()); assert!(body.get("runtime_id").is_none()); @@ -1310,7 +1330,7 @@ mod tests { "operation": "command_start", "result": "command-1" }))])); - let session = WorkspaceAttachedWorkdirSession::handle(client.clone()); + let session = WorkspaceAttachedWorkdirSession::handle(client.clone(), "workdir"); let handle = session .start_command(CommandRequest { @@ -1348,7 +1368,7 @@ mod tests { workdir::http::WorkdirTransportErrorCode::InvalidRequest, "Workdir operation request is invalid", )])); - let session = WorkspaceAttachedWorkdirSession::handle(client); + let session = WorkspaceAttachedWorkdirSession::handle(client, "workdir"); let error = session .glob(workdir::GlobRequest { @@ -1368,7 +1388,7 @@ mod tests { status: 502, body: "secret token and /host/private/path".to_string(), }])); - let session = WorkspaceAttachedWorkdirSession::handle(client); + let session = WorkspaceAttachedWorkdirSession::handle(client, "workdir"); let error = session .stat(StatRequest { @@ -1395,6 +1415,7 @@ mod tests { ])); let broker = workdir::WorkdirToolBroker::new(WorkspaceAttachedWorkdirSession::handle( client.clone(), + "workdir", )); let scoped = broker .scope(workdir::WorkdirToolScope { @@ -1495,7 +1516,7 @@ mod tests { router .attach( workdir::WorkdirAttachmentAlias::new("checkout").unwrap(), - WorkspaceAttachedWorkdirSession::handle(client.clone()), + WorkspaceAttachedWorkdirSession::handle(client.clone(), "workdir"), ) .unwrap(); let cleanup_calls = Arc::new(AtomicUsize::new(0)); @@ -1533,7 +1554,7 @@ mod tests { router .attach( workdir::WorkdirAttachmentAlias::new("checkout").unwrap(), - WorkspaceAttachedWorkdirSession::handle(client.clone()), + WorkspaceAttachedWorkdirSession::handle(client.clone(), "workdir"), ) .unwrap(); let before_release: BeforeWorkdirRelease = diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index 2c7366cd1..0b50042ce 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -772,7 +772,7 @@ pub(crate) async fn prepare_internal_worker_session( visibility: InternalWorkerVisibility, child_registry: Option>, on_turn_end: Option>, - command_event_broker: Option, + command_event_broker: Option, ) -> Result { let (event_tx, _event_rx) = broadcast::channel(256); let sink = worker.sink(); @@ -780,7 +780,12 @@ pub(crate) async fn prepare_internal_worker_session( let alerter = Alerter::new(event_tx.clone()); let in_flight = InFlightEvents::new(event_tx.clone()); if let Some(broker) = command_event_broker.as_ref() { - wire_workdir_command_events(&broker.tool_session(), &in_flight); + let sessions = broker.sessions(); + for alias in sessions.aliases() { + if let Ok(selected) = sessions.resolve(Some(alias.as_str())) { + wire_workdir_command_events(&selected.session, &in_flight); + } + } } else if let Some(session) = worker.workdir_session() { wire_workdir_command_events(&session, &in_flight); } diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 7c2606f74..8cac5d7e5 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -27,7 +27,7 @@ use session_store::{ }; use tokio::sync::{Notify, broadcast}; use tracing::warn; -use workdir::WorkdirScopeLease; +use workdir::WorkdirScopeLeaseSet; use crate::internal_worker::{InternalWorkerSessionHandle, InternalWorkerVisibility}; use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord}; @@ -70,7 +70,7 @@ pub(crate) struct SubWorkerStopSummary { pub(crate) struct InternalSpawnedWorkerRecord { pub worker_name: String, pub scope_delegated: Vec, - pub workdir_tool_scope: Arc, + pub workdir_tool_scope: Arc, #[cfg(test)] pub installed_tools: Arc<[String]>, pub session: InternalWorkerSessionHandle, @@ -89,7 +89,7 @@ impl InternalSpawnedWorkerRecord { pub(crate) fn new( worker_name: String, scope_delegated: Vec, - workdir_tool_scope: WorkdirScopeLease, + workdir_tool_scope: WorkdirScopeLeaseSet, #[cfg(test)] installed_tools: Vec, session: InternalWorkerSessionHandle, child_registry: Arc, @@ -1179,7 +1179,11 @@ mod tests { InternalSpawnedWorkerRecord::new( name.into(), Vec::new(), - delegation, + workdir::WorkdirScopeLeaseSet::from_single( + workdir::WorkdirAttachmentAlias::new("workdir").unwrap(), + delegation, + ) + .unwrap(), Vec::new(), session, registry(), diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index 60efb8c68..18dfa2a79 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -5,6 +5,7 @@ //! in-process Internal Worker session actor. No Runtime Worker record, OS process, PID, Unix socket, //! or machine-wide child allocation is created. +use std::collections::BTreeMap; use std::path::{Path, PathBuf}; use std::sync::Arc; @@ -22,7 +23,7 @@ use manifest::{ use serde::Deserialize; use tokio::sync::mpsc; use workdir::{ - WorkdirToolBroker, WorkdirToolScope, WorkdirToolScopePermission, WorkdirToolScopeRule, + WorkdirToolBrokerRouter, WorkdirToolScope, WorkdirToolScopePermission, WorkdirToolScopeRule, }; use crate::PromptCatalogSource; @@ -81,6 +82,10 @@ struct ReviewerHandoffInput { #[derive(Debug, Deserialize, schemars::JsonSchema)] struct ScopeRuleInput { + /// Worker-local alias of the parent Workdir attachment. May be omitted only + /// when the parent has exactly one attachment. + #[serde(default)] + target_workdir: Option, /// Logical Workdir-relative target such as `.` or `src`. Absolute host /// paths and parent traversal are rejected. target: String, @@ -310,7 +315,7 @@ pub struct SubWorkerSpawnTool { /// Directory the spawned SubWorker's tools should use when the LLM did not /// override it. Defaults to the spawner's cwd. /// Parent-owned broker for scoped Workdir tool execution. - workdir_tool_broker: Option, + 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 @@ -337,7 +342,7 @@ impl SubWorkerSpawnTool { runtime_base: PathBuf, bash_output_dir: PathBuf, workspace_root: PathBuf, - workdir_tool_broker: Option, + workdir_tool_broker: Option, registry: Arc, spawner_manifest: WorkerManifest, prompt_loader: PromptCatalogSource, @@ -428,14 +433,19 @@ impl Tool for SubWorkerSpawnTool { )) })?; 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 tool_scopes = routed_workdir_tool_scopes( + workdir_tool_broker, + input.cwd.as_deref(), + workdir_rules, + input.command, + )?; let workdir_scope = workdir_tool_broker - .scope(tool_scope) + .scope(tool_scopes) .await .map_err(|error| { ToolError::InvalidArgument(format!("scope parent-owned Workdir tools: {error}")) })?; - let child_workdir_tool_broker = workdir_scope.broker(); + let child_workdir_tool_broker = workdir_scope.broker_router(); let spawn_selector = parse_spawn_profile_selector(input.profile.as_deref()).map_err(|msg| { @@ -537,6 +547,7 @@ impl Tool for SubWorkerSpawnTool { )) })?; let child_scope = child.scope().clone(); + child.bind_workdir_sessions(workdir_scope.session_router()); let child_registry = SpawnedWorkerRegistry::new_internal(input.name.clone(), child_scope); register_worker_tools( &mut child, @@ -717,22 +728,27 @@ fn logical_workdir_path(value: &str, field: &str) -> Result { }) } -fn parse_workdir_scope(rules: &[ScopeRuleInput]) -> Result, ToolError> { +fn parse_workdir_scope( + rules: &[ScopeRuleInput], +) -> Result, WorkdirToolScopeRule)>, ToolError> { if rules.is_empty() { return Err(ToolError::InvalidArgument("scope must not be empty".into())); } rules .iter() .map(|rule| { - Ok(WorkdirToolScopeRule { - target: logical_workdir_path(&rule.target, "scope.target")?, - permission: match rule.permission { - PermissionInput::Read => WorkdirToolScopePermission::Read, - PermissionInput::Write => WorkdirToolScopePermission::Write, + Ok(( + rule.target_workdir.clone(), + WorkdirToolScopeRule { + target: logical_workdir_path(&rule.target, "scope.target")?, + permission: match rule.permission { + PermissionInput::Read => WorkdirToolScopePermission::Read, + PermissionInput::Write => WorkdirToolScopePermission::Write, + }, + recursive: rule.recursive, + symlink_policy: rule.symlink_policy.into(), }, - recursive: rule.recursive, - symlink_policy: rule.symlink_policy.into(), - }) + )) }) .collect() } @@ -749,9 +765,28 @@ fn workdir_tool_scope( }) } +fn routed_workdir_tool_scopes( + broker: &WorkdirToolBrokerRouter, + cwd: Option<&str>, + rules: Vec<(Option, WorkdirToolScopeRule)>, + command: bool, +) -> Result, ToolError> { + let mut grouped = BTreeMap::>::new(); + for (target_workdir, rule) in rules { + let alias = broker + .resolve_alias(target_workdir.as_deref()) + .map_err(|error| ToolError::InvalidArgument(error.to_string()))?; + grouped.entry(alias).or_default().push(rule); + } + grouped + .into_iter() + .map(|(alias, rules)| workdir_tool_scope(cwd, rules, command).map(|scope| (alias, scope))) + .collect() +} + fn require_workdir_tool_broker( - broker: Option<&WorkdirToolBroker>, -) -> Result<&WorkdirToolBroker, ToolError> { + broker: Option<&WorkdirToolBrokerRouter>, +) -> Result<&WorkdirToolBrokerRouter, ToolError> { broker.ok_or_else(|| { ToolError::InvalidArgument( "SubWorkerSpawn requires parent-owned Workdir tools; attach a Workdir before granting filesystem access" @@ -992,7 +1027,7 @@ pub(crate) fn sub_worker_spawn_tool( runtime_base: PathBuf, bash_output_dir: PathBuf, workspace_root: PathBuf, - workdir_tool_broker: Option, + workdir_tool_broker: Option, registry: Arc, spawner_manifest: WorkerManifest, prompts: Arc>, @@ -1018,7 +1053,7 @@ fn sub_worker_spawn_tool_impl( runtime_base: PathBuf, bash_output_dir: PathBuf, workspace_root: PathBuf, - workdir_tool_broker: Option, + workdir_tool_broker: Option, registry: Arc, spawner_manifest: WorkerManifest, prompts: Arc>, @@ -1096,12 +1131,14 @@ mod tests { fn workdir_scope_uses_logical_relative_paths() { let rules = parse_workdir_scope(&[ ScopeRuleInput { + target_workdir: None, target: ".".to_string(), permission: PermissionInput::Read, recursive: true, symlink_policy: Default::default(), }, ScopeRuleInput { + target_workdir: Some("checkout".to_string()), target: "src".to_string(), permission: PermissionInput::Write, recursive: false, @@ -1109,12 +1146,14 @@ mod tests { }, ]) .unwrap(); - assert_eq!(rules[0].target.as_str(), ""); - assert_eq!(rules[1].target.as_str(), "src"); - assert_eq!(rules[0].symlink_policy, SymlinkPolicy::Resolved); - assert_eq!(rules[1].symlink_policy, SymlinkPolicy::Logical); + assert_eq!(rules[0].1.target.as_str(), ""); + assert_eq!(rules[1].0.as_deref(), Some("checkout")); + assert_eq!(rules[1].1.target.as_str(), "src"); + assert_eq!(rules[0].1.symlink_policy, SymlinkPolicy::Resolved); + assert_eq!(rules[1].1.symlink_policy, SymlinkPolicy::Logical); for target in ["/host/path", "../escape"] { let error = parse_workdir_scope(&[ScopeRuleInput { + target_workdir: None, target: target.to_string(), permission: PermissionInput::Read, recursive: true, @@ -1278,15 +1317,17 @@ 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 workdir_tool_broker = workdir::WorkdirToolBroker::new(Arc::new( - workdir::LocalWorkdirSession::materialized_bound( + let workdir_tool_broker = workdir::WorkdirToolBrokerRouter::from_single( + workdir::WorkdirAttachmentAlias::new("workdir").unwrap(), + Arc::new(workdir::LocalWorkdirSession::materialized_bound( workdir::Workdir::new("test-workdir"), workspace_root.clone(), workspace_root.clone(), spawner_scope.clone(), workdir::WorkdirSessionCapabilities::ALL, - ), - )); + )), + ) + .unwrap(); let tool = SubWorkerSpawnTool::new( "parent".into(), workspace_context, @@ -1528,9 +1569,11 @@ enabled = false Arc::new(AvailableWorkspaceClient), ); let remote_client = Arc::new(StrictRemoteWorkdirWorkspaceClient::default()); - let workdir_tool_broker = workdir::WorkdirToolBroker::new( - WorkspaceAttachedWorkdirSession::handle(remote_client.clone()), - ); + let workdir_tool_broker = workdir::WorkdirToolBrokerRouter::from_single( + workdir::WorkdirAttachmentAlias::new("workdir").unwrap(), + WorkspaceAttachedWorkdirSession::handle(remote_client.clone(), "workdir"), + ) + .unwrap(); let calls = Arc::new(AtomicUsize::new(0)); let (parent_method_tx, _parent_method_rx) = mpsc::channel(8); let tool = SubWorkerSpawnTool::new( @@ -1607,6 +1650,7 @@ enabled = false assert!(properties.contains_key("command"), "schema: {schema}"); let schema_text = serde_json::to_string(&schema).unwrap(); assert!(schema_text.contains("symlink_policy"), "schema: {schema}"); + assert!(schema_text.contains("target_workdir"), "schema: {schema}"); assert!(schema_text.contains("logical"), "schema: {schema}"); let required = schema .get("required") diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index e1a90d770..9a60ab229 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -2950,11 +2950,7 @@ impl Worker { location.session_id, location.segment_id, ); - let read_only_tools = self - .workdir_session() - .map(|source| Arc::new(ReadOnlyWorkdirSession::new(source)) as WorkdirSessionHandle) - .map(tools::read_only_builtin_tools) - .unwrap_or_default(); + let read_only_tools = tools::read_only_routed_builtin_tools(self.workdir_sessions()); let client = self .engine .as_ref() @@ -5401,11 +5397,14 @@ impl Worker { // Default references: the N most-recently-touched files in the // session, surfaced so the compact worker can inspect them and // decide which (if any) the next session needs. - let default_refs: Vec = self - .tracker - .as_ref() - .map(|t| t.recent_files(manifest::defaults::COMPACT_DEFAULT_REFERENCE_COUNT)) - .unwrap_or_default(); + let default_refs: Vec = if self.workdir_sessions.len() == 1 { + self.tracker + .as_ref() + .map(|t| t.recent_files(manifest::defaults::COMPACT_DEFAULT_REFERENCE_COUNT)) + .unwrap_or_default() + } else { + Vec::new() + }; // Input text fed to the compact worker. Includes the default // references, current TaskStore snapshot, current TaskStore snapshot, and the (pruned) conversation text. diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 912c30295..e050a3b77 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -406,10 +406,6 @@ struct WorkdirSessionRegistry { } impl WorkdirSessionRegistry { - fn attachment(&self, worker: &RuntimeWorkerRef) -> Option { - self.attachments.get(worker).cloned() - } - fn insert_attachment( &mut self, worker: RuntimeWorkerRef, @@ -8416,57 +8412,6 @@ fn current_worker_identity( authenticate_worker_mutation_source(api, workspace_id, headers) } -fn current_worker_active_attachment( - api: &WorkspaceApi, - worker: &RuntimeWorkerRef, -) -> ApiResult { - let active_links = api - .store - .list_worker_workdir_links(&api.config.workspace_id, worker)? - .into_iter() - .filter(|link| link.unlinked_at.is_none()) - .collect::>(); - if active_links.len() == 1 { - return Ok(active_links.into_iter().next().expect("one active link")); - } - if active_links.len() > 1 { - return Err(Error::WorkdirAttachmentConflict(format!( - "Worker {}:{} has multiple Workdir attachments; implicit operation routing is unavailable", - worker.runtime_id, worker.worker_id - )) - .into()); - } - - if api - .store - .worker_workdir_link_history_exists(&api.config.workspace_id, worker)? - { - return Err(Error::WorkdirAttachmentConflict(format!( - "Worker {}:{} has no active Workdir attachment", - worker.runtime_id, worker.worker_id - )) - .into()); - } - - // A Runtime can start the Worker immediately after reserving its local binding, before the - // outer spawn handler has projected that binding into the Backend registry. Import the same - // binding transactionally on the first identity-bound operation so initial input cannot race - // attachment authority. - let observed_worker = api - .runtime - .worker(worker) - .map_err(|error| error.into_error())?; - if !observed_worker.workdir_attachments.is_empty() { - sync_worker_observation(api, &observed_worker)?; - return current_worker_active_attachment(api, worker); - } - Err(Error::WorkdirAttachmentConflict(format!( - "Worker {}:{} has no active Workdir attachment", - worker.runtime_id, worker.worker_id - )) - .into()) -} - fn current_worker_session_lock( api: &WorkspaceApi, worker: &RuntimeWorkerRef, @@ -8529,13 +8474,11 @@ async fn open_current_worker_workdir_session_locked( api.runtime .authorize_working_directory_repository_access(&workdir.runtime_id, access) .map_err(RuntimeRegistryError::into_error)?; - } else if let Some(session) = api - .workdir_sessions - .lock() - .expect("Workdir session registry lock poisoned") - .attachment(worker) - { - return Ok(session); + } else { + // The registry is an operation-session cache, not routing authority. + // Close any prior alias session before opening the explicitly selected + // attachment so multiple aliases can never share one provider namespace. + close_current_worker_attachment_session_locked(api, worker).await?; } let owner_worker_id = runtime_local_owner_worker_id(worker, &workdir.runtime_id); let session = api @@ -8843,8 +8786,43 @@ async fn scoped_detach_current_worker_workdir( fn validated_current_worker_attachment( api: &WorkspaceApi, worker: &RuntimeWorkerRef, + target_workdir: &str, ) -> ApiResult { - current_worker_active_attachment(api, worker) + let find = || -> Result> { + Ok(api + .store + .list_worker_workdir_links(&api.config.workspace_id, worker)? + .into_iter() + .find(|link| link.unlinked_at.is_none() && link.alias == target_workdir)) + }; + if let Some(link) = find()? { + return Ok(link); + } + + // A Runtime can start the Worker immediately after reserving its local + // bindings, before the outer spawn handler projects them into the Backend. + // Import that exact alias set once; never substitute another attachment. + if !api + .store + .worker_workdir_link_history_exists(&api.config.workspace_id, worker)? + { + let observed_worker = api + .runtime + .worker(worker) + .map_err(|error| error.into_error())?; + if !observed_worker.workdir_attachments.is_empty() { + sync_worker_observation(api, &observed_worker)?; + if let Some(link) = find()? { + return Ok(link); + } + } + } + + Err(Error::WorkdirAttachmentConflict(format!( + "Worker {}:{} has no Workdir attachment alias `{target_workdir}`", + worker.runtime_id, worker.worker_id + )) + .into()) } #[derive(Debug)] @@ -8899,11 +8877,12 @@ 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 target_workdir = request.target_workdir; 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)?; + let link = validated_current_worker_attachment(&api, &worker, &target_workdir)?; let source = open_current_worker_workdir_session_locked(&api, &worker, &link).await?; let provider_handle = source .start_command(command) @@ -8929,7 +8908,7 @@ async fn scoped_execute_current_worker_workdir_operation( } WorkdirSessionOperation::CommandStatus(external_handle) => { let (session, provider_handle) = - current_worker_command_session(&api, &worker, &external_handle)?; + current_worker_command_session(&api, &worker, &target_workdir, &external_handle)?; session .command_status(provider_handle) .await @@ -8938,7 +8917,7 @@ async fn scoped_execute_current_worker_workdir_operation( } WorkdirSessionOperation::CommandOutput(mut output) => { let (session, provider_handle) = - current_worker_command_session(&api, &worker, &output.handle)?; + current_worker_command_session(&api, &worker, &target_workdir, &output.handle)?; output.handle = provider_handle; session .command_output(output) @@ -8948,7 +8927,7 @@ async fn scoped_execute_current_worker_workdir_operation( } WorkdirSessionOperation::CommandCancel(external_handle) => { let (session, provider_handle) = - current_worker_command_session(&api, &worker, &external_handle)?; + current_worker_command_session(&api, &worker, &target_workdir, &external_handle)?; session .cancel_command(provider_handle) .await @@ -8966,7 +8945,7 @@ 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)?; + let link = validated_current_worker_attachment(&api, &worker, &target_workdir)?; let source = open_current_worker_workdir_session_locked(&api, &worker, &link).await?; execute_workdir_session_operation(&source, operation) .await @@ -8979,9 +8958,10 @@ async fn scoped_execute_current_worker_workdir_operation( fn current_worker_command_session( api: &WorkspaceApi, worker: &RuntimeWorkerRef, + target_workdir: &str, external_handle: &CommandHandle, ) -> std::result::Result<(WorkdirSessionHandle, CommandHandle), WorkdirOperationApiError> { - let _link = validated_current_worker_attachment(api, worker)?; + let _link = validated_current_worker_attachment(api, worker, target_workdir)?; let command = api .workdir_sessions .lock()