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

fix: complete multi-workdir isolation and launch UI

Browse Source
This commit is contained in:
Hare
2026-09-20 22:07:59 +09:00
parent c102119412
commit 74fbadb1cd
11 changed files with 804 additions and 277 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/tools/src/edit.rs
+3 -1
View File
@@ -52,7 +52,9 @@ impl Tool for EditTool {
params.target_workdir.as_deref(),
workdir::WorkdirSessionCapability::Edit,
)?;
let tracker = self.tracker.scoped(selected.alias.as_str());
let tracker = self
.tracker
.scoped_attachment(&selected.alias, selected.generation);
let path = WorkdirPath::new(&params.file_path).map_err(ToolsError::from)?;
tracing::debug!(path = %path, replace_all = params.replace_all, "Edit");
crates/tools/src/read.rs
+5 -3
View File
@@ -53,7 +53,9 @@ impl Tool for ReadTool {
params.target_workdir.as_deref(),
workdir::WorkdirSessionCapability::Read,
)?;
let tracker = self.tracker.scoped(selected.alias.as_str());
let tracker = self
.tracker
.scoped_attachment(&selected.alias, selected.generation);
let offset = params.offset.unwrap_or(0);
let limit = params.limit.unwrap_or(DEFAULT_LIMIT).max(1);
@@ -216,7 +218,7 @@ mod tests {
// History recorded
assert!(
tracker
.scoped("workdir")
.scoped_attachment(&workdir::WorkdirAttachmentAlias::new("workdir").unwrap(), 0,)
.expected_workdir_hash(&WorkdirPath::new("a.txt").unwrap())
.is_ok()
);
@@ -247,7 +249,7 @@ mod tests {
#[tokio::test]
async fn read_tool_missing_file() {
let (dir, fs, tracker) = setup();
let (_dir, fs, tracker) = setup();
let def = read_tool(fs, tracker);
let (_, tool) = def();
let input = serde_json::json!({
crates/tools/src/tracker.rs
+11 -2
View File
@@ -161,8 +161,9 @@ impl Tracker {
}
/// 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.
/// attachment incarnation (the Worker-local alias plus router generation),
/// 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(),
@@ -171,6 +172,14 @@ impl Tracker {
}
}
pub(crate) fn scoped_attachment(
&self,
alias: &workdir::WorkdirAttachmentAlias,
generation: u64,
) -> Self {
self.scoped(format!("{}#{generation}", alias.as_str()))
}
fn key(&self, path: PathBuf) -> TrackedPath {
TrackedPath {
namespace: self.namespace.clone(),
crates/tools/src/write.rs
+3 -1
View File
@@ -49,7 +49,9 @@ impl Tool for WriteTool {
params.target_workdir.as_deref(),
workdir::WorkdirSessionCapability::Write,
)?;
let tracker = self.tracker.scoped(selected.alias.as_str());
let tracker = self
.tracker
.scoped_attachment(&selected.alias, selected.generation);
let path = WorkdirPath::new(&params.file_path).map_err(ToolsError::from)?;
tracing::debug!(path = %path, bytes = params.content.len(), "Write");
crates/tools/tests/integration.rs
+60
View File
@@ -339,6 +339,66 @@ async fn routed_tools_observe_detach_after_registration() {
router.close_all().await.unwrap();
}
#[tokio::test]
async fn detach_and_reattach_same_alias_requires_a_fresh_read() {
let first = TempDir::new().unwrap();
let second = TempDir::new().unwrap();
let spill = TempDir::new().unwrap();
std::fs::write(first.path().join("same.txt"), "same bytes\n").unwrap();
std::fs::write(second.path().join("same.txt"), "same bytes\n").unwrap();
let router = Arc::new(WorkdirSessionRouter::new());
let alias = WorkdirAttachmentAlias::new("checkout").unwrap();
router
.attach(
alias.clone(),
Arc::new(LocalWorkdirSession::new(
scope_with_spill(first.path(), spill.path()),
first.path().to_path_buf(),
)),
)
.unwrap();
let reg = Registry::new(routed_builtin_tools(
router.clone(),
Tracker::new(),
spill.path().to_path_buf(),
));
call(
&reg.get("Read"),
json!({ "target_workdir": "checkout", "file_path": "same.txt" }),
)
.await;
router.detach(&alias).await.unwrap();
router
.attach(
alias.clone(),
Arc::new(LocalWorkdirSession::new(
scope_with_spill(second.path(), spill.path()),
second.path().to_path_buf(),
)),
)
.unwrap();
let error = call_err(
&reg.get("Edit"),
json!({
"target_workdir": "checkout",
"file_path": "same.txt",
"old_string": "same",
"new_string": "changed"
}),
)
.await;
assert!(error.to_string().contains("has not been read"), "{error}");
assert_eq!(
std::fs::read_to_string(second.path().join("same.txt")).unwrap(),
"same bytes\n"
);
router.close_all().await.unwrap();
}
#[tokio::test]
async fn view_image_reads_scoped_bytes_into_durable_tool_detail() {
let dir = TempDir::new().unwrap();
crates/workspace-server/src/server.rs
+163 -14
View File
@@ -383,6 +383,7 @@ static EMBEDDED_RUNTIME_REQUEST_IDENTITY: std::sync::LazyLock<
#[derive(Clone)]
struct WorkdirCommandSession {
attachment_alias: String,
source: WorkdirSessionHandle,
provider_handle: CommandHandle,
}
@@ -421,6 +422,7 @@ impl WorkdirSessionRegistry {
fn register_command(
&mut self,
worker: RuntimeWorkerRef,
attachment_alias: String,
source: WorkdirSessionHandle,
provider_handle: CommandHandle,
) -> CommandHandle {
@@ -436,6 +438,7 @@ impl WorkdirSessionRegistry {
self.commands.insert(
(worker, external_handle.clone()),
WorkdirCommandSession {
attachment_alias,
source,
provider_handle,
},
@@ -446,13 +449,42 @@ impl WorkdirSessionRegistry {
fn command(
&self,
worker: &RuntimeWorkerRef,
attachment_alias: &str,
external_handle: &CommandHandle,
) -> Option<WorkdirCommandSession> {
self.commands
.get(&(worker.clone(), external_handle.clone()))
.filter(|session| session.attachment_alias == attachment_alias)
.cloned()
}
fn take_worker_commands_for_alias(
&mut self,
worker: &RuntimeWorkerRef,
attachment_alias: &str,
) -> Vec<RegisteredWorkdirSession> {
let command_handles = self
.commands
.iter()
.filter(|((owner, _), session)| {
owner == worker && session.attachment_alias == attachment_alias
})
.map(|((_, handle), _)| handle.clone())
.collect::<Vec<_>>();
command_handles
.into_iter()
.filter_map(|external_handle| {
self.commands
.remove(&(worker.clone(), external_handle.clone()))
.map(|session| RegisteredWorkdirSession::Command {
worker: worker.clone(),
external_handle,
session,
})
})
.collect()
}
fn take_worker(&mut self, worker: &RuntimeWorkerRef) -> Vec<RegisteredWorkdirSession> {
let mut sessions = Vec::new();
if let Some(session) = self.attachments.remove(worker) {
@@ -498,14 +530,10 @@ impl WorkdirSessionRegistry {
}
}
async fn close_worker_workdir_sessions(
async fn close_registered_workdir_sessions(
registry: &Arc<Mutex<WorkdirSessionRegistry>>,
worker: &RuntimeWorkerRef,
registered: Vec<RegisteredWorkdirSession>,
) -> std::result::Result<(), String> {
let registered = registry
.lock()
.map_err(|_| "Workdir session registry was poisoned".to_string())?
.take_worker(worker);
let mut failed = Vec::new();
let mut errors = Vec::new();
for item in registered {
@@ -533,6 +561,29 @@ async fn close_worker_workdir_sessions(
}
}
async fn close_worker_workdir_sessions(
registry: &Arc<Mutex<WorkdirSessionRegistry>>,
worker: &RuntimeWorkerRef,
) -> std::result::Result<(), String> {
let registered = registry
.lock()
.map_err(|_| "Workdir session registry was poisoned".to_string())?
.take_worker(worker);
close_registered_workdir_sessions(registry, registered).await
}
async fn close_worker_workdir_command_sessions_for_alias(
registry: &Arc<Mutex<WorkdirSessionRegistry>>,
worker: &RuntimeWorkerRef,
attachment_alias: &str,
) -> std::result::Result<(), String> {
let registered = registry
.lock()
.map_err(|_| "Workdir session registry was poisoned".to_string())?
.take_worker_commands_for_alias(worker, attachment_alias);
close_registered_workdir_sessions(registry, registered).await
}
const ATTACHMENT_UPLOAD_GRANT_TTL_MS: u64 = 5 * 60 * 1_000;
#[derive(Clone)]
@@ -8529,6 +8580,20 @@ async fn close_current_worker_attachment_session_locked(
Ok(())
}
async fn close_current_worker_command_sessions_for_alias_locked(
api: &WorkspaceApi,
worker: &RuntimeWorkerRef,
attachment_alias: &str,
) -> Result<()> {
close_worker_workdir_command_sessions_for_alias(&api.workdir_sessions, worker, attachment_alias)
.await
.map_err(|message| Error::RuntimeOperationFailed {
runtime_id: worker.runtime_id.clone(),
code: "workdir_session_close_failed".to_string(),
message,
})
}
async fn close_current_worker_session_locked(
api: &WorkspaceApi,
worker: &RuntimeWorkerRef,
@@ -8583,7 +8648,7 @@ async fn refresh_current_worker_session_locked(
api: &WorkspaceApi,
worker: &RuntimeWorkerRef,
) -> Result<()> {
close_current_worker_session_locked(api, worker).await?;
close_current_worker_attachment_session_locked(api, worker).await?;
let remaining = api
.store
.list_worker_workdir_links(&api.config.workspace_id, worker)?
@@ -8653,7 +8718,7 @@ async fn scoped_attach_current_worker_workdir(
attached: true,
}));
}
close_current_worker_session_locked(&api, &worker).await?;
close_current_worker_attachment_session_locked(&api, &worker).await?;
let link = api.store.attach_worker_workdir(&WorkerWorkdirLinkRecord {
workspace_id: api.config.workspace_id.clone(),
worker: worker.clone(),
@@ -8734,7 +8799,8 @@ async fn scoped_detach_current_worker_workdir(
worker.runtime_id, worker.worker_id
))
})?;
close_current_worker_session_locked(&api, &worker).await?;
close_current_worker_attachment_session_locked(&api, &worker).await?;
close_current_worker_command_sessions_for_alias_locked(&api, &worker, alias.as_str()).await?;
api.store.detach_worker_workdir(
&api.config.workspace_id,
&worker,
@@ -8903,7 +8969,12 @@ async fn scoped_execute_current_worker_workdir_operation(
.workdir_sessions
.lock()
.expect("Workdir session registry lock poisoned")
.register_command(worker.clone(), registered_source, provider_handle);
.register_command(
worker.clone(),
target_workdir.clone(),
registered_source,
provider_handle,
);
WorkdirSessionOperationResult::CommandStart(external_handle)
}
WorkdirSessionOperation::CommandStatus(external_handle) => {
@@ -8966,7 +9037,7 @@ fn current_worker_command_session(
.workdir_sessions
.lock()
.expect("Workdir session registry lock poisoned")
.command(worker, external_handle)
.command(worker, target_workdir, external_handle)
.ok_or_else(|| {
WorkdirOperationApiError::Provider(current_worker_workdir_operation_error(
worker,
@@ -20100,8 +20171,12 @@ mod tests {
let mut registry = WorkdirSessionRegistry::default();
registry.insert_attachment(worker.clone(), source.clone());
let registered_source = registry.remove_attachment(&worker).unwrap();
let external_handle =
registry.register_command(worker.clone(), registered_source, provider_handle.clone());
let external_handle = registry.register_command(
worker.clone(),
"checkout".to_string(),
registered_source,
provider_handle.clone(),
);
assert_ne!(external_handle, provider_handle);
let refreshed: WorkdirSessionHandle = Arc::new(workdir::LocalWorkdirSession::new(
@@ -20116,7 +20191,14 @@ mod tests {
.await
.unwrap();
let command = registry.command(&worker, &external_handle).unwrap();
let command = registry
.command(&worker, "checkout", &external_handle)
.unwrap();
assert!(
registry
.command(&worker, "other", &external_handle)
.is_none()
);
let output = command
.source
.command_output(workdir::CommandOutputRequest {
@@ -20176,6 +20258,73 @@ mod tests {
);
}
#[tokio::test]
async fn attachment_scoped_command_cleanup_preserves_other_aliases() {
let left_directory = tempfile::tempdir().unwrap();
let right_directory = tempfile::tempdir().unwrap();
let worker = RuntimeWorkerRef::new("runtime-command-alias", "worker-command-alias");
let make_session = |path: &std::path::Path| -> WorkdirSessionHandle {
Arc::new(workdir::LocalWorkdirSession::new(
manifest::Scope::writable(path).unwrap(),
path.to_path_buf(),
))
};
let left = make_session(left_directory.path());
let right = make_session(right_directory.path());
let request = || workdir::CommandRequest {
command: "sleep 30".to_string(),
timeout_secs: 60,
output_limit: 4096,
cwd: None,
spill_dir: None,
tool_call_id: None,
};
let left_provider = left.start_command(request()).await.unwrap();
let right_provider = right.start_command(request()).await.unwrap();
let registry = Arc::new(Mutex::new(WorkdirSessionRegistry::default()));
let (left_external, right_external) = {
let mut sessions = registry.lock().unwrap();
let left_external =
sessions.register_command(worker.clone(), "left".to_string(), left, left_provider);
let right_external = sessions.register_command(
worker.clone(),
"right".to_string(),
right,
right_provider,
);
assert!(sessions.command(&worker, "right", &left_external).is_none());
assert!(sessions.command(&worker, "left", &right_external).is_none());
(left_external, right_external)
};
close_worker_workdir_command_sessions_for_alias(&registry, &worker, "right")
.await
.unwrap();
let left_command = {
let sessions = registry.lock().unwrap();
assert!(
sessions
.command(&worker, "right", &right_external)
.is_none()
);
sessions.command(&worker, "left", &left_external).unwrap()
};
assert_eq!(
left_command
.source
.command_status(left_command.provider_handle.clone())
.await
.unwrap(),
workdir::CommandStatus::Running
);
close_worker_workdir_sessions(&registry, &worker)
.await
.unwrap();
}
fn seed_test_api_token(store: &dyn ControlPlaneStore, suffix: &str) -> String {
let account_id = format!("account-{suffix}");
let user_id = format!("user-{suffix}");