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/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}");