feat: support multiple worker workdir attachments

This commit is contained in:
2026-09-19 23:27:57 +09:00
parent 308ae19f5e
commit 37df3ec545
52 changed files with 3715 additions and 1404 deletions
+4 -7
View File
@@ -145,7 +145,7 @@ where
.map_err(WorkerBootstrapError::Worker)?;
if let Some(workdir_session) = self.workdir_session {
worker.bind_workdir_session(Some(workdir_session));
worker.bind_single_workdir_session(Some(workdir_session));
}
Ok(PreparedWorker::new(worker, self.layout, self.transport))
}
@@ -168,7 +168,7 @@ where
.map_err(WorkerBootstrapError::Worker)?;
if let Some(workdir_session) = self.workdir_session {
worker.bind_workdir_session(Some(workdir_session));
worker.bind_single_workdir_session(Some(workdir_session));
}
Ok(PreparedWorker::new(worker, self.layout, self.transport))
}
@@ -222,7 +222,7 @@ where
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
{
let cleanup_session = worker.workdir_session().cloned();
let cleanup_sessions = worker.workdir_sessions();
let controller = match layout {
WorkerBootstrapLayout::Direct {
runtime_base,
@@ -257,10 +257,7 @@ where
controller_task,
}),
Err(source) => {
let cleanup_failed = match cleanup_session {
Some(session) => session.close().await.is_err(),
None => false,
};
let cleanup_failed = cleanup_sessions.close_all().await.is_err();
Err(WorkerBootstrapError::Controller {
source,
cleanup_failed,
+15 -19
View File
@@ -617,7 +617,7 @@ impl WorkerController {
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
{
let session = worker.workdir_session().cloned();
let sessions = worker.workdir_sessions();
let result = Self::spawn_initialized(
worker,
runtime_base,
@@ -628,10 +628,9 @@ impl WorkerController {
)
.await;
if result.is_err()
&& let Some(session) = session
&& let Err(error) = session.close().await
&& let Err(error) = sessions.close_all().await
{
tracing::warn!(%error, "Workdir session close after controller startup failure failed");
tracing::warn!(%error, "Workdir sessions close after controller startup failure failed");
}
result
}
@@ -1154,20 +1153,17 @@ where
// below so the worker borrow doesn't conflict with reads on `worker`.
let feature_config = worker.manifest().feature.clone();
let mut workdir_tool_broker = inherited_workdir_tool_broker;
if feature_config.manage_workdir.enabled && worker.workdir_session().is_none() {
if feature_config.manage_workdir.enabled {
let workspace_client = worker.workspace_client_handle();
let broker = workdir::WorkdirToolBroker::new(
workdir_tool_broker = Some(workdir::WorkdirToolBroker::new(
crate::feature::builtin::manage_workdir::WorkspaceAttachedWorkdirSession::handle(
workspace_client,
),
);
worker.bind_workdir_session(Some(broker.tool_session()));
workdir_tool_broker = Some(broker);
));
} else if workdir_tool_broker.is_none()
&& let Some(existing) = worker.workdir_session().cloned()
&& let Some(existing) = worker.workdir_session()
{
let broker = workdir::WorkdirToolBroker::new(existing);
worker.bind_workdir_session(Some(broker.tool_session()));
workdir_tool_broker = Some(broker);
}
let worker_workdir = workdir_tool_broker
@@ -1344,6 +1340,7 @@ where
feature_registry.add_module(
crate::feature::builtin::manage_workdir::ManageWorkdirFeature::with_child_lifecycle(
workspace_client,
worker.workdir_sessions(),
Arc::new(move || {
let child_registry = shutdown_registry.clone();
Box::pin(async move { child_registry.shutdown_internal().await })
@@ -2610,14 +2607,13 @@ async fn controller_loop<C, St>(
}
}
if let Some(session) = worker.workdir_session() {
loop {
match session.close().await {
Ok(()) => break,
Err(error) => {
tracing::warn!(%error, "Workdir session close failed; retrying");
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
let sessions = worker.workdir_sessions();
while !sessions.is_empty() {
match sessions.close_all().await {
Ok(()) => break,
Err(error) => {
tracing::warn!(%error, "Workdir session set close failed; retrying");
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
}
}
@@ -52,9 +52,9 @@ const DETACH_TOOL: &str = "WorkdirDetach";
const DELETE_TOOL: &str = "WorkdirDelete";
const LIST_DESCRIPTION: &str = "List persistent Workdirs in the current Workspace through Backend Workspace API authority. The result contains safe summaries and diagnostics, never host paths or Runtime connection details.";
const CREATE_DESCRIPTION: &str = "Materialize a persistent Workdir on a selected Runtime from a Workspace repository and optional selector. This does not change this Worker's attachment; use WorkdirAttach explicitly after creation.";
const ATTACH_DESCRIPTION: &str = "Attach this Worker to one existing Workdir. The Backend enforces one active Workdir per Worker and one active Worker per Workdir, then opens an ephemeral operation session.";
const DETACH_DESCRIPTION: &str = "Detach this Worker from its active Workdir and release Workdir occupancy. Any ephemeral operation session is closed.";
const CREATE_DESCRIPTION: &str = "Materialize a persistent Workdir with an optional human-facing display name on a selected Runtime from a Workspace repository and optional selector. The display name is independent of attachment aliases. This does not change this Worker's attachment; use WorkdirAttach explicitly after creation.";
const ATTACH_DESCRIPTION: &str = "Attach one existing Workdir to this Worker under a stable Worker-local alias. The Backend preserves one active Worker per Workdir while allowing this Worker to own multiple aliases.";
const DETACH_DESCRIPTION: &str = "Detach one Workdir attachment by its Worker-local alias. Active operations, commands, or Internal SubWorkers cause safe rejection rather than implicit authority loss.";
pub(crate) type BeforeWorkdirRelease =
Arc<dyn Fn() -> Pin<Box<dyn Future<Output = std::io::Result<()>> + Send>> + Send + Sync>;
pub(crate) type AfterWorkdirAttach = Arc<dyn Fn() + Send + Sync>;
@@ -64,6 +64,7 @@ const DELETE_DESCRIPTION: &str = "Request removal of one persistent Workdir by i
#[derive(Clone)]
pub struct ManageWorkdirFeature {
client: Arc<dyn WorkspaceClient>,
session_router: Arc<workdir::WorkdirSessionRouter>,
before_workdir_release: Option<BeforeWorkdirRelease>,
after_workdir_attach: Option<AfterWorkdirAttach>,
}
@@ -82,6 +83,7 @@ impl ManageWorkdirFeature {
pub fn new(client: Arc<dyn WorkspaceClient>) -> Self {
Self {
client,
session_router: Arc::new(workdir::WorkdirSessionRouter::new()),
before_workdir_release: None,
after_workdir_attach: None,
}
@@ -89,11 +91,13 @@ impl ManageWorkdirFeature {
pub(crate) fn with_child_lifecycle(
client: Arc<dyn WorkspaceClient>,
session_router: Arc<workdir::WorkdirSessionRouter>,
before_workdir_release: BeforeWorkdirRelease,
after_workdir_attach: AfterWorkdirAttach,
) -> Self {
Self {
client,
session_router,
before_workdir_release: Some(before_workdir_release),
after_workdir_attach: Some(after_workdir_attach),
}
@@ -116,10 +120,12 @@ impl FeatureModule for ManageWorkdirFeature {
}
fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> {
let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone()).with_child_lifecycle(
self.before_workdir_release.clone(),
self.after_workdir_attach.clone(),
);
let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone())
.with_session_router(self.session_router.clone())
.with_child_lifecycle(
self.before_workdir_release.clone(),
self.after_workdir_attach.clone(),
);
for (name, definition) in [
(
LIST_TOOL,
@@ -183,6 +189,8 @@ impl FeatureModule for ManageWorkdirFeature {
#[derive(Clone)]
struct WorkspaceHttpWorkdirBackend {
client: Arc<dyn WorkspaceClient>,
session_router: Arc<workdir::WorkdirSessionRouter>,
attachment_mutation_lock: Arc<tokio::sync::Mutex<()>>,
before_workdir_release: Option<BeforeWorkdirRelease>,
after_workdir_attach: Option<AfterWorkdirAttach>,
}
@@ -399,11 +407,18 @@ impl WorkspaceHttpWorkdirBackend {
fn new(client: Arc<dyn WorkspaceClient>) -> Self {
Self {
client,
session_router: Arc::new(workdir::WorkdirSessionRouter::new()),
attachment_mutation_lock: Arc::new(tokio::sync::Mutex::new(())),
before_workdir_release: None,
after_workdir_attach: None,
}
}
fn with_session_router(mut self, session_router: Arc<workdir::WorkdirSessionRouter>) -> Self {
self.session_router = session_router;
self
}
fn with_child_lifecycle(
mut self,
before_workdir_release: Option<BeforeWorkdirRelease>,
@@ -449,10 +464,12 @@ impl WorkspaceHttpWorkdirBackend {
.transpose()?;
let repository_key =
validate_identity(&input.repository_key, CREATE_TOOL, "repository_key")?;
let display_name = validate_optional_display_name(input.display_name)?;
let selector = validate_optional_selector(input.selector)?;
let workspace_id = encode_path_segment(self.workspace_id()?);
let request = WorkdirCreateRequest {
runtime_id: runtime_id.map(str::to_string),
display_name,
repository_key: repository_key.to_string(),
selector,
operation_id: Some(operation_id),
@@ -472,32 +489,67 @@ impl WorkspaceHttpWorkdirBackend {
}
fn attach(&self, input: WorkdirAttachInput) -> Result<ToolOutput, ToolError> {
let workdir_id = validate_identity(&input.workdir_id, ATTACH_TOOL, "workdir_id")?;
let response = self.attach_response(workdir_id)?;
workdir_output(format!("Attached to Workdir {workdir_id}"), &response)
let alias = validate_identity(&input.alias, ATTACH_TOOL, "alias")?;
workdir::WorkdirAttachmentAlias::new(alias)
.map_err(|error| ToolError::InvalidArgument(error.to_string()))?;
let working_directory_id = validate_identity(
&input.working_directory_id,
ATTACH_TOOL,
"working_directory_id",
)?;
let alias_key = workdir::WorkdirAttachmentAlias::new(alias)
.map_err(|error| ToolError::InvalidArgument(error.to_string()))?;
if self.session_router.session(&alias_key).is_some() {
return Err(ToolError::InvalidArgument(format!(
"attachment alias `{alias}` already exists"
)));
}
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()),
) {
let _ = self.detach(alias);
return Err(ToolError::ExecutionFailed(format!(
"attach Workdir session route: {error}"
)));
}
workdir_output(
format!("Attached Workdir {working_directory_id} as {alias}"),
&response,
)
}
fn attach_response(&self, workdir_id: &str) -> Result<WorkdirAttachmentResponse, ToolError> {
fn attach_response(
&self,
alias: &str,
working_directory_id: &str,
) -> Result<WorkdirAttachmentResponse, ToolError> {
let workspace_id = encode_path_segment(self.workspace_id()?);
self.execute_json::<WorkdirAttachmentResponse>(WorkspaceRequest::json(
WorkspaceRequestMethod::Post,
format!("/api/w/{workspace_id}/workers/self/workdir-attachment"),
format!("/api/w/{workspace_id}/workers/self/workdir-attachments"),
serde_json::to_string(&WorkdirAttachRequest {
workdir_id: workdir_id.to_string(),
alias: alias.to_string(),
working_directory_id: working_directory_id.to_string(),
})
.map_err(decode_error)?,
))
}
fn detach(&self) -> Result<ToolOutput, ToolError> {
fn detach(&self, alias: &str) -> Result<ToolOutput, ToolError> {
let workspace_id = encode_path_segment(self.workspace_id()?);
let alias_path = encode_path_segment(alias);
let response = self.execute_json::<WorkdirAttachmentResponse>(WorkspaceRequest {
method: WorkspaceRequestMethod::Delete,
path: format!("/api/w/{workspace_id}/workers/self/workdir-attachment"),
path: format!("/api/w/{workspace_id}/workers/self/workdir-attachments/{alias_path}"),
body: None,
})?;
workdir_output(
format!("Detached from Workdir {}", response.workdir_id),
format!(
"Detached Workdir {} ({})",
response.working_directory_id, response.alias
),
&response,
)
}
@@ -590,6 +642,7 @@ impl Tool for WorkspaceHttpWorkdirTool {
ctx.call_id.to_string(),
),
WorkdirOperation::Attach => {
let _mutation_guard = self.backend.attachment_mutation_lock.lock().await;
let result = self
.backend
.attach(parse_input::<WorkdirAttachInput>(input_json)?);
@@ -601,15 +654,41 @@ impl Tool for WorkspaceHttpWorkdirTool {
result
}
WorkdirOperation::Detach => {
let _input = parse_input::<WorkdirDetachInput>(input_json)?;
if let Some(before_release) = &self.backend.before_workdir_release {
before_release().await.map_err(|error| {
let _mutation_guard = self.backend.attachment_mutation_lock.lock().await;
let input = parse_input::<WorkdirDetachInput>(input_json)?;
let alias = validate_identity(&input.alias, DETACH_TOOL, "alias")?;
let alias_key = workdir::WorkdirAttachmentAlias::new(alias)
.map_err(|error| ToolError::InvalidArgument(error.to_string()))?;
self.backend
.session_router
.begin_detach(&alias_key)
.await
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?;
if let Some(before_release) = &self.backend.before_workdir_release
&& let Err(error) = before_release().await
{
self.backend.session_router.cancel_detach(&alias_key);
return Err(ToolError::ExecutionFailed(format!(
"stop Internal SubWorkers before Workdir detach: {error}"
)));
}
let result = match self.backend.detach(alias) {
Ok(result) => result,
Err(error) => {
self.backend.session_router.cancel_detach(&alias_key);
return Err(error);
}
};
self.backend
.session_router
.finish_detach(&alias_key)
.await
.map_err(|error| {
ToolError::ExecutionFailed(format!(
"stop Internal SubWorkers before Workdir detach: {error}"
"detach Workdir session route after Backend release: {error}"
))
})?;
}
self.backend.detach()
Ok(result)
}
WorkdirOperation::Delete => self
.backend
@@ -672,6 +751,25 @@ fn validate_delete_reason(reason: &str) -> Result<&str, ToolError> {
Ok(reason)
}
fn validate_optional_display_name(
display_name: Option<String>,
) -> Result<Option<String>, ToolError> {
let Some(display_name) = display_name else {
return Ok(None);
};
let display_name = display_name.trim();
if display_name.is_empty()
|| display_name.len() > 80
|| display_name.chars().any(char::is_control)
{
return Err(ToolError::InvalidArgument(
"WorkdirCreate display_name must be non-empty, contain no control characters, and be at most 80 bytes"
.to_string(),
));
}
Ok(Some(display_name.to_string()))
}
fn validate_optional_selector(selector: Option<String>) -> Result<Option<String>, ToolError> {
let Some(selector) = selector else {
return Ok(None);
@@ -725,6 +823,7 @@ fn create_schema() -> serde_json::Value {
"required": ["repository_key"],
"properties": {
"runtime_id": {"type": ["string", "null"], "minLength": 1},
"display_name": {"type": ["string", "null"], "minLength": 1, "maxLength": 80},
"repository_key": {"type": "string", "minLength": 1},
"selector": {"type": ["string", "null"], "minLength": 1}
}
@@ -735,15 +834,23 @@ fn attach_schema() -> serde_json::Value {
json!({
"type": "object",
"additionalProperties": false,
"required": ["workdir_id"],
"required": ["alias", "working_directory_id"],
"properties": {
"workdir_id": {"type": "string", "minLength": 1}
"alias": {"type": "string", "minLength": 1, "maxLength": 64},
"working_directory_id": {"type": "string", "minLength": 1}
}
})
}
fn detach_schema() -> serde_json::Value {
list_schema()
json!({
"type": "object",
"additionalProperties": false,
"required": ["alias"],
"properties": {
"alias": {"type": "string", "minLength": 1, "maxLength": 64}
}
})
}
fn delete_schema() -> serde_json::Value {
@@ -767,6 +874,8 @@ struct WorkdirListInput {}
struct WorkdirCreateInput {
#[serde(default)]
runtime_id: Option<String>,
#[serde(default)]
display_name: Option<String>,
repository_key: String,
#[serde(default)]
selector: Option<String>,
@@ -775,22 +884,27 @@ struct WorkdirCreateInput {
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct WorkdirAttachInput {
workdir_id: String,
alias: String,
working_directory_id: String,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct WorkdirDetachInput {}
struct WorkdirDetachInput {
alias: String,
}
#[derive(Debug, Serialize)]
struct WorkdirAttachRequest {
workdir_id: String,
alias: String,
working_directory_id: String,
}
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
struct WorkdirAttachmentResponse {
workspace_id: String,
workdir_id: String,
alias: String,
working_directory_id: String,
attached: bool,
}
@@ -1021,7 +1135,10 @@ mod tests {
);
assert!(create["properties"].get("path").is_none());
assert!(create["properties"].get("session_id").is_none());
assert_eq!(attach_schema()["required"], json!(["workdir_id"]));
assert_eq!(
attach_schema()["required"],
json!(["alias", "working_directory_id"])
);
assert!(attach_schema()["properties"].get("session_id").is_none());
assert_eq!(
delete_schema()["required"],
@@ -1057,12 +1174,14 @@ mod tests {
})),
response(json!({
"workspace_id": "workspace/test",
"workdir_id": "wd-created",
"alias": "checkout",
"working_directory_id": "wd-created",
"attached": true
})),
response(json!({
"workspace_id": "workspace/test",
"workdir_id": "wd-created",
"alias": "checkout",
"working_directory_id": "wd-created",
"attached": false
})),
response(json!({
@@ -1089,6 +1208,7 @@ mod tests {
.create(
WorkdirCreateInput {
runtime_id: Some("runtime/one".to_string()),
display_name: Some("Review checkout".to_string()),
repository_key: "main".to_string(),
selector: Some("refs/heads/topic".to_string()),
},
@@ -1106,10 +1226,11 @@ mod tests {
assert_eq!(client.requests().len(), 2);
backend
.attach(WorkdirAttachInput {
workdir_id: "wd-created".to_string(),
alias: "checkout".to_string(),
working_directory_id: "wd-created".to_string(),
})
.unwrap();
backend.detach().unwrap();
backend.detach("checkout").unwrap();
backend
.delete(WorkdirDeleteInput {
working_directory_id: "wd-created".to_string(),
@@ -1132,16 +1253,17 @@ mod tests {
serde_json::from_str(requests[1].body.as_deref().unwrap()).unwrap();
assert_eq!(body["repository_key"], "main");
assert_eq!(body["runtime_id"], "runtime/one");
assert_eq!(body["display_name"], "Review checkout");
assert_eq!(body["operation_id"], "call-create-1");
assert_eq!(body["selector"], "refs/heads/topic");
assert_eq!(
requests[2].path,
"/api/w/workspace%2Ftest/workers/self/workdir-attachment"
"/api/w/workspace%2Ftest/workers/self/workdir-attachments"
);
assert_eq!(requests[2].method, WorkspaceRequestMethod::Post);
assert_eq!(
requests[3].path,
"/api/w/workspace%2Ftest/workers/self/workdir-attachment"
"/api/w/workspace%2Ftest/workers/self/workdir-attachments/checkout"
);
assert_eq!(requests[3].method, WorkspaceRequestMethod::Delete);
assert_eq!(
@@ -1323,6 +1445,7 @@ mod tests {
.create(
WorkdirCreateInput {
runtime_id: None,
display_name: None,
repository_key: "main".to_string(),
selector: None,
},
@@ -1348,6 +1471,7 @@ mod tests {
.create(
WorkdirCreateInput {
runtime_id: Some(" ".to_string()),
display_name: None,
repository_key: "main".to_string(),
selector: None,
},
@@ -1363,9 +1487,17 @@ mod tests {
async fn detach_stops_internal_subworkers_before_backend_release() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![response(json!({
"workspace_id": "workspace/test",
"workdir_id": "wd-attached",
"alias": "checkout",
"working_directory_id": "wd-attached",
"attached": false
}))]));
let router = Arc::new(workdir::WorkdirSessionRouter::new());
router
.attach(
workdir::WorkdirAttachmentAlias::new("checkout").unwrap(),
WorkspaceAttachedWorkdirSession::handle(client.clone()),
)
.unwrap();
let cleanup_calls = Arc::new(AtomicUsize::new(0));
let cleanup_calls_for_guard = cleanup_calls.clone();
let before_release: BeforeWorkdirRelease = Arc::new(move || {
@@ -1377,11 +1509,12 @@ mod tests {
});
let tool = WorkspaceHttpWorkdirTool {
backend: WorkspaceHttpWorkdirBackend::new(client.clone())
.with_session_router(router)
.with_child_lifecycle(Some(before_release), None),
operation: WorkdirOperation::Detach,
};
tool.execute("{}", ToolExecutionContext::default())
tool.execute(r#"{"alias":"checkout"}"#, ToolExecutionContext::default())
.await
.unwrap();
@@ -1389,23 +1522,31 @@ mod tests {
assert_eq!(client.requests().len(), 1);
assert_eq!(
client.requests()[0].path,
"/api/w/workspace%2Ftest/workers/self/workdir-attachment"
"/api/w/workspace%2Ftest/workers/self/workdir-attachments/checkout"
);
}
#[tokio::test]
async fn detach_does_not_release_backend_when_child_cleanup_fails() {
let client = Arc::new(RecordingWorkspaceClient::new(Vec::new()));
let router = Arc::new(workdir::WorkdirSessionRouter::new());
router
.attach(
workdir::WorkdirAttachmentAlias::new("checkout").unwrap(),
WorkspaceAttachedWorkdirSession::handle(client.clone()),
)
.unwrap();
let before_release: BeforeWorkdirRelease =
Arc::new(|| Box::pin(async { Err(std::io::Error::other("child cleanup failed")) }));
let tool = WorkspaceHttpWorkdirTool {
backend: WorkspaceHttpWorkdirBackend::new(client.clone())
.with_session_router(router)
.with_child_lifecycle(Some(before_release), None),
operation: WorkdirOperation::Detach,
};
let error = tool
.execute("{}", ToolExecutionContext::default())
.execute(r#"{"alias":"checkout"}"#, ToolExecutionContext::default())
.await
.unwrap_err();
@@ -1413,11 +1554,77 @@ mod tests {
assert!(client.requests().is_empty());
}
#[tokio::test]
async fn detach_reconciles_naturally_completed_workspace_command() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![
response(json!({
"workspace_id": "workspace/test",
"alias": "checkout",
"working_directory_id": "wd-attached",
"attached": true
})),
response(json!({
"operation": "command_start",
"result": "command-1"
})),
response(json!({
"operation": "command_status",
"result": "completed"
})),
response(json!({
"workspace_id": "workspace/test",
"alias": "checkout",
"working_directory_id": "wd-attached",
"attached": false
})),
]));
let backend = WorkspaceHttpWorkdirBackend::new(client.clone());
backend
.attach(WorkdirAttachInput {
alias: "checkout".to_string(),
working_directory_id: "wd-attached".to_string(),
})
.unwrap();
let alias = workdir::WorkdirAttachmentAlias::new("checkout").unwrap();
backend
.session_router
.session(&alias)
.unwrap()
.start_command(CommandRequest {
command: "printf done".to_string(),
timeout_secs: 5,
output_limit: 1024,
cwd: None,
spill_dir: None,
tool_call_id: None,
})
.await
.unwrap();
let tool = WorkspaceHttpWorkdirTool {
backend,
operation: WorkdirOperation::Detach,
};
tool.execute(r#"{"alias":"checkout"}"#, ToolExecutionContext::default())
.await
.unwrap();
let requests = client.requests();
assert_eq!(requests.len(), 4);
assert!(
requests[2]
.body
.as_deref()
.is_some_and(|body| body.contains("command_status"))
);
}
#[tokio::test]
async fn successful_attach_reopens_internal_subworker_admission() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![response(json!({
"workspace_id": "workspace/test",
"workdir_id": "wd-attached",
"alias": "checkout",
"working_directory_id": "wd-attached",
"attached": true
}))]));
let reopen_calls = Arc::new(AtomicUsize::new(0));
@@ -1432,7 +1639,7 @@ mod tests {
};
tool.execute(
r#"{"workdir_id":"wd-attached"}"#,
r#"{"alias":"checkout","working_directory_id":"wd-attached"}"#,
ToolExecutionContext::default(),
)
.await
+4 -4
View File
@@ -165,7 +165,7 @@ where
})?;
worker.disable_manifest_lifecycle_features();
if let Some(session) = inherited_workdir_session {
worker.bind_workdir_session(Some(session));
worker.bind_single_workdir_session(Some(session));
}
let install_report = worker.install_features(features);
@@ -608,7 +608,7 @@ pub(crate) async fn spawn_internal_worker_session(
})?;
worker.disable_manifest_lifecycle_features();
if let Some(session) = inherited_workdir_session {
worker.bind_workdir_session(Some(session));
worker.bind_single_workdir_session(Some(session));
}
let install_report = worker.install_features(features);
let installed_tools = install_report.installed_tool_names();
@@ -701,7 +701,7 @@ pub(crate) fn prepare_internal_worker_from_spec(
})?;
worker.disable_manifest_lifecycle_features();
if let Some(session) = inherited_workdir_session {
worker.bind_workdir_session(Some(session));
worker.bind_single_workdir_session(Some(session));
}
let install_report = worker.install_features(features);
let installed_tools = install_report.installed_tool_names();
@@ -782,7 +782,7 @@ pub(crate) async fn prepare_internal_worker_session(
if let Some(broker) = command_event_broker.as_ref() {
wire_workdir_command_events(&broker.tool_session(), &in_flight);
} else if let Some(session) = worker.workdir_session() {
wire_workdir_command_events(session, &in_flight);
wire_workdir_command_events(&session, &in_flight);
}
let actor_in_flight = in_flight.clone();
worker.attach_alerter(alerter.clone());
+64 -31
View File
@@ -374,7 +374,8 @@ use protocol::{
use tokio::net::UnixStream;
use tokio::sync::broadcast;
use workdir::{
LocalWorkdirSession, ReadOnlyWorkdirSession, WorkdirSessionCapabilities, WorkdirSessionHandle,
LocalWorkdirSession, ReadOnlyWorkdirSession, WorkdirAttachmentAlias,
WorkdirSessionCapabilities, WorkdirSessionHandle, WorkdirSessionRouter,
};
const RESTORE_RECONCILIATION_REACHABILITY_TIMEOUT: Duration = Duration::from_millis(500);
@@ -2128,9 +2129,9 @@ pub struct Worker<C: LlmClient, St: Store> {
/// Explicit local filesystem authority, or `None` for Workers with no
/// local cwd and no filesystem/Bash tool surface.
filesystem_authority: WorkerFilesystemAuthority,
/// Live WorkdirSession provider derived once from the Worker–Workdir binding.
/// Local tools, file views, and compaction workers clone this handle.
workdir_session: Option<WorkdirSessionHandle>,
/// Alias-keyed live Workdir sessions. Provider transport remains hidden
/// behind each handle and no attachment is designated primary.
workdir_sessions: Arc<WorkdirSessionRouter>,
/// Path-free workspace identity/client context injected by Runtime/host.
/// This never grants local filesystem authority.
workspace_context: WorkerWorkspaceContext,
@@ -2464,7 +2465,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
let prompts = Arc::new(ArcSwap::from(PromptCatalog::builtins_only()?));
DelegationScope::from_config(&manifest.delegation_scope).map_err(WorkerError::Scope)?;
let scope = SharedScope::new(scope);
let workdir_session = workdir_session_from_authority(&filesystem_authority, &scope);
let workdir_sessions = workdir_sessions_from_authority(&filesystem_authority, &scope);
let mut worker = Self {
manifest,
engine: Some(worker),
@@ -2475,7 +2476,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
worker_metadata_segment_cas: None,
segment_state: SegmentState::new(session_id, segment_id, 0),
filesystem_authority,
workdir_session,
workdir_sessions,
workspace_context,
flow_runtime_state: Arc::new(Mutex::new(None)),
flow_feature_enabled: false,
@@ -2652,15 +2653,36 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
self.filesystem_authority.as_local()
}
pub fn workdir_session(&self) -> Option<&WorkdirSessionHandle> {
self.workdir_session.as_ref()
/// Compatibility projection for tool schemas that do not yet carry a
/// target alias. It is available only when the attachment set has exactly
/// one member, so multiple attachments never acquire an implicit primary.
pub fn workdir_session(&self) -> Option<WorkdirSessionHandle> {
self.workdir_sessions.only_session()
}
/// Replace the constructor fallback with the provider binding resolved by
/// the owning Runtime. Runtime calls this before the Worker controller is
/// spawned, so tools only ever observe the Runtime-bound handle.
pub fn bind_workdir_session(&mut self, workdir_session: Option<WorkdirSessionHandle>) {
self.workdir_session = workdir_session;
pub fn workdir_sessions(&self) -> Arc<WorkdirSessionRouter> {
self.workdir_sessions.clone()
}
/// Replace the constructor fallback with the complete alias-keyed provider
/// set resolved by the owning Runtime before controller startup.
pub fn bind_workdir_sessions(&mut self, workdir_sessions: Arc<WorkdirSessionRouter>) {
self.workdir_sessions = workdir_sessions;
}
/// Bind one session for direct/Internal Worker construction. Runtime Worker
/// creation uses [`Self::bind_workdir_sessions`] with the complete set.
pub fn bind_single_workdir_session(&mut self, workdir_session: Option<WorkdirSessionHandle>) {
let router = Arc::new(WorkdirSessionRouter::new());
if let Some(session) = workdir_session {
router
.attach(
WorkdirAttachmentAlias::new("workdir").expect("static alias is valid"),
session,
)
.expect("fresh Workdir router accepts its only attachment");
}
self.workdir_sessions = router;
}
/// Path-free workspace identity, if Runtime/host associated this Worker
@@ -2929,8 +2951,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
location.segment_id,
);
let read_only_tools = self
.workdir_session
.clone()
.workdir_session()
.map(|source| Arc::new(ReadOnlyWorkdirSession::new(source)) as WorkdirSessionHandle)
.map(tools::read_only_builtin_tools)
.unwrap_or_default();
@@ -4061,7 +4082,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// unresolved placeholder stays in the flattened user message so the LLM
/// still sees the intent.
async fn resolve_file_refs(&self, segments: &[Segment]) -> Vec<SystemItem> {
let Some(workdir) = self.workdir_session.clone() else {
let Some(workdir) = self.workdir_session() else {
for seg in segments {
if let Segment::FileRef { path } = seg {
self.alert(
@@ -5428,7 +5449,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
// Build a normal parent-owned Internal Worker over a pinned immutable
// capture. Only SessionExplore and compaction output tools are installed.
let workdir = self.workdir_session.clone();
let workdir = self.workdir_session();
let read_only_workdir = workdir.clone().map(|session| {
Arc::new(ReadOnlyWorkdirSession::new(session)) as workdir::WorkdirSessionHandle
});
@@ -6139,7 +6160,8 @@ where
let worker_metadata_writer = Some(worker_metadata_writer_for_store(&store));
let worker_metadata_segment_cas = Some(worker_metadata_segment_cas_for_store(&store));
let scope = SharedScope::new(common.scope);
let workdir_session = workdir_session_from_authority(&common.filesystem_authority, &scope);
let workdir_sessions =
workdir_sessions_from_authority(&common.filesystem_authority, &scope);
let mut worker = Self {
manifest,
@@ -6151,7 +6173,7 @@ where
worker_metadata_segment_cas,
segment_state: SegmentState::new(session_id, segment_id, 0),
filesystem_authority: common.filesystem_authority,
workdir_session,
workdir_sessions,
workspace_context: common.workspace_context,
flow_runtime_state: Arc::new(Mutex::new(None)),
flow_feature_enabled: false,
@@ -6228,7 +6250,8 @@ where
apply_worker_manifest(&mut engine, &manifest.engine);
engine.set_cache_key(Some(segment_id.to_string()));
let scope = SharedScope::new(common.scope);
let workdir_session = workdir_session_from_authority(&common.filesystem_authority, &scope);
let workdir_sessions =
workdir_sessions_from_authority(&common.filesystem_authority, &scope);
let mut worker = Self {
manifest,
engine: Some(engine),
@@ -6239,7 +6262,7 @@ where
worker_metadata_segment_cas: None,
segment_state: SegmentState::new(session_id, segment_id, 0),
filesystem_authority: common.filesystem_authority,
workdir_session,
workdir_sessions,
workspace_context: common.workspace_context,
flow_runtime_state: Arc::new(Mutex::new(None)),
flow_feature_enabled: false,
@@ -6350,7 +6373,8 @@ where
let worker_metadata_writer = Some(worker_metadata_writer_for_store(&store));
let worker_metadata_segment_cas = Some(worker_metadata_segment_cas_for_store(&store));
let scope = SharedScope::new(common.scope);
let workdir_session = workdir_session_from_authority(&common.filesystem_authority, &scope);
let workdir_sessions =
workdir_sessions_from_authority(&common.filesystem_authority, &scope);
let mut worker = Self {
manifest,
@@ -6362,7 +6386,7 @@ where
worker_metadata_segment_cas,
segment_state: SegmentState::new(session_id, segment_id, 0),
filesystem_authority: common.filesystem_authority,
workdir_session,
workdir_sessions,
workspace_context: common.workspace_context,
flow_runtime_state: Arc::new(Mutex::new(None)),
flow_feature_enabled: false,
@@ -6726,7 +6750,8 @@ where
let worker_metadata_writer = Some(worker_metadata_writer_for_store(&store));
let worker_metadata_segment_cas = Some(worker_metadata_segment_cas_for_store(&store));
let scope = SharedScope::new(common.scope);
let workdir_session = workdir_session_from_authority(&common.filesystem_authority, &scope);
let workdir_sessions =
workdir_sessions_from_authority(&common.filesystem_authority, &scope);
let mut worker = Self {
manifest,
@@ -6738,7 +6763,7 @@ where
worker_metadata_segment_cas,
segment_state: SegmentState::new(session_id, segment_id, state.entries_count),
filesystem_authority: common.filesystem_authority,
workdir_session,
workdir_sessions,
workspace_context: common.workspace_context,
flow_runtime_state: Arc::new(Mutex::new(restored_flow_runtime_state(
&state.extensions,
@@ -7580,18 +7605,26 @@ pub enum WorkerError {
},
}
fn workdir_session_from_authority(
fn workdir_sessions_from_authority(
authority: &WorkerFilesystemAuthority,
scope: &SharedScope,
) -> Option<WorkdirSessionHandle> {
authority.as_local().map(|local| {
Arc::new(LocalWorkdirSession::materialized(
) -> Arc<WorkdirSessionRouter> {
let router = Arc::new(WorkdirSessionRouter::new());
if let Some(local) = authority.as_local() {
let session = Arc::new(LocalWorkdirSession::materialized(
local.root.clone(),
local.cwd.clone(),
scope.clone(),
WorkdirSessionCapabilities::ALL,
)) as WorkdirSessionHandle
})
)) as WorkdirSessionHandle;
router
.attach(
WorkdirAttachmentAlias::new("workdir").expect("static alias is valid"),
session,
)
.expect("fresh Workdir router accepts its only attachment");
}
router
}
/// Bundle of resources that every high-level Worker constructor needs:
+4 -4
View File
@@ -349,7 +349,7 @@ async fn shutdown_closes_bound_workdir_session() {
})
.await
.unwrap();
worker.bind_workdir_session(Some(Arc::clone(&session)));
worker.bind_single_workdir_session(Some(Arc::clone(&session)));
let runtime_base = tempfile::tempdir().unwrap();
let bash_output_dir = runtime_base.path().join("bash-output");
@@ -384,7 +384,7 @@ async fn controller_projects_workdir_command_events_and_snapshot_state() {
worker.scope().clone(),
WorkdirSessionCapabilities::ALL,
));
worker.bind_workdir_session(Some(Arc::clone(&session)));
worker.bind_single_workdir_session(Some(Arc::clone(&session)));
let handle = spawn_controller(worker).await;
let mut events = handle.subscribe();
@@ -495,7 +495,7 @@ async fn controller_refreshes_command_snapshot_after_high_output_provider_lag()
worker.scope().clone(),
WorkdirSessionCapabilities::ALL,
));
worker.bind_workdir_session(Some(Arc::clone(&session)));
worker.bind_single_workdir_session(Some(Arc::clone(&session)));
let handle = spawn_controller(worker).await;
// Local command telemetry uses 8 KiB chunks and a 256-event channel. One
@@ -572,7 +572,7 @@ async fn controller_startup_failure_closes_bound_workdir_session() {
worker.scope().clone(),
WorkdirSessionCapabilities::ALL,
));
worker.bind_workdir_session(Some(Arc::clone(&session)));
worker.bind_single_workdir_session(Some(Arc::clone(&session)));
let runtime_base = tempfile::tempdir().unwrap();
let invalid_runtime_base = runtime_base.path().join("not-a-directory");
std::fs::write(&invalid_runtime_base, "file").unwrap();