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

feat: route tools across workdir attachments

Browse Source
This commit is contained in:
Hare
2026-09-20 19:38:20 +09:00
parent c492cd9a94
commit c102119412
22 changed files with 1237 additions and 278 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/tools/src/bash.rs
+38 -21
View File
@@ -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<String>,
command: String,
#[serde(default)]
timeout: Option<u64>,
}
pub(crate) struct BashTool {
session: WorkdirSessionHandle,
router: Arc<WorkdirSessionRouter>,
output_dir: PathBuf,
state: Arc<Mutex<BashExecutionState>>,
}
@@ -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<ToolOutput, ToolError> {
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::<Vec<_>>()
};
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<WorkdirSessionRouter>,
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<dyn Tool> = Arc::new(BashTool {
session: session.clone(),
router: router.clone(),
output_dir: output_dir.clone(),
state: Arc::new(Mutex::new(BashExecutionState::default())),
});
crates/tools/src/edit.rs
+27 -11
View File
@@ -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<String>,
/// 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<WorkdirSessionRouter>,
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(&params.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<WorkdirSessionRouter>,
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<dyn Tool> = Arc::new(EditTool {
session: session.clone(),
router: router.clone(),
tracker: tracker.clone(),
});
(meta, tool)
crates/tools/src/glob.rs
+17 -5
View File
@@ -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<String>,
/// 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<WorkdirSessionRouter>,
}
#[async_trait]
@@ -32,13 +35,18 @@ impl Tool for GlobTool {
) -> Result<ToolOutput, ToolError> {
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<WorkdirSessionRouter>) -> 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<dyn Tool> = Arc::new(GlobTool {
session: session.clone(),
router: router.clone(),
});
(meta, tool)
})
crates/tools/src/grep.rs
+19 -5
View File
@@ -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<String>,
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<WorkdirSessionRouter>,
}
#[async_trait]
@@ -60,6 +65,11 @@ impl Tool for GrepTool {
) -> Result<ToolOutput, ToolError> {
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<WorkdirSessionRouter>) -> 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<dyn Tool> = Arc::new(GrepTool {
session: session.clone(),
router: router.clone(),
});
(meta, tool)
})
crates/tools/src/lib.rs
+83 -8
View File
@@ -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<workdir::WorkdirSessionRouter>,
tracker: Tracker,
bash_output_dir: std::path::PathBuf,
) -> Vec<agen::tool::ToolDefinition> {
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<workdir::WorkdirSessionRouter>,
tracker: Tracker,
bash_output_dir: std::path::PathBuf,
) -> Vec<agen::tool::ToolDefinition> {
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<workdir::WorkdirSessionRouter>,
) -> Vec<agen::tool::ToolDefinition> {
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<manifest::WebConfig>,
) -> Vec<agen::tool::ToolDefinition> {
crates/tools/src/read.rs
+26 -8
View File
@@ -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<String>,
/// 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<WorkdirSessionRouter>,
tracker: Tracker,
}
@@ -44,13 +48,19 @@ impl Tool for ReadTool {
) -> Result<ToolOutput, ToolError> {
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(&params.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<WorkdirSessionRouter>,
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<dyn Tool> = 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()
);
crates/tools/src/routing.rs
+32
View File
@@ -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<WorkdirSessionRouter> {
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<ResolvedWorkdirSession, ToolError> {
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)
}
crates/tools/src/tracker.rs
+66 -18
View File
@@ -65,9 +65,15 @@ fn hash_bytes(bytes: &[u8]) -> ContentHash {
hasher.finalize().into()
}
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
struct TrackedPath {
namespace: Option<Arc<str>>,
path: PathBuf,
}
#[derive(Debug, Clone, Default)]
struct FileMutationCoordinator {
locks: Arc<tokio::sync::Mutex<HashMap<PathBuf, Arc<tokio::sync::Mutex<()>>>>>,
locks: Arc<tokio::sync::Mutex<HashMap<TrackedPath, Arc<tokio::sync::Mutex<()>>>>>,
}
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<PathBuf, ContentHash>,
/// Hash of each file's last observed contents, keyed by attachment and canonical path.
hashes: HashMap<TrackedPath, ContentHash>,
/// Line count paired with observations that included the file content.
line_counts: HashMap<PathBuf, usize>,
line_counts: HashMap<TrackedPath, usize>,
/// LRU list of touched files. Front = most recently touched.
recency: VecDeque<PathBuf>,
recency: VecDeque<TrackedPath>,
/// Successful Write/Edit mutations attributed to this session's tools.
change_stat: ChangeStat,
}
@@ -146,6 +151,7 @@ struct Inner {
pub struct Tracker {
inner: Arc<Mutex<Inner>>,
mutations: FileMutationCoordinator,
namespace: Option<Arc<str>>,
}
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<String>) -> 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<usize> {
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<workdir::ContentHash, ToolsError> {
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<PathBuf> {
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();
crates/tools/src/view_image.rs
+19 -6
View File
@@ -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<String>,
/// Logical path relative to the bound Workdir root.
path: String,
}
struct ViewImageTool {
session: WorkdirSessionHandle,
router: Arc<WorkdirSessionRouter>,
}
#[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<WorkdirSessionRouter>) -> 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<dyn Tool> = Arc::new(ViewImageTool {
session: session.clone(),
router: router.clone(),
});
(meta, tool)
})
crates/tools/src/write.rs
+37 -17
View File
@@ -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<String>,
/// 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<WorkdirSessionRouter>,
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(&params.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<WorkdirSessionRouter>,
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<dyn Tool> = Arc::new(WriteTool {
session: session.clone(),
router: router.clone(),
tracker: tracker.clone(),
});
(meta, tool)
crates/tools/tests/integration.rs
+237 -2
View File
@@ -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<WorkdirSessionRouter>,
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<dyn Tool>, 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 &reg.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(&reg.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(
&reg.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(&reg.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(
&reg.get("Glob"),
json!({ "target_workdir": "left", "pattern": "*.log" }),
)
.await;
assert!(glob.content.unwrap().contains("only-left.log"));
let grep = call(
&reg.get("Grep"),
json!({
"target_workdir": "right",
"pattern": "RIGHT-NEEDLE",
"output_mode": "content"
}),
)
.await;
assert!(grep.content.unwrap().contains("only-right.log"));
let bash = call(
&reg.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(
&reg.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(
&reg.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(
&reg.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();