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

worker: unify workdir attachments

Browse Source
This commit is contained in:
Hare
2026-08-05 05:30:58 +09:00
parent ba009c0a20
commit fb115fbb7e
6 changed files with 1439 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+12 -1
View File
@@ -599,6 +599,18 @@ where
// Worker-immutable snapshots taken before the mutable worker borrow
// below so the worker borrow doesn't conflict with reads on `worker`.
let scope_handle = worker.scope().clone();
let feature_config = worker.manifest().feature.clone();
if feature_config.manage_workdir.enabled {
if let Some(existing) = worker.workdir_session().cloned() {
existing.close().await.map_err(std::io::Error::other)?;
}
let workspace_client = worker.workspace_client_handle();
worker.bind_workdir_session(Some(
crate::feature::builtin::manage_workdir::WorkspaceAttachedWorkdirSession::handle(
workspace_client,
),
));
}
let worker_workdir = worker.workdir_session().cloned();
let local_filesystem = worker.local_working_directory().cloned();
let local_workspace_root = local_filesystem.as_ref().map(|local| local.root.clone());
@@ -606,7 +618,6 @@ where
let memory_config = worker.manifest().memory.clone();
let web_config = worker.manifest().web.clone();
let mcp_config = worker.manifest().mcp.clone();
let feature_config = worker.manifest().feature.clone();
let spawner_name = worker.manifest().worker.name.clone();
let spawner_manifest = worker.manifest().clone();
let prompts = worker.prompts().clone();
crates/worker/src/feature/builtin/manage_workdir.rs
+366 -9
View File
@@ -13,6 +13,13 @@ use llm_engine::tool::{
};
use serde::{Deserialize, Serialize};
use serde_json::json;
use workdir::http::{WorkdirSessionOperation, WorkdirSessionOperationResult};
use workdir::{
CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandStatus, EditRequest,
EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult,
ReadRequest, ReadResult, StatRequest, StatResult, Workdir, WorkdirError, WorkdirSession,
WorkdirSessionCapabilities, WorkdirSessionHandle, WriteRequest, WriteResult,
};
use crate::feature::{
FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureModule, ToolContribution,
@@ -27,10 +34,14 @@ const FEATURE_DESCRIPTION: &str =
const LIST_TOOL: &str = "WorkdirList";
const CREATE_TOOL: &str = "WorkdirCreate";
const ATTACH_TOOL: &str = "WorkdirAttach";
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. Repository resolution and materialization remain Backend/Runtime authority.";
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 DELETE_DESCRIPTION: &str = "Delete one persistent Workdir by id through Backend Workspace API authority. Occupied, blocked, or dirty Workdirs requiring confirmation are rejected.";
#[derive(Clone, Debug)]
@@ -54,6 +65,8 @@ impl FeatureModule for ManageWorkdirFeature {
.with_description(FEATURE_DESCRIPTION)
.with_tool(ToolDeclaration::new(LIST_TOOL, LIST_DESCRIPTION))
.with_tool(ToolDeclaration::new(CREATE_TOOL, CREATE_DESCRIPTION))
.with_tool(ToolDeclaration::new(ATTACH_TOOL, ATTACH_DESCRIPTION))
.with_tool(ToolDeclaration::new(DETACH_TOOL, DETACH_DESCRIPTION))
.with_tool(ToolDeclaration::new(DELETE_TOOL, DELETE_DESCRIPTION))
}
@@ -80,6 +93,26 @@ impl FeatureModule for ManageWorkdirFeature {
WorkdirOperation::Create,
),
),
(
ATTACH_TOOL,
workdir_tool(
ATTACH_TOOL,
ATTACH_DESCRIPTION,
attach_schema(),
backend.clone(),
WorkdirOperation::Attach,
),
),
(
DETACH_TOOL,
workdir_tool(
DETACH_TOOL,
DETACH_DESCRIPTION,
detach_schema(),
backend.clone(),
WorkdirOperation::Detach,
),
),
(
DELETE_TOOL,
workdir_tool(
@@ -104,6 +137,164 @@ struct WorkspaceHttpWorkdirBackend {
client: Arc<dyn WorkspaceClient>,
}
/// Worker-local Workdir handle whose operation authority remains in the Workspace Backend.
///
/// The Backend resolves the caller from the identity-bound [`WorkspaceClient`] and routes each
/// operation to that Worker's active attachment. Runtime endpoints, credentials, and ephemeral
/// session ids remain outside model-visible tool contracts.
#[derive(Debug)]
pub struct WorkspaceAttachedWorkdirSession {
client: Arc<dyn WorkspaceClient>,
workdir: Workdir,
}
impl WorkspaceAttachedWorkdirSession {
pub fn handle(client: Arc<dyn WorkspaceClient>) -> WorkdirSessionHandle {
Arc::new(Self {
client,
workdir: Workdir::new("workspace-attachment"),
})
}
fn operate(
&self,
operation: WorkdirSessionOperation,
) -> Result<WorkdirSessionOperationResult, WorkdirError> {
let workspace_id = self.client.workspace_id().ok_or_else(|| {
WorkdirError::Unavailable("Workspace identity is unavailable".to_string())
})?;
let request = WorkspaceRequest::json(
WorkspaceRequestMethod::Post,
format!(
"/api/w/{}/workers/self/workdir-session/operations",
encode_path_segment(workspace_id)
),
serde_json::to_string(&operation).map_err(|error| {
WorkdirError::Unavailable(format!(
"failed to encode Workspace Workdir operation: {error}"
))
})?,
);
let response = self
.client
.execute(request)
.map_err(|error| WorkdirError::Unavailable(error.to_string()))?;
if !response.is_success() {
return Err(WorkdirError::Unavailable(format!(
"Workspace Workdir API returned HTTP {}: {}",
response.status,
bounded_error_body(&response.body)
)));
}
serde_json::from_str(&response.body).map_err(|error| {
WorkdirError::Unavailable(format!(
"failed to decode Workspace Workdir operation result: {error}"
))
})
}
fn mismatch(expected: &str) -> WorkdirError {
WorkdirError::Unavailable(format!(
"Workspace Backend returned a mismatched Workdir operation result; expected {expected}"
))
}
}
#[async_trait]
impl WorkdirSession for WorkspaceAttachedWorkdirSession {
fn workdir(&self) -> &Workdir {
&self.workdir
}
fn capabilities(&self) -> WorkdirSessionCapabilities {
WorkdirSessionCapabilities::ALL
}
async fn stat(&self, request: StatRequest) -> Result<StatResult, WorkdirError> {
match self.operate(WorkdirSessionOperation::Stat(request))? {
WorkdirSessionOperationResult::Stat(result) => Ok(result),
_ => Err(Self::mismatch("stat")),
}
}
async fn read(&self, request: ReadRequest) -> Result<ReadResult, WorkdirError> {
match self.operate(WorkdirSessionOperation::Read(request))? {
WorkdirSessionOperationResult::Read(result) => Ok(result),
_ => Err(Self::mismatch("read")),
}
}
async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError> {
match self.operate(WorkdirSessionOperation::Write(request))? {
WorkdirSessionOperationResult::Write(result) => Ok(result),
_ => Err(Self::mismatch("write")),
}
}
async fn edit(&self, request: EditRequest) -> Result<EditResult, WorkdirError> {
match self.operate(WorkdirSessionOperation::Edit(request))? {
WorkdirSessionOperationResult::Edit(result) => Ok(result),
_ => Err(Self::mismatch("edit")),
}
}
async fn list(&self, request: ListRequest) -> Result<ListResult, WorkdirError> {
match self.operate(WorkdirSessionOperation::List(request))? {
WorkdirSessionOperationResult::List(result) => Ok(result),
_ => Err(Self::mismatch("list")),
}
}
async fn glob(&self, request: GlobRequest) -> Result<GlobResult, WorkdirError> {
match self.operate(WorkdirSessionOperation::Glob(request))? {
WorkdirSessionOperationResult::Glob(result) => Ok(result),
_ => Err(Self::mismatch("glob")),
}
}
async fn grep(&self, request: GrepRequest) -> Result<GrepResult, WorkdirError> {
match self.operate(WorkdirSessionOperation::Grep(request))? {
WorkdirSessionOperationResult::Grep(result) => Ok(result),
_ => Err(Self::mismatch("grep")),
}
}
async fn start_command(&self, request: CommandRequest) -> Result<CommandHandle, WorkdirError> {
match self.operate(WorkdirSessionOperation::CommandStart(request))? {
WorkdirSessionOperationResult::CommandStart(result) => Ok(result),
_ => Err(Self::mismatch("command_start")),
}
}
async fn command_status(&self, handle: CommandHandle) -> Result<CommandStatus, WorkdirError> {
match self.operate(WorkdirSessionOperation::CommandStatus(handle))? {
WorkdirSessionOperationResult::CommandStatus(result) => Ok(result),
_ => Err(Self::mismatch("command_status")),
}
}
async fn command_output(
&self,
request: CommandOutputRequest,
) -> Result<CommandOutput, WorkdirError> {
match self.operate(WorkdirSessionOperation::CommandOutput(request))? {
WorkdirSessionOperationResult::CommandOutput(result) => Ok(result),
_ => Err(Self::mismatch("command_output")),
}
}
async fn cancel_command(&self, handle: CommandHandle) -> Result<(), WorkdirError> {
match self.operate(WorkdirSessionOperation::CommandCancel(handle))? {
WorkdirSessionOperationResult::CommandCancel => Ok(()),
_ => Err(Self::mismatch("command_cancel")),
}
}
async fn close(&self) -> Result<(), WorkdirError> {
Ok(())
}
}
impl WorkspaceHttpWorkdirBackend {
fn new(client: Arc<dyn WorkspaceClient>) -> Self {
Self { client }
@@ -154,6 +345,37 @@ 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)
}
fn attach_response(&self, workdir_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"),
serde_json::to_string(&WorkdirAttachRequest {
workdir_id: workdir_id.to_string(),
})
.map_err(decode_error)?,
))
}
fn detach(&self) -> Result<ToolOutput, ToolError> {
let workspace_id = encode_path_segment(self.workspace_id()?);
let response = self.execute_json::<WorkdirAttachmentResponse>(WorkspaceRequest {
method: WorkspaceRequestMethod::Delete,
path: format!("/api/w/{workspace_id}/workers/self/workdir-attachment"),
body: None,
})?;
workdir_output(
format!("Detached from Workdir {}", response.workdir_id),
&response,
)
}
fn delete(&self, input: WorkdirDeleteInput) -> Result<ToolOutput, ToolError> {
let workdir_id = validate_identity(
&input.working_directory_id,
@@ -185,6 +407,8 @@ impl WorkspaceHttpWorkdirBackend {
enum WorkdirOperation {
List,
Create,
Attach,
Detach,
Delete,
}
@@ -229,6 +453,13 @@ impl Tool for WorkspaceHttpWorkdirTool {
WorkdirOperation::Create => self
.backend
.create(parse_input::<WorkdirCreateInput>(input_json)?),
WorkdirOperation::Attach => self
.backend
.attach(parse_input::<WorkdirAttachInput>(input_json)?),
WorkdirOperation::Detach => {
let _input = parse_input::<WorkdirDetachInput>(input_json)?;
self.backend.detach()
}
WorkdirOperation::Delete => self
.backend
.delete(parse_input::<WorkdirDeleteInput>(input_json)?),
@@ -337,6 +568,21 @@ fn create_schema() -> serde_json::Value {
})
}
fn attach_schema() -> serde_json::Value {
json!({
"type": "object",
"additionalProperties": false,
"required": ["workdir_id"],
"properties": {
"workdir_id": {"type": "string", "minLength": 1}
}
})
}
fn detach_schema() -> serde_json::Value {
list_schema()
}
fn delete_schema() -> serde_json::Value {
json!({
"type": "object",
@@ -369,6 +615,28 @@ struct WorkdirCreateRequest {
selector: Option<String>,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct WorkdirAttachInput {
workdir_id: String,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct WorkdirDetachInput {}
#[derive(Debug, Serialize)]
struct WorkdirAttachRequest {
workdir_id: String,
}
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
struct WorkdirAttachmentResponse {
workspace_id: String,
workdir_id: String,
attached: bool,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct WorkdirDeleteInput {
@@ -407,8 +675,12 @@ struct WorkdirSummary {
status: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
cleanliness: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
primary_worker_id: Option<u64>,
#[serde(
default,
alias = "primary_worker_id",
skip_serializing_if = "Option::is_none"
)]
attached_worker_id: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
occupied_by: Option<WorkdirOccupancy>,
}
@@ -525,7 +797,16 @@ mod tests {
.collect::<Vec<_>>();
assert_eq!(descriptor.id.as_str(), "builtin:manage-workdir");
assert_eq!(names, [LIST_TOOL, CREATE_TOOL, DELETE_TOOL]);
assert_eq!(
names,
[
LIST_TOOL,
CREATE_TOOL,
ATTACH_TOOL,
DETACH_TOOL,
DELETE_TOOL
]
);
}
#[test]
@@ -544,9 +825,24 @@ mod tests {
assert!(report.reports[0].installed);
assert_eq!(
report.installed_tool_names(),
[LIST_TOOL, CREATE_TOOL, DELETE_TOOL]
[
LIST_TOOL,
CREATE_TOOL,
ATTACH_TOOL,
DETACH_TOOL,
DELETE_TOOL
]
);
assert_eq!(
names,
[
LIST_TOOL,
CREATE_TOOL,
ATTACH_TOOL,
DETACH_TOOL,
DELETE_TOOL
]
);
assert_eq!(names, [LIST_TOOL, CREATE_TOOL, DELETE_TOOL]);
}
#[test]
@@ -587,11 +883,13 @@ mod tests {
assert_eq!(create["required"], json!(["runtime_id", "repository_id"]));
assert!(create["properties"].get("path").is_none());
assert!(create["properties"].get("session_id").is_none());
assert_eq!(attach_schema()["required"], json!(["workdir_id"]));
assert!(attach_schema()["properties"].get("session_id").is_none());
assert_eq!(delete_schema()["required"], json!(["working_directory_id"]));
}
#[test]
fn list_create_and_delete_use_scoped_workspace_authority_paths() {
fn list_create_explicit_attach_detach_and_delete_use_scoped_workspace_authority_paths() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![
response(json!({
"workspace_id": "workspace/test",
@@ -603,6 +901,16 @@ mod tests {
"item": workdir_json("wd-created"),
"diagnostics": []
})),
response(json!({
"workspace_id": "workspace/test",
"workdir_id": "wd-created",
"attached": true
})),
response(json!({
"workspace_id": "workspace/test",
"workdir_id": "wd-created",
"attached": false
})),
response(json!({
"workspace_id": "workspace/test",
"item": {
@@ -617,13 +925,25 @@ mod tests {
let backend = WorkspaceHttpWorkdirBackend::new(client.clone());
backend.list().unwrap();
backend
let created = backend
.create(WorkdirCreateInput {
runtime_id: "runtime/one".to_string(),
repository_id: "main".to_string(),
selector: Some("refs/heads/topic".to_string()),
})
.unwrap();
assert_eq!(created.summary, "Created Workdir wd-created");
let created: serde_json::Value =
serde_json::from_str(created.content.as_deref().unwrap()).unwrap();
assert_eq!(created["item"]["working_directory_id"], "wd-created");
assert!(created.get("attachment").is_none());
assert_eq!(client.requests().len(), 2);
backend
.attach(WorkdirAttachInput {
workdir_id: "wd-created".to_string(),
})
.unwrap();
backend.detach().unwrap();
backend
.delete(WorkdirDeleteInput {
working_directory_id: "wd-created".to_string(),
@@ -647,9 +967,46 @@ mod tests {
assert_eq!(body["selector"], "refs/heads/topic");
assert_eq!(
requests[2].path,
"/api/w/workspace%2Ftest/workers/self/workdir-attachment"
);
assert_eq!(requests[2].method, WorkspaceRequestMethod::Post);
assert_eq!(
requests[3].path,
"/api/w/workspace%2Ftest/workers/self/workdir-attachment"
);
assert_eq!(requests[3].method, WorkspaceRequestMethod::Delete);
assert_eq!(
requests[4].path,
"/api/w/workspace%2Ftest/working-directories/wd-created"
);
assert_eq!(requests[2].method, WorkspaceRequestMethod::Delete);
assert_eq!(requests[4].method, WorkspaceRequestMethod::Delete);
}
#[tokio::test]
async fn attached_session_proxies_operations_without_runtime_or_session_arguments() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![response(json!({
"operation": "stat",
"result": {"path": "visible.txt", "kind": "file", "size": 8}
}))]));
let session = WorkspaceAttachedWorkdirSession::handle(client.clone());
let result = session
.stat(StatRequest {
path: workdir::WorkdirPath::new("visible.txt").unwrap(),
})
.await
.unwrap();
assert_eq!(result.size, 8);
let requests = client.requests();
assert_eq!(requests.len(), 1);
assert_eq!(
requests[0].path,
"/api/w/workspace%2Ftest/workers/self/workdir-session/operations"
);
let body: serde_json::Value =
serde_json::from_str(requests[0].body.as_deref().unwrap()).unwrap();
assert_eq!(body["operation"], "stat");
assert!(body.get("runtime_id").is_none());
assert!(body.get("session_id").is_none());
}
#[test]
crates/workspace-server/src/hosts.rs
+21 -1
View File
@@ -16,7 +16,7 @@ use std::{
time::Duration,
};
use workdir::{
Workdir, WorkdirError,
Workdir, WorkdirError, WorkdirSessionHandle,
http::{OpenWorkdirSessionRequest, RemoteWorkdirSession, WorkdirHttpAuthorization},
};
use worker_runtime::RuntimeWorkspaceScope;
@@ -1177,6 +1177,26 @@ impl RuntimeRegistry {
Ok(runtime.working_directory(working_directory_id))
}
pub async fn open_workdir_session(
&self,
runtime_id: &str,
working_directory_id: &str,
owner_worker_id: &str,
) -> Result<WorkdirSessionHandle, RuntimeRegistryError> {
validate_backend_identifier("runtime_id", runtime_id)?;
validate_backend_identifier("working_directory_id", working_directory_id)?;
validate_backend_identifier("owner_worker_id", owner_worker_id)?;
let runtime = self.runtime(runtime_id)?;
runtime
.open_workdir_session(working_directory_id, Some(owner_worker_id))
.await
.map_err(|error| RuntimeRegistryError::RuntimeOperationFailed {
runtime_id: runtime_id.to_string(),
code: "workdir_session_open_failed".to_string(),
message: error.to_string(),
})
}
pub fn cleanup_working_directory(
&self,
runtime_id: &str,
crates/workspace-server/src/lib.rs
+2
View File
@@ -91,6 +91,8 @@ pub enum Error {
WorkspaceIdMismatch,
#[error("Ticket assignment conflict: {0}")]
TicketAssignmentConflict(String),
#[error("Workdir attachment conflict: {0}")]
WorkdirAttachmentConflict(String),
#[error("Worker source identity is invalid: {0}")]
WorkerSourceIdentity(String),
#[error("workspace identity error: {0}")]
crates/workspace-server/src/server.rs
+545 -24
View File
@@ -39,6 +39,8 @@ use webauthn_rs::prelude::{
PublicKeyCredential, RegisterPublicKeyCredential, RequestChallengeResponse, Webauthn,
WebauthnBuilder,
};
use workdir::WorkdirSessionHandle;
use workdir::http::{WorkdirSessionOperation, WorkdirSessionOperationResult};
use worker_runtime::resource::{BackendResourceError, BackendResourceFetchRequest};
use worker_runtime::worker_backend::{ProfileRuntimeWorkerFactory, WorkerRuntimeExecutionBackend};
@@ -257,6 +259,8 @@ pub struct WorkspaceApi {
observation_proxy: BackendObservationProxy,
runtime_subscription_broker: RuntimeSubscriptionBroker,
resource_broker: BackendResourceBroker,
workdir_sessions: Arc<Mutex<HashMap<(String, u64), WorkdirSessionHandle>>>,
workdir_session_locks: Arc<Mutex<HashMap<(String, u64), Arc<tokio::sync::Mutex<()>>>>>,
}
impl WorkspaceApi {
@@ -358,6 +362,8 @@ impl WorkspaceApi {
observation_proxy,
runtime_subscription_broker,
resource_broker,
workdir_sessions: Arc::new(Mutex::new(HashMap::new())),
workdir_session_locks: Arc::new(Mutex::new(HashMap::new())),
})
}
@@ -390,11 +396,45 @@ impl WorkspaceApi {
) -> ApiResult<WorkerSpawnResult> {
let workspace_api = self.workspace_api_ref(runtime_id);
request.resolved_workspace_api = Some(workspace_api.clone());
let result = self
.runtime
.spawn_worker(runtime_id, request)
.map_err(|error| error.into_error())?;
let attachment_reservation =
request
.resolved_working_directory
.as_ref()
.map(|working_directory| {
(
working_directory.working_directory_id.clone(),
Uuid::new_v4().to_string(),
)
});
if let Some((workdir_id, reservation_id)) = attachment_reservation.as_ref() {
self.store.reserve_worker_workdir_attachment(
&self.config.workspace_id,
workdir_id,
reservation_id,
&now_registry_timestamp(),
)?;
}
let result = match self.runtime.spawn_worker(runtime_id, request) {
Ok(result) => result,
Err(error) => {
if let Some((workdir_id, reservation_id)) = attachment_reservation.as_ref() {
let _ = self.store.release_worker_workdir_attachment_reservation(
&self.config.workspace_id,
workdir_id,
reservation_id,
);
}
return Err(error.into_error().into());
}
};
let Some(worker) = result.worker.as_ref() else {
if let Some((workdir_id, reservation_id)) = attachment_reservation.as_ref() {
self.store.release_worker_workdir_attachment_reservation(
&self.config.workspace_id,
workdir_id,
reservation_id,
)?;
}
return Ok(result);
};
let replacement = match self.runtime.replace_worker_workspace_api(
@@ -405,11 +445,25 @@ impl WorkspaceApi {
Ok(replacement) => replacement,
Err(error) => {
let _ = self.runtime.delete_worker(runtime_id, &worker.worker_id);
if let Some((workdir_id, reservation_id)) = attachment_reservation.as_ref() {
let _ = self.store.release_worker_workdir_attachment_reservation(
&self.config.workspace_id,
workdir_id,
reservation_id,
);
}
return Err(error.into_error().into());
}
};
if replacement.state != WorkerOperationState::Accepted {
let _ = self.runtime.delete_worker(runtime_id, &worker.worker_id);
if let Some((workdir_id, reservation_id)) = attachment_reservation.as_ref() {
let _ = self.store.release_worker_workdir_attachment_reservation(
&self.config.workspace_id,
workdir_id,
reservation_id,
);
}
return Err(Error::RuntimeOperationFailed {
runtime_id: runtime_id.to_string(),
code: "worker_workspace_api_replace_failed".to_string(),
@@ -423,6 +477,41 @@ impl WorkspaceApi {
}
.into());
}
if let Some((workdir_id, reservation_id)) = attachment_reservation.as_ref() {
let runtime_worker_id = match parse_runtime_worker_id_for_registry(&worker.worker_id) {
Ok(worker_id) => worker_id,
Err(error) => {
let _ = self.runtime.delete_worker(runtime_id, &worker.worker_id);
let _ = self.store.release_worker_workdir_attachment_reservation(
&self.config.workspace_id,
workdir_id,
reservation_id,
);
return Err(error);
}
};
let attachment = WorkerWorkdirLinkRecord {
workspace_id: self.config.workspace_id.clone(),
runtime_id: runtime_id.to_string(),
runtime_worker_id,
workdir_id: workdir_id.clone(),
role: "attachment".to_string(),
linked_at: now_registry_timestamp(),
unlinked_at: None,
};
if let Err(error) = self
.store
.finalize_reserved_worker_workdir_attachment(&attachment, reservation_id)
{
let _ = self.runtime.delete_worker(runtime_id, &worker.worker_id);
let _ = self.store.release_worker_workdir_attachment_reservation(
&self.config.workspace_id,
workdir_id,
reservation_id,
);
return Err(error.into());
}
}
Ok(result)
}
@@ -772,6 +861,15 @@ pub fn build_router(api: WorkspaceApi) -> Router {
"/api/w/{workspace_id}/runtimes/{runtime_id}/working-directories/{working_directory_id}",
get(scoped_runtime_working_directory_detail).delete(scoped_cleanup_runtime_working_directory),
)
.route(
"/api/w/{workspace_id}/workers/self/workdir-attachment",
post(scoped_attach_current_worker_workdir)
.delete(scoped_detach_current_worker_workdir),
)
.route(
"/api/w/{workspace_id}/workers/self/workdir-session/operations",
post(scoped_execute_current_worker_workdir_operation),
)
.route(
"/api/w/{workspace_id}/working-directories",
get(scoped_list_working_directories).post(scoped_create_working_directory),
@@ -1473,6 +1571,19 @@ struct ScopedWorkspacePath {
workspace_id: String,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct AttachCurrentWorkerWorkdirRequest {
workdir_id: String,
}
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
struct CurrentWorkerWorkdirAttachmentResponse {
workspace_id: String,
workdir_id: String,
attached: bool,
}
#[derive(Debug, Deserialize)]
struct ScopedRecordPath {
workspace_id: String,
@@ -2977,6 +3088,331 @@ fn authenticate_worker_mutation_source(
})
}
fn current_worker_identity(
api: &WorkspaceApi,
workspace_id: &str,
headers: &HeaderMap,
) -> Result<(String, u64)> {
let source = authenticate_worker_mutation_source(api, workspace_id, headers)?;
let worker_id = source.worker_id.parse::<u64>().map_err(|_| {
Error::WorkerSourceIdentity(format!(
"Runtime-bound Worker id must be numeric, got `{}`",
source.worker_id
))
})?;
Ok((source.runtime_id, worker_id))
}
fn current_worker_active_attachment(
api: &WorkspaceApi,
runtime_id: &str,
worker_id: u64,
) -> ApiResult<WorkerWorkdirLinkRecord> {
if let Some(link) = api
.store
.list_worker_workdir_links(&api.config.workspace_id, runtime_id, worker_id)?
.into_iter()
.next()
{
return Ok(link);
}
if api.store.worker_workdir_link_history_exists(
&api.config.workspace_id,
runtime_id,
worker_id,
)? {
return Err(Error::WorkdirAttachmentConflict(format!(
"Worker {runtime_id}:{worker_id} has no active Workdir attachment"
))
.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 worker = api
.runtime
.worker(runtime_id, &worker_id.to_string())
.map_err(|error| error.into_error())?;
if worker.working_directory.is_some() {
sync_worker_observation(api, &worker)?;
if let Some(link) = api
.store
.list_worker_workdir_links(&api.config.workspace_id, runtime_id, worker_id)?
.into_iter()
.next()
{
return Ok(link);
}
}
Err(Error::WorkdirAttachmentConflict(format!(
"Worker {runtime_id}:{worker_id} has no active Workdir attachment"
))
.into())
}
fn current_worker_session_lock(
api: &WorkspaceApi,
runtime_id: &str,
worker_id: u64,
) -> Arc<tokio::sync::Mutex<()>> {
api.workdir_session_locks
.lock()
.expect("Workdir session lock registry poisoned")
.entry((runtime_id.to_string(), worker_id))
.or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
.clone()
}
async fn open_current_worker_workdir_session_locked(
api: &WorkspaceApi,
runtime_id: &str,
worker_id: u64,
link: &WorkerWorkdirLinkRecord,
) -> Result<WorkdirSessionHandle> {
if let Some(session) = api
.workdir_sessions
.lock()
.expect("Workdir session registry lock poisoned")
.get(&(runtime_id.to_string(), worker_id))
.cloned()
{
return Ok(session);
}
let workdir = api
.store
.list_workdir_registry(&api.config.workspace_id, 10_000)?
.into_iter()
.find(|workdir| workdir.workdir_id == link.workdir_id)
.ok_or_else(|| Error::RuntimeOperationFailed {
runtime_id: runtime_id.to_string(),
code: "working_directory_not_found".to_string(),
message: format!(
"attached Workdir {} is not registered in this Workspace",
link.workdir_id
),
})?;
let owner_worker_id = format!("{runtime_id}-{worker_id}");
let session = api
.runtime
.open_workdir_session(&workdir.runtime_id, &workdir.workdir_id, &owner_worker_id)
.await
.map_err(|error| error.into_error())?;
let key = (runtime_id.to_string(), worker_id);
let (selected, unused) = {
let mut sessions = api
.workdir_sessions
.lock()
.expect("Workdir session registry lock poisoned");
match sessions.entry(key) {
std::collections::hash_map::Entry::Occupied(entry) => {
(entry.get().clone(), Some(session))
}
std::collections::hash_map::Entry::Vacant(entry) => {
entry.insert(session.clone());
(session, None)
}
}
};
if let Some(unused) = unused {
unused
.close()
.await
.map_err(|error| Error::RuntimeOperationFailed {
runtime_id: runtime_id.to_string(),
code: "duplicate_workdir_session_close_failed".to_string(),
message: error.to_string(),
})?;
}
Ok(selected)
}
async fn close_current_worker_session_locked(
api: &WorkspaceApi,
runtime_id: &str,
worker_id: u64,
) -> Result<()> {
let key = (runtime_id.to_string(), worker_id);
let session = api
.workdir_sessions
.lock()
.expect("Workdir session registry lock poisoned")
.get(&key)
.cloned();
if let Some(session) = session {
session
.close()
.await
.map_err(|error| Error::RuntimeOperationFailed {
runtime_id: runtime_id.to_string(),
code: "workdir_session_close_failed".to_string(),
message: error.to_string(),
})?;
api.workdir_sessions
.lock()
.expect("Workdir session registry lock poisoned")
.remove(&key);
}
Ok(())
}
async fn scoped_attach_current_worker_workdir(
State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedWorkspacePath>,
headers: HeaderMap,
Json(request): Json<AttachCurrentWorkerWorkdirRequest>,
) -> ApiResult<Json<CurrentWorkerWorkdirAttachmentResponse>> {
validate_workspace_scope(&api, &path.workspace_id)?;
let (runtime_id, worker_id) = current_worker_identity(&api, &path.workspace_id, &headers)?;
let workdir_id = request.workdir_id.trim();
if workdir_id.is_empty() || workdir_id.chars().any(char::is_control) {
return Err(Error::InvalidRecordId(request.workdir_id).into());
}
if !api
.store
.list_workdir_registry(&api.config.workspace_id, 10_000)?
.iter()
.any(|workdir| workdir.workdir_id == workdir_id)
{
return Err(Error::RuntimeOperationFailed {
runtime_id: runtime_id.clone(),
code: "working_directory_not_found".to_string(),
message: format!("unknown Workdir `{workdir_id}`"),
}
.into());
}
let session_lock = current_worker_session_lock(&api, &runtime_id, worker_id);
let _session_guard = session_lock.lock().await;
let link = api.store.attach_worker_workdir(&WorkerWorkdirLinkRecord {
workspace_id: api.config.workspace_id.clone(),
runtime_id: runtime_id.clone(),
runtime_worker_id: worker_id,
workdir_id: workdir_id.to_string(),
role: "attachment".to_string(),
linked_at: Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true),
unlinked_at: None,
})?;
if let Err(error) =
open_current_worker_workdir_session_locked(&api, &runtime_id, worker_id, &link).await
{
let _ = api.store.detach_worker_workdir(
&api.config.workspace_id,
&runtime_id,
worker_id,
Some(workdir_id),
&Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true),
);
return Err(error.into());
}
Ok(Json(CurrentWorkerWorkdirAttachmentResponse {
workspace_id: api.config.workspace_id.clone(),
workdir_id: workdir_id.to_string(),
attached: true,
}))
}
async fn scoped_detach_current_worker_workdir(
State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedWorkspacePath>,
headers: HeaderMap,
) -> ApiResult<Json<CurrentWorkerWorkdirAttachmentResponse>> {
validate_workspace_scope(&api, &path.workspace_id)?;
let (runtime_id, worker_id) = current_worker_identity(&api, &path.workspace_id, &headers)?;
let session_lock = current_worker_session_lock(&api, &runtime_id, worker_id);
let _session_guard = session_lock.lock().await;
let link = current_worker_active_attachment(&api, &runtime_id, worker_id)?;
close_current_worker_session_locked(&api, &runtime_id, worker_id).await?;
api.store.detach_worker_workdir(
&api.config.workspace_id,
&runtime_id,
worker_id,
Some(&link.workdir_id),
&Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true),
)?;
Ok(Json(CurrentWorkerWorkdirAttachmentResponse {
workspace_id: api.config.workspace_id.clone(),
workdir_id: link.workdir_id,
attached: false,
}))
}
async fn scoped_execute_current_worker_workdir_operation(
State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedWorkspacePath>,
headers: HeaderMap,
Json(operation): Json<WorkdirSessionOperation>,
) -> ApiResult<Json<WorkdirSessionOperationResult>> {
validate_workspace_scope(&api, &path.workspace_id)?;
let (runtime_id, worker_id) = current_worker_identity(&api, &path.workspace_id, &headers)?;
let session_lock = current_worker_session_lock(&api, &runtime_id, worker_id);
let _session_guard = session_lock.lock().await;
let link = current_worker_active_attachment(&api, &runtime_id, worker_id)?;
let session =
open_current_worker_workdir_session_locked(&api, &runtime_id, worker_id, &link).await?;
let result = execute_workdir_session_operation(&session, operation)
.await
.map_err(|error| Error::RuntimeOperationFailed {
runtime_id,
code: "workdir_session_operation_failed".to_string(),
message: error.to_string(),
})?;
Ok(Json(result))
}
async fn execute_workdir_session_operation(
session: &WorkdirSessionHandle,
operation: WorkdirSessionOperation,
) -> std::result::Result<WorkdirSessionOperationResult, workdir::WorkdirError> {
match operation {
WorkdirSessionOperation::Stat(request) => session
.stat(request)
.await
.map(WorkdirSessionOperationResult::Stat),
WorkdirSessionOperation::Read(request) => session
.read(request)
.await
.map(WorkdirSessionOperationResult::Read),
WorkdirSessionOperation::Write(request) => session
.write(request)
.await
.map(WorkdirSessionOperationResult::Write),
WorkdirSessionOperation::Edit(request) => session
.edit(request)
.await
.map(WorkdirSessionOperationResult::Edit),
WorkdirSessionOperation::List(request) => session
.list(request)
.await
.map(WorkdirSessionOperationResult::List),
WorkdirSessionOperation::Glob(request) => session
.glob(request)
.await
.map(WorkdirSessionOperationResult::Glob),
WorkdirSessionOperation::Grep(request) => session
.grep(request)
.await
.map(WorkdirSessionOperationResult::Grep),
WorkdirSessionOperation::CommandStart(request) => session
.start_command(request)
.await
.map(WorkdirSessionOperationResult::CommandStart),
WorkdirSessionOperation::CommandStatus(handle) => session
.command_status(handle)
.await
.map(WorkdirSessionOperationResult::CommandStatus),
WorkdirSessionOperation::CommandOutput(request) => session
.command_output(request)
.await
.map(WorkdirSessionOperationResult::CommandOutput),
WorkdirSessionOperation::CommandCancel(handle) => session
.cancel_command(handle)
.await
.map(|()| WorkdirSessionOperationResult::CommandCancel),
}
}
async fn run_orchestrator_turn_end_hook(api: WorkspaceApi) {
let Ok(mut subscription) = api.runtime_subscription_broker.subscribe(
EMBEDDED_WORKER_RUNTIME_ID,
@@ -4331,7 +4767,7 @@ fn cleanup_plan_digest(
Ok(format!("sha256:{digest}"))
}
fn execute_runtime_cleanup(
async fn execute_runtime_cleanup(
api: &WorkspaceApi,
runtime_id: &str,
request: ExecuteRuntimeCleanupRequest,
@@ -4375,12 +4811,20 @@ fn execute_runtime_cleanup(
));
}
let runtime_worker_id = parse_runtime_worker_id_for_registry(&candidate.runtime_worker_id)?;
let session_lock = current_worker_session_lock(api, runtime_id, runtime_worker_id);
let _session_guard = session_lock.lock().await;
close_current_worker_session_locked(api, runtime_id, runtime_worker_id).await?;
cleanup_runtime_worker_for_execution(api, runtime_id, candidate)?;
api.store.delete_worker_registry(
&api.config.workspace_id,
candidate.runtime_id.as_str(),
runtime_worker_id,
)?;
drop(_session_guard);
api.workdir_session_locks
.lock()
.expect("Workdir session lock registry poisoned")
.remove(&(runtime_id.to_string(), runtime_worker_id));
results.push(RuntimeCleanupExecutionResult {
target_id: candidate.target_id.clone(),
action: candidate.action.clone(),
@@ -4813,7 +5257,7 @@ async fn scoped_execute_runtime_cleanup(
Json(request): Json<ExecuteRuntimeCleanupRequest>,
) -> ApiResult<Json<RuntimeCleanupExecutionResponse>> {
validate_workspace_scope(&api, &path.workspace_id)?;
let response = execute_runtime_cleanup(&api, path.runtime_id.as_str(), request)?;
let response = execute_runtime_cleanup(&api, path.runtime_id.as_str(), request).await?;
Ok(Json(response))
}
@@ -6200,7 +6644,12 @@ fn record_browser_worker_spawn(
let workdir_record =
workdir_record_from_summary(api, worker.runtime_id.as_str(), working_directory);
api.store.upsert_workdir_registry(&workdir_record)?;
link_worker_to_workdir(api, &worker_record, &working_directory.working_directory_id)?;
link_worker_to_workdir(
api,
&worker_record,
&working_directory.working_directory_id,
None,
)?;
}
if let Some(workdir_id) = selected_working_directory_id.as_deref() {
if api
@@ -6228,7 +6677,7 @@ fn record_browser_worker_spawn(
.get_workdir_registry(&api.config.workspace_id, workdir_id)?
.is_some()
{
link_worker_to_workdir(api, &worker_record, workdir_id)?;
link_worker_to_workdir(api, &worker_record, workdir_id, None)?;
}
}
let runtime_id = worker.runtime_id.clone();
@@ -6533,7 +6982,7 @@ async fn create_runtime_worker(
.get_workdir_registry(&api.config.workspace_id, workdir_id)?
.is_some()
{
link_worker_to_workdir(&api, &record, workdir_id)?;
link_worker_to_workdir(&api, &record, workdir_id, None)?;
}
}
}
@@ -6614,6 +7063,9 @@ async fn stop_runtime_worker(
.stop_worker(&runtime_id, &worker_id, request)
.map_err(|err| err.into_error())?;
let runtime_worker_id = parse_runtime_worker_id_for_registry(&worker_id)?;
let session_lock = current_worker_session_lock(&api, &runtime_id, runtime_worker_id);
let _session_guard = session_lock.lock().await;
close_current_worker_session_locked(&api, &runtime_id, runtime_worker_id).await?;
if let Some(record) =
api.store
.get_worker_registry(&api.config.workspace_id, &runtime_id, runtime_worker_id)?
@@ -8033,7 +8485,7 @@ fn sync_worker_observation(
let workdir_record =
workdir_record_from_summary(api, worker.runtime_id.as_str(), working_directory);
api.store.upsert_workdir_registry(&workdir_record)?;
link_worker_to_workdir(api, &record, &working_directory.working_directory_id)?;
link_worker_to_workdir(api, &record, &working_directory.working_directory_id, None)?;
}
Ok(record)
}
@@ -8326,18 +8778,24 @@ fn link_worker_to_workdir(
api: &WorkspaceApi,
worker_record: &WorkerRegistryRecord,
workdir_id: &str,
reservation_id: Option<&str>,
) -> ApiResult<()> {
let timestamp = now_registry_timestamp();
api.store
.upsert_worker_workdir_link(&WorkerWorkdirLinkRecord {
let record = WorkerWorkdirLinkRecord {
workspace_id: api.config.workspace_id.clone(),
runtime_id: worker_record.runtime_id.clone(),
runtime_worker_id: worker_record.runtime_worker_id.clone(),
runtime_worker_id: worker_record.runtime_worker_id,
workdir_id: workdir_id.to_string(),
role: "primary_cwd".to_string(),
role: "attachment".to_string(),
linked_at: timestamp,
unlinked_at: None,
})?;
};
if let Some(reservation_id) = reservation_id {
api.store
.finalize_reserved_worker_workdir_attachment(&record, reservation_id)?;
} else {
api.store.attach_worker_workdir(&record)?;
}
Ok(())
}
@@ -8740,7 +9198,9 @@ impl ApiError {
impl IntoResponse for ApiError {
fn into_response(self) -> Response {
let status = match &self.error {
Error::TicketAssignmentConflict(_) => StatusCode::CONFLICT,
Error::TicketAssignmentConflict(_) | Error::WorkdirAttachmentConflict(_) => {
StatusCode::CONFLICT
}
Error::WorkerSourceIdentity(_) => StatusCode::BAD_REQUEST,
Error::InvalidRuntimeIdentifier { .. } | Error::ReservedWorkerName(_) => {
StatusCode::BAD_REQUEST
@@ -8931,7 +9391,7 @@ mod tests {
runtime_id: worker.runtime_id.clone(),
runtime_worker_id: worker.runtime_worker_id.clone(),
workdir_id: workdir.workdir_id.clone(),
role: "primary_cwd".to_string(),
role: "attachment".to_string(),
linked_at: "4".to_string(),
unlinked_at: None,
};
@@ -9085,12 +9545,12 @@ mod tests {
})
.unwrap();
api.store
.upsert_worker_workdir_link(&WorkerWorkdirLinkRecord {
.attach_worker_workdir(&WorkerWorkdirLinkRecord {
workspace_id: TEST_WORKSPACE_ID.to_string(),
runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(),
runtime_worker_id: 7,
workdir_id: "managed".to_string(),
role: "primary_cwd".to_string(),
role: "attachment".to_string(),
linked_at: "3".to_string(),
unlinked_at: None,
})
@@ -10741,12 +11201,12 @@ mod tests {
fn seed_cleanup_link(api: &WorkspaceApi, runtime_worker_id: &str, workdir_id: &str) {
let runtime_worker_id = runtime_worker_id.parse::<u64>().unwrap();
api.store
.upsert_worker_workdir_link(&WorkerWorkdirLinkRecord {
.attach_worker_workdir(&WorkerWorkdirLinkRecord {
workspace_id: api.config.workspace_id.clone(),
runtime_id: "runtime-test".to_string(),
runtime_worker_id,
workdir_id: workdir_id.to_string(),
role: "primary".to_string(),
role: "attachment".to_string(),
linked_at: now_registry_timestamp(),
unlinked_at: None,
})
@@ -10766,6 +11226,54 @@ mod tests {
);
}
#[tokio::test]
async fn current_worker_workdir_routes_require_a_live_runtime_worker_identity() {
let workspace = tempfile::tempdir().unwrap();
init_clean_git_workspace(workspace.path());
let api = test_api(workspace.path()).await;
let response = build_router(api)
.oneshot(
Request::builder()
.method("POST")
.uri(format!(
"/api/w/{TEST_WORKSPACE_ID}/workers/self/workdir-attachment"
))
.header("content-type", "application/json")
.body(Body::from(json!({"workdir_id": "wd"}).to_string()))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
}
#[tokio::test]
async fn backend_workdir_session_proxy_executes_typed_operations() {
use manifest::Scope;
use workdir::LocalWorkdirSession;
let root = tempfile::tempdir().unwrap();
fs::write(root.path().join("visible.txt"), "attached").unwrap();
let session: WorkdirSessionHandle = Arc::new(LocalWorkdirSession::new(
Scope::writable(root.path()).unwrap(),
root.path().to_path_buf(),
));
let result = execute_workdir_session_operation(
&session,
WorkdirSessionOperation::List(workdir::ListRequest {
path: workdir::WorkdirPath::root(),
limit: 10,
}),
)
.await
.unwrap();
assert!(matches!(
result,
WorkdirSessionOperationResult::List(ref result)
if result.entries.iter().any(|entry| entry.path.as_str().ends_with("visible.txt"))
));
}
#[tokio::test]
async fn simple_workdir_cleanup_rejects_dirty_and_blocked_candidates() {
let workspace = tempfile::tempdir().unwrap();
@@ -10884,7 +11392,11 @@ mod tests {
workdir_target_ids: Vec::new(),
confirm_dirty_discard_target_ids: Vec::new(),
};
assert!(execute_runtime_cleanup(&api, "runtime-test", stale).is_err());
assert!(
execute_runtime_cleanup(&api, "runtime-test", stale)
.await
.is_err()
);
let pinned = ExecuteRuntimeCleanupRequest {
expected_plan_revision: plan.revision,
expected_plan_digest: plan.digest,
@@ -10892,7 +11404,11 @@ mod tests {
workdir_target_ids: Vec::new(),
confirm_dirty_discard_target_ids: Vec::new(),
};
assert!(execute_runtime_cleanup(&api, "runtime-test", pinned).is_err());
assert!(
execute_runtime_cleanup(&api, "runtime-test", pinned)
.await
.is_err()
);
}
#[tokio::test]
@@ -10925,7 +11441,11 @@ mod tests {
workdir_target_ids: vec![dirty_target],
confirm_dirty_discard_target_ids: Vec::new(),
};
assert!(execute_runtime_cleanup(&api, "runtime-test", missing_confirmation).is_err());
assert!(
execute_runtime_cleanup(&api, "runtime-test", missing_confirmation)
.await
.is_err()
);
let delete_removed = ExecuteRuntimeCleanupRequest {
expected_plan_revision: plan.revision,
expected_plan_digest: plan.digest,
@@ -10934,6 +11454,7 @@ mod tests {
confirm_dirty_discard_target_ids: Vec::new(),
};
let response = execute_runtime_cleanup(&api, "runtime-test", delete_removed)
.await
.unwrap_or_else(|err| panic!("cleanup execution: {}", err.error));
assert_eq!(response.results[0].status, "deleted");
assert!(
crates/workspace-server/src/store.rs
+488 -18
View File
@@ -127,6 +127,11 @@ const MIGRATIONS: &[Migration] = &[
name: "drop Ticket notification outbox",
apply: drop_ticket_notification_tables,
},
Migration {
version: 23,
name: "enforce exclusive Worker Workdir attachments and spawn reservations",
apply: enforce_exclusive_worker_workdir_attachments,
},
];
struct Migration {
@@ -618,7 +623,42 @@ pub trait ControlPlaneStore: Send + Sync {
) -> Result<Vec<WorkdirRegistryRecord>>;
fn delete_workdir_registry(&self, workspace_id: &str, workdir_id: &str) -> Result<bool>;
fn upsert_worker_workdir_link(&self, record: &WorkerWorkdirLinkRecord) -> Result<()>;
fn reserve_worker_workdir_attachment(
&self,
workspace_id: &str,
workdir_id: &str,
reservation_id: &str,
reserved_at: &str,
) -> Result<()>;
fn release_worker_workdir_attachment_reservation(
&self,
workspace_id: &str,
workdir_id: &str,
reservation_id: &str,
) -> Result<()>;
fn finalize_reserved_worker_workdir_attachment(
&self,
record: &WorkerWorkdirLinkRecord,
reservation_id: &str,
) -> Result<WorkerWorkdirLinkRecord>;
fn attach_worker_workdir(
&self,
record: &WorkerWorkdirLinkRecord,
) -> Result<WorkerWorkdirLinkRecord>;
fn detach_worker_workdir(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
expected_workdir_id: Option<&str>,
unlinked_at: &str,
) -> Result<Option<WorkerWorkdirLinkRecord>>;
fn worker_workdir_link_history_exists(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
) -> Result<bool>;
fn list_worker_workdir_links(
&self,
workspace_id: &str,
@@ -1680,10 +1720,18 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
runtime_worker_id: u64,
) -> Result<bool> {
self.with_conn(|conn| {
let changed = conn.execute(
let tx = conn.unchecked_transaction()?;
tx.execute(
r#"UPDATE worker_workdir_links
SET unlinked_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
WHERE workspace_id = ?1 AND runtime_id = ?2 AND runtime_worker_id = ?3 AND unlinked_at IS NULL"#,
params![workspace_id, runtime_id, runtime_worker_id],
)?;
let changed = tx.execute(
"DELETE FROM worker_registry WHERE workspace_id = ?1 AND runtime_id = ?2 AND runtime_worker_id = ?3",
params![workspace_id, runtime_id, runtime_worker_id],
)?;
tx.commit()?;
Ok(changed > 0)
})
}
@@ -2173,23 +2221,147 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
fn delete_workdir_registry(&self, workspace_id: &str, workdir_id: &str) -> Result<bool> {
self.with_conn(|conn| {
let changed = conn.execute(
let tx = rusqlite::Transaction::new_unchecked(
conn,
rusqlite::TransactionBehavior::Immediate,
)?;
let blocked: bool = tx.query_row(
r#"SELECT EXISTS(
SELECT 1 FROM worker_workdir_links
WHERE workspace_id = ?1 AND workdir_id = ?2 AND unlinked_at IS NULL
UNION ALL
SELECT 1 FROM worker_workdir_attachment_reservations
WHERE workspace_id = ?1 AND workdir_id = ?2
)"#,
params![workspace_id, workdir_id],
|row| row.get(0),
)?;
if blocked {
return Err(Error::WorkdirAttachmentConflict(format!(
"Workdir {workdir_id} has an active or pending Worker attachment"
)));
}
let changed = tx.execute(
"DELETE FROM workdir_registry WHERE workspace_id = ?1 AND workdir_id = ?2",
params![workspace_id, workdir_id],
)?;
tx.commit()?;
Ok(changed > 0)
})
}
fn upsert_worker_workdir_link(&self, record: &WorkerWorkdirLinkRecord) -> Result<()> {
fn reserve_worker_workdir_attachment(
&self,
workspace_id: &str,
workdir_id: &str,
reservation_id: &str,
reserved_at: &str,
) -> Result<()> {
self.with_conn(|conn| {
let tx = rusqlite::Transaction::new_unchecked(
conn,
rusqlite::TransactionBehavior::Immediate,
)?;
let registered: bool = tx.query_row(
r#"SELECT EXISTS(
SELECT 1 FROM workdir_registry
WHERE workspace_id = ?1 AND workdir_id = ?2
)"#,
params![workspace_id, workdir_id],
|row| row.get(0),
)?;
if !registered {
return Err(Error::WorkdirAttachmentConflict(format!(
"Workdir {workdir_id} is not registered in Workspace {workspace_id}"
)));
}
let occupied: bool = tx.query_row(
r#"SELECT EXISTS(
SELECT 1 FROM worker_workdir_links
WHERE workspace_id = ?1 AND workdir_id = ?2 AND unlinked_at IS NULL
)"#,
params![workspace_id, workdir_id],
|row| row.get(0),
)?;
if occupied {
return Err(Error::WorkdirAttachmentConflict(format!(
"Workdir {workdir_id} already has an active attachment"
)));
}
tx.execute(
r#"INSERT INTO worker_workdir_attachment_reservations (
workspace_id, workdir_id, reservation_id, reserved_at
) VALUES (?1, ?2, ?3, ?4)"#,
params![workspace_id, workdir_id, reservation_id, reserved_at],
)
.map_err(|error| {
if matches!(error, rusqlite::Error::SqliteFailure(ref code, _) if code.extended_code == rusqlite::ffi::SQLITE_CONSTRAINT_PRIMARYKEY || code.extended_code == rusqlite::ffi::SQLITE_CONSTRAINT_UNIQUE)
{
Error::WorkdirAttachmentConflict(format!(
"Workdir {workdir_id} already has a pending attachment reservation"
))
} else {
error.into()
}
})?;
tx.commit()?;
Ok(())
})
}
fn release_worker_workdir_attachment_reservation(
&self,
workspace_id: &str,
workdir_id: &str,
reservation_id: &str,
) -> Result<()> {
self.with_conn(|conn| {
conn.execute(
r#"DELETE FROM worker_workdir_attachment_reservations
WHERE workspace_id = ?1 AND workdir_id = ?2 AND reservation_id = ?3"#,
params![workspace_id, workdir_id, reservation_id],
)?;
Ok(())
})
}
fn finalize_reserved_worker_workdir_attachment(
&self,
record: &WorkerWorkdirLinkRecord,
reservation_id: &str,
) -> Result<WorkerWorkdirLinkRecord> {
if record.role != "attachment" || record.unlinked_at.is_some() {
return Err(Error::WorkdirAttachmentConflict(
"reserved attachment finalization requires an active canonical attachment"
.to_string(),
));
}
self.with_conn(|conn| {
let tx = rusqlite::Transaction::new_unchecked(
conn,
rusqlite::TransactionBehavior::Immediate,
)?;
let owns_reservation: bool = tx.query_row(
r#"SELECT EXISTS(
SELECT 1 FROM worker_workdir_attachment_reservations
WHERE workspace_id = ?1 AND workdir_id = ?2 AND reservation_id = ?3
)"#,
params![record.workspace_id, record.workdir_id, reservation_id],
|row| row.get(0),
)?;
if !owns_reservation {
return Err(Error::WorkdirAttachmentConflict(format!(
"Workdir {} attachment reservation is missing or owned by another spawn",
record.workdir_id
)));
}
tx.execute(
r#"INSERT INTO worker_workdir_links (
workspace_id, runtime_id, runtime_worker_id, workdir_id, role, linked_at, unlinked_at
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, NULL)
ON CONFLICT(workspace_id, runtime_id, runtime_worker_id, workdir_id, role) DO UPDATE SET
linked_at = excluded.linked_at,
unlinked_at = excluded.unlinked_at"#,
unlinked_at = NULL"#,
params![
record.workspace_id,
record.runtime_id,
@@ -2197,10 +2369,213 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
record.workdir_id,
record.role,
record.linked_at,
record.unlinked_at,
],
)
.map_err(|error| {
if matches!(error, rusqlite::Error::SqliteFailure(ref code, _) if code.extended_code == rusqlite::ffi::SQLITE_CONSTRAINT_UNIQUE)
{
Error::WorkdirAttachmentConflict(
"Worker or Workdir acquired another active attachment".to_string(),
)
} else {
error.into()
}
})?;
tx.execute(
r#"DELETE FROM worker_workdir_attachment_reservations
WHERE workspace_id = ?1 AND workdir_id = ?2 AND reservation_id = ?3"#,
params![record.workspace_id, record.workdir_id, reservation_id],
)?;
Ok(())
tx.commit()?;
Ok(record.clone())
})
}
fn attach_worker_workdir(
&self,
record: &WorkerWorkdirLinkRecord,
) -> Result<WorkerWorkdirLinkRecord> {
if record.role != "attachment" {
return Err(Error::WorkdirAttachmentConflict(format!(
"unsupported Workdir attachment role `{}`",
record.role
)));
}
if record.unlinked_at.is_some() {
return Err(Error::WorkdirAttachmentConflict(
"a new attachment cannot already be unlinked".to_string(),
));
}
self.with_conn(|conn| {
let tx = rusqlite::Transaction::new_unchecked(
conn,
rusqlite::TransactionBehavior::Immediate,
)?;
let registered: bool = tx.query_row(
r#"SELECT EXISTS(
SELECT 1 FROM workdir_registry
WHERE workspace_id = ?1 AND workdir_id = ?2
)"#,
params![record.workspace_id, record.workdir_id],
|row| row.get(0),
)?;
if !registered {
return Err(Error::WorkdirAttachmentConflict(format!(
"Workdir {} is not registered in Workspace {}",
record.workdir_id, record.workspace_id
)));
}
let reserved: bool = tx.query_row(
r#"SELECT EXISTS(
SELECT 1 FROM worker_workdir_attachment_reservations
WHERE workspace_id = ?1 AND workdir_id = ?2
)"#,
params![record.workspace_id, record.workdir_id],
|row| row.get(0),
)?;
if reserved {
return Err(Error::WorkdirAttachmentConflict(format!(
"Workdir {} has a pending attachment reservation",
record.workdir_id
)));
}
let active_for_worker = tx
.query_row(
r#"SELECT workspace_id, runtime_id, runtime_worker_id, workdir_id, role, linked_at, unlinked_at
FROM worker_workdir_links
WHERE workspace_id = ?1 AND runtime_id = ?2 AND runtime_worker_id = ?3 AND unlinked_at IS NULL"#,
params![
record.workspace_id,
record.runtime_id,
record.runtime_worker_id,
],
read_worker_workdir_link_record,
)
.optional()?;
if let Some(active) = active_for_worker {
if active.workdir_id == record.workdir_id {
tx.commit()?;
return Ok(active);
}
return Err(Error::WorkdirAttachmentConflict(format!(
"Worker {}:{} is already attached to Workdir {}",
record.runtime_id, record.runtime_worker_id, active.workdir_id
)));
}
let active_for_workdir = tx
.query_row(
r#"SELECT workspace_id, runtime_id, runtime_worker_id, workdir_id, role, linked_at, unlinked_at
FROM worker_workdir_links
WHERE workspace_id = ?1 AND workdir_id = ?2 AND unlinked_at IS NULL"#,
params![record.workspace_id, record.workdir_id],
read_worker_workdir_link_record,
)
.optional()?;
if let Some(active) = active_for_workdir {
return Err(Error::WorkdirAttachmentConflict(format!(
"Workdir {} is already attached to Worker {}:{}",
record.workdir_id, active.runtime_id, active.runtime_worker_id
)));
}
let write = tx.execute(
r#"INSERT INTO worker_workdir_links (
workspace_id, runtime_id, runtime_worker_id, workdir_id, role, linked_at, unlinked_at
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, NULL)
ON CONFLICT(workspace_id, runtime_id, runtime_worker_id, workdir_id, role) DO UPDATE SET
linked_at = excluded.linked_at,
unlinked_at = NULL"#,
params![
record.workspace_id,
record.runtime_id,
record.runtime_worker_id,
record.workdir_id,
record.role,
record.linked_at,
],
);
if let Err(error) = write {
if matches!(error, rusqlite::Error::SqliteFailure(ref code, _) if code.extended_code == rusqlite::ffi::SQLITE_CONSTRAINT_UNIQUE)
{
return Err(Error::WorkdirAttachmentConflict(
"Worker or Workdir acquired another active attachment".to_string(),
));
}
return Err(error.into());
}
tx.commit()?;
Ok(record.clone())
})
}
fn detach_worker_workdir(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
expected_workdir_id: Option<&str>,
unlinked_at: &str,
) -> Result<Option<WorkerWorkdirLinkRecord>> {
self.with_conn(|conn| {
let tx = rusqlite::Transaction::new_unchecked(
conn,
rusqlite::TransactionBehavior::Immediate,
)?;
let active = tx
.query_row(
r#"SELECT workspace_id, runtime_id, runtime_worker_id, workdir_id, role, linked_at, unlinked_at
FROM worker_workdir_links
WHERE workspace_id = ?1 AND runtime_id = ?2 AND runtime_worker_id = ?3 AND unlinked_at IS NULL"#,
params![workspace_id, runtime_id, runtime_worker_id],
read_worker_workdir_link_record,
)
.optional()?;
let Some(active) = active else {
tx.commit()?;
return Ok(None);
};
if let Some(expected_workdir_id) = expected_workdir_id {
if active.workdir_id != expected_workdir_id {
return Err(Error::WorkdirAttachmentConflict(format!(
"Worker {runtime_id}:{runtime_worker_id} is attached to Workdir {}, not {expected_workdir_id}",
active.workdir_id
)));
}
}
let changed = tx.execute(
r#"UPDATE worker_workdir_links
SET unlinked_at = ?4
WHERE workspace_id = ?1 AND runtime_id = ?2 AND runtime_worker_id = ?3 AND unlinked_at IS NULL"#,
params![workspace_id, runtime_id, runtime_worker_id, unlinked_at],
)?;
if changed != 1 {
return Err(Error::WorkdirAttachmentConflict(format!(
"Worker {runtime_id}:{runtime_worker_id} attachment changed during detach"
)));
}
tx.commit()?;
Ok(Some(WorkerWorkdirLinkRecord {
unlinked_at: Some(unlinked_at.to_string()),
..active
}))
})
}
fn worker_workdir_link_history_exists(
&self,
workspace_id: &str,
runtime_id: &str,
runtime_worker_id: u64,
) -> Result<bool> {
self.with_conn(|conn| {
let exists = conn.query_row(
r#"SELECT EXISTS(
SELECT 1 FROM worker_workdir_links
WHERE workspace_id = ?1 AND runtime_id = ?2 AND runtime_worker_id = ?3
)"#,
params![workspace_id, runtime_id, runtime_worker_id],
|row| row.get(0),
)?;
Ok(exists)
})
}
@@ -3037,6 +3412,29 @@ DROP TABLE IF EXISTS ticket_notification_outbox;
Ok(())
}
fn enforce_exclusive_worker_workdir_attachments(conn: &Connection) -> Result<()> {
conn.execute_batch(
r#"
CREATE TABLE IF NOT EXISTS worker_workdir_attachment_reservations (
workspace_id TEXT NOT NULL,
workdir_id TEXT NOT NULL,
reservation_id TEXT NOT NULL,
reserved_at TEXT NOT NULL,
PRIMARY KEY (workspace_id, workdir_id)
);
CREATE UNIQUE INDEX IF NOT EXISTS ux_worker_workdir_attachment_reservation_id
ON worker_workdir_attachment_reservations(workspace_id, reservation_id);
CREATE UNIQUE INDEX IF NOT EXISTS ux_worker_workdir_links_active_worker
ON worker_workdir_links(workspace_id, runtime_id, runtime_worker_id)
WHERE unlinked_at IS NULL;
CREATE UNIQUE INDEX IF NOT EXISTS ux_worker_workdir_links_active_workdir
ON worker_workdir_links(workspace_id, workdir_id)
WHERE unlinked_at IS NULL;
"#,
)?;
Ok(())
}
fn create_objective_event_tables(conn: &Connection) -> Result<()> {
conn.execute_batch(
r#"
@@ -3722,7 +4120,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT);
let db = dir.path().join("control-plane.sqlite");
let store = SqliteWorkspaceStore::open(&db).unwrap();
assert_eq!(store.schema_version().await.unwrap(), 22);
assert_eq!(store.schema_version().await.unwrap(), 23);
assert!(
!store
.with_conn(|conn| table_exists(conn, "worker_workspace_credentials"))
@@ -3739,7 +4137,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT);
store.upsert_workspace(&record).await.unwrap();
let reopened = SqliteWorkspaceStore::open(&db).unwrap();
assert_eq!(reopened.schema_version().await.unwrap(), 22);
assert_eq!(reopened.schema_version().await.unwrap(), 23);
assert_eq!(
reopened.get_workspace("local-dev").await.unwrap(),
Some(record)
@@ -4190,7 +4588,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT);
.unwrap();
let store = SqliteWorkspaceStore::from_connection(conn).unwrap();
assert_eq!(store.schema_version().await.unwrap(), 22);
assert_eq!(store.schema_version().await.unwrap(), 23);
store
.with_conn(|conn| {
@@ -4379,7 +4777,7 @@ CREATE TABLE ticket_assignment_operations (
#[tokio::test]
async fn repository_records_round_trip() {
let store = SqliteWorkspaceStore::in_memory().unwrap();
assert_eq!(store.schema_version().await.unwrap(), 22);
assert_eq!(store.schema_version().await.unwrap(), 23);
let workspace = WorkspaceRecord {
workspace_id: "local-dev".to_string(),
owner_account_id: None,
@@ -4417,7 +4815,7 @@ CREATE TABLE ticket_assignment_operations (
#[tokio::test]
async fn memory_authority_records_round_trip_and_close_staging() {
let store = SqliteWorkspaceStore::in_memory().unwrap();
assert_eq!(store.schema_version().await.unwrap(), 22);
assert_eq!(store.schema_version().await.unwrap(), 23);
let workspace = WorkspaceRecord {
workspace_id: "local-dev".to_string(),
owner_account_id: None,
@@ -4560,13 +4958,39 @@ CREATE TABLE ticket_assignment_operations (
let link = WorkerWorkdirLinkRecord {
workspace_id: "local-dev".to_string(),
runtime_id: worker.runtime_id.clone(),
runtime_worker_id: worker.runtime_worker_id.clone(),
runtime_worker_id: worker.runtime_worker_id,
workdir_id: workdir.workdir_id.clone(),
role: "primary_cwd".to_string(),
role: "attachment".to_string(),
linked_at: "4".to_string(),
unlinked_at: None,
};
store.upsert_worker_workdir_link(&link).unwrap();
store
.reserve_worker_workdir_attachment("local-dev", &workdir.workdir_id, "spawn-1", "4")
.unwrap();
assert!(matches!(
store.reserve_worker_workdir_attachment(
"local-dev",
&workdir.workdir_id,
"spawn-2",
"4"
),
Err(Error::WorkdirAttachmentConflict(_))
));
assert!(matches!(
store.delete_workdir_registry("local-dev", &workdir.workdir_id),
Err(Error::WorkdirAttachmentConflict(_))
));
assert!(matches!(
store.attach_worker_workdir(&link),
Err(Error::WorkdirAttachmentConflict(_))
));
assert_eq!(
store
.finalize_reserved_worker_workdir_attachment(&link, "spawn-1")
.unwrap(),
link
);
assert_eq!(store.attach_worker_workdir(&link).unwrap(), link);
assert_eq!(
store
@@ -4588,14 +5012,60 @@ CREATE TABLE ticket_assignment_operations (
store
.list_worker_workdir_links("local-dev", "embedded", 1)
.unwrap(),
vec![link]
vec![link.clone()]
);
let worker_conflict = WorkerWorkdirLinkRecord {
workdir_id: unmanaged_workdir.workdir_id.clone(),
linked_at: "5".to_string(),
..link.clone()
};
assert!(matches!(
store.attach_worker_workdir(&worker_conflict),
Err(Error::WorkdirAttachmentConflict(_))
));
let second_worker = WorkerRegistryRecord {
runtime_worker_id: 2,
display_name: "Browser 2".to_string(),
created_at: "5".to_string(),
updated_at: "5".to_string(),
..worker.clone()
};
store.upsert_worker_registry(&second_worker).unwrap();
let workdir_conflict = WorkerWorkdirLinkRecord {
runtime_worker_id: second_worker.runtime_worker_id,
linked_at: "5".to_string(),
..link.clone()
};
assert!(matches!(
store.attach_worker_workdir(&workdir_conflict),
Err(Error::WorkdirAttachmentConflict(_))
));
assert!(matches!(
store.detach_worker_workdir("local-dev", "embedded", 1, Some("wrong-workdir"), "6"),
Err(Error::WorkdirAttachmentConflict(_))
));
let detached = store
.detach_worker_workdir("local-dev", "embedded", 1, Some(&workdir.workdir_id), "6")
.unwrap()
.unwrap();
assert_eq!(detached.unlinked_at.as_deref(), Some("6"));
assert!(
store
.worker_workdir_link_history_exists("local-dev", "embedded", 1)
.unwrap()
);
assert_eq!(
store.attach_worker_workdir(&workdir_conflict).unwrap(),
workdir_conflict
);
}
#[tokio::test]
async fn account_and_login_records_round_trip() {
let store = SqliteWorkspaceStore::in_memory().unwrap();
assert_eq!(store.schema_version().await.unwrap(), 22);
assert_eq!(store.schema_version().await.unwrap(), 23);
let now = "2026-07-22T00:00:00Z".to_string();
let account = AccountRecord {
account_id: "acct-user-alice".to_string(),