diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 424bd110..298a63d2 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -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(); diff --git a/crates/worker/src/feature/builtin/manage_workdir.rs b/crates/worker/src/feature/builtin/manage_workdir.rs index 07eb0b11..ed717cfb 100644 --- a/crates/worker/src/feature/builtin/manage_workdir.rs +++ b/crates/worker/src/feature/builtin/manage_workdir.rs @@ -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, } +/// 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, + workdir: Workdir, +} + +impl WorkspaceAttachedWorkdirSession { + pub fn handle(client: Arc) -> WorkdirSessionHandle { + Arc::new(Self { + client, + workdir: Workdir::new("workspace-attachment"), + }) + } + + fn operate( + &self, + operation: WorkdirSessionOperation, + ) -> Result { + 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 { + match self.operate(WorkdirSessionOperation::Stat(request))? { + WorkdirSessionOperationResult::Stat(result) => Ok(result), + _ => Err(Self::mismatch("stat")), + } + } + + async fn read(&self, request: ReadRequest) -> Result { + match self.operate(WorkdirSessionOperation::Read(request))? { + WorkdirSessionOperationResult::Read(result) => Ok(result), + _ => Err(Self::mismatch("read")), + } + } + + async fn write(&self, request: WriteRequest) -> Result { + match self.operate(WorkdirSessionOperation::Write(request))? { + WorkdirSessionOperationResult::Write(result) => Ok(result), + _ => Err(Self::mismatch("write")), + } + } + + async fn edit(&self, request: EditRequest) -> Result { + match self.operate(WorkdirSessionOperation::Edit(request))? { + WorkdirSessionOperationResult::Edit(result) => Ok(result), + _ => Err(Self::mismatch("edit")), + } + } + + async fn list(&self, request: ListRequest) -> Result { + match self.operate(WorkdirSessionOperation::List(request))? { + WorkdirSessionOperationResult::List(result) => Ok(result), + _ => Err(Self::mismatch("list")), + } + } + + async fn glob(&self, request: GlobRequest) -> Result { + match self.operate(WorkdirSessionOperation::Glob(request))? { + WorkdirSessionOperationResult::Glob(result) => Ok(result), + _ => Err(Self::mismatch("glob")), + } + } + + async fn grep(&self, request: GrepRequest) -> Result { + match self.operate(WorkdirSessionOperation::Grep(request))? { + WorkdirSessionOperationResult::Grep(result) => Ok(result), + _ => Err(Self::mismatch("grep")), + } + } + + async fn start_command(&self, request: CommandRequest) -> Result { + match self.operate(WorkdirSessionOperation::CommandStart(request))? { + WorkdirSessionOperationResult::CommandStart(result) => Ok(result), + _ => Err(Self::mismatch("command_start")), + } + } + + async fn command_status(&self, handle: CommandHandle) -> Result { + match self.operate(WorkdirSessionOperation::CommandStatus(handle))? { + WorkdirSessionOperationResult::CommandStatus(result) => Ok(result), + _ => Err(Self::mismatch("command_status")), + } + } + + async fn command_output( + &self, + request: CommandOutputRequest, + ) -> Result { + 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) -> Self { Self { client } @@ -154,6 +345,37 @@ impl WorkspaceHttpWorkdirBackend { ) } + fn attach(&self, input: WorkdirAttachInput) -> Result { + 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 { + let workspace_id = encode_path_segment(self.workspace_id()?); + self.execute_json::(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 { + let workspace_id = encode_path_segment(self.workspace_id()?); + let response = self.execute_json::(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 { 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::(input_json)?), + WorkdirOperation::Attach => self + .backend + .attach(parse_input::(input_json)?), + WorkdirOperation::Detach => { + let _input = parse_input::(input_json)?; + self.backend.detach() + } WorkdirOperation::Delete => self .backend .delete(parse_input::(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, } +#[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, - #[serde(default, skip_serializing_if = "Option::is_none")] - primary_worker_id: Option, + #[serde( + default, + alias = "primary_worker_id", + skip_serializing_if = "Option::is_none" + )] + attached_worker_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] occupied_by: Option, } @@ -525,7 +797,16 @@ mod tests { .collect::>(); 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] diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index 4e10c56d..37e6a5b0 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -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 { + 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, diff --git a/crates/workspace-server/src/lib.rs b/crates/workspace-server/src/lib.rs index f9cff884..ff111ff6 100644 --- a/crates/workspace-server/src/lib.rs +++ b/crates/workspace-server/src/lib.rs @@ -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}")] diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index f0aca945..0d6ac88b 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -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>>, + workdir_session_locks: Arc>>>>, } 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 { 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::().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 { + 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> { + 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 { + 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, + AxumPath(path): AxumPath, + headers: HeaderMap, + Json(request): Json, +) -> ApiResult> { + 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, + AxumPath(path): AxumPath, + headers: HeaderMap, +) -> ApiResult> { + 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, + AxumPath(path): AxumPath, + headers: HeaderMap, + Json(operation): Json, +) -> ApiResult> { + 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 { + 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, ) -> ApiResult> { 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 { - workspace_id: api.config.workspace_id.clone(), - runtime_id: worker_record.runtime_id.clone(), - runtime_worker_id: worker_record.runtime_worker_id.clone(), - workdir_id: workdir_id.to_string(), - role: "primary_cwd".to_string(), - linked_at: timestamp, - unlinked_at: None, - })?; + 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, + workdir_id: workdir_id.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::().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!( diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index 69acbcc3..e1e2c0a4 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -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>; fn delete_workdir_registry(&self, workspace_id: &str, workdir_id: &str) -> Result; - 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; + fn attach_worker_workdir( + &self, + record: &WorkerWorkdirLinkRecord, + ) -> Result; + fn detach_worker_workdir( + &self, + workspace_id: &str, + runtime_id: &str, + runtime_worker_id: u64, + expected_workdir_id: Option<&str>, + unlinked_at: &str, + ) -> Result>; + fn worker_workdir_link_history_exists( + &self, + workspace_id: &str, + runtime_id: &str, + runtime_worker_id: u64, + ) -> Result; fn list_worker_workdir_links( &self, workspace_id: &str, @@ -1680,10 +1720,18 @@ impl ControlPlaneStore for SqliteWorkspaceStore { runtime_worker_id: u64, ) -> Result { 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 { 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 { + 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 { + 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> { + 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 { + 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(),