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
+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.