feat: route tools across workdir attachments

This commit is contained in:
2026-09-20 19:38:20 +09:00
parent c492cd9a94
commit c102119412
22 changed files with 1237 additions and 278 deletions
+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())),
});
+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)
+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)
})
+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)
})
+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> {
+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()
);
+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)
}
+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();
+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)
})
+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)
+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();
+12 -3
View File
@@ -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,
+146 -1
View File
@@ -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<String>,
pub available_aliases: Vec<String>,
}
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<HashSet<CommandHandle>>,
@@ -299,6 +346,7 @@ impl WorkdirSession for RoutedWorkdirSession {
#[derive(Debug, Default)]
pub struct WorkdirSessionRouter {
sessions: RwLock<BTreeMap<WorkdirAttachmentAlias, Arc<RoutedSessionState>>>,
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<ResolvedWorkdirSession, WorkdirRouteError> {
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::<Vec<_>>();
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();
+262 -1
View File
@@ -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<crate::WorkdirSessionRouter>,
brokers: Arc<Mutex<BTreeMap<crate::WorkdirAttachmentAlias, (u64, WorkdirToolBroker)>>>,
}
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<Self, WorkdirError> {
let sessions = Arc::new(crate::WorkdirSessionRouter::new());
sessions.attach(alias, session)?;
Ok(Self::new(sessions))
}
pub fn new(sessions: Arc<crate::WorkdirSessionRouter>) -> Self {
Self {
sessions,
brokers: Arc::new(Mutex::new(BTreeMap::new())),
}
}
pub fn sessions(&self) -> Arc<crate::WorkdirSessionRouter> {
self.sessions.clone()
}
pub fn resolve_alias(
&self,
target_workdir: Option<&str>,
) -> Result<crate::WorkdirAttachmentAlias, crate::WorkdirRouteError> {
self.sessions
.resolve(target_workdir)
.map(|selected| selected.alias)
}
fn broker(
&self,
alias: &crate::WorkdirAttachmentAlias,
) -> Result<WorkdirToolBroker, WorkdirError> {
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<WorkdirScopeLeaseSet, WorkdirError> {
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::<Vec<_>>(),
)
.field("active", &self.is_active())
.finish()
}
}
impl WorkdirScopeLeaseSet {
pub fn from_single(
alias: crate::WorkdirAttachmentAlias,
lease: WorkdirScopeLease,
) -> Result<Self, WorkdirError> {
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<crate::WorkdirSessionRouter> {
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();
+2
View File
@@ -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,
}
+30 -41
View File
@@ -1143,7 +1143,7 @@ pub(crate) async fn register_worker_tools<C, St>(
runtime_base: PathBuf,
spawned_registry: Arc<SpawnedWorkerRegistry>,
parent_method_tx: Option<mpsc::WeakSender<Method>>,
inherited_workdir_tool_broker: Option<workdir::WorkdirToolBroker>,
inherited_workdir_tool_broker: Option<workdir::WorkdirToolBrokerRouter>,
) -> std::io::Result<Option<workdir::WorkdirSessionHandle>>
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)
}
@@ -213,14 +213,27 @@ impl std::fmt::Debug for WorkspaceHttpWorkdirBackend {
#[derive(Debug)]
pub struct WorkspaceAttachedWorkdirSession {
client: Arc<dyn WorkspaceClient>,
target_workdir: String,
workdir: Workdir,
}
impl WorkspaceAttachedWorkdirSession {
pub fn handle(client: Arc<dyn WorkspaceClient>) -> WorkdirSessionHandle {
pub fn handle(
client: Arc<dyn WorkspaceClient>,
target_workdir: impl Into<String>,
) -> WorkdirSessionHandle {
Self::handle_for_workdir(client, target_workdir, "workspace-attachment")
}
pub fn handle_for_workdir(
client: Arc<dyn WorkspaceClient>,
target_workdir: impl Into<String>,
workdir_id: impl Into<String>,
) -> 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 =
+7 -2
View File
@@ -772,7 +772,7 @@ pub(crate) async fn prepare_internal_worker_session(
visibility: InternalWorkerVisibility,
child_registry: Option<Arc<SpawnedWorkerRegistry>>,
on_turn_end: Option<Arc<dyn Fn(InternalWorkerSessionStatus) + Send + Sync>>,
command_event_broker: Option<workdir::WorkdirToolBroker>,
command_event_broker: Option<workdir::WorkdirToolBrokerRouter>,
) -> Result<InternalWorkerSessionHandle, InternalWorkerSessionError> {
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);
}
+8 -4
View File
@@ -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<ScopeRule>,
pub workdir_tool_scope: Arc<WorkdirScopeLease>,
pub workdir_tool_scope: Arc<WorkdirScopeLeaseSet>,
#[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<ScopeRule>,
workdir_tool_scope: WorkdirScopeLease,
workdir_tool_scope: WorkdirScopeLeaseSet,
#[cfg(test)] installed_tools: Vec<String>,
session: InternalWorkerSessionHandle,
child_registry: Arc<SpawnedWorkerRegistry>,
@@ -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(),
+74 -30
View File
@@ -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<String>,
/// 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<WorkdirToolBroker>,
workdir_tool_broker: Option<WorkdirToolBrokerRouter>,
/// Parent-owned in-memory registry shared by the five SubWorker tools.
registry: Arc<SpawnedWorkerRegistry>,
/// 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<WorkdirToolBroker>,
workdir_tool_broker: Option<WorkdirToolBrokerRouter>,
registry: Arc<SpawnedWorkerRegistry>,
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<FsPath, ToolError> {
})
}
fn parse_workdir_scope(rules: &[ScopeRuleInput]) -> Result<Vec<WorkdirToolScopeRule>, ToolError> {
fn parse_workdir_scope(
rules: &[ScopeRuleInput],
) -> Result<Vec<(Option<String>, 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<String>, WorkdirToolScopeRule)>,
command: bool,
) -> Result<Vec<(workdir::WorkdirAttachmentAlias, WorkdirToolScope)>, ToolError> {
let mut grouped = BTreeMap::<workdir::WorkdirAttachmentAlias, Vec<WorkdirToolScopeRule>>::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<WorkdirToolBroker>,
workdir_tool_broker: Option<WorkdirToolBrokerRouter>,
registry: Arc<SpawnedWorkerRegistry>,
spawner_manifest: WorkerManifest,
prompts: Arc<ArcSwap<PromptCatalog>>,
@@ -1018,7 +1053,7 @@ fn sub_worker_spawn_tool_impl(
runtime_base: PathBuf,
bash_output_dir: PathBuf,
workspace_root: PathBuf,
workdir_tool_broker: Option<WorkdirToolBroker>,
workdir_tool_broker: Option<WorkdirToolBrokerRouter>,
registry: Arc<SpawnedWorkerRegistry>,
spawner_manifest: WorkerManifest,
prompts: Arc<ArcSwap<PromptCatalog>>,
@@ -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")
+9 -10
View File
@@ -2950,11 +2950,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
// 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<PathBuf> = self
.tracker
.as_ref()
.map(|t| t.recent_files(manifest::defaults::COMPACT_DEFAULT_REFERENCE_COUNT))
.unwrap_or_default();
let default_refs: Vec<PathBuf> = 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.
+49 -69
View File
@@ -406,10 +406,6 @@ struct WorkdirSessionRegistry {
}
impl WorkdirSessionRegistry {
fn attachment(&self, worker: &RuntimeWorkerRef) -> Option<WorkdirSessionHandle> {
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<WorkerWorkdirLinkRecord> {
let active_links = api
.store
.list_worker_workdir_links(&api.config.workspace_id, worker)?
.into_iter()
.filter(|link| link.unlinked_at.is_none())
.collect::<Vec<_>>();
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<WorkerWorkdirLinkRecord> {
current_worker_active_attachment(api, worker)
let find = || -> Result<Option<WorkerWorkdirLinkRecord>> {
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<Json<WorkdirSessionOperationResult>, 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()