diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index da3cead9d..2235614bc 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -2275,6 +2275,9 @@ impl Engine { info!("Paused during response stream"); } self.timeline.abort_current_block(); + // Preserve any UsageEvent already received for billing, but + // do not treat this paused request as a normal completion. + self.timeline.flush_usage(); return Err(EngineError::PauseRequested); } cancel = self.cancel_rx.recv() => { diff --git a/crates/agen/tests/engine_state_test.rs b/crates/agen/tests/engine_state_test.rs index 5060015d5..051f94394 100644 --- a/crates/agen/tests/engine_state_test.rs +++ b/crates/agen/tests/engine_state_test.rs @@ -7,6 +7,7 @@ mod common; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; +use std::time::Duration; use agen::Item; use agen::interceptor::{ @@ -18,16 +19,67 @@ use agen::interceptor::{ }; use agen::llm_client::{ ClientError, LlmClient, Request, ResponseStream, - event::{Event, ResponseStatus, StatusEvent}, + event::{Event, ResponseStatus, StatusEvent, UsageEvent}, }; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::{Engine, EngineError, EngineRunExit, History, RunInterruptionReason}; use async_trait::async_trait; use common::MockLlmClient; +use futures::StreamExt; -// ============================================================================= -// Mutable State Tests -// ============================================================================= +#[derive(Clone)] +struct UsageThenPendingClient; + +#[async_trait] +impl LlmClient for UsageThenPendingClient { + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + async fn stream(&self, _request: Request) -> Result { + let usage = Event::Usage(UsageEvent { + input_tokens: Some(90), + output_tokens: Some(2), + total_tokens: Some(92), + cache_read_input_tokens: Some(0), + cache_creation_input_tokens: Some(0), + }); + Ok(Box::pin( + futures::stream::once(async move { Ok(usage) }) + .chain(futures::stream::pending::>()), + )) + } +} + +#[tokio::test] +async fn pause_during_response_stream_flushes_observed_usage() { + let observed = Arc::new(Mutex::new(Vec::new())); + let observed_for_callback = Arc::clone(&observed); + let mut engine = Engine::new(UsageThenPendingClient); + engine.on_usage(move |usage| { + observed_for_callback.lock().unwrap().push(usage.clone()); + }); + let pause = engine.pause_sender(); + let mut history = History::new(); + let mut engine = engine.lock(&history); + + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(20)).await; + pause.send(()).await.unwrap(); + }); + + let exit = tokio::time::timeout( + Duration::from_secs(2), + engine.run(&mut history, "pause after usage"), + ) + .await + .expect("run should observe pause"); + assert!(matches!(exit, EngineRunExit::Paused)); + let observed = observed.lock().unwrap(); + assert_eq!(observed.len(), 1); + assert_eq!(observed[0].input_tokens, Some(90)); + assert_eq!(observed[0].output_tokens, Some(2)); +} /// Verify that system prompt can be set in Mutable state #[test] diff --git a/crates/client/src/backend_runtime.rs b/crates/client/src/backend_runtime.rs index 77c352827..a3e78e3dc 100644 --- a/crates/client/src/backend_runtime.rs +++ b/crates/client/src/backend_runtime.rs @@ -943,7 +943,7 @@ mod tests { profile: Some("builtin:coder".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory: None, + workdir_attachments: Vec::new(), control_operation_id: None, }; @@ -964,7 +964,7 @@ mod tests { assert_eq!(body["display_name"], "Coder one"); assert_eq!(body["profile"], "builtin:coder"); assert_eq!(body["initial_submit"], serde_json::json!([])); - assert_eq!(body["working_directory"], serde_json::Value::Null); + assert_eq!(body["workdir_attachments"], serde_json::json!([])); } #[test] @@ -1014,29 +1014,35 @@ mod tests { "state": "idle", "implementation": {"kind": "worker", "display_hint": "Coder"}, "capabilities": {"can_stop": true, "can_spawn_followup": false}, - "working_directory": { - "working_directory_id": "wd-1", - "repository_key": "main", - "materializer_kind": "runtime_git_clone", - "status": "active", - "occupied_by": { - "runtime_id": "arcadia", - "worker_id": "worker-opaque-64", - "display_name": "Coder", - "linked_at": "2026-08-12T00:00:00Z" + "workdir_attachments": [{ + "alias": "checkout", + "working_directory": { + "working_directory_id": "wd-1", + "repository_key": "main", + "materializer_kind": "runtime_git_clone", + "status": "active", + "occupied_by": { + "runtime_id": "arcadia", + "worker_id": "worker-opaque-64", + "display_name": "Coder", + "linked_at": "2026-08-12T00:00:00Z" + } } - } + }] }); let worker: BackendWorkerSummary = serde_json::from_value(payload.clone()).unwrap(); - let workdir = worker.working_directory.unwrap(); + let attachment = worker.workdir_attachments.into_iter().next().unwrap(); + assert_eq!(attachment.alias, "checkout"); + let workdir = attachment.working_directory; assert_eq!(workdir.repository_key, "main"); let occupied_by = workdir.occupied_by.expect("occupied Workdir"); assert_eq!(occupied_by.runtime_id, "arcadia"); assert_eq!(occupied_by.worker_id, "worker-opaque-64"); let mut stale = payload; - stale["working_directory"]["occupied_by"]["runtime_worker_id"] = serde_json::json!(64); + stale["workdir_attachments"][0]["working_directory"]["occupied_by"]["runtime_worker_id"] = + serde_json::json!(64); assert!(serde_json::from_value::(stale).is_err()); } diff --git a/crates/client/src/workspace_product.rs b/crates/client/src/workspace_product.rs index 4ba758f17..f69c61927 100644 --- a/crates/client/src/workspace_product.rs +++ b/crates/client/src/workspace_product.rs @@ -312,7 +312,7 @@ impl BackendWorkspaceProductClient { initial_submit: vec![protocol::Segment::Text { content: format!("Please handle intake for Ticket {ticket_id}."), }], - working_directory: None, + workdir_attachments: Vec::new(), control_operation_id: None, }; let response: BrowserCreateWorkerResponse = diff --git a/crates/protocol/src/subscription.rs b/crates/protocol/src/subscription.rs index 6d25c6d04..3f017b640 100644 --- a/crates/protocol/src/subscription.rs +++ b/crates/protocol/src/subscription.rs @@ -572,6 +572,41 @@ pub enum SubscriptionWorkerAvailability { Unavailable, } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +pub struct SubscriptionWorkerWorkdirAttachment { + /// Stable Worker-local routing alias. + pub alias: String, + /// Runtime-internal Repository id. Workspace-facing TypeScript contracts + /// omit this field and require `repository_key` from the Server projection. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "typescript", ts(skip))] + pub repository_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub repository_key: Option, + pub working_directory_id: SubscriptionWorkdirId, +} + +impl SubscriptionWorkerWorkdirAttachment { + pub fn validate(&self) -> Result<(), SubscriptionValidationError> { + validate_identifier( + "workdir_attachment_alias", + &self.alias, + MAX_RESOURCE_ID_BYTES, + )?; + self.working_directory_id.validate()?; + match (&self.repository_id, &self.repository_key) { + (Some(repository_id), None) => { + validate_identifier("repository_id", repository_id, MAX_RESOURCE_ID_BYTES) + } + (None, Some(repository_key)) => validate_repository_key(repository_key), + _ => Err(SubscriptionValidationError::InvalidIdentifier { + field: "repository_authority", + }), + } + } +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[cfg_attr(feature = "typescript", derive(ts_rs::TS))] pub struct SubscriptionWorker { @@ -604,15 +639,8 @@ pub struct SubscriptionWorker { pub display_name: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub profile: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - #[cfg_attr(feature = "typescript", ts(skip))] - pub repository_id: Option, - /// Workspace-facing Repository key. Runtime producers leave this unset and - /// Workspace Server projections replace `repository_id` with this field. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub repository_key: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub working_directory_id: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub workdir_attachments: Vec, } impl SubscriptionWorker { @@ -624,19 +652,20 @@ impl SubscriptionWorker { if let Some(resource_key) = &self.resource_key { validate_identifier("resource_key", resource_key, MAX_RESOURCE_ID_BYTES)?; } - if let Some(repository_id) = &self.repository_id { - validate_identifier("repository_id", repository_id, MAX_RESOURCE_ID_BYTES)?; - } - if let Some(repository_key) = &self.repository_key { - validate_repository_key(repository_key)?; - } - if self.repository_id.is_some() && self.repository_key.is_some() { - return Err(SubscriptionValidationError::InvalidIdentifier { - field: "repository_authority", - }); - } - if let Some(working_directory_id) = &self.working_directory_id { - working_directory_id.validate()?; + let mut aliases = HashSet::new(); + let mut workdir_ids = HashSet::new(); + for attachment in &self.workdir_attachments { + attachment.validate()?; + if !aliases.insert(attachment.alias.as_str()) { + return Err(SubscriptionValidationError::InvalidIdentifier { + field: "workdir_attachment_alias", + }); + } + if !workdir_ids.insert(attachment.working_directory_id.as_str()) { + return Err(SubscriptionValidationError::InvalidIdentifier { + field: "working_directory_id", + }); + } } Ok(()) } @@ -654,8 +683,6 @@ pub struct SubscriptionWorkdir { #[serde(default, skip_serializing_if = "Option::is_none")] pub repository_key: Option, pub state: String, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub primary_worker_id: Option, } impl SubscriptionWorkdir { @@ -673,9 +700,6 @@ impl SubscriptionWorkdir { } } validate_identifier("workdir_state", &self.state, MAX_RESOURCE_ID_BYTES)?; - if let Some(worker_id) = &self.primary_worker_id { - worker_id.validate()?; - } Ok(()) } } @@ -688,8 +712,6 @@ pub struct WorkspaceSubscriptionWorkdir { pub working_directory_id: SubscriptionWorkdirId, pub repository_key: String, pub state: String, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub primary_worker_id: Option, } impl WorkspaceSubscriptionWorkdir { @@ -697,9 +719,6 @@ impl WorkspaceSubscriptionWorkdir { self.working_directory_id.validate()?; validate_repository_key(&self.repository_key)?; validate_identifier("workdir_state", &self.state, MAX_RESOURCE_ID_BYTES)?; - if let Some(worker_id) = &self.primary_worker_id { - worker_id.validate()?; - } Ok(()) } } @@ -903,36 +922,58 @@ mod tests { workspace_id: Some("workspace-1".to_string()), display_name: Some(format!("Worker {value}")), profile: Some("builtin:coder".to_string()), - repository_id: None, - repository_key: None, - working_directory_id: None, + workdir_attachments: Vec::new(), } } #[test] fn runtime_and_workspace_repository_identity_projections_do_not_alias() { let mut runtime_worker = worker("worker-1"); - runtime_worker.repository_id = Some("01890f47-3c22-7cc0-98c4-dc0c0c07398f".to_string()); + runtime_worker + .workdir_attachments + .push(SubscriptionWorkerWorkdirAttachment { + alias: "checkout".to_string(), + repository_id: Some("01890f47-3c22-7cc0-98c4-dc0c0c07398f".to_string()), + repository_key: None, + working_directory_id: SubscriptionWorkdirId::new("workdir-1").unwrap(), + }); runtime_worker.validate().unwrap(); let runtime_json = serde_json::to_value(&runtime_worker).unwrap(); assert_eq!( - runtime_json["repository_id"], + runtime_json["workdir_attachments"][0]["repository_id"], "01890f47-3c22-7cc0-98c4-dc0c0c07398f" ); - assert!(runtime_json.get("repository_key").is_none()); + assert!( + runtime_json["workdir_attachments"][0] + .get("repository_key") + .is_none() + ); let mut workspace_worker = worker("worker-1"); - workspace_worker.repository_key = Some("main".to_string()); + workspace_worker + .workdir_attachments + .push(SubscriptionWorkerWorkdirAttachment { + alias: "checkout".to_string(), + repository_id: None, + repository_key: Some("main".to_string()), + working_directory_id: SubscriptionWorkdirId::new("workdir-1").unwrap(), + }); workspace_worker.validate().unwrap(); let workspace_json = serde_json::to_value(&workspace_worker).unwrap(); - assert_eq!(workspace_json["repository_key"], "main"); - assert!(workspace_json.get("repository_id").is_none()); + assert_eq!( + workspace_json["workdir_attachments"][0]["repository_key"], + "main" + ); + assert!( + workspace_json["workdir_attachments"][0] + .get("repository_id") + .is_none() + ); let workspace_workdir = WorkspaceSubscriptionWorkdir { working_directory_id: SubscriptionWorkdirId::new("workdir-1").unwrap(), repository_key: "main".to_string(), state: "active".to_string(), - primary_worker_id: Some(worker_id("worker-1")), }; workspace_workdir.validate().unwrap(); let workdir_json = serde_json::to_value(&workspace_workdir).unwrap(); diff --git a/crates/protocol/src/typescript.rs b/crates/protocol/src/typescript.rs index 9ea536cd9..8ceb9ccd7 100644 --- a/crates/protocol/src/typescript.rs +++ b/crates/protocol/src/typescript.rs @@ -22,7 +22,8 @@ use crate::{ SubscriptionRequestId, SubscriptionResponse, SubscriptionSnapshot, SubscriptionTerminationCode, SubscriptionWorkdirId, SubscriptionWorker, SubscriptionWorkerAvailability, SubscriptionWorkerId, SubscriptionWorkerIds, - SubscriptionWorkerProtocolMethod, SubscriptionWorkerState, WorkspaceSubscriptionWorkdir, + SubscriptionWorkerProtocolMethod, SubscriptionWorkerState, + SubscriptionWorkerWorkdirAttachment, WorkspaceSubscriptionWorkdir, }, }; @@ -114,6 +115,7 @@ pub fn generated_protocol_types() -> String { push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); diff --git a/crates/runtime-api/src/lib.rs b/crates/runtime-api/src/lib.rs index c6ab08b3a..c15b04a0a 100644 --- a/crates/runtime-api/src/lib.rs +++ b/crates/runtime-api/src/lib.rs @@ -335,6 +335,8 @@ pub struct RepositoryMaterializationContext { #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] pub struct WorkingDirectoryRequest { pub repository: WorkingDirectoryRepository, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub display_name: Option, #[serde(default)] pub materializer: MaterializerKind, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -344,12 +346,31 @@ pub struct WorkingDirectoryRequest { } #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] -pub struct WorkingDirectoryClaim { +pub struct WorkingDirectoryAttachmentClaim { + /// Stable Worker-local routing key. This is not a Workdir id or display name. + pub alias: workdir::WorkdirAttachmentAlias, pub working_directory_id: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub relative_cwd: Option, } +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct WorkerWorkdirAttachmentsRequest { + pub workdir_attachments: Vec, +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct WorkingDirectoryAttachmentRequest { + pub alias: workdir::WorkdirAttachmentAlias, + pub working_directory: WorkingDirectoryRequest, +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct WorkingDirectoryAttachmentStatus { + pub alias: workdir::WorkdirAttachmentAlias, + pub working_directory: WorkingDirectoryStatus, +} + #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] pub struct WorkingDirectoryStatus { pub summary: RuntimeWorkingDirectorySummary, @@ -387,10 +408,10 @@ pub struct CreateWorkerRequest { pub config_bundle: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub initial_input: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub working_directory_request: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub working_directory: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub workdir_attachment_requests: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub workdir_attachments: Vec, #[serde(default, skip_serializing_if = "std::ops::Not::not")] pub worker_observation_enabled: bool, #[serde(default, skip_serializing_if = "Vec::is_empty")] @@ -413,8 +434,8 @@ pub struct WorkerSummary { pub worker_state: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub workspace_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub working_directory: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub workdir_attachments: Vec, pub profile: ProfileSelector, #[serde(default, skip_serializing_if = "Option::is_none")] pub display_name: Option, @@ -726,6 +747,17 @@ pub trait RuntimeApi { #[body] request: WorkerWorkspaceApiRequest, ) -> Result; + #[post( + "/v1/workers/{worker_id}/workdir-attachments", + status = 200, + error_status = 400 + )] + async fn replace_worker_workdir_attachments( + &self, + #[path] worker_id: String, + #[body] request: WorkerWorkdirAttachmentsRequest, + ) -> Result; + #[post( "/v1/workers/{worker_id}/completions", status = 200, @@ -1002,6 +1034,13 @@ mod tests { ) -> Result { Err(test_error(501)) } + async fn replace_worker_workdir_attachments( + &self, + _worker_id: String, + _request: WorkerWorkdirAttachmentsRequest, + ) -> Result { + Err(test_error(501)) + } async fn complete_worker_arguments( &self, _worker_id: String, @@ -1072,7 +1111,7 @@ mod tests { #[test] fn contract_inventory_is_complete_and_unique() { let operations = RuntimeApiMetadata::OPERATIONS; - assert_eq!(operations.len(), 15); + assert_eq!(operations.len(), 16); let mut routes = operations .iter() .map(|operation| (format!("{:?}", operation.method), operation.path)) diff --git a/crates/server-api/src/lib.rs b/crates/server-api/src/lib.rs index ebe0d0ccf..62f51a1c3 100644 --- a/crates/server-api/src/lib.rs +++ b/crates/server-api/src/lib.rs @@ -1583,6 +1583,9 @@ pub struct RuntimeWorkingDirectoryCleanupTarget { #[serde(deny_unknown_fields)] pub struct RuntimeWorkingDirectorySummary { pub working_directory_id: String, + /// Optional human-facing label, never a routing key. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub display_name: Option, pub repository_id: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub creation_selector: Option, @@ -1606,8 +1609,6 @@ pub struct RuntimeWorkingDirectorySummary { #[serde(default, skip_serializing_if = "Option::is_none")] pub cleanliness: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - pub primary_worker_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] pub occupied_by: Option, } @@ -1618,6 +1619,9 @@ pub struct RuntimeWorkingDirectorySummary { #[serde(deny_unknown_fields)] pub struct WorkingDirectorySummary { pub working_directory_id: String, + /// Optional human-facing label, never a routing key. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub display_name: Option, pub repository_key: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub creation_selector: Option, @@ -1641,8 +1645,6 @@ pub struct WorkingDirectorySummary { #[serde(default, skip_serializing_if = "Option::is_none")] pub cleanliness: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - pub primary_worker_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] pub occupied_by: Option, } @@ -1713,6 +1715,7 @@ impl From for RuntimeWorking fn from(value: workdir::workspace::RuntimeWorkingDirectorySummary) -> Self { Self { working_directory_id: value.working_directory_id, + display_name: value.display_name, repository_id: value.repository_id, creation_selector: value.creation_selector, creation_ref: value.creation_ref, @@ -1725,7 +1728,6 @@ impl From for RuntimeWorking cleanup_target: value.cleanup_target.map(Into::into), status: value.status.into(), cleanliness: value.cleanliness, - primary_worker_id: value.primary_worker_id, occupied_by: value.occupied_by.map(Into::into), } } @@ -1735,6 +1737,7 @@ impl From for WorkingDirectorySumma fn from(value: workdir::workspace::WorkingDirectorySummary) -> Self { Self { working_directory_id: value.working_directory_id, + display_name: value.display_name, repository_key: value.repository_key, creation_selector: value.creation_selector, creation_ref: value.creation_ref, @@ -1747,7 +1750,6 @@ impl From for WorkingDirectorySumma cleanup_target: value.cleanup_target.map(Into::into), status: value.status.into(), cleanliness: value.cleanliness, - primary_worker_id: value.primary_worker_id, occupied_by: value.occupied_by.map(Into::into), } } @@ -1765,6 +1767,9 @@ impl From for WorkingDirectorySumma pub struct WorkingDirectoryCreateRequest { #[serde(default, skip_serializing_if = "Option::is_none")] pub runtime_id: Option, + /// Optional human-facing label, independent of every Worker attachment alias. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub display_name: Option, pub repository_key: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub selector: Option, @@ -2445,6 +2450,24 @@ pub struct WorkspaceWorkerDiscoveryPage { pub next_cursor: Option, } +/// One Workspace Worker Workdir attachment. `alias` is the stable Worker-local +/// routing key; the nested Workdir id and display name are inventory metadata. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(deny_unknown_fields)] +pub struct WorkerWorkdirAttachmentSummary { + pub alias: String, + pub working_directory: WorkingDirectorySummary, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(deny_unknown_fields)] +pub struct RuntimeWorkerWorkdirAttachmentSummary { + pub alias: String, + pub working_directory: RuntimeWorkingDirectorySummary, +} + /// Workspace-authoritative Worker projection. /// /// `resource_key` is required here even though Runtime-internal Worker summaries @@ -2478,8 +2501,8 @@ pub struct WorkerSummary { pub retention_state: String, pub implementation: WorkerImplementationSummary, pub capabilities: WorkerCapabilitySummary, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub working_directory: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub workdir_attachments: Vec, #[serde(default)] pub diagnostics: Vec, } @@ -2507,9 +2530,8 @@ pub struct WorkerLaunchWorkerSummary { pub retention_state: String, pub implementation: WorkerImplementationSummary, pub capabilities: WorkerCapabilitySummary, - #[serde(default, skip_serializing_if = "Option::is_none")] - #[cfg_attr(feature = "typescript", ts(optional = nullable))] - pub working_directory: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub workdir_attachments: Vec, #[serde(default)] pub diagnostics: Vec, } @@ -2563,6 +2585,8 @@ pub struct WorkingDirectoryRepositoryOption { #[cfg_attr(feature = "typescript", derive(ts_rs::TS))] #[serde(deny_unknown_fields)] pub struct BrowserWorkerWorkingDirectorySelection { + /// Stable Worker-local routing alias. + pub alias: String, pub working_directory_id: String, #[serde(default)] pub relative_cwd: Option, @@ -2589,7 +2613,7 @@ pub struct CreateWorkspaceWorkerRequest { #[serde(default)] pub initial_submit: Vec, #[serde(default)] - pub working_directory: Option, + pub workdir_attachments: Vec, /// Backend idempotency key used only for authenticated Worker-owned spawn/control. #[serde(default)] pub control_operation_id: Option, @@ -3663,6 +3687,7 @@ pub fn worker_launch_api_typescript() -> String { WorkerWorkspaceSummary::decl(&config), WorkerImplementationSummary::decl(&config), WorkerCapabilitySummary::decl(&config), + RuntimeWorkerWorkdirAttachmentSummary::decl(&config), WorkerLaunchWorkerSummary::decl(&config), WorkerLaunchRuntimeOption::decl(&config), WorkerLaunchProfileCandidate::decl(&config), @@ -4100,7 +4125,7 @@ mod tests { can_stop: true, can_spawn_followup: false, }, - working_directory: None, + workdir_attachments: Vec::new(), diagnostics: Vec::new(), } } @@ -4178,7 +4203,7 @@ mod tests { profile: None, ticket_assignment: None, initial_submit: Vec::new(), - working_directory: None, + workdir_attachments: Vec::new(), control_operation_id: None, }) .unwrap(); @@ -4190,7 +4215,7 @@ mod tests { "profile": null, "ticket_assignment": null, "initial_submit": [], - "working_directory": null, + "workdir_attachments": [], "control_operation_id": null, }) ); @@ -5097,6 +5122,7 @@ mod tests { fn workdir_summary_omits_absent_optional_fields_on_the_wire() { let value = serde_json::to_value(WorkingDirectorySummary { working_directory_id: "workdir-1".into(), + display_name: None, repository_key: "main".into(), creation_selector: None, creation_ref: None, @@ -5109,7 +5135,6 @@ mod tests { cleanup_target: None, status: WorkingDirectoryStatusKind::Active, cleanliness: None, - primary_worker_id: None, occupied_by: None, }) .expect("serialize Workdir summary"); @@ -5125,7 +5150,6 @@ mod tests { "observed_at_epoch_seconds", "cleanup_target", "cleanliness", - "primary_worker_id", "occupied_by", ] { assert!( diff --git a/crates/tools/src/bash.rs b/crates/tools/src/bash.rs index 34e979b56..a213a30eb 100644 --- a/crates/tools/src/bash.rs +++ b/crates/tools/src/bash.rs @@ -6,7 +6,9 @@ use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use async_trait::async_trait; use schemars::JsonSchema; use serde::Deserialize; -use workdir::{CommandHandle, CommandOutputRequest, CommandRequest, WorkdirSessionHandle}; +use workdir::{ + CommandHandle, CommandOutputRequest, CommandRequest, WorkdirSessionHandle, WorkdirSessionRouter, +}; const DEFAULT_TIMEOUT_SECS: u64 = 120; const MAX_TIMEOUT_SECS: u64 = 600; @@ -14,13 +16,16 @@ const INLINE_BYTE_BUDGET: usize = 12 * 1024; #[derive(Debug, Deserialize, JsonSchema)] struct BashParams { + /// Worker-local alias of the Workdir attachment to use. + #[serde(default)] + target_workdir: Option, command: String, #[serde(default)] timeout: Option, } pub(crate) struct BashTool { - session: WorkdirSessionHandle, + router: Arc, output_dir: PathBuf, state: Arc>, } @@ -29,6 +34,7 @@ pub(crate) struct BashTool { struct ActiveCommand { call_id: String, execution_nonce: u64, + session: WorkdirSessionHandle, handle: CommandHandle, } @@ -93,6 +99,12 @@ impl Tool for BashTool { ) -> Result { let params: BashParams = serde_json::from_str(input_json) .map_err(|error| ToolError::InvalidArgument(format!("invalid Bash input: {error}")))?; + let selected = crate::routing::resolve_session( + &self.router, + params.target_workdir.as_deref(), + workdir::WorkdirSessionCapability::Command, + )?; + let session = selected.session; let timeout_secs = params .timeout .unwrap_or(DEFAULT_TIMEOUT_SECS) @@ -106,14 +118,13 @@ impl Tool for BashTool { state.next_execution_nonce }; let mut guard = CommandGuard { - session: self.session.clone(), + session: session.clone(), state: self.state.clone(), execution_id: execution_id.clone(), execution_nonce, handle: None, }; - let handle = self - .session + let handle = session .start_command(CommandRequest { command: params.command, timeout_secs, @@ -131,6 +142,7 @@ impl Tool for BashTool { ActiveCommand { call_id: call_id.clone(), execution_nonce, + session: session.clone(), handle: handle.clone(), }, ); @@ -139,13 +151,12 @@ impl Tool for BashTool { }; guard.handle = Some(handle.clone()); if cancel_after_start { - self.session + session .cancel_command(handle.clone()) .await .map_err(crate::ToolsError::from)?; } - let output = self - .session + let output = session .command_output(CommandOutputRequest { handle, cursor: 0, @@ -222,12 +233,13 @@ impl Tool for BashTool { .active .values() .filter(|active| active.call_id == call_id) - .map(|active| active.handle.clone()) + .cloned() .collect::>() }; - for handle in handles { - self.session - .cancel_command(handle) + for active in handles { + active + .session + .cancel_command(active.handle) .await .map_err(crate::ToolsError::from)?; } @@ -242,14 +254,12 @@ impl Tool for BashTool { let handle = { let mut state = self.state.lock().unwrap(); state.cancellation_requested.insert(execution_id.clone()); - state - .active - .get(&execution_id) - .map(|active| active.handle.clone()) + state.active.get(&execution_id).cloned() }; - if let Some(handle) = handle { - self.session - .cancel_command(handle) + if let Some(active) = handle { + active + .session + .cancel_command(active.handle) .await .map_err(crate::ToolsError::from)?; } @@ -268,13 +278,20 @@ fn truncate_for_summary(command: &str) -> String { } pub fn bash_tool(session: WorkdirSessionHandle, output_dir: PathBuf) -> ToolDefinition { + routed_bash_tool(crate::routing::singleton_router(session), output_dir) +} + +pub(crate) fn routed_bash_tool( + router: Arc, + output_dir: PathBuf, +) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(BashParams); let meta = ToolMeta::new("Bash") - .description("Execute a shell command in the bound Workdir. Process start, bounded inline output, full-output spill, timeout and cancellation are owned by the WorkdirSession provider. This is not a sandbox.") + .description("Execute a shell command in one selected Workdir. Process start, bounded inline output, full-output spill, timeout and cancellation are owned by that WorkdirSession provider. This is not a sandbox.") .input_schema(serde_json::to_value(schema).expect("Bash schema serialization")); let tool: Arc = Arc::new(BashTool { - session: session.clone(), + router: router.clone(), output_dir: output_dir.clone(), state: Arc::new(Mutex::new(BashExecutionState::default())), }); diff --git a/crates/tools/src/edit.rs b/crates/tools/src/edit.rs index a303628fb..4a22cb218 100644 --- a/crates/tools/src/edit.rs +++ b/crates/tools/src/edit.rs @@ -9,15 +9,18 @@ use serde::Deserialize; use crate::error::ToolsError; use crate::tracker::Tracker; -use workdir::{EditRequest, WorkdirPath, WorkdirSessionHandle}; +use workdir::{EditRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter}; -const DESCRIPTION: &str = "Replace a substring in an existing file. By default \ +const DESCRIPTION: &str = "Replace a substring in an existing file in the selected Workdir attachment. By default \ `old_string` must be unique in the file; set `replace_all: true` to replace \ -every occurrence. The file must have been read first (via the Read tool) in \ -this session. Paths are relative to the bound Workdir."; +every occurrence. The file must have been read first through the same attachment. \ +When exactly one Workdir is attached, target_workdir may be omitted."; #[derive(Debug, Deserialize, schemars::JsonSchema)] pub(crate) struct EditParams { + /// Worker-local alias of the Workdir attachment to use. + #[serde(default)] + pub target_workdir: Option, /// Logical path relative to the bound Workdir root. pub file_path: String, /// String to replace. Must be unique in the file unless `replace_all` is true. @@ -30,7 +33,7 @@ pub(crate) struct EditParams { } pub(crate) struct EditTool { - session: WorkdirSessionHandle, + router: Arc, tracker: Tracker, } @@ -44,6 +47,14 @@ impl Tool for EditTool { let params: EditParams = serde_json::from_str(input_json) .map_err(|e| ToolError::InvalidArgument(format!("invalid Edit input: {e}")))?; + let selected = crate::routing::resolve_session( + &self.router, + params.target_workdir.as_deref(), + workdir::WorkdirSessionCapability::Edit, + )?; + let tracker = self + .tracker + .scoped_attachment(&selected.alias, selected.generation); let path = WorkdirPath::new(¶ms.file_path).map_err(ToolsError::from)?; tracing::debug!(path = %path, replace_all = params.replace_all, "Edit"); @@ -59,9 +70,9 @@ impl Tool for EditTool { } let mutation_key = PathBuf::from(path.as_str()); - let _mutation_permit = self.tracker.acquire_mutation(&mutation_key, &ctx).await; - let expected_hash = self.tracker.expected_workdir_hash(&path)?; - let result = self + let _mutation_permit = tracker.acquire_mutation(&mutation_key, &ctx).await; + let expected_hash = tracker.expected_workdir_hash(&path)?; + let result = selected .session .edit(EditRequest { path: path.clone(), @@ -73,7 +84,7 @@ impl Tool for EditTool { .await .map_err(ToolsError::from)?; let replacements = result.replacements; - self.tracker.record_workdir_edit( + tracker.record_workdir_edit( &path, result.content_hash, replacements, @@ -122,8 +133,15 @@ fn make_preview(text: &str, needle: &str) -> String { out } -/// Factory for the `Edit` tool. +/// Factory for the `Edit` tool bound to one compatibility session. pub fn edit_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinition { + routed_edit_tool(crate::routing::singleton_router(session), tracker) +} + +pub(crate) fn routed_edit_tool( + router: Arc, + tracker: Tracker, +) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(EditParams); let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({})); @@ -131,7 +149,7 @@ pub fn edit_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinit .description(DESCRIPTION) .input_schema(schema_value); let tool: Arc = Arc::new(EditTool { - session: session.clone(), + router: router.clone(), tracker: tracker.clone(), }); (meta, tool) diff --git a/crates/tools/src/glob.rs b/crates/tools/src/glob.rs index f2e8616e7..0c2914dfd 100644 --- a/crates/tools/src/glob.rs +++ b/crates/tools/src/glob.rs @@ -4,7 +4,7 @@ use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use async_trait::async_trait; use schemars::JsonSchema; use serde::Deserialize; -use workdir::{GlobRequest, WorkdirPath, WorkdirSessionHandle}; +use workdir::{GlobRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter}; use crate::ToolsError; @@ -12,6 +12,9 @@ const RESULT_LIMIT: usize = 1000; #[derive(Debug, Deserialize, JsonSchema)] struct GlobParams { + /// Worker-local alias of the Workdir attachment to use. + #[serde(default)] + target_workdir: Option, /// Glob pattern, for example `**/*.rs` or `src/**/test_*.py`. pattern: String, /// Logical Workdir-relative directory. Defaults to the Workdir root. @@ -20,7 +23,7 @@ struct GlobParams { } struct GlobTool { - session: WorkdirSessionHandle, + router: Arc, } #[async_trait] @@ -32,13 +35,18 @@ impl Tool for GlobTool { ) -> Result { let params: GlobParams = serde_json::from_str(input_json) .map_err(|error| ToolError::InvalidArgument(format!("invalid Glob input: {error}")))?; + let selected = crate::routing::resolve_session( + &self.router, + params.target_workdir.as_deref(), + workdir::WorkdirSessionCapability::Glob, + )?; let path = match params.path { Some(path) => WorkdirPath::new(&path).map_err(ToolsError::from)?, None => WorkdirPath::root(), }; let pattern = params.pattern; tracing::debug!(%pattern, %path, "Glob"); - let result = self + let result = selected .session .glob(GlobRequest { pattern: pattern.clone(), @@ -75,13 +83,17 @@ impl Tool for GlobTool { } pub fn glob_tool(session: WorkdirSessionHandle) -> ToolDefinition { + routed_glob_tool(crate::routing::singleton_router(session)) +} + +pub(crate) fn routed_glob_tool(router: Arc) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(GlobParams); let meta = ToolMeta::new("Glob") - .description("Find files matching a glob pattern inside the bound Workdir. Results are sorted and capped at 1000 entries. Paths are Workdir-relative.") + .description("Find files matching a glob pattern inside the selected Workdir. Results are sorted and capped at 1000 entries. Paths are Workdir-relative.") .input_schema(serde_json::to_value(schema).expect("Glob schema serialization")); let tool: Arc = Arc::new(GlobTool { - session: session.clone(), + router: router.clone(), }); (meta, tool) }) diff --git a/crates/tools/src/grep.rs b/crates/tools/src/grep.rs index 4774e6ffa..190b8590e 100644 --- a/crates/tools/src/grep.rs +++ b/crates/tools/src/grep.rs @@ -4,7 +4,9 @@ use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use async_trait::async_trait; use schemars::JsonSchema; use serde::Deserialize; -use workdir::{GrepOutputMode, GrepRequest, WorkdirPath, WorkdirSessionHandle}; +use workdir::{ + GrepOutputMode, GrepRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter, +}; use crate::ToolsError; @@ -21,6 +23,9 @@ enum OutputMode { #[derive(Debug, Deserialize, JsonSchema)] struct GrepParams { + /// Worker-local alias of the Workdir attachment to use. + #[serde(default)] + target_workdir: Option, pattern: String, /// Workdir-relative path, or an absolute path covered by readable scope. Defaults to the Workdir root. #[serde(default)] @@ -48,7 +53,7 @@ struct GrepParams { } struct GrepTool { - session: WorkdirSessionHandle, + router: Arc, } #[async_trait] @@ -60,6 +65,11 @@ impl Tool for GrepTool { ) -> Result { let params: GrepParams = serde_json::from_str(input_json) .map_err(|error| ToolError::InvalidArgument(format!("invalid Grep input: {error}")))?; + let selected = crate::routing::resolve_session( + &self.router, + params.target_workdir.as_deref(), + workdir::WorkdirSessionCapability::Grep, + )?; let path = match params.path { Some(path) => WorkdirPath::new_scoped(&path).map_err(ToolsError::from)?, None => WorkdirPath::root(), @@ -74,7 +84,7 @@ impl Tool for GrepTool { .map(|context| (context, context)) .unwrap_or((params.before.unwrap_or(0), params.after.unwrap_or(0))); let head_limit = params.head_limit.unwrap_or(DEFAULT_HEAD_LIMIT); - let result = self + let result = selected .session .grep(GrepRequest { pattern: params.pattern, @@ -126,13 +136,17 @@ impl Tool for GrepTool { } pub fn grep_tool(session: WorkdirSessionHandle) -> ToolDefinition { + routed_grep_tool(crate::routing::singleton_router(session)) +} + +pub(crate) fn routed_grep_tool(router: Arc) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(GrepParams); let meta = ToolMeta::new("Grep") - .description("Search a Workdir file or directory with a regex. Content results group lines by file; `>` marks matching lines and unmarked lines are context. Directory traversal executes inside the WorkdirSession provider. Results are bounded and Workdir-relative.") + .description("Search a file or directory in the selected Workdir with a regex. Content results group lines by file; `>` marks matching lines and unmarked lines are context. Directory traversal executes inside the WorkdirSession provider. Results are bounded and Workdir-relative.") .input_schema(serde_json::to_value(schema).expect("Grep schema serialization")); let tool: Arc = Arc::new(GrepTool { - session: session.clone(), + router: router.clone(), }); (meta, tool) }) diff --git a/crates/tools/src/lib.rs b/crates/tools/src/lib.rs index 7b838dadd..0aaabe8d6 100644 --- a/crates/tools/src/lib.rs +++ b/crates/tools/src/lib.rs @@ -17,6 +17,7 @@ mod edit; mod glob; mod grep; mod read; +mod routing; mod view_image; mod web; mod write; @@ -28,11 +29,11 @@ pub use glob::glob_tool; pub use grep::grep_tool; pub use read::read_tool; pub use tracker::{ChangeStat, Tracker}; -pub use view_image::view_image_tool; +pub use view_image::{routed_view_image_tool, view_image_tool}; pub use web::{web_fetch_tool, web_search_tool}; pub use write::write_tool; -/// Build the local filesystem/command tool surface implemented by a WorkdirSession. +/// Build the filesystem/command tool surface for a single Workdir session. /// Profile/manifest policy may narrow this set further in the Engine. pub fn core_builtin_tools( session: workdir::WorkdirSessionHandle, @@ -42,24 +43,75 @@ pub fn core_builtin_tools( use workdir::WorkdirSessionCapability; let capabilities = session.capabilities(); + let router = routing::singleton_router(session); let mut tools = Vec::with_capacity(6); if capabilities.supports(WorkdirSessionCapability::Read) { - tools.push(read_tool(session.clone(), tracker.clone())); + tools.push(read::routed_read_tool(router.clone(), tracker.clone())); } if capabilities.supports(WorkdirSessionCapability::Write) { - tools.push(write_tool(session.clone(), tracker.clone())); + tools.push(write::routed_write_tool(router.clone(), tracker.clone())); } if capabilities.supports(WorkdirSessionCapability::Edit) { - tools.push(edit_tool(session.clone(), tracker)); + tools.push(edit::routed_edit_tool(router.clone(), tracker)); } if capabilities.supports(WorkdirSessionCapability::Glob) { - tools.push(glob_tool(session.clone())); + tools.push(glob::routed_glob_tool(router.clone())); } if capabilities.supports(WorkdirSessionCapability::Grep) { - tools.push(grep_tool(session.clone())); + tools.push(grep::routed_grep_tool(router.clone())); } if capabilities.supports(WorkdirSessionCapability::Command) { - tools.push(bash_tool(session, bash_output_dir)); + tools.push(bash::routed_bash_tool(router, bash_output_dir)); + } + tools +} + +/// Build the stable routed tool surface for a Worker's live attachment set. +/// +/// The definitions remain installed while the set is empty or changes at runtime; +/// each invocation resolves `target_workdir` against the router and then lets the +/// selected provider enforce its capabilities. +pub fn routed_builtin_tools( + router: std::sync::Arc, + tracker: Tracker, + bash_output_dir: std::path::PathBuf, +) -> Vec { + vec![ + read::routed_read_tool(router.clone(), tracker.clone()), + write::routed_write_tool(router.clone(), tracker.clone()), + edit::routed_edit_tool(router.clone(), tracker), + glob::routed_glob_tool(router.clone()), + grep::routed_grep_tool(router.clone()), + bash::routed_bash_tool(router, bash_output_dir), + ] +} + +pub fn routed_builtin_tools_for_current_capabilities( + router: std::sync::Arc, + tracker: Tracker, + bash_output_dir: std::path::PathBuf, +) -> Vec { + use workdir::WorkdirSessionCapability; + + let capabilities = router.capabilities(); + let mut tools = Vec::with_capacity(6); + if capabilities.supports(WorkdirSessionCapability::Read) { + tools.push(read::routed_read_tool(router.clone(), tracker.clone())); + } + if capabilities.supports(WorkdirSessionCapability::Write) { + tools.push(write::routed_write_tool(router.clone(), tracker.clone())); + } + if capabilities.supports(WorkdirSessionCapability::Edit) { + tools.push(edit::routed_edit_tool(router.clone(), tracker)); + } + if capabilities.supports(WorkdirSessionCapability::Glob) { + tools.push(glob::routed_glob_tool(router.clone())); + } + if capabilities.supports(WorkdirSessionCapability::Grep) { + tools.push(grep::routed_grep_tool(router.clone())); + } + if capabilities.supports(WorkdirSessionCapability::Command) { + tools.push(bash::routed_bash_tool(router, bash_output_dir)); } tools } @@ -75,6 +127,29 @@ pub fn read_only_builtin_tools( core_builtin_tools(session, Tracker::new(), std::path::PathBuf::new()) } +pub fn read_only_routed_builtin_tools( + source: std::sync::Arc, +) -> Vec { + let router = std::sync::Arc::new(workdir::WorkdirSessionRouter::new()); + for alias in source.aliases() { + let Ok(selected) = source.resolve(Some(alias.as_str())) else { + continue; + }; + router + .attach( + alias, + std::sync::Arc::new(workdir::ReadOnlyWorkdirSession::new(selected.session)), + ) + .expect("read-only projection preserves unique aliases"); + } + let tracker = Tracker::new(); + vec![ + read::routed_read_tool(router.clone(), tracker), + glob::routed_glob_tool(router.clone()), + grep::routed_grep_tool(router), + ] +} + pub fn web_builtin_tools( web_config: Option, ) -> Vec { diff --git a/crates/tools/src/read.rs b/crates/tools/src/read.rs index 668f4a909..5a6800b02 100644 --- a/crates/tools/src/read.rs +++ b/crates/tools/src/read.rs @@ -8,18 +8,22 @@ use serde::Deserialize; use crate::error::ToolsError; use crate::tracker::Tracker; -use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle}; +use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter}; -const DESCRIPTION: &str = "Read a text file from the local filesystem. \ +const DESCRIPTION: &str = "Read a text file from a Workdir attachment selected by its Worker-local alias. \ Supports offset/limit for large files. Returns line-numbered output (1-based). \ Directories cannot be read. The file must be read before Write or Edit can \ -modify it. Paths are Workdir-relative unless an absolute path is explicitly readable."; +modify it. Paths are Workdir-relative unless an absolute path is explicitly readable. \ +When exactly one Workdir is attached, target_workdir may be omitted."; const DEFAULT_LIMIT: usize = 2000; const PROVIDER_BYTE_LIMIT: usize = 256 * 1024; #[derive(Debug, Deserialize, schemars::JsonSchema)] pub(crate) struct ReadParams { + /// Worker-local alias of the Workdir attachment to use. + #[serde(default)] + pub target_workdir: Option, /// Workdir-relative path, or an absolute path covered by readable scope. pub file_path: String, /// 0-based line offset from the start. Defaults to 0. @@ -31,7 +35,7 @@ pub(crate) struct ReadParams { } pub(crate) struct ReadTool { - session: WorkdirSessionHandle, + router: Arc, tracker: Tracker, } @@ -44,13 +48,21 @@ impl Tool for ReadTool { ) -> Result { let params: ReadParams = serde_json::from_str(input_json) .map_err(|e| ToolError::InvalidArgument(format!("invalid Read input: {e}")))?; + let selected = crate::routing::resolve_session( + &self.router, + params.target_workdir.as_deref(), + workdir::WorkdirSessionCapability::Read, + )?; + let tracker = self + .tracker + .scoped_attachment(&selected.alias, selected.generation); let offset = params.offset.unwrap_or(0); let limit = params.limit.unwrap_or(DEFAULT_LIMIT).max(1); let path = WorkdirPath::new_scoped(¶ms.file_path).map_err(ToolsError::from)?; tracing::debug!(path = %path, offset, limit, "Read"); - let result = self + let result = selected .session .read(ReadRequest { path: path.clone(), @@ -60,7 +72,7 @@ impl Tool for ReadTool { }) .await .map_err(ToolsError::from)?; - self.tracker.record_workdir_hash(&path, result.content_hash); + tracker.record_workdir_hash(&path, result.content_hash); let text = String::from_utf8_lossy(&result.bytes).into_owned(); let rendered = render_provider_read( @@ -144,8 +156,15 @@ fn render_numbered(text: &str, offset: usize, limit: usize) -> Rendered { } } -/// Factory for the `Read` tool. +/// Factory for the `Read` tool bound to one compatibility session. pub fn read_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinition { + routed_read_tool(crate::routing::singleton_router(session), tracker) +} + +pub(crate) fn routed_read_tool( + router: Arc, + tracker: Tracker, +) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(ReadParams); let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({})); @@ -153,7 +172,7 @@ pub fn read_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinit .description(DESCRIPTION) .input_schema(schema_value); let tool: Arc = Arc::new(ReadTool { - session: session.clone(), + router: router.clone(), tracker: tracker.clone(), }); (meta, tool) @@ -199,6 +218,7 @@ mod tests { // History recorded assert!( tracker + .scoped_attachment(&workdir::WorkdirAttachmentAlias::new("workdir").unwrap(), 0,) .expected_workdir_hash(&WorkdirPath::new("a.txt").unwrap()) .is_ok() ); @@ -229,7 +249,7 @@ mod tests { #[tokio::test] async fn read_tool_missing_file() { - let (dir, fs, tracker) = setup(); + let (_dir, fs, tracker) = setup(); let def = read_tool(fs, tracker); let (_, tool) = def(); let input = serde_json::json!({ diff --git a/crates/tools/src/routing.rs b/crates/tools/src/routing.rs new file mode 100644 index 000000000..e586da4f8 --- /dev/null +++ b/crates/tools/src/routing.rs @@ -0,0 +1,32 @@ +use std::sync::Arc; + +use agen::tool::ToolError; +use workdir::{ + ResolvedWorkdirSession, WorkdirAttachmentAlias, WorkdirSessionCapability, WorkdirSessionHandle, + WorkdirSessionRouter, +}; + +pub(crate) fn singleton_router(session: WorkdirSessionHandle) -> Arc { + let router = Arc::new(WorkdirSessionRouter::new()); + router + .attach( + WorkdirAttachmentAlias::new("workdir").expect("static attachment alias is valid"), + session, + ) + .expect("fresh singleton router accepts its session"); + router +} + +pub(crate) fn resolve_session( + router: &WorkdirSessionRouter, + target_workdir: Option<&str>, + required: WorkdirSessionCapability, +) -> Result { + let selected = router + .resolve(target_workdir) + .map_err(|error| ToolError::InvalidArgument(error.to_string()))?; + if !selected.session.capabilities().supports(required) { + return Err(crate::ToolsError::from(workdir::WorkdirError::Unsupported(required)).into()); + } + Ok(selected) +} diff --git a/crates/tools/src/tracker.rs b/crates/tools/src/tracker.rs index a544aa67c..feaaf30ff 100644 --- a/crates/tools/src/tracker.rs +++ b/crates/tools/src/tracker.rs @@ -65,9 +65,15 @@ fn hash_bytes(bytes: &[u8]) -> ContentHash { hasher.finalize().into() } +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct TrackedPath { + namespace: Option>, + path: PathBuf, +} + #[derive(Debug, Clone, Default)] struct FileMutationCoordinator { - locks: Arc>>>>, + locks: Arc>>>>, } pub(crate) struct FileMutationPermit { @@ -75,8 +81,7 @@ pub(crate) struct FileMutationPermit { } impl FileMutationCoordinator { - async fn acquire(&self, path: &Path) -> FileMutationPermit { - let key = file_mutation_key(path); + async fn acquire(&self, key: TrackedPath) -> FileMutationPermit { let lock = { let mut locks = self.locks.lock().await; locks @@ -127,12 +132,12 @@ pub struct ChangeStat { #[derive(Debug, Default)] struct Inner { - /// Hash of each file's last observed contents, keyed by canonical path. - hashes: HashMap, + /// Hash of each file's last observed contents, keyed by attachment and canonical path. + hashes: HashMap, /// Line count paired with observations that included the file content. - line_counts: HashMap, + line_counts: HashMap, /// LRU list of touched files. Front = most recently touched. - recency: VecDeque, + recency: VecDeque, /// Successful Write/Edit mutations attributed to this session's tools. change_stat: ChangeStat, } @@ -146,6 +151,7 @@ struct Inner { pub struct Tracker { inner: Arc>, mutations: FileMutationCoordinator, + namespace: Option>, } impl Tracker { @@ -154,6 +160,33 @@ impl Tracker { Self::default() } + /// Return a view whose read-before-mutation state is isolated to one + /// attachment incarnation (the Worker-local alias plus router generation), + /// while aggregate recency/change statistics remain shared with the Worker + /// tracker. + pub fn scoped(&self, namespace: impl Into) -> Self { + Self { + inner: self.inner.clone(), + mutations: self.mutations.clone(), + namespace: Some(Arc::from(namespace.into())), + } + } + + pub(crate) fn scoped_attachment( + &self, + alias: &workdir::WorkdirAttachmentAlias, + generation: u64, + ) -> Self { + self.scoped(format!("{}#{generation}", alias.as_str())) + } + + fn key(&self, path: PathBuf) -> TrackedPath { + TrackedPath { + namespace: self.namespace.clone(), + path, + } + } + /// Acquire the per-target-file mutation guard shared by `Write` and `Edit`. /// /// The guard is keyed by canonical target path where possible so equivalent @@ -170,7 +203,9 @@ impl Tracker { call_index = ctx.call_index, "acquire file mutation guard" ); - self.mutations.acquire(path).await + self.mutations + .acquire(self.key(file_mutation_key(path))) + .await } /// Record that `path` has been observed with the given content bytes. @@ -182,7 +217,7 @@ impl Tracker { /// Also bumps `path` to the front of the recency list. If the list /// grows past [`RECENCY_CAPACITY`], the oldest entry is evicted. pub fn record(&self, path: &Path, bytes: &[u8]) { - let key = canonicalize_or_owned(path); + let key = self.key(canonicalize_or_owned(path)); let hash = hash_bytes(bytes); let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); inner.hashes.insert(key.clone(), hash); @@ -198,7 +233,7 @@ impl Tracker { } pub fn record_workdir_content(&self, path: &workdir::WorkdirPath, content: &[u8]) { - let key = PathBuf::from(path.as_str()); + let key = self.key(PathBuf::from(path.as_str())); let hash = hash_bytes(content); let line_count = String::from_utf8_lossy(content).lines().count(); let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); @@ -212,16 +247,17 @@ impl Tracker { } pub fn observed_workdir_line_count(&self, path: &workdir::WorkdirPath) -> Option { + let key = self.key(PathBuf::from(path.as_str())); self.inner .lock() .unwrap_or_else(|e| e.into_inner()) .line_counts - .get(Path::new(path.as_str())) + .get(&key) .copied() } pub fn record_workdir_hash(&self, path: &workdir::WorkdirPath, hash: workdir::ContentHash) { - let key = PathBuf::from(path.as_str()); + let key = self.key(PathBuf::from(path.as_str())); let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); inner.hashes.insert(key.clone(), hash); inner.recency.retain(|candidate| candidate != &key); @@ -253,7 +289,7 @@ impl Tracker { ) { let added = added_lines_per_replacement.saturating_mul(replacements); let deleted = deleted_lines_per_replacement.saturating_mul(replacements); - let key = PathBuf::from(path.as_str()); + let key = self.key(PathBuf::from(path.as_str())); let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); inner.change_stat.added = inner.change_stat.added.saturating_add(added as u64); inner.change_stat.deleted = inner.change_stat.deleted.saturating_add(deleted as u64); @@ -279,14 +315,14 @@ impl Tracker { &self, path: &workdir::WorkdirPath, ) -> Result { - let key = PathBuf::from(path.as_str()); + let key = self.key(PathBuf::from(path.as_str())); self.inner .lock() .unwrap_or_else(|error| error.into_inner()) .hashes .get(&key) .copied() - .ok_or_else(|| ToolsError::NotRead(key)) + .ok_or_else(|| ToolsError::NotRead(key.path.clone())) } /// Verify that `path` was previously recorded and its current bytes @@ -296,7 +332,7 @@ impl Tracker { /// - If the current content hashes differ from the recorded value, /// returns [`ToolsError::ExternallyModified`]. pub fn verify(&self, path: &Path, current_bytes: &[u8]) -> Result<(), ToolsError> { - let key = canonicalize_or_owned(path); + let key = self.key(canonicalize_or_owned(path)); let guard = self.inner.lock().unwrap_or_else(|e| e.into_inner()); let recorded = guard .hashes @@ -316,13 +352,18 @@ impl Tracker { /// can pass them as default references to the compaction worker. pub fn recent_files(&self, n: usize) -> Vec { let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - inner.recency.iter().take(n).cloned().collect() + inner + .recency + .iter() + .take(n) + .map(|key| key.path.clone()) + .collect() } /// Returns true if `path` has a history entry. Test-only. #[cfg(test)] pub(crate) fn has(&self, path: &Path) -> bool { - let key = canonicalize_or_owned(path); + let key = self.key(canonicalize_or_owned(path)); self.inner .lock() .unwrap_or_else(|e| e.into_inner()) @@ -420,6 +461,22 @@ mod tests { } } + #[test] + fn attachment_namespaces_do_not_share_read_state() { + let tracker = Tracker::new(); + let left = tracker.scoped("left"); + let right = tracker.scoped("right"); + let path = workdir::WorkdirPath::new("same.txt").unwrap(); + + left.record_workdir_content(&path, b"left"); + + assert!(left.expected_workdir_hash(&path).is_ok()); + assert!(matches!( + right.expected_workdir_hash(&path), + Err(ToolsError::NotRead(_)) + )); + } + #[test] fn clone_shares_state() { let dir = TempDir::new().unwrap(); diff --git a/crates/tools/src/view_image.rs b/crates/tools/src/view_image.rs index 5c9adb520..1fc86ca44 100644 --- a/crates/tools/src/view_image.rs +++ b/crates/tools/src/view_image.rs @@ -7,24 +7,28 @@ use agen::tool::{ }; use async_trait::async_trait; use serde::Deserialize; -use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle}; +use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter}; use crate::error::ToolsError; /// Maximum image body accepted for one model request. pub const MAX_IMAGE_BYTES: usize = 10 * 1024 * 1024; -const DESCRIPTION: &str = "Attach an image from the bound Workdir to the next model request. \ -The path must be logical and Workdir-relative. Supported formats: PNG, JPEG, GIF, and WebP."; +const DESCRIPTION: &str = "Attach an image from the selected Workdir attachment to the next model request. \ +The path must be logical and Workdir-relative. Supported formats: PNG, JPEG, GIF, and WebP. \ +When exactly one Workdir is attached, target_workdir may be omitted."; #[derive(Debug, Deserialize, schemars::JsonSchema)] struct ViewImageParams { + /// Worker-local alias of the Workdir attachment to use. + #[serde(default)] + target_workdir: Option, /// Logical path relative to the bound Workdir root. path: String, } struct ViewImageTool { - session: WorkdirSessionHandle, + router: Arc, } #[async_trait] @@ -37,8 +41,13 @@ impl Tool for ViewImageTool { let input: ViewImageParams = serde_json::from_str(input_json).map_err(|error| { ToolError::InvalidArgument(format!("invalid ViewImage input: {error}")) })?; + let selected = crate::routing::resolve_session( + &self.router, + input.target_workdir.as_deref(), + workdir::WorkdirSessionCapability::Read, + )?; let path = WorkdirPath::new(&input.path).map_err(ToolsError::from)?; - let result = self + let result = selected .session .read(ReadRequest { path: path.clone(), @@ -75,6 +84,10 @@ impl Tool for ViewImageTool { } pub fn view_image_tool(session: WorkdirSessionHandle) -> ToolDefinition { + routed_view_image_tool(crate::routing::singleton_router(session)) +} + +pub fn routed_view_image_tool(router: Arc) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(ViewImageParams); let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({})); @@ -82,7 +95,7 @@ pub fn view_image_tool(session: WorkdirSessionHandle) -> ToolDefinition { .description(DESCRIPTION) .input_schema(schema_value); let tool: Arc = Arc::new(ViewImageTool { - session: session.clone(), + router: router.clone(), }); (meta, tool) }) diff --git a/crates/tools/src/write.rs b/crates/tools/src/write.rs index 4eb567556..8479de806 100644 --- a/crates/tools/src/write.rs +++ b/crates/tools/src/write.rs @@ -9,15 +9,20 @@ use serde::Deserialize; use crate::error::ToolsError; use crate::tracker::Tracker; -use workdir::{StatRequest, WorkdirError, WorkdirPath, WorkdirSessionHandle, WriteRequest}; +use workdir::{ + StatRequest, WorkdirError, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter, + WriteRequest, +}; -const DESCRIPTION: &str = "Create a new file or overwrite an existing one with \ -the given content. Missing parent directories within scope are created \ -automatically. Existing files must have been read first (via the Read tool) \ -in this session. Paths are relative to the bound Workdir."; +const DESCRIPTION: &str = "Create a new file or overwrite an existing one in the selected Workdir attachment. \ +Missing parent directories within scope are created automatically. Existing files must have been read first \ +through the same attachment. When exactly one Workdir is attached, target_workdir may be omitted."; #[derive(Debug, Deserialize, schemars::JsonSchema)] pub(crate) struct WriteParams { + /// Worker-local alias of the Workdir attachment to use. + #[serde(default)] + pub target_workdir: Option, /// Logical path relative to the bound Workdir root. pub file_path: String, /// Full content to write. Overwrites any existing content. @@ -25,7 +30,7 @@ pub(crate) struct WriteParams { } pub(crate) struct WriteTool { - session: WorkdirSessionHandle, + router: Arc, tracker: Tracker, } @@ -39,19 +44,31 @@ impl Tool for WriteTool { let params: WriteParams = serde_json::from_str(input_json) .map_err(|e| ToolError::InvalidArgument(format!("invalid Write input: {e}")))?; + let selected = crate::routing::resolve_session( + &self.router, + params.target_workdir.as_deref(), + workdir::WorkdirSessionCapability::Write, + )?; + let tracker = self + .tracker + .scoped_attachment(&selected.alias, selected.generation); let path = WorkdirPath::new(¶ms.file_path).map_err(ToolsError::from)?; tracing::debug!(path = %path, bytes = params.content.len(), "Write"); let mutation_key = PathBuf::from(path.as_str()); - let _mutation_permit = self.tracker.acquire_mutation(&mutation_key, &ctx).await; - let expected_hash = match self.session.stat(StatRequest { path: path.clone() }).await { - Ok(_) => Some(self.tracker.expected_workdir_hash(&path)?), + let _mutation_permit = tracker.acquire_mutation(&mutation_key, &ctx).await; + let expected_hash = match selected + .session + .stat(StatRequest { path: path.clone() }) + .await + { + Ok(_) => Some(tracker.expected_workdir_hash(&path)?), Err(WorkdirError::NotFound(_)) => None, Err(error) => return Err(ToolsError::from(error).into()), }; - let old_line_count = self.tracker.observed_workdir_line_count(&path).unwrap_or(0); - let outcome = self + let old_line_count = tracker.observed_workdir_line_count(&path).unwrap_or(0); + let outcome = selected .session .write(WriteRequest { path: path.clone(), @@ -61,10 +78,8 @@ impl Tool for WriteTool { .await .map_err(ToolsError::from)?; - self.tracker - .record_change(params.content.lines().count(), old_line_count); - self.tracker - .record_workdir_content(&path, params.content.as_bytes()); + tracker.record_change(params.content.lines().count(), old_line_count); + tracker.record_workdir_content(&path, params.content.as_bytes()); let summary = format!( "{} {} ({} bytes)", @@ -84,8 +99,15 @@ impl Tool for WriteTool { } } -/// Factory for the `Write` tool. +/// Factory for the `Write` tool bound to one compatibility session. pub fn write_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinition { + routed_write_tool(crate::routing::singleton_router(session), tracker) +} + +pub(crate) fn routed_write_tool( + router: Arc, + tracker: Tracker, +) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(WriteParams); let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({})); @@ -93,7 +115,7 @@ pub fn write_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefini .description(DESCRIPTION) .input_schema(schema_value); let tool: Arc = Arc::new(WriteTool { - session: session.clone(), + router: router.clone(), tracker: tracker.clone(), }); (meta, tool) diff --git a/crates/tools/tests/integration.rs b/crates/tools/tests/integration.rs index 02b8c986c..f1d346e60 100644 --- a/crates/tools/tests/integration.rs +++ b/crates/tools/tests/integration.rs @@ -14,8 +14,12 @@ use agen::tool::{ use manifest::{Permission, Scope, ScopeConfig, ScopeRule}; use serde_json::json; use tempfile::TempDir; -use tools::{Tracker, core_builtin_tools, view_image_tool}; -use workdir::{LocalWorkdirSession, WorkdirSessionHandle}; +use tools::{ + Tracker, core_builtin_tools, routed_builtin_tools, routed_view_image_tool, view_image_tool, +}; +use workdir::{ + LocalWorkdirSession, WorkdirAttachmentAlias, WorkdirSessionHandle, WorkdirSessionRouter, +}; fn scope_with_spill(workspace: &Path, spill: &Path) -> Scope { let base = Scope::writable(workspace).unwrap(); @@ -66,6 +70,30 @@ fn setup() -> (TempDir, TempDir, Registry) { (dir, spill, reg) } +fn setup_routed() -> ( + TempDir, + TempDir, + TempDir, + Arc, + Registry, +) { + let left = TempDir::new().unwrap(); + let right = TempDir::new().unwrap(); + let spill = TempDir::new().unwrap(); + let router = Arc::new(WorkdirSessionRouter::new()); + for (alias, dir) in [("left", &left), ("right", &right)] { + let scope = scope_with_spill(dir.path(), spill.path()); + let session: WorkdirSessionHandle = + Arc::new(LocalWorkdirSession::new(scope, dir.path().to_path_buf())); + router + .attach(WorkdirAttachmentAlias::new(alias).unwrap(), session) + .unwrap(); + } + let definitions = + routed_builtin_tools(router.clone(), Tracker::new(), spill.path().to_path_buf()); + (left, right, spill, router, Registry::new(definitions)) +} + async fn call(tool: &Arc, input: serde_json::Value) -> agen::tool::ToolOutput { tool.execute(&input.to_string(), Default::default()) .await @@ -104,6 +132,273 @@ fn meta_has_description_and_schema() { } } +#[test] +fn routed_tool_schemas_share_optional_target_workdir() { + let (_left, _right, _spill, router, mut reg) = setup_routed(); + reg.entries.push(routed_view_image_tool(router)()); + for (meta, _) in ®.entries { + let target = &meta.input_schema["properties"]["target_workdir"]; + assert!(target.is_object(), "{} lacks target_workdir", meta.name); + assert_eq!(target["type"], json!(["string", "null"])); + } +} + +#[tokio::test] +async fn routed_tools_require_alias_only_when_attachment_set_is_ambiguous() { + let (_left, _right, _spill, router, reg) = setup_routed(); + let error = call_err(®.get("Read"), json!({ "file_path": "same.txt" })).await; + let message = error.to_string(); + assert!(message.contains("target_workdir_required"), "{message}"); + assert!(message.contains("left"), "{message}"); + assert!(message.contains("right"), "{message}"); + + let unknown = call_err( + ®.get("Glob"), + json!({ "target_workdir": "missing", "pattern": "*" }), + ) + .await + .to_string(); + assert!(unknown.contains("unknown_target_workdir"), "{unknown}"); + assert!(unknown.contains("available_aliases"), "{unknown}"); + + router.close_all().await.unwrap(); +} + +#[tokio::test] +async fn empty_routed_surface_returns_no_attachment_error() { + let router = Arc::new(WorkdirSessionRouter::new()); + let reg = Registry::new(routed_builtin_tools( + router, + Tracker::new(), + Path::new("unused").to_path_buf(), + )); + let message = call_err(®.get("Read"), json!({ "file_path": "file.txt" })) + .await + .to_string(); + assert!(message.contains("no_workdir_attached"), "{message}"); +} + +#[tokio::test] +async fn target_alias_routes_each_tool_and_isolates_read_before_edit() { + let (left, right, _spill, router, mut reg) = setup_routed(); + std::fs::write(left.path().join("same.txt"), "left value\n").unwrap(); + std::fs::write(right.path().join("same.txt"), "right value\n").unwrap(); + std::fs::write(left.path().join("only-left.log"), "LEFT-NEEDLE\n").unwrap(); + std::fs::write(right.path().join("only-right.log"), "RIGHT-NEEDLE\n").unwrap(); + + let read = reg.get("Read"); + let left_read = call( + &read, + json!({ "target_workdir": "left", "file_path": "same.txt" }), + ) + .await; + assert!(left_read.content.unwrap().contains("left value")); + + let edit = reg.get("Edit"); + let isolated = call_err( + &edit, + json!({ + "target_workdir": "right", + "file_path": "same.txt", + "old_string": "right", + "new_string": "edited" + }), + ) + .await; + assert!(isolated.to_string().contains("has not been read")); + call( + &read, + json!({ "target_workdir": "right", "file_path": "same.txt" }), + ) + .await; + call( + &edit, + json!({ + "target_workdir": "right", + "file_path": "same.txt", + "old_string": "right", + "new_string": "edited" + }), + ) + .await; + assert_eq!( + std::fs::read_to_string(right.path().join("same.txt")).unwrap(), + "edited value\n" + ); + assert_eq!( + std::fs::read_to_string(left.path().join("same.txt")).unwrap(), + "left value\n" + ); + + let glob = call( + ®.get("Glob"), + json!({ "target_workdir": "left", "pattern": "*.log" }), + ) + .await; + assert!(glob.content.unwrap().contains("only-left.log")); + let grep = call( + ®.get("Grep"), + json!({ + "target_workdir": "right", + "pattern": "RIGHT-NEEDLE", + "output_mode": "content" + }), + ) + .await; + assert!(grep.content.unwrap().contains("only-right.log")); + + let bash = call( + ®.get("Bash"), + json!({ "target_workdir": "right", "command": "pwd" }), + ) + .await; + assert_eq!( + std::fs::canonicalize(bash.content.unwrap().trim()).unwrap(), + std::fs::canonicalize(right.path()).unwrap() + ); + + let png = b"\x89PNG\r\n\x1a\nrouted"; + std::fs::write(left.path().join("image.png"), png).unwrap(); + reg.entries.push(routed_view_image_tool(router.clone())()); + let image = call( + ®.get("ViewImage"), + json!({ "target_workdir": "left", "path": "image.png" }), + ) + .await; + let agen::tool::Attachment::Image(image) = &image.attachments[0]; + assert_eq!(image.data(), png); + + router.close_all().await.unwrap(); +} + +#[tokio::test] +async fn selected_attachment_capability_is_checked_without_fallback() { + let left = TempDir::new().unwrap(); + let right = TempDir::new().unwrap(); + let spill = TempDir::new().unwrap(); + let router = Arc::new(WorkdirSessionRouter::new()); + router + .attach( + WorkdirAttachmentAlias::new("writable").unwrap(), + Arc::new(LocalWorkdirSession::new( + scope_with_spill(left.path(), spill.path()), + left.path().to_path_buf(), + )), + ) + .unwrap(); + router + .attach( + WorkdirAttachmentAlias::new("readonly").unwrap(), + Arc::new(LocalWorkdirSession::materialized_bound( + workdir::Workdir::new("readonly"), + right.path().to_path_buf(), + right.path().to_path_buf(), + manifest::SharedScope::new(scope_with_spill(right.path(), spill.path())), + workdir::WorkdirSessionCapabilities::READ_ONLY, + )), + ) + .unwrap(); + let reg = Registry::new(routed_builtin_tools( + router.clone(), + Tracker::new(), + spill.path().to_path_buf(), + )); + + let message = call_err( + ®.get("Write"), + json!({ + "target_workdir": "readonly", + "file_path": "created.txt", + "content": "must not be routed elsewhere" + }), + ) + .await + .to_string(); + assert!(message.contains("does not support Write"), "{message}"); + assert!(!left.path().join("created.txt").exists()); + assert!(!right.path().join("created.txt").exists()); + + router.close_all().await.unwrap(); +} + +#[tokio::test] +async fn routed_tools_observe_detach_after_registration() { + let (_left, _right, _spill, router, reg) = setup_routed(); + let alias = WorkdirAttachmentAlias::new("left").unwrap(); + router.detach(&alias).await.unwrap(); + + let message = call_err( + ®.get("Glob"), + json!({ "target_workdir": "left", "pattern": "*" }), + ) + .await + .to_string(); + assert!(message.contains("unknown_target_workdir"), "{message}"); + assert!(message.contains("right"), "{message}"); + + router.close_all().await.unwrap(); +} + +#[tokio::test] +async fn detach_and_reattach_same_alias_requires_a_fresh_read() { + let first = TempDir::new().unwrap(); + let second = TempDir::new().unwrap(); + let spill = TempDir::new().unwrap(); + std::fs::write(first.path().join("same.txt"), "same bytes\n").unwrap(); + std::fs::write(second.path().join("same.txt"), "same bytes\n").unwrap(); + + let router = Arc::new(WorkdirSessionRouter::new()); + let alias = WorkdirAttachmentAlias::new("checkout").unwrap(); + router + .attach( + alias.clone(), + Arc::new(LocalWorkdirSession::new( + scope_with_spill(first.path(), spill.path()), + first.path().to_path_buf(), + )), + ) + .unwrap(); + let reg = Registry::new(routed_builtin_tools( + router.clone(), + Tracker::new(), + spill.path().to_path_buf(), + )); + call( + ®.get("Read"), + json!({ "target_workdir": "checkout", "file_path": "same.txt" }), + ) + .await; + + router.detach(&alias).await.unwrap(); + router + .attach( + alias.clone(), + Arc::new(LocalWorkdirSession::new( + scope_with_spill(second.path(), spill.path()), + second.path().to_path_buf(), + )), + ) + .unwrap(); + + let error = call_err( + ®.get("Edit"), + json!({ + "target_workdir": "checkout", + "file_path": "same.txt", + "old_string": "same", + "new_string": "changed" + }), + ) + .await; + assert!(error.to_string().contains("has not been read"), "{error}"); + assert_eq!( + std::fs::read_to_string(second.path().join("same.txt")).unwrap(), + "same bytes\n" + ); + + router.close_all().await.unwrap(); +} + #[tokio::test] async fn view_image_reads_scoped_bytes_into_durable_tool_detail() { let dir = TempDir::new().unwrap(); diff --git a/crates/tui/src/backend_spawn.rs b/crates/tui/src/backend_spawn.rs index 4c4203d07..f6d440d62 100644 --- a/crates/tui/src/backend_spawn.rs +++ b/crates/tui/src/backend_spawn.rs @@ -164,7 +164,7 @@ fn request_from_selection(selection: Selection) -> BackendCreateWorkerRequest { display_name: selection.display_name, profile: Some(selection.profile), initial_submit: Vec::new(), - working_directory: None, + workdir_attachments: Vec::new(), ticket_assignment: None, control_operation_id: None, } @@ -477,7 +477,7 @@ mod tests { assert_eq!(request.display_name, "Coder one"); assert_eq!(request.profile.as_deref(), Some("builtin:coder")); assert!(request.initial_submit.is_empty()); - assert!(request.working_directory.is_none()); + assert!(request.workdir_attachments.is_empty()); assert!(request.ticket_assignment.is_none()); } } diff --git a/crates/tui/src/backend_worker_picker.rs b/crates/tui/src/backend_worker_picker.rs index 1768334b2..6f94ffc29 100644 --- a/crates/tui/src/backend_worker_picker.rs +++ b/crates/tui/src/backend_worker_picker.rs @@ -408,10 +408,22 @@ fn short_worker_id(worker: &BackendWorkerSummary) -> String { } fn working_directory_text(worker: &BackendWorkerSummary) -> String { - let Some(wd) = worker.working_directory.as_ref() else { + if worker.workdir_attachments.is_empty() { return "wd:—".to_string(); - }; - format!("wd:{}・{}", wd.repository_key, wd.working_directory_id) + } + if worker.workdir_attachments.len() > 1 { + return format!("wd:{} attachments", worker.workdir_attachments.len()); + } + let attachment = &worker.workdir_attachments[0]; + let directory = &attachment.working_directory; + let label = directory + .display_name + .as_deref() + .unwrap_or(&directory.repository_key); + format!( + "wd:{}:{}・{}", + attachment.alias, label, directory.working_directory_id + ) } #[cfg(test)] @@ -456,7 +468,7 @@ mod tests { can_stop: true, can_spawn_followup: false, }, - working_directory: None, + workdir_attachments: Vec::new(), diagnostics: Vec::new(), } } @@ -529,22 +541,26 @@ mod tests { worker.label = "Coder · T-585".to_string(); worker.state = "stopped".to_string(); worker.worker_state = None; - worker.working_directory = Some( + worker.workdir_attachments = vec![ serde_json::from_value(serde_json::json!({ - "working_directory_id": "001a06a9f0202000000", - "repository_key": "main", - "materializer_kind": "runtime_git_clone", - "status": "active", - "cleanliness": "clean" + "alias": "checkout", + "working_directory": { + "working_directory_id": "001a06a9f0202000000", + "display_name": "Checkout", + "repository_key": "main", + "materializer_kind": "runtime_git_clone", + "status": "active", + "cleanliness": "clean" + } })) .unwrap(), - ); + ]; let widths = WorkerColumnWidths::from_workers(std::slice::from_ref(&worker)); let text = row_text(&worker, &widths); assert_eq!( text, - " W-90 Coder · T-585 [stopped] wd:main・001a06a9f0202000000" + " W-90 Coder · T-585 [stopped] wd:checkout:Checkout・001a06a9f0202000000" ); assert!(!text.contains("profile:")); assert!(!text.contains("active clean")); @@ -570,15 +586,18 @@ mod tests { long.worker_state = None; for worker in [&mut short, &mut long] { - worker.working_directory = Some( + worker.workdir_attachments = vec![ serde_json::from_value(serde_json::json!({ - "working_directory_id": "workdir-1", - "repository_key": "main", - "materializer_kind": "runtime_git_clone", - "status": "active" + "alias": "checkout", + "working_directory": { + "working_directory_id": "workdir-1", + "repository_key": "main", + "materializer_kind": "runtime_git_clone", + "status": "active" + } })) .unwrap(), - ); + ]; } let workers = vec![short, long]; @@ -595,8 +614,8 @@ mod tests { display_column(&second, "[stopped]") ); assert_eq!( - display_column(&first, "wd:main"), - display_column(&second, "wd:main") + display_column(&first, "wd:checkout:main"), + display_column(&second, "wd:checkout:main") ); } diff --git a/crates/workdir/src/lib.rs b/crates/workdir/src/lib.rs index 0356936d4..ba1070541 100644 --- a/crates/workdir/src/lib.rs +++ b/crates/workdir/src/lib.rs @@ -8,6 +8,7 @@ pub mod http; mod local; mod operation; +mod router; mod scope; pub mod workspace; @@ -27,10 +28,14 @@ pub use local::{ LocalWorkdirSession, SymlinkInfo, WorkdirSessionResource, direct_symlink, first_symlink, }; pub use operation::*; +pub use router::{ + ResolvedWorkdirSession, RoutedWorkdirSession, WorkdirAttachmentAlias, WorkdirRouteError, + WorkdirRouteErrorCode, WorkdirSessionRouter, +}; pub use scope::{ ReadOnlyWorkdirSession, WorkdirScopeAuthorizationRequest, WorkdirScopeLease, - WorkdirScopeOverlapRequest, WorkdirToolBroker, WorkdirToolScope, WorkdirToolScopePermission, - WorkdirToolScopeRule, + WorkdirScopeLeaseSet, WorkdirScopeOverlapRequest, WorkdirToolBroker, WorkdirToolBrokerRouter, + WorkdirToolScope, WorkdirToolScopePermission, WorkdirToolScopeRule, }; /// Persistent, opaque identity of one materialized Workdir. @@ -122,6 +127,12 @@ impl WorkdirSessionCapabilities { bits: Self::READ | Self::GLOB | Self::GREP, }; + pub const fn union(self, other: Self) -> Self { + Self { + bits: self.bits | other.bits, + } + } + pub const fn supports(self, capability: WorkdirSessionCapability) -> bool { let bit = match capability { WorkdirSessionCapability::Read => Self::READ, diff --git a/crates/workdir/src/router.rs b/crates/workdir/src/router.rs new file mode 100644 index 000000000..d8e2c9a20 --- /dev/null +++ b/crates/workdir/src/router.rs @@ -0,0 +1,796 @@ +//! Alias-keyed Worker Workdir session routing. +//! +//! The router owns live provider sessions without exposing whether a session is +//! local or transported over HTTP. Attachment aliases are Worker-local stable +//! identities; provider Workdir ids and human display names are deliberately not +//! routing keys. + +use std::collections::{BTreeMap, HashSet}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, RwLock}; + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use tokio::sync::broadcast; + +use crate::{ + CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, + CommandSnapshot, CommandStatus, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest, + GrepResult, ListRequest, ListResult, ReadRequest, ReadResult, StatRequest, StatResult, Workdir, + WorkdirError, WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, + WorkdirSessionCapabilities, WorkdirSessionHandle, WriteRequest, WriteResult, +}; + +/// Stable Worker-local routing key for one attached Workdir. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)] +#[serde(transparent)] +pub struct WorkdirAttachmentAlias(String); + +impl WorkdirAttachmentAlias { + pub fn new(value: impl Into) -> Result { + let value = value.into(); + let valid = !value.is_empty() + && value.len() <= 64 + && value.bytes().enumerate().all(|(index, byte)| match byte { + b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' => true, + b'_' | b'-' | b'.' => index > 0, + _ => false, + }); + if !valid { + return Err(WorkdirError::InvalidArgument( + "attachment alias must be 1-64 ASCII letters, digits, '.', '_' or '-', and start with a letter or digit" + .to_string(), + )); + } + Ok(Self(value)) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl<'de> Deserialize<'de> for WorkdirAttachmentAlias { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let value = String::deserialize(deserializer)?; + Self::new(value).map_err(serde::de::Error::custom) + } +} + +impl std::fmt::Display for WorkdirAttachmentAlias { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(&self.0) + } +} + +/// Stable machine-readable reason why a tool target could not be resolved. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum WorkdirRouteErrorCode { + NoWorkdirAttached, + TargetWorkdirRequired, + UnknownTargetWorkdir, +} + +/// Provider-neutral attachment routing failure returned before an operation starts. +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct WorkdirRouteError { + pub code: WorkdirRouteErrorCode, + #[serde(skip_serializing_if = "Option::is_none")] + pub target_workdir: Option, + pub available_aliases: Vec, +} + +impl std::fmt::Display for WorkdirRouteError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let encoded = serde_json::to_string(self).map_err(|_| std::fmt::Error)?; + formatter.write_str(&encoded) + } +} + +impl std::error::Error for WorkdirRouteError {} + +/// One attachment selected for a single tool invocation. +#[derive(Clone)] +pub struct ResolvedWorkdirSession { + pub alias: WorkdirAttachmentAlias, + pub generation: u64, + pub session: WorkdirSessionHandle, +} + +impl std::fmt::Debug for ResolvedWorkdirSession { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ResolvedWorkdirSession") + .field("alias", &self.alias) + .field("generation", &self.generation) + .field("workdir", self.session.workdir()) + .finish_non_exhaustive() + } +} + +#[derive(Debug)] +struct RoutedSessionState { + session: WorkdirSessionHandle, + generation: u64, + detached: AtomicBool, + active_operations: AtomicUsize, + active_commands: Mutex>, +} + +impl RoutedSessionState { + fn enter(&self) -> Result, WorkdirError> { + if self.detached.load(Ordering::Acquire) { + return Err(WorkdirError::Unavailable( + "Workdir attachment is detached".to_string(), + )); + } + self.active_operations.fetch_add(1, Ordering::AcqRel); + if self.detached.load(Ordering::Acquire) { + self.active_operations.fetch_sub(1, Ordering::AcqRel); + return Err(WorkdirError::Unavailable( + "Workdir attachment is detaching".to_string(), + )); + } + Ok(ActiveOperation(self)) + } + + fn is_busy(&self) -> bool { + if self.active_operations.load(Ordering::Acquire) != 0 { + return true; + } + self.reconcile_terminal_commands(&self.session.command_snapshot()); + !self + .active_commands + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_empty() + } + + async fn is_busy_for_detach(&self) -> bool { + if self.active_operations.load(Ordering::Acquire) != 0 { + return true; + } + self.reconcile_terminal_commands(&self.session.command_snapshot()); + let handles = self + .active_commands + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .iter() + .cloned() + .collect::>(); + for handle in handles { + match self.session.command_status(handle.clone()).await { + Ok(status) => self.mark_terminal(&handle, status), + Err(WorkdirError::UnknownCommand(_)) => { + self.active_commands + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&handle); + } + Err(_) => {} + } + } + self.active_operations.load(Ordering::Acquire) != 0 + || !self + .active_commands + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_empty() + } + + fn reconcile_terminal_commands(&self, snapshots: &[CommandSnapshot]) { + let terminal = snapshots + .iter() + .filter(|snapshot| snapshot.status != CommandStatus::Running) + .map(|snapshot| snapshot.command_id.as_str()) + .collect::>(); + if terminal.is_empty() { + return; + } + self.active_commands + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .retain(|handle| !terminal.contains(handle.0.as_str())); + } + + fn mark_terminal(&self, handle: &CommandHandle, status: CommandStatus) { + if status != CommandStatus::Running { + self.active_commands + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(handle); + } + } +} + +struct ActiveOperation<'a>(&'a RoutedSessionState); + +impl Drop for ActiveOperation<'_> { + fn drop(&mut self) { + self.0.active_operations.fetch_sub(1, Ordering::AcqRel); + } +} + +/// A session handle bound to one alias in a [`WorkdirSessionRouter`]. +#[derive(Debug, Clone)] +pub struct RoutedWorkdirSession { + state: Arc, +} + +#[async_trait] +impl WorkdirSession for RoutedWorkdirSession { + fn workdir(&self) -> &Workdir { + self.state.session.workdir() + } + + fn capabilities(&self) -> WorkdirSessionCapabilities { + self.state.session.capabilities() + } + + async fn authorize_scope_path( + &self, + request: WorkdirScopeAuthorizationRequest, + ) -> Result<(), WorkdirError> { + let _active = self.state.enter()?; + self.state.session.authorize_scope_path(request).await + } + + async fn scope_rules_overlap( + &self, + request: WorkdirScopeOverlapRequest, + ) -> Result { + let _active = self.state.enter()?; + self.state.session.scope_rules_overlap(request).await + } + + async fn stat(&self, request: StatRequest) -> Result { + let _active = self.state.enter()?; + self.state.session.stat(request).await + } + + async fn read(&self, request: ReadRequest) -> Result { + let _active = self.state.enter()?; + self.state.session.read(request).await + } + + async fn write(&self, request: WriteRequest) -> Result { + let _active = self.state.enter()?; + self.state.session.write(request).await + } + + async fn edit(&self, request: EditRequest) -> Result { + let _active = self.state.enter()?; + self.state.session.edit(request).await + } + + async fn list(&self, request: ListRequest) -> Result { + let _active = self.state.enter()?; + self.state.session.list(request).await + } + + async fn glob(&self, request: GlobRequest) -> Result { + let _active = self.state.enter()?; + self.state.session.glob(request).await + } + + async fn grep(&self, request: GrepRequest) -> Result { + let _active = self.state.enter()?; + self.state.session.grep(request).await + } + + async fn start_command(&self, request: CommandRequest) -> Result { + let _active = self.state.enter()?; + let handle = self.state.session.start_command(request).await?; + self.state + .active_commands + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(handle.clone()); + Ok(handle) + } + + async fn command_status(&self, handle: CommandHandle) -> Result { + let _active = self.state.enter()?; + let status = self.state.session.command_status(handle.clone()).await?; + self.state.mark_terminal(&handle, status); + Ok(status) + } + + async fn command_output( + &self, + request: CommandOutputRequest, + ) -> Result { + let _active = self.state.enter()?; + let handle = request.handle.clone(); + let output = self.state.session.command_output(request).await?; + self.state.mark_terminal(&handle, output.status); + Ok(output) + } + + async fn cancel_command(&self, handle: CommandHandle) -> Result<(), WorkdirError> { + let _active = self.state.enter()?; + self.state.session.cancel_command(handle.clone()).await?; + self.state + .active_commands + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&handle); + Ok(()) + } + + fn subscribe_command_events(&self) -> Option> { + self.state.session.subscribe_command_events() + } + + fn command_snapshot(&self) -> Vec { + let snapshots = self.state.session.command_snapshot(); + self.state.reconcile_terminal_commands(&snapshots); + snapshots + } + + async fn close(&self) -> Result<(), WorkdirError> { + if self.state.is_busy() { + return Err(WorkdirError::Conflict( + "Workdir attachment has active operations or commands".to_string(), + )); + } + self.state.session.close().await + } +} + +/// Worker-owned map of stable attachment aliases to provider-neutral sessions. +#[derive(Debug, Default)] +pub struct WorkdirSessionRouter { + sessions: RwLock>>, + next_generation: AtomicU64, +} + +impl WorkdirSessionRouter { + pub fn new() -> Self { + Self::default() + } + + pub fn aliases(&self) -> Vec { + self.sessions + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .keys() + .cloned() + .collect() + } + + pub fn capabilities(&self) -> WorkdirSessionCapabilities { + self.sessions + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .values() + .fold(WorkdirSessionCapabilities::EMPTY, |capabilities, state| { + capabilities.union(state.session.capabilities()) + }) + } + + pub fn len(&self) -> usize { + self.sessions + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .len() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Resolve one model-visible target without exposing provider identity. + /// + /// Omission is accepted only for the unambiguous single-attachment case. + /// Every error includes the current alias snapshot so callers can recover + /// without consulting Workdir ids, display names, or transport metadata. + pub fn resolve( + &self, + target_workdir: Option<&str>, + ) -> Result { + let sessions = self + .sessions + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let available_aliases = sessions + .keys() + .map(|alias| alias.as_str().to_owned()) + .collect::>(); + let selected = match target_workdir { + Some(target) => sessions + .iter() + .find(|(alias, _)| alias.as_str() == target) + .map(|(alias, state)| (alias.clone(), state.clone())) + .ok_or_else(|| WorkdirRouteError { + code: WorkdirRouteErrorCode::UnknownTargetWorkdir, + target_workdir: Some(target.to_owned()), + available_aliases: available_aliases.clone(), + })?, + None if sessions.is_empty() => { + return Err(WorkdirRouteError { + code: WorkdirRouteErrorCode::NoWorkdirAttached, + target_workdir: None, + available_aliases, + }); + } + None if sessions.len() > 1 => { + return Err(WorkdirRouteError { + code: WorkdirRouteErrorCode::TargetWorkdirRequired, + target_workdir: None, + available_aliases, + }); + } + None => { + let (alias, state) = sessions + .first_key_value() + .expect("single attachment exists after cardinality checks"); + (alias.clone(), state.clone()) + } + }; + Ok(ResolvedWorkdirSession { + alias: selected.0, + generation: selected.1.generation, + session: Arc::new(RoutedWorkdirSession { state: selected.1 }), + }) + } + + pub fn attach( + &self, + alias: WorkdirAttachmentAlias, + session: WorkdirSessionHandle, + ) -> Result { + let mut sessions = self.sessions.write().map_err(|_| { + WorkdirError::Unavailable("Workdir attachment router is poisoned".to_string()) + })?; + if sessions.contains_key(&alias) { + return Err(WorkdirError::Conflict(format!( + "Workdir attachment alias `{alias}` already exists" + ))); + } + let state = Arc::new(RoutedSessionState { + session, + generation: self.next_generation.fetch_add(1, Ordering::AcqRel), + detached: AtomicBool::new(false), + active_operations: AtomicUsize::new(0), + active_commands: Mutex::new(HashSet::new()), + }); + sessions.insert(alias, state.clone()); + Ok(Arc::new(RoutedWorkdirSession { state })) + } + + pub fn session(&self, alias: &WorkdirAttachmentAlias) -> Option { + self.sessions + .read() + .ok()? + .get(alias) + .cloned() + .map(|state| Arc::new(RoutedWorkdirSession { state }) as WorkdirSessionHandle) + } + + /// Compatibility projection used only while filesystem tool schemas have no + /// `target_workdir`: exactly one attachment is unambiguous; zero or multiple + /// attachments expose no implicit primary. + pub fn only_session(&self) -> Option { + let sessions = self.sessions.read().ok()?; + if sessions.len() != 1 { + return None; + } + sessions + .values() + .next() + .cloned() + .map(|state| Arc::new(RoutedWorkdirSession { state }) as WorkdirSessionHandle) + } + + pub fn can_detach(&self, alias: &WorkdirAttachmentAlias) -> Result<(), WorkdirError> { + let state = self + .sessions + .read() + .map_err(|_| { + WorkdirError::Unavailable("Workdir attachment router is poisoned".to_string()) + })? + .get(alias) + .cloned() + .ok_or_else(|| WorkdirError::NotFound(alias.as_str().into()))?; + if state.detached.load(Ordering::Acquire) { + return Err(WorkdirError::Unavailable( + "Workdir attachment is already detaching".to_string(), + )); + } + if state.is_busy() { + return Err(WorkdirError::Conflict(format!( + "Workdir attachment `{alias}` has active operations or commands" + ))); + } + Ok(()) + } + + /// Atomically fence new operations before an external detach transaction begins. + pub async fn begin_detach(&self, alias: &WorkdirAttachmentAlias) -> Result<(), WorkdirError> { + let state = self + .sessions + .read() + .map_err(|_| { + WorkdirError::Unavailable("Workdir attachment router is poisoned".to_string()) + })? + .get(alias) + .cloned() + .ok_or_else(|| WorkdirError::NotFound(alias.as_str().into()))?; + if state.detached.swap(true, Ordering::AcqRel) { + return Err(WorkdirError::Unavailable( + "Workdir attachment is already detaching".to_string(), + )); + } + if state.is_busy_for_detach().await { + state.detached.store(false, Ordering::Release); + return Err(WorkdirError::Conflict(format!( + "Workdir attachment `{alias}` has active operations or commands" + ))); + } + Ok(()) + } + + /// Reopen a route after the external detach transaction failed before commit. + pub fn cancel_detach(&self, alias: &WorkdirAttachmentAlias) { + if let Ok(sessions) = self.sessions.read() + && let Some(state) = sessions.get(alias) + { + state.detached.store(false, Ordering::Release); + } + } + + pub async fn finish_detach( + &self, + alias: &WorkdirAttachmentAlias, + ) -> Result { + let state = { + let sessions = self.sessions.read().map_err(|_| { + WorkdirError::Unavailable("Workdir attachment router is poisoned".to_string()) + })?; + sessions + .get(alias) + .cloned() + .ok_or_else(|| WorkdirError::NotFound(alias.as_str().into()))? + }; + if !state.detached.load(Ordering::Acquire) { + return Err(WorkdirError::Conflict(format!( + "Workdir attachment `{alias}` was not prepared for detach" + ))); + } + if state.is_busy_for_detach().await { + state.detached.store(false, Ordering::Release); + return Err(WorkdirError::Conflict(format!( + "Workdir attachment `{alias}` has active operations or commands" + ))); + } + if let Err(error) = state.session.close().await { + state.detached.store(false, Ordering::Release); + return Err(error); + } + let mut sessions = self.sessions.write().map_err(|_| { + WorkdirError::Unavailable("Workdir attachment router is poisoned".to_string()) + })?; + if sessions + .get(alias) + .is_some_and(|current| Arc::ptr_eq(current, &state)) + { + sessions.remove(alias); + } + Ok(state.session.clone()) + } + + pub async fn detach( + &self, + alias: &WorkdirAttachmentAlias, + ) -> Result { + self.begin_detach(alias).await?; + self.finish_detach(alias).await + } + + /// Close every attachment during Worker shutdown or startup rollback. + /// + /// Unlike [`Self::detach`], lifecycle teardown is authoritative and delegates + /// cancellation or draining of active commands to each provider session. + pub async fn close_all(&self) -> Result<(), WorkdirError> { + let aliases = self.aliases(); + for alias in aliases { + let state = { + let sessions = self.sessions.read().map_err(|_| { + WorkdirError::Unavailable("Workdir attachment router is poisoned".to_string()) + })?; + sessions.get(&alias).cloned() + }; + let Some(state) = state else { + continue; + }; + state.detached.store(true, Ordering::Release); + if let Err(error) = state.session.close().await { + state.detached.store(false, Ordering::Release); + return Err(error); + } + let mut sessions = self.sessions.write().map_err(|_| { + WorkdirError::Unavailable("Workdir attachment router is poisoned".to_string()) + })?; + if sessions + .get(&alias) + .is_some_and(|current| Arc::ptr_eq(current, &state)) + { + sessions.remove(&alias); + } + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::LocalWorkdirSession; + use manifest::Scope; + use tempfile::TempDir; + + fn session(id: &str, dir: &TempDir) -> WorkdirSessionHandle { + Arc::new(LocalWorkdirSession::materialized_bound( + Workdir::new(id), + dir.path().to_path_buf(), + dir.path().to_path_buf(), + manifest::SharedScope::new(Scope::writable(dir.path()).unwrap()), + WorkdirSessionCapabilities::ALL, + )) + } + + #[tokio::test] + async fn router_holds_multiple_aliases_without_implicit_primary() { + let left = TempDir::new().unwrap(); + let right = TempDir::new().unwrap(); + let router = WorkdirSessionRouter::new(); + router + .attach( + WorkdirAttachmentAlias::new("left").unwrap(), + session("wd-left", &left), + ) + .unwrap(); + router + .attach( + WorkdirAttachmentAlias::new("right").unwrap(), + session("wd-right", &right), + ) + .unwrap(); + assert_eq!( + router + .aliases() + .iter() + .map(WorkdirAttachmentAlias::as_str) + .collect::>(), + ["left", "right"] + ); + assert!(router.only_session().is_none()); + let missing = router.resolve(None).unwrap_err(); + assert_eq!(missing.code, WorkdirRouteErrorCode::TargetWorkdirRequired); + assert_eq!(missing.available_aliases, ["left", "right"]); + assert_eq!( + router + .resolve(Some("right")) + .unwrap() + .session + .workdir() + .id() + .as_str(), + "wd-right" + ); + let unknown = router.resolve(Some("display-name")).unwrap_err(); + assert_eq!(unknown.code, WorkdirRouteErrorCode::UnknownTargetWorkdir); + assert_eq!(unknown.target_workdir.as_deref(), Some("display-name")); + assert_eq!(unknown.available_aliases, ["left", "right"]); + router.close_all().await.unwrap(); + } + + #[test] + fn empty_router_returns_structured_no_attachment_error() { + let error = WorkdirSessionRouter::new().resolve(None).unwrap_err(); + assert_eq!(error.code, WorkdirRouteErrorCode::NoWorkdirAttached); + assert!(error.available_aliases.is_empty()); + assert_eq!( + serde_json::to_value(error).unwrap(), + serde_json::json!({ + "code": "no_workdir_attached", + "available_aliases": [] + }) + ); + } + + #[tokio::test] + async fn detach_allows_naturally_completed_command_without_final_poll() { + let dir = TempDir::new().unwrap(); + let router = WorkdirSessionRouter::new(); + let alias = WorkdirAttachmentAlias::new("checkout").unwrap(); + let underlying = session("wd", &dir); + let routed = router + .attach(alias.clone(), Arc::clone(&underlying)) + .unwrap(); + let handle = routed + .start_command(CommandRequest { + command: "printf done".to_string(), + timeout_secs: 5, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: None, + }) + .await + .unwrap(); + let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2); + loop { + let completed = underlying.command_snapshot().iter().any(|snapshot| { + snapshot.command_id == handle.0 && snapshot.status != CommandStatus::Running + }); + if completed { + break; + } + assert!(tokio::time::Instant::now() < deadline); + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + + router.detach(&alias).await.unwrap(); + assert!(router.is_empty()); + } + + #[tokio::test] + async fn lifecycle_close_all_closes_sessions_with_active_commands() { + let dir = TempDir::new().unwrap(); + let router = WorkdirSessionRouter::new(); + let alias = WorkdirAttachmentAlias::new("checkout").unwrap(); + let underlying = session("wd", &dir); + let routed = router.attach(alias, Arc::clone(&underlying)).unwrap(); + let handle = routed + .start_command(CommandRequest { + command: "sleep 30".to_string(), + timeout_secs: 60, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: None, + }) + .await + .unwrap(); + + router.close_all().await.unwrap(); + + assert!(router.is_empty()); + assert!(matches!( + underlying.command_status(handle).await, + Err(WorkdirError::Unavailable(_)) + )); + } + + #[tokio::test] + async fn detach_rejects_active_command_until_cancelled() { + let dir = TempDir::new().unwrap(); + let router = WorkdirSessionRouter::new(); + let alias = WorkdirAttachmentAlias::new("checkout").unwrap(); + let routed = router.attach(alias.clone(), session("wd", &dir)).unwrap(); + let handle = routed + .start_command(CommandRequest { + command: "sleep 30".to_string(), + timeout_secs: 60, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: None, + }) + .await + .unwrap(); + assert!(matches!( + router.detach(&alias).await, + Err(WorkdirError::Conflict(_)) + )); + routed.cancel_command(handle).await.unwrap(); + router.detach(&alias).await.unwrap(); + assert!(router.is_empty()); + } +} diff --git a/crates/workdir/src/scope.rs b/crates/workdir/src/scope.rs index 5e158d231..7eadad4e9 100644 --- a/crates/workdir/src/scope.rs +++ b/crates/workdir/src/scope.rs @@ -1,4 +1,4 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::path::Path; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex, Weak}; @@ -284,6 +284,204 @@ impl Drop for WorkdirScopeLease { } } +/// Alias-keyed parent authority used to attenuate one or more Workdir +/// attachments for an Internal SubWorker. +#[derive(Clone)] +pub struct WorkdirToolBrokerRouter { + sessions: Arc, + brokers: Arc>>, +} + +impl std::fmt::Debug for WorkdirToolBrokerRouter { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkdirToolBrokerRouter") + .field("aliases", &self.sessions.aliases()) + .finish_non_exhaustive() + } +} + +impl WorkdirToolBrokerRouter { + pub fn from_single( + alias: crate::WorkdirAttachmentAlias, + session: WorkdirSessionHandle, + ) -> Result { + let sessions = Arc::new(crate::WorkdirSessionRouter::new()); + sessions.attach(alias, session)?; + Ok(Self::new(sessions)) + } + + pub fn new(sessions: Arc) -> Self { + Self { + sessions, + brokers: Arc::new(Mutex::new(BTreeMap::new())), + } + } + + pub fn sessions(&self) -> Arc { + self.sessions.clone() + } + + pub fn resolve_alias( + &self, + target_workdir: Option<&str>, + ) -> Result { + self.sessions + .resolve(target_workdir) + .map(|selected| selected.alias) + } + + fn broker( + &self, + alias: &crate::WorkdirAttachmentAlias, + ) -> Result { + let selected = self + .sessions + .resolve(Some(alias.as_str())) + .map_err(|error| WorkdirError::InvalidArgument(error.to_string()))?; + let mut brokers = self.brokers.lock().map_err(|_| { + WorkdirError::Unavailable("Workdir tool broker router is poisoned".to_string()) + })?; + if let Some((generation, broker)) = brokers.get(alias) + && *generation == selected.generation + { + return Ok(broker.clone()); + } + let broker = WorkdirToolBroker::new(selected.session); + brokers.insert(alias.clone(), (selected.generation, broker.clone())); + Ok(broker) + } + + pub async fn scope( + &self, + targets: Vec<(crate::WorkdirAttachmentAlias, WorkdirToolScope)>, + ) -> Result { + if targets.is_empty() { + return Err(WorkdirError::InvalidArgument( + "at least one Workdir attachment scope is required".to_string(), + )); + } + let child_sessions = Arc::new(crate::WorkdirSessionRouter::new()); + let mut child_brokers = BTreeMap::new(); + let mut leases = Vec::with_capacity(targets.len()); + for (alias, request) in targets { + if leases.iter().any( + |(existing, _): &(crate::WorkdirAttachmentAlias, WorkdirScopeLease)| { + existing == &alias + }, + ) { + for (_, lease) in &leases { + let _ = lease.close().await; + } + return Err(WorkdirError::InvalidArgument(format!( + "duplicate Workdir attachment scope `{alias}`" + ))); + } + let lease = match self.broker(&alias)?.scope(request).await { + Ok(lease) => lease, + Err(error) => { + for (_, lease) in &leases { + let _ = lease.close().await; + } + return Err(error); + } + }; + if let Err(error) = child_sessions.attach(alias.clone(), lease.tool_session()) { + let _ = lease.close().await; + for (_, existing) in &leases { + let _ = existing.close().await; + } + return Err(error); + } + let child_generation = child_sessions + .resolve(Some(alias.as_str())) + .expect("newly attached child route resolves") + .generation; + child_brokers.insert(alias.clone(), (child_generation, lease.broker())); + leases.push((alias, lease)); + } + Ok(WorkdirScopeLeaseSet { + broker_router: WorkdirToolBrokerRouter { + sessions: child_sessions, + brokers: Arc::new(Mutex::new(child_brokers)), + }, + leases, + }) + } +} + +pub struct WorkdirScopeLeaseSet { + broker_router: WorkdirToolBrokerRouter, + leases: Vec<(crate::WorkdirAttachmentAlias, WorkdirScopeLease)>, +} + +impl std::fmt::Debug for WorkdirScopeLeaseSet { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkdirScopeLeaseSet") + .field( + "aliases", + &self + .leases + .iter() + .map(|(alias, _)| alias) + .collect::>(), + ) + .field("active", &self.is_active()) + .finish() + } +} + +impl WorkdirScopeLeaseSet { + pub fn from_single( + alias: crate::WorkdirAttachmentAlias, + lease: WorkdirScopeLease, + ) -> Result { + let sessions = Arc::new(crate::WorkdirSessionRouter::new()); + sessions.attach(alias.clone(), lease.tool_session())?; + let generation = sessions + .resolve(Some(alias.as_str())) + .expect("newly attached single route resolves") + .generation; + let brokers = BTreeMap::from([(alias.clone(), (generation, lease.broker()))]); + Ok(Self { + broker_router: WorkdirToolBrokerRouter { + sessions, + brokers: Arc::new(Mutex::new(brokers)), + }, + leases: vec![(alias, lease)], + }) + } + + pub fn broker_router(&self) -> WorkdirToolBrokerRouter { + self.broker_router.clone() + } + + pub fn session_router(&self) -> Arc { + self.broker_router.sessions() + } + + pub async fn close(&self) -> Result<(), WorkdirError> { + let mut first_error = None; + for (_, lease) in &self.leases { + if let Err(error) = lease.close().await { + first_error.get_or_insert(error); + } + } + first_error.map_or(Ok(()), Err) + } + + pub fn is_active(&self) -> bool { + self.leases.iter().all(|(_, lease)| lease.is_active()) + } + + pub fn revoke(&self) { + for (_, lease) in &self.leases { + lease.revoke(); + } + } +} + #[derive(Debug)] struct SessionValidity { active: AtomicBool, @@ -2303,6 +2501,69 @@ mod tests { child.close().await.unwrap(); } + #[tokio::test] + async fn broker_router_delegates_multiple_aliases_without_namespace_mixing() { + let left = TempDir::new().unwrap(); + let right = TempDir::new().unwrap(); + fs::write(left.path().join("same.txt"), "left").unwrap(); + fs::write(right.path().join("same.txt"), "right").unwrap(); + let sessions = Arc::new(crate::WorkdirSessionRouter::new()); + sessions + .attach( + crate::WorkdirAttachmentAlias::new("left").unwrap(), + Arc::new(crate::LocalWorkdirSession::new( + manifest::Scope::writable(left.path()).unwrap(), + left.path().to_path_buf(), + )), + ) + .unwrap(); + sessions + .attach( + crate::WorkdirAttachmentAlias::new("right").unwrap(), + Arc::new(crate::LocalWorkdirSession::new( + manifest::Scope::writable(right.path()).unwrap(), + right.path().to_path_buf(), + )), + ) + .unwrap(); + let brokers = WorkdirToolBrokerRouter::new(sessions); + let leases = brokers + .scope(vec![ + ( + crate::WorkdirAttachmentAlias::new("left").unwrap(), + request("", WorkdirToolScopePermission::Read), + ), + ( + crate::WorkdirAttachmentAlias::new("right").unwrap(), + request("", WorkdirToolScopePermission::Read), + ), + ]) + .await + .unwrap(); + + let child = leases.session_router(); + let left_content = child + .resolve(Some("left")) + .unwrap() + .session + .read(read("same.txt")) + .await + .unwrap(); + let right_content = child + .resolve(Some("right")) + .unwrap() + .session + .read(read("same.txt")) + .await + .unwrap(); + assert_eq!(left_content.bytes, b"left"); + assert_eq!(right_content.bytes, b"right"); + assert!(child.resolve(None).is_err()); + + leases.close().await.unwrap(); + assert!(!leases.is_active()); + } + #[tokio::test] async fn closing_parent_invalidates_scoped_tools() { let root = TempDir::new().unwrap(); diff --git a/crates/workdir/src/workspace.rs b/crates/workdir/src/workspace.rs index 76990f398..3bec80329 100644 --- a/crates/workdir/src/workspace.rs +++ b/crates/workdir/src/workspace.rs @@ -70,6 +70,9 @@ pub struct RuntimeWorkingDirectoryCleanupTarget { #[serde(deny_unknown_fields)] pub struct RuntimeWorkingDirectorySummary { pub working_directory_id: String, + /// Optional human-facing label. Never used for routing or identity. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub display_name: Option, pub repository_id: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub creation_selector: Option, @@ -92,8 +95,6 @@ pub struct RuntimeWorkingDirectorySummary { #[serde(default, skip_serializing_if = "Option::is_none")] pub cleanliness: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - pub primary_worker_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] pub occupied_by: Option, } @@ -101,6 +102,9 @@ pub struct RuntimeWorkingDirectorySummary { #[serde(deny_unknown_fields)] pub struct WorkingDirectorySummary { pub working_directory_id: String, + /// Optional human-facing label. Never used for attachment routing. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub display_name: Option, pub repository_key: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub creation_selector: Option, @@ -123,8 +127,6 @@ pub struct WorkingDirectorySummary { #[serde(default, skip_serializing_if = "Option::is_none")] pub cleanliness: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - pub primary_worker_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] pub occupied_by: Option, } @@ -184,8 +186,6 @@ pub struct WorkingDirectoryCurrentObservation { #[serde(default, skip_serializing_if = "Option::is_none")] pub cleanliness: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - pub primary_worker_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] pub occupied_by: Option, } @@ -214,6 +214,7 @@ mod tests { fn working_directory_summary_preserves_server_projection_wire_shape() { let summary = WorkingDirectorySummary { working_directory_id: "wd-1".to_string(), + display_name: Some("Checkout".to_string()), repository_key: "repo-1".to_string(), creation_selector: None, creation_ref: None, @@ -226,7 +227,6 @@ mod tests { cleanup_target: None, status: WorkingDirectoryStatusKind::Active, cleanliness: None, - primary_worker_id: None, occupied_by: None, }; let value = serde_json::to_value(summary).unwrap(); @@ -240,5 +240,7 @@ mod tests { #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct WorkspaceWorkdirSessionOperationRequest { + /// Worker-local attachment alias selected by the calling tool router. + pub target_workdir: String, pub operation: crate::http::WorkdirSessionOperation, } diff --git a/crates/worker-runtime/src/catalog.rs b/crates/worker-runtime/src/catalog.rs index ff99a2d9a..51a4b64e3 100644 --- a/crates/worker-runtime/src/catalog.rs +++ b/crates/worker-runtime/src/catalog.rs @@ -161,6 +161,8 @@ pub struct WorkingDirectoryRepositoryAccessRequest { #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct WorkingDirectoryRequest { pub repository: WorkingDirectoryRepository, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub display_name: Option, #[serde(default)] pub materializer: MaterializerKind, /// Backend-assigned stable Workdir id. Runtimes use this when present so the @@ -197,12 +199,25 @@ pub struct RepositoryRefObservation { } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct WorkingDirectoryClaim { +pub struct WorkingDirectoryAttachmentClaim { + pub alias: workdir::WorkdirAttachmentAlias, pub working_directory_id: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub relative_cwd: Option, } +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct WorkingDirectoryAttachmentRequest { + pub alias: workdir::WorkdirAttachmentAlias, + pub working_directory: WorkingDirectoryRequest, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct WorkingDirectoryAttachmentStatus { + pub alias: workdir::WorkdirAttachmentAlias, + pub working_directory: WorkingDirectoryStatus, +} + #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct WorkingDirectoryStatus { pub summary: WorkingDirectorySummary, @@ -247,10 +262,10 @@ pub struct CreateWorkerRequest { pub config_bundle: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub initial_input: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub working_directory_request: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub working_directory: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub workdir_attachment_requests: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub workdir_attachments: Vec, /// Backend-only feature enablement. Grants still define local Runtime peers; /// the Workspace provider reauthorizes its dynamic set per operation. #[serde(default, skip_serializing_if = "is_false")] @@ -308,8 +323,8 @@ pub struct WorkerSummary { pub worker_state: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub workspace_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub working_directory: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub workdir_attachments: Vec, pub profile: ProfileSelector, #[serde(default, skip_serializing_if = "Option::is_none")] pub display_name: Option, @@ -333,8 +348,8 @@ pub struct WorkerDetail { pub worker_state: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub workspace_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub working_directory: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub workdir_attachments: Vec, pub profile: ProfileSelector, #[serde(default, skip_serializing_if = "Option::is_none")] pub display_name: Option, diff --git a/crates/worker-runtime/src/execution.rs b/crates/worker-runtime/src/execution.rs index 54934181a..257dab2b2 100644 --- a/crates/worker-runtime/src/execution.rs +++ b/crates/worker-runtime/src/execution.rs @@ -12,6 +12,7 @@ use crate::observation::WorkerObservationEvent; use crate::working_directory::{WorkingDirectoryBinding, WorkingDirectoryDiagnostic}; use protocol::{Method, UploadedFileRef}; use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; use std::fmt; use std::sync::Arc; use workdir::WorkdirSessionHandle; @@ -246,7 +247,7 @@ pub struct WorkerExecutionSpawnRequest { pub request: crate::catalog::CreateWorkerRequest, pub workspace_scope: Option, pub context: WorkerExecutionContext, - pub working_directory: Option, + pub workdir_attachments: BTreeMap, pub config_bundle: Option, } @@ -257,8 +258,8 @@ pub struct WorkerExecutionRestoreRequest { pub request: crate::catalog::CreateWorkerRequest, pub workspace_scope: Option, pub context: WorkerExecutionContext, - pub previous_working_directory: Option, - pub working_directory: Option, + pub previous_workdir_attachments: Vec, + pub workdir_attachments: BTreeMap, pub config_bundle: Option, } @@ -284,7 +285,7 @@ pub enum WorkerExecutionSpawnResult { Connected { handle: WorkerExecutionHandle, worker_state: protocol::WorkerStateSnapshot, - working_directory: Option, + workdir_attachments: Vec, }, /// A read-only preflight rejected the operation before live work started. Rejected(WorkerExecutionResult), @@ -297,7 +298,7 @@ pub enum WorkerExecutionSpawnResult { result: WorkerExecutionResult, handle: Option, worker_state: Option, - working_directory: Option, + workdir_attachments: Vec, }, /// Legacy spawn failure. Restore implementations must use `RolledBack` or /// `ReconciliationRequired` after crossing their live side-effect boundary. @@ -308,12 +309,12 @@ impl WorkerExecutionSpawnResult { pub fn connected( handle: WorkerExecutionHandle, worker_state: protocol::WorkerStateSnapshot, - working_directory: Option, + workdir_attachments: Vec, ) -> Self { Self::Connected { handle, worker_state, - working_directory, + workdir_attachments, } } } diff --git a/crates/worker-runtime/src/fs_store.rs b/crates/worker-runtime/src/fs_store.rs index d72842b82..cddf440e1 100644 --- a/crates/worker-runtime/src/fs_store.rs +++ b/crates/worker-runtime/src/fs_store.rs @@ -1,6 +1,7 @@ use crate::catalog::{ ConfigBundleRef, CreateWorkerRequest, ProfileSelector, WorkerRestoreIntent, WorkerStatus, - WorkingDirectoryStatus, + WorkingDirectoryAttachmentClaim, WorkingDirectoryAttachmentRequest, + WorkingDirectoryAttachmentStatus, WorkingDirectoryRequest, }; use crate::config_bundle::ConfigBundle; use crate::diagnostics::{DiagnosticSeverity, RuntimeDiagnostic}; @@ -442,7 +443,7 @@ pub(crate) struct PersistedWorkerRecord { pub(crate) status: WorkerStatus, pub(crate) execution_state: PersistedWorkerExecutionState, pub(crate) workspace_id: Option, - pub(crate) working_directory: Option, + pub(crate) workdir_attachments: Vec, } fn runtime_io_error(operation: &'static str, path: &Path, source: std::io::Error) -> RuntimeError { @@ -1685,11 +1686,18 @@ struct WorkerIdentityRecord { status: WorkerStatus, #[serde(default, skip_serializing_if = "Option::is_none")] workspace_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - working_directory: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + workdir_attachments: Vec, + /// Schema-v8 compatibility input. New records serialize only `workdir_attachments`. + #[serde( + default, + rename = "working_directory", + skip_serializing_if = "Option::is_none" + )] + legacy_working_directory: Option, } -#[derive(Clone, Debug, Serialize, Deserialize)] +#[derive(Clone, Debug, Serialize)] #[serde(deny_unknown_fields)] struct WorkerExecutionRecord { schema_version: u32, @@ -1699,6 +1707,72 @@ struct WorkerExecutionRecord { restore_intent: WorkerRestoreIntent, } +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RawWorkerExecutionRecord { + schema_version: u32, + request: serde_json::Value, + #[serde(default)] + binding: Option, + restore_intent: WorkerRestoreIntent, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct LegacyWorkingDirectoryClaim { + working_directory_id: String, + #[serde(default)] + relative_cwd: Option, +} + +impl<'de> Deserialize<'de> for WorkerExecutionRecord { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let raw = RawWorkerExecutionRecord::deserialize(deserializer)?; + let mut request_value = raw.request; + let request_object = request_value + .as_object_mut() + .ok_or_else(|| serde::de::Error::custom("Worker create request must be an object"))?; + let legacy_request = request_object.remove("working_directory_request"); + let legacy_claim = request_object.remove("working_directory"); + let mut request: CreateWorkerRequest = + serde_json::from_value(request_value).map_err(serde::de::Error::custom)?; + if request.workdir_attachment_requests.is_empty() && request.workdir_attachments.is_empty() + { + let alias = workdir::WorkdirAttachmentAlias::new("workdir") + .expect("static compatibility alias is valid"); + if let Some(value) = legacy_request.filter(|value| !value.is_null()) { + let working_directory: WorkingDirectoryRequest = + serde_json::from_value(value).map_err(serde::de::Error::custom)?; + request + .workdir_attachment_requests + .push(WorkingDirectoryAttachmentRequest { + alias, + working_directory, + }); + } else if let Some(value) = legacy_claim.filter(|value| !value.is_null()) { + let claim: LegacyWorkingDirectoryClaim = + serde_json::from_value(value).map_err(serde::de::Error::custom)?; + request + .workdir_attachments + .push(WorkingDirectoryAttachmentClaim { + alias, + working_directory_id: claim.working_directory_id, + relative_cwd: claim.relative_cwd, + }); + } + } + Ok(Self { + schema_version: raw.schema_version, + request, + binding: raw.binding, + restore_intent: raw.restore_intent, + }) + } +} + impl WorkerIdentityRecord { fn from_persisted(worker: &PersistedWorkerRecord) -> Self { Self { @@ -1712,7 +1786,8 @@ impl WorkerIdentityRecord { created_at_ms: worker.created_at_ms, status: worker.status, workspace_id: worker.workspace_id.clone(), - working_directory: worker.working_directory.clone(), + workdir_attachments: worker.workdir_attachments.clone(), + legacy_working_directory: None, } } @@ -1760,6 +1835,19 @@ impl WorkerIdentityRecord { self, execution_state: PersistedWorkerExecutionState, ) -> PersistedWorkerRecord { + let workdir_attachments = if self.workdir_attachments.is_empty() { + self.legacy_working_directory + .map(|working_directory| { + vec![WorkingDirectoryAttachmentStatus { + alias: workdir::WorkdirAttachmentAlias::new("workdir") + .expect("static compatibility alias is valid"), + working_directory, + }] + }) + .unwrap_or_default() + } else { + self.workdir_attachments + }; PersistedWorkerRecord { worker_ref: self.worker_ref, worker_id: self.worker_id, @@ -1771,7 +1859,7 @@ impl WorkerIdentityRecord { status: self.status, execution_state, workspace_id: self.workspace_id, - working_directory: self.working_directory, + workdir_attachments, } } } @@ -2215,6 +2303,72 @@ mod tests { document } + #[test] + fn schema_v8_singular_workdir_fields_normalize_to_one_attachment() { + let worker_id = WorkerId::now_v7(); + let workdir_status = serde_json::json!({ + "summary": { + "working_directory_id": "wd-legacy", + "repository_id": "repo", + "materializer_kind": "runtime_git_clone", + "status": "active" + } + }); + let identity: WorkerIdentityRecord = serde_json::from_value(serde_json::json!({ + "schema_version": SCHEMA_VERSION, + "worker_ref": { "worker_id": worker_id }, + "worker_id": worker_id, + "profile": { "kind": "builtin", "value": "builtin:coder" }, + "display_name": null, + "profile_source": { + "id": "archive-1", + "digest": "sha256:archive", + "size_bytes": 0, + "source_graph": { + "source_count": 0, + "total_source_bytes": 0, + "entrypoints": {}, + "import_count": 0 + } + }, + "config_bundle": null, + "status": "stopped", + "working_directory": workdir_status + })) + .unwrap(); + let persisted = identity.into_persisted(PersistedWorkerExecutionState::Unavailable); + assert_eq!(persisted.workdir_attachments.len(), 1); + assert_eq!(persisted.workdir_attachments[0].alias.as_str(), "workdir"); + assert_eq!( + persisted.workdir_attachments[0] + .working_directory + .summary + .working_directory_id, + "wd-legacy" + ); + + let mut request = schema_v7_worker_document(worker_id)["request"].clone(); + request["working_directory"] = serde_json::json!({ + "working_directory_id": "wd-legacy", + "relative_cwd": "crates/yoi" + }); + let execution: WorkerExecutionRecord = serde_json::from_value(serde_json::json!({ + "schema_version": SCHEMA_VERSION, + "request": request, + "binding": null, + "restore_intent": "explicit" + })) + .unwrap(); + assert!(execution.request.workdir_attachment_requests.is_empty()); + assert_eq!(execution.request.workdir_attachments.len(), 1); + let attachment = &execution.request.workdir_attachments[0]; + assert_eq!(attachment.alias.as_str(), "workdir"); + assert_eq!(attachment.working_directory_id, "wd-legacy"); + assert_eq!(attachment.relative_cwd.as_deref(), Some("crates/yoi")); + let serialized = serde_json::to_value(execution).unwrap(); + assert!(serialized["request"].get("working_directory").is_none()); + } + #[test] fn startup_recovers_staging_created_before_source_backup() { let temp = tempfile::tempdir().unwrap(); diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index 404071938..98bde5397 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -2282,7 +2282,11 @@ fn required_runtime_permission(method: &Method, path: &str) -> Option<&'static s if path.ends_with("/workspace-api") { return Some("workers:create"); } - if path.ends_with("/input") || path.ends_with("/restore") || path.contains("/attachments") { + if path.ends_with("/input") + || path.ends_with("/restore") + || path.contains("/attachments") + || path.contains("/workdir-attachments") + { return Some("workers:input"); } if path.ends_with("/stop") || path.ends_with("/cancel") { @@ -3072,8 +3076,8 @@ mod tests { digest: bundle.metadata.digest, }), initial_input: None, - working_directory_request: None, - working_directory: None, + workdir_attachment_requests: Vec::new(), + workdir_attachments: Vec::new(), worker_observation_enabled: false, worker_observation_grants: Vec::new(), workspace_api: None, @@ -3319,10 +3323,16 @@ mod tests { worker_state: protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, - working_directory: request - .working_directory - .as_ref() - .map(|binding| binding.status()), + workdir_attachments: request + .workdir_attachments + .iter() + .map( + |(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus { + alias: alias.clone(), + working_directory: binding.status(), + }, + ) + .collect(), } } @@ -3335,7 +3345,7 @@ mod tests { worker_state: protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, - working_directory: request.previous_working_directory, + workdir_attachments: request.previous_workdir_attachments, } } @@ -3907,10 +3917,16 @@ mod ws_tests { worker_state: protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, - working_directory: request - .working_directory - .as_ref() - .map(|binding| binding.status()), + workdir_attachments: request + .workdir_attachments + .iter() + .map( + |(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus { + alias: alias.clone(), + working_directory: binding.status(), + }, + ) + .collect(), } } @@ -4020,8 +4036,8 @@ mod ws_tests { digest: bundle.metadata.digest, }), initial_input: None, - working_directory_request: None, - working_directory: None, + workdir_attachment_requests: Vec::new(), + workdir_attachments: Vec::new(), worker_observation_enabled: false, worker_observation_grants: Vec::new(), workspace_api: None, diff --git a/crates/worker-runtime/src/http_server/runtime_management_api.rs b/crates/worker-runtime/src/http_server/runtime_management_api.rs index 8d29fe188..51b65d5ca 100644 --- a/crates/worker-runtime/src/http_server/runtime_management_api.rs +++ b/crates/worker-runtime/src/http_server/runtime_management_api.rs @@ -423,6 +423,31 @@ impl runtime_api::RuntimeApi for RuntimeManagementApi { Ok(runtime_api::WorkerResponse { worker }) } + async fn replace_worker_workdir_attachments( + &self, + worker_id: String, + value: runtime_api::WorkerWorkdirAttachmentsRequest, + ) -> Result { + let attachments = request(value.workdir_attachments)?; + let worker_ref = worker_ref_for(&self.state.runtime, worker_id).map_err(api_error)?; + let worker = match auth_workspace_scope(&self.state, auth_extension().as_ref()) + .map_err(api_error)? + { + Some(scope) => self + .state + .runtime + .replace_worker_workdir_attachments_scoped(&scope, &worker_ref, attachments), + None => self + .state + .runtime + .replace_worker_workdir_attachments(&worker_ref, attachments), + } + .map_err(RuntimeHttpRestError::runtime) + .map_err(api_error)?; + let worker = response(worker)?; + Ok(runtime_api::WorkerResponse { worker }) + } + async fn complete_worker_arguments( &self, worker_id: String, diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 6b440706a..1f383e9c8 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -1,7 +1,8 @@ use crate::catalog::{ ConfigBundleRef, CreateWorkerRequest, ProfileSelector, ProfileSourceArchiveSource, RepositoryRefObservation, RepositoryRefObservationRequest, WorkerDetail, WorkerLifecycleAck, - WorkerRestoreIntent, WorkerStatus, WorkerSummary, WorkingDirectoryRepositoryAccessRequest, + WorkerRestoreIntent, WorkerStatus, WorkerSummary, WorkingDirectoryAttachmentClaim, + WorkingDirectoryAttachmentStatus, WorkingDirectoryRepositoryAccessRequest, WorkingDirectoryRequest, WorkingDirectoryStatus as CatalogWorkingDirectoryStatus, WorkspaceApiRef, }; @@ -43,7 +44,7 @@ use crate::retention::{ use protocol::subscription::{ EventSubscriptionSelector, SubscriptionEventPayload, SubscriptionSnapshot, SubscriptionValidationError, SubscriptionWorkdirId, SubscriptionWorker, SubscriptionWorkerId, - SubscriptionWorkerState, + SubscriptionWorkerState, SubscriptionWorkerWorkdirAttachment, }; use protocol::{Event, Method}; use server_api::WorkerRestoreState; @@ -729,8 +730,9 @@ impl Runtime { .get(&worker_ref.worker_id) .is_some_and(|worker| { worker.belongs_to_workspace(&scope.workspace_id) - && worker.working_directory.as_ref().is_some_and(|status| { - status.summary.working_directory_id == working_directory_id + && worker.workdir_attachments.iter().any(|attachment| { + attachment.working_directory.summary.working_directory_id + == working_directory_id }) }); if !owns_workdir { @@ -764,7 +766,7 @@ impl Runtime { let backend = { let state = self.lock()?; state.ensure_running()?; - if let Some(worker_id) = state.primary_worker_id_for_workdir(working_directory_id) { + if let Some(worker_id) = state.worker_id_for_workdir(working_directory_id) { return Err(RuntimeError::InvalidRequest(format!( "working directory {working_directory_id} is assigned to worker {worker_id}" ))); @@ -792,12 +794,8 @@ impl Runtime { fn annotate_working_directory_status( &self, - mut status: CatalogWorkingDirectoryStatus, + status: CatalogWorkingDirectoryStatus, ) -> Result { - let state = self.lock()?; - status.summary.primary_worker_id = state - .primary_worker_id_for_workdir(status.summary.working_directory_id.as_str()) - .map(|worker_id| worker_id.to_string()); Ok(status) } @@ -1023,10 +1021,8 @@ impl Runtime { return Ok(existing.detail()); } state.validate_worker_config_boundary(&request)?; - if let Some(working_directory_id) = requested_primary_workdir_id(&request) { - if let Some(owner_worker_id) = - state.primary_worker_id_for_workdir(working_directory_id) - { + for working_directory_id in requested_workdir_ids(&request) { + if let Some(owner_worker_id) = state.worker_id_for_workdir(working_directory_id) { return Err(RuntimeError::InvalidRequest(format!( "working directory {working_directory_id} is already assigned to worker {owner_worker_id}" ))); @@ -1064,7 +1060,7 @@ impl Runtime { execution_metadata_available: true, execution_bound: true, restore_intent: WorkerRestoreIntent::Explicit, - working_directory: None, + workdir_attachments: Vec::new(), execution_handle: None, internal_workers: InternalWorkerActivityProjection::default(), }; @@ -1076,19 +1072,19 @@ impl Runtime { request, workspace_scope: scope.cloned(), context: self.execution_context(worker_ref.clone()), - working_directory: None, + workdir_attachments: BTreeMap::new(), config_bundle, }; (backend, worker_ref, spawn_request) }; let spawn_result = backend.spawn_worker(spawn_request); - let (handle, initial_worker_state, working_directory) = match spawn_result { + let (handle, initial_worker_state, workdir_attachments) = match spawn_result { WorkerExecutionSpawnResult::Connected { handle, worker_state, - working_directory, - } => (handle, worker_state, working_directory), + workdir_attachments, + } => (handle, worker_state, workdir_attachments), WorkerExecutionSpawnResult::Rejected(result) | WorkerExecutionSpawnResult::RolledBack(result) | WorkerExecutionSpawnResult::Errored(result) => { @@ -1101,8 +1097,30 @@ impl Runtime { result, }); } - WorkerExecutionSpawnResult::ReconciliationRequired { result, .. } => { - self.rollback_failed_create(&worker_ref)?; + WorkerExecutionSpawnResult::ReconciliationRequired { + result, + handle, + worker_state, + workdir_attachments, + } => { + let retain_result = + if let (Some(handle), Some(worker_state)) = (handle, worker_state) { + self.retain_restore_execution_evidence( + &worker_ref, + handle, + worker_state, + workdir_attachments, + ) + } else { + self.retain_create_reconciliation_pending(&worker_ref, workdir_attachments) + }; + if let Err(error) = retain_result { + tracing::error!( + worker_id = %worker_ref.worker_id, + %error, + "failed to persist Worker create reconciliation evidence" + ); + } return Err(RuntimeError::WorkerExecutionRejected { worker_id: worker_ref.worker_id.clone(), operation: result.operation, @@ -1160,7 +1178,7 @@ impl Runtime { &worker_ref, handle.clone(), initial_worker_state.clone(), - working_directory, + workdir_attachments, dispatch_result, ) { Ok(detail) => detail, @@ -1179,7 +1197,7 @@ impl Runtime { &worker_ref, handle.clone(), initial_worker_state, - working_directory, + workdir_attachments, WorkerExecutionResult::accepted(WorkerExecutionOperation::Spawn), ) { Ok(detail) => Ok(detail), @@ -1425,6 +1443,119 @@ impl Runtime { Ok(state.worker(worker_ref)?.detail()) } + pub fn replace_worker_workdir_attachments_scoped( + &self, + scope: &RuntimeWorkspaceScope, + worker_ref: &WorkerRef, + attachments: Vec, + ) -> Result { + self.ensure_worker_in_workspace(scope, worker_ref)?; + self.replace_worker_workdir_attachments(worker_ref, attachments) + } + + /// Persist the exact current attachment set used by live mutation and future restore. + pub fn replace_worker_workdir_attachments( + &self, + worker_ref: &WorkerRef, + mut attachments: Vec, + ) -> Result { + let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; + let _operation_guard = operation_lock + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + + let backend = { + let state = self.lock()?; + let worker = state.worker(worker_ref)?; + if worker.execution_handle.is_none() || !worker.status.is_active() { + return Err(RuntimeError::WorkerExecutionUnavailable { + worker_id: worker.worker_id, + message: "Workdir attachments can only be changed for a live Worker" + .to_string(), + }); + } + let request = worker.request.as_ref().ok_or_else(|| { + RuntimeError::WorkerExecutionUnavailable { + worker_id: worker.worker_id, + message: "persisted Worker restore request is unavailable".to_string(), + } + })?; + for attachment in &mut attachments { + if attachment.relative_cwd.is_none() + && let Some(previous) = request.workdir_attachments.iter().find(|previous| { + previous.alias == attachment.alias + && previous.working_directory_id == attachment.working_directory_id + }) + { + attachment.relative_cwd = previous.relative_cwd.clone(); + } + } + let mut candidate = request.clone(); + candidate.workdir_attachment_requests.clear(); + candidate.workdir_attachments = attachments.clone(); + validate_create_worker_request(&candidate)?; + state.execution_backend.clone().ok_or_else(|| { + RuntimeError::ExecutionBackendUnavailable { + message: "Workdir attachment mutation requires an execution backend" + .to_string(), + } + })? + }; + + let statuses = attachments + .iter() + .map(|attachment| { + backend + .working_directory(&attachment.working_directory_id) + .map(|working_directory| WorkingDirectoryAttachmentStatus { + alias: attachment.alias.clone(), + working_directory, + }) + .map_err(RuntimeError::from) + }) + .collect::, RuntimeError>>()?; + if let Some(inactive) = statuses.iter().find(|attachment| { + attachment.working_directory.summary.status + != workdir::workspace::WorkingDirectoryStatusKind::Active + }) { + return Err(RuntimeError::InvalidRequest(format!( + "working directory {} is not active ({})", + inactive.working_directory.summary.working_directory_id, + inactive.working_directory.summary.status + ))); + } + + let mut state = self.lock()?; + for attachment in &attachments { + if let Some(owner_worker_id) = + state.worker_id_for_workdir(&attachment.working_directory_id) + && owner_worker_id != worker_ref.worker_id + { + return Err(RuntimeError::InvalidRequest(format!( + "working directory {} is already assigned to worker {owner_worker_id}", + attachment.working_directory_id + ))); + } + } + let previous = state.worker(worker_ref)?.clone(); + { + let worker = state.worker_mut(worker_ref)?; + let request = worker + .request + .as_mut() + .expect("restore request checked above"); + request.workdir_attachment_requests.clear(); + request.workdir_attachments = attachments; + worker.workdir_attachments = statuses; + } + if let Err(error) = state.persist_worker(&worker_ref.worker_id) { + state.workers.insert(worker_ref.worker_id, previous); + return Err(error); + } + state.publish_worker_upsert(worker_ref.worker_id)?; + Ok(state.worker(worker_ref)?.detail()) + } + /// Observe a Worker's live protocol availability or its retained Session. /// The same lifecycle lock used by stop/restore fences the decision and read. pub fn worker_session_scoped( @@ -1570,8 +1701,8 @@ impl Runtime { request: worker_request, workspace_scope, context: self.execution_context(worker_ref.clone()), - previous_working_directory: worker.working_directory.clone(), - working_directory: None, + previous_workdir_attachments: worker.workdir_attachments.clone(), + workdir_attachments: BTreeMap::new(), config_bundle: None, }; (backend, request) @@ -1589,14 +1720,14 @@ impl Runtime { WorkerExecutionSpawnResult::Connected { handle, worker_state, - working_directory, + workdir_attachments, } => { let commit = self.commit_restored_worker_execution( worker_ref, handle.clone(), worker_state.clone(), WorkerStatus::Idle, - working_directory.clone(), + workdir_attachments.clone(), ); match commit { Ok(worker) => Ok(RuntimeWorkerRestoreResult::accepted(worker)), @@ -1625,7 +1756,7 @@ impl Runtime { worker_ref, handle, worker_state, - working_directory, + workdir_attachments, ) { tracing::error!( worker_id = %worker_ref.worker_id, @@ -1663,14 +1794,14 @@ impl Runtime { result: _result, handle, worker_state, - working_directory, + workdir_attachments, } => { if let (Some(handle), Some(worker_state)) = (handle, worker_state) { if let Err(retain_error) = self.retain_restore_execution_evidence( worker_ref, handle, worker_state, - working_directory, + workdir_attachments, ) { tracing::error!( worker_id = %worker_ref.worker_id, @@ -2079,7 +2210,7 @@ impl Runtime { worker_ref: &WorkerRef, handle: WorkerExecutionHandle, initial_worker_state: protocol::WorkerStateSnapshot, - working_directory: Option, + workdir_attachments: Vec, result: WorkerExecutionResult, ) -> Result { let mut state = self.lock()?; @@ -2094,7 +2225,7 @@ impl Runtime { let _ = worker.apply_worker_state(snapshot); } worker.restore_intent = restore_intent_for_status(worker.status); - worker.working_directory = working_directory; + worker.workdir_attachments = workdir_attachments; worker.detail() }; state.publish_worker_upsert(worker_ref.worker_id)?; @@ -2367,9 +2498,13 @@ impl Runtime { state.ensure_running()?; state.ensure_worker_ref(worker_ref)?; let worker = state.worker(worker_ref)?; - if worker.execution_handle.is_some() && worker.status.is_active() { + if worker.status.is_active() + && (worker.execution_handle.is_some() + || (!worker.execution_metadata_available + && !worker.workdir_attachments.is_empty())) + { return Err(RuntimeError::InvalidRequest(format!( - "worker {} is running and must be stopped before deletion", + "worker {} is active or requires reconciliation and must be stopped before deletion", worker_ref.worker_id ))); } @@ -2398,9 +2533,12 @@ impl Runtime { state.ensure_running()?; state.ensure_worker_ref(worker_ref)?; let worker = state.worker(worker_ref)?; - if worker.execution_handle.is_some() && worker.status.is_active() { + if worker.status.is_active() + && (worker.execution_handle.is_some() + || (!worker.execution_metadata_available && !worker.workdir_attachments.is_empty())) + { return Err(RuntimeError::InvalidRequest(format!( - "worker {} became active before deletion", + "worker {} became active or reconciliation-required before deletion", worker_ref.worker_id ))); } @@ -2662,7 +2800,7 @@ impl Runtime { handle: WorkerExecutionHandle, worker_state: protocol::WorkerStateSnapshot, status: WorkerStatus, - working_directory: Option, + workdir_attachments: Vec, ) -> Result { let mut state = self.lock()?; state.ensure_worker_ref(worker_ref)?; @@ -2679,7 +2817,7 @@ impl Runtime { candidate.status = status; let _ = candidate.apply_worker_state(&worker_state); candidate.restore_intent = restore_intent_for_status(candidate.status); - candidate.working_directory = working_directory; + candidate.workdir_attachments = workdir_attachments; // Validate the exact subscription projection before persistence. The // subsequent publish uses the same pure projection while the Worker @@ -2708,7 +2846,7 @@ impl Runtime { worker_ref: &WorkerRef, handle: WorkerExecutionHandle, worker_state: protocol::WorkerStateSnapshot, - working_directory: Option, + workdir_attachments: Vec, ) -> Result<(), RuntimeError> { let mut state = self.lock()?; let mut candidate = state.worker(worker_ref)?.clone(); @@ -2718,7 +2856,26 @@ impl Runtime { candidate.status = WorkerStatus::Idle; let _ = candidate.apply_worker_state(&worker_state); candidate.restore_intent = WorkerRestoreIntent::Automatic; - candidate.working_directory = working_directory; + candidate.workdir_attachments = workdir_attachments; + state.workers.insert(worker_ref.worker_id, candidate); + state.publish_worker_upsert(worker_ref.worker_id)?; + state.persist_worker(&worker_ref.worker_id)?; + Ok(()) + } + + fn retain_create_reconciliation_pending( + &self, + worker_ref: &WorkerRef, + workdir_attachments: Vec, + ) -> Result<(), RuntimeError> { + let mut state = self.lock()?; + let mut candidate = state.worker(worker_ref)?.clone(); + candidate.execution_handle = None; + candidate.execution_metadata_available = false; + candidate.execution_bound = true; + candidate.status = WorkerStatus::Idle; + candidate.restore_intent = WorkerRestoreIntent::Explicit; + candidate.workdir_attachments = workdir_attachments; state.workers.insert(worker_ref.worker_id, candidate); state.publish_worker_upsert(worker_ref.worker_id)?; state.persist_worker(&worker_ref.worker_id)?; @@ -3078,7 +3235,7 @@ impl RuntimeState { execution_metadata_available, execution_bound, restore_intent, - working_directory: worker.working_directory, + workdir_attachments: worker.workdir_attachments, execution_handle: None, internal_workers: InternalWorkerActivityProjection::default(), }, @@ -3394,18 +3551,25 @@ impl RuntimeState { ) -> Result { let worker_id = SubscriptionWorkerId::new(worker.worker_id.to_string()) .map_err(subscription_validation_error)?; - let repository_id = worker - .working_directory - .as_ref() - .map(|working_directory| working_directory.summary.repository_id.clone()); - let working_directory_id = worker - .working_directory - .as_ref() - .map(|working_directory| { - SubscriptionWorkdirId::new(working_directory.summary.working_directory_id.clone()) - .map_err(subscription_validation_error) + let workdir_attachments = worker + .workdir_attachments + .iter() + .map(|attachment| { + Ok(SubscriptionWorkerWorkdirAttachment { + alias: attachment.alias.to_string(), + repository_id: Some(attachment.working_directory.summary.repository_id.clone()), + repository_key: None, + working_directory_id: SubscriptionWorkdirId::new( + attachment + .working_directory + .summary + .working_directory_id + .clone(), + ) + .map_err(subscription_validation_error)?, + }) }) - .transpose()?; + .collect::, RuntimeError>>()?; let profile = match &worker.profile { ProfileSelector::Builtin(name) | ProfileSelector::Named(name) => Some(name.clone()), }; @@ -3425,9 +3589,7 @@ impl RuntimeState { workspace_id: worker.workspace_id.clone(), display_name: worker.display_name.clone(), profile, - repository_id, - repository_key: None, - working_directory_id, + workdir_attachments, }) } @@ -3525,22 +3687,15 @@ impl RuntimeState { } } - fn primary_worker_id_for_workdir(&self, working_directory_id: &str) -> Option { + fn worker_id_for_workdir(&self, working_directory_id: &str) -> Option { self.workers.values().find_map(|worker| { - if worker - .working_directory - .as_ref() - .is_some_and(|binding| binding.summary.working_directory_id == working_directory_id) - || worker - .request - .as_ref() - .and_then(requested_primary_workdir_id) - == Some(working_directory_id) - { - Some(worker.worker_id) - } else { - None - } + let has_status = worker.workdir_attachments.iter().any(|attachment| { + attachment.working_directory.summary.working_directory_id == working_directory_id + }); + let has_request = worker.request.as_ref().is_some_and(|request| { + requested_workdir_ids(request).contains(&working_directory_id) + }); + (has_status || has_request).then_some(worker.worker_id) }) } @@ -3976,7 +4131,7 @@ struct WorkerRecord { execution_metadata_available: bool, execution_bound: bool, restore_intent: WorkerRestoreIntent, - working_directory: Option, + workdir_attachments: Vec, execution_handle: Option, internal_workers: InternalWorkerActivityProjection, } @@ -3999,7 +4154,7 @@ impl WorkerRecord { execution_metadata_available: self.execution_metadata_available, worker_state: self.worker_state.clone(), workspace_id: self.workspace_id.clone(), - working_directory: self.working_directory.clone(), + workdir_attachments: self.workdir_attachments.clone(), profile: self.profile.clone(), display_name: self.display_name.clone(), profile_source: self.profile_source.clone(), @@ -4016,7 +4171,7 @@ impl WorkerRecord { execution_metadata_available: self.execution_metadata_available, worker_state: self.worker_state.clone(), workspace_id: self.workspace_id.clone(), - working_directory: self.working_directory.clone(), + workdir_attachments: self.workdir_attachments.clone(), profile: self.profile.clone(), display_name: self.display_name.clone(), profile_source: self.profile_source.clone(), @@ -4049,7 +4204,7 @@ impl WorkerRecord { status: self.status, execution_state, workspace_id: self.workspace_id.clone(), - working_directory: self.working_directory.clone(), + workdir_attachments: self.workdir_attachments.clone(), } } } @@ -4100,25 +4255,28 @@ fn repository_resource_error(error: BackendResourceError) -> RuntimeError { fn durable_create_worker_request(request: &CreateWorkerRequest) -> CreateWorkerRequest { let mut durable = request.clone(); - if let Some(working_directory) = durable.working_directory_request.as_mut() - && let Some(materialization) = working_directory.materialization.as_mut() - { - materialization.ssh = None; + for attachment in &mut durable.workdir_attachment_requests { + if let Some(materialization) = attachment.working_directory.materialization.as_mut() { + materialization.ssh = None; + } } durable } -fn requested_primary_workdir_id(request: &CreateWorkerRequest) -> Option<&str> { +fn requested_workdir_ids(request: &CreateWorkerRequest) -> Vec<&str> { request - .working_directory - .as_ref() + .workdir_attachments + .iter() .map(|claim| claim.working_directory_id.as_str()) - .or_else(|| { + .chain( request - .working_directory_request - .as_ref() - .and_then(|request| request.backend_workdir_id.as_deref()) - }) + .workdir_attachment_requests + .iter() + .filter_map(|attachment| { + attachment.working_directory.backend_workdir_id.as_deref() + }), + ) + .collect() } fn validate_create_worker_request(request: &CreateWorkerRequest) -> Result<(), RuntimeError> { @@ -4127,6 +4285,42 @@ fn validate_create_worker_request(request: &CreateWorkerRequest) -> Result<(), R "create_fingerprint must not be empty".to_string(), )); } + let mut aliases = std::collections::BTreeSet::new(); + let mut workdir_ids = std::collections::BTreeSet::new(); + for (alias, workdir_id) in request + .workdir_attachments + .iter() + .map(|attachment| { + ( + &attachment.alias, + Some(attachment.working_directory_id.as_str()), + ) + }) + .chain( + request + .workdir_attachment_requests + .iter() + .map(|attachment| { + ( + &attachment.alias, + attachment.working_directory.backend_workdir_id.as_deref(), + ) + }), + ) + { + if !aliases.insert(alias.as_str()) { + return Err(RuntimeError::InvalidRequest(format!( + "duplicate Workdir attachment alias `{alias}`" + ))); + } + if let Some(workdir_id) = workdir_id + && !workdir_ids.insert(workdir_id) + { + return Err(RuntimeError::InvalidRequest(format!( + "working directory `{workdir_id}` cannot be attached more than once" + ))); + } + } match &request.profile_source { crate::catalog::ProfileSourceArchiveSource::Embedded { archive } => { archive.verify().map_err(|err| { @@ -4420,8 +4614,8 @@ mod tests { use crate::catalog::{ ConfigBundleRef, MaterializerKind, ProfileSelector, RepositoryMaterializationContext, RepositorySshCredentialCandidate, RepositorySshMaterializationAccess, SensitiveString, - WorkingDirectoryClaim, WorkingDirectoryRepository, WorkingDirectoryRequest, - WorkspaceApiRef, + WorkingDirectoryAttachmentClaim, WorkingDirectoryAttachmentRequest, + WorkingDirectoryRepository, WorkingDirectoryRequest, WorkspaceApiRef, }; use crate::config_bundle::{ ConfigBundle, ConfigBundleMetadata, ConfigBundleProvenance, ConfigDeclaration, @@ -4913,8 +5107,8 @@ mod tests { digest: bundle.metadata.digest, }), initial_input: None, - working_directory_request: None, - working_directory: None, + workdir_attachment_requests: Vec::new(), + workdir_attachments: Vec::new(), worker_observation_enabled: false, worker_observation_grants: Vec::new(), workspace_api: None, @@ -4929,60 +5123,64 @@ mod tests { #[test] fn durable_worker_request_omits_repository_credentials() { let mut request = task_request("worker-secret-redaction"); - request.working_directory_request = Some(WorkingDirectoryRequest { - repository: WorkingDirectoryRepository { - id: "repository-1".to_string(), - provider: "git".to_string(), - source: server_api::RepositorySource { - kind: server_api::RepositorySourceKind::Ssh, - uri: "ssh://git@example.test/repo.git".to_string(), + request.workdir_attachment_requests = vec![WorkingDirectoryAttachmentRequest { + alias: workdir::WorkdirAttachmentAlias::new("checkout").unwrap(), + working_directory: WorkingDirectoryRequest { + repository: WorkingDirectoryRepository { + id: "repository-1".to_string(), + provider: "git".to_string(), + source: server_api::RepositorySource { + kind: server_api::RepositorySourceKind::Ssh, + uri: "ssh://git@example.test/repo.git".to_string(), + }, + source_revision: 1, + source_fingerprint: "sha256:source".to_string(), + selector: None, }, - source_revision: 1, - source_fingerprint: "sha256:source".to_string(), - selector: None, - }, - materializer: MaterializerKind::RuntimeGitClone, - backend_workdir_id: Some("working-directory-1".to_string()), - materialization: Some(RepositoryMaterializationContext { - workspace_id: "workspace-1".to_string(), - runtime_id: "runtime-1".to_string(), - operation_id: "operation-1".to_string(), - config_revision: 1, - config_projection_digest: "sha256:projection".to_string(), - ssh: Some(RepositorySshMaterializationAccess { - credential_candidates: vec![RepositorySshCredentialCandidate { - credential_id: "credential-1".to_string(), - credential_revision: 1, - private_key: SensitiveString::new("private-key-bytes"), - }], - host_trust_id: "host-trust-1".to_string(), - host_trust_revision: 1, - access: server_api::RepositoryAccessMode::ReadOnly, - expires_at_epoch_seconds: u64::MAX, - repository_id: "repository-1".to_string(), - repository_source_fingerprint: "sha256:source".to_string(), - repository_uri: "ssh://git@example.test/repo.git".to_string(), - secret_resource: repository_resource_handle(), - known_hosts_entry: SensitiveString::new("known-hosts-entry"), + display_name: None, + materializer: MaterializerKind::RuntimeGitClone, + backend_workdir_id: Some("working-directory-1".to_string()), + materialization: Some(RepositoryMaterializationContext { + workspace_id: "workspace-1".to_string(), + runtime_id: "runtime-1".to_string(), + operation_id: "operation-1".to_string(), + config_revision: 1, + config_projection_digest: "sha256:projection".to_string(), + ssh: Some(RepositorySshMaterializationAccess { + credential_candidates: vec![RepositorySshCredentialCandidate { + credential_id: "credential-1".to_string(), + credential_revision: 1, + private_key: SensitiveString::new("private-key-bytes"), + }], + host_trust_id: "host-trust-1".to_string(), + host_trust_revision: 1, + access: server_api::RepositoryAccessMode::ReadOnly, + expires_at_epoch_seconds: u64::MAX, + repository_id: "repository-1".to_string(), + repository_source_fingerprint: "sha256:source".to_string(), + repository_uri: "ssh://git@example.test/repo.git".to_string(), + secret_resource: repository_resource_handle(), + known_hosts_entry: SensitiveString::new("known-hosts-entry"), + }), }), - }), - }); + }, + }]; let durable = durable_create_worker_request(&request); assert!( request - .working_directory_request - .as_ref() - .and_then(|working_directory| working_directory.materialization.as_ref()) + .workdir_attachment_requests + .first() + .and_then(|attachment| attachment.working_directory.materialization.as_ref()) .and_then(|materialization| materialization.ssh.as_ref()) .is_some() ); assert!( durable - .working_directory_request - .as_ref() - .and_then(|working_directory| working_directory.materialization.as_ref()) + .workdir_attachment_requests + .first() + .and_then(|attachment| attachment.working_directory.materialization.as_ref()) .and_then(|materialization| materialization.ssh.as_ref()) .is_none() ); @@ -5172,6 +5370,7 @@ mod tests { source_fingerprint: "sha256:source".to_string(), selector: None, }, + display_name: None, materializer: MaterializerKind::RuntimeGitClone, backend_workdir_id: Some("working-directory-1".to_string()), materialization: Some(RepositoryMaterializationContext { @@ -5342,6 +5541,7 @@ mod tests { #[derive(Default)] struct TestExecutionBackend { + spawn_result: Mutex>, dispatch_result: Mutex>, stop_result: Mutex>, stop_gate: Mutex>>, @@ -5447,6 +5647,9 @@ mod tests { } fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult { + if let Some(result) = self.spawn_result.lock().unwrap().take() { + return result; + } self.config_bundles .lock() .unwrap() @@ -5460,10 +5663,16 @@ mod tests { worker_state: protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, - working_directory: request - .working_directory - .as_ref() - .map(|binding| binding.status()), + workdir_attachments: request + .workdir_attachments + .iter() + .map( + |(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus { + alias: alias.clone(), + working_directory: binding.status(), + }, + ) + .collect(), } } @@ -5502,10 +5711,16 @@ mod tests { worker_state: protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, - working_directory: request - .working_directory - .as_ref() - .map(|binding| binding.status()), + workdir_attachments: request + .workdir_attachments + .iter() + .map( + |(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus { + alias: alias.clone(), + working_directory: binding.status(), + }, + ) + .collect(), } } @@ -6263,10 +6478,11 @@ mod tests { let runtime = runtime_with_backend(); let mut request = task_request("idempotent"); request.create_fingerprint = "sha256:input-1".to_string(); - request.working_directory = Some(WorkingDirectoryClaim { + request.workdir_attachments = vec![WorkingDirectoryAttachmentClaim { + alias: workdir::WorkdirAttachmentAlias::new("workdir").unwrap(), working_directory_id: "workdir-idempotent".to_string(), relative_cwd: None, - }); + }]; let first = runtime.create_worker(request.clone()).unwrap(); let workdir_count_after_first = runtime.list_working_directories().unwrap().len(); @@ -6631,6 +6847,49 @@ mod tests { assert!(runtime.list_workers().unwrap().is_empty()); } + #[test] + fn reconciliation_required_create_retains_durable_worker_evidence() { + let (runtime, backend) = runtime_and_backend(); + runtime.store_config_bundle(test_bundle()).unwrap(); + *backend.spawn_result.lock().unwrap() = + Some(WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::errored( + WorkerExecutionOperation::Spawn, + "Workdir cleanup could not be proven", + ), + handle: None, + worker_state: None, + workdir_attachments: vec![WorkingDirectoryAttachmentStatus { + alias: workdir::WorkdirAttachmentAlias::new("checkout").unwrap(), + working_directory: serde_json::from_value(serde_json::json!({ + "summary": { + "working_directory_id": "wd-uncertain", + "repository_id": "repo", + "materializer_kind": "runtime_git_clone", + "status": "active" + } + })) + .unwrap(), + }], + }); + let request = task_request("uncertain create cleanup"); + let worker_ref = WorkerRef::new(request.worker_id); + + let error = runtime.create_worker(request).unwrap_err(); + + assert!(matches!( + error, + RuntimeError::WorkerExecutionRejected { .. } + )); + let retained = runtime.worker_detail(&worker_ref).unwrap(); + assert_eq!(retained.status, WorkerStatus::Idle); + assert!(!retained.execution_metadata_available); + assert!(matches!( + runtime.delete_worker(&worker_ref), + Err(RuntimeError::InvalidRequest(_)) + )); + } + #[test] fn create_worker_without_execution_backend_is_rejected_and_not_persisted() { let runtime = Runtime::new_memory(); @@ -6804,10 +7063,16 @@ mod tests { worker_state: protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, - working_directory: request - .working_directory - .as_ref() - .map(|binding| binding.status()), + workdir_attachments: request + .workdir_attachments + .iter() + .map( + |(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus { + alias: alias.clone(), + working_directory: binding.status(), + }, + ) + .collect(), } } @@ -7667,7 +7932,7 @@ mod tests { ), handle: None, worker_state: None, - working_directory: None, + workdir_attachments: Vec::new(), }, ); diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 13807a326..edfc77409 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -7,7 +7,7 @@ //! socket paths, session paths, manifests, credentials, or handles leave this //! module. -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::future::Future; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; @@ -17,8 +17,9 @@ use std::time::Duration; use crate::auth::{BACKEND_RESOURCE_FETCH_PERMISSION, RuntimeIdentityMaterial}; use crate::catalog::{ CreateWorkerRequest, ProfileSourceArchiveSource, RepositoryRefObservation, - RepositoryRefObservationRequest, WorkingDirectoryRepositoryAccessRequest, - WorkingDirectoryRequest, WorkingDirectoryStatus, WorkingDirectoryStatusKind, + RepositoryRefObservationRequest, WorkingDirectoryAttachmentStatus, + WorkingDirectoryRepositoryAccessRequest, WorkingDirectoryRequest, WorkingDirectoryStatus, + WorkingDirectoryStatusKind, }; use crate::config_bundle::{ConfigBundle, workspace_config_etag}; use crate::execution::{ @@ -44,6 +45,45 @@ use protocol::{Event, Method, Segment, WorkerCommandEnvelope}; static NEXT_INTERNAL_COMMAND_ID: AtomicU64 = AtomicU64::new(1); +fn rollback_materialized_spawn_workdirs( + materializer: &dyn WorkingDirectoryMaterializer, + bindings: &BTreeMap, + workdir_ids: &[String], + failure_context: &str, +) -> Result<(), WorkerExecutionSpawnResult> { + let mut cleanup_failures = Vec::new(); + let mut uncertain_attachments = Vec::new(); + for workdir_id in workdir_ids.iter().rev() { + if let Err(error) = materializer.cleanup_working_directory(workdir_id) { + cleanup_failures.push(format!("{workdir_id}: {error}")); + if let Some((alias, binding)) = bindings + .iter() + .find(|(_, binding)| binding.working_directory.id == *workdir_id) + { + uncertain_attachments.push(WorkingDirectoryAttachmentStatus { + alias: alias.clone(), + working_directory: binding.status(), + }); + } + } + } + if cleanup_failures.is_empty() { + return Ok(()); + } + Err(WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::errored( + WorkerExecutionOperation::Spawn, + format!( + "{failure_context}; Workdir cleanup could not be proven: {}", + cleanup_failures.join("; ") + ), + ), + handle: None, + worker_state: None, + workdir_attachments: uncertain_attachments, + }) +} + fn next_internal_command( state: &RwLock, ) -> Result { @@ -68,7 +108,10 @@ use session_store::{FsStore, FsWorkerStore}; use tokio::runtime::Runtime; #[cfg(feature = "ws-server")] use tokio::sync::broadcast; -use workdir::{LocalWorkdirSession, Workdir, WorkdirSessionCapabilities, WorkdirSessionHandle}; +use workdir::{ + LocalWorkdirSession, Workdir, WorkdirAttachmentAlias, WorkdirSessionCapabilities, + WorkdirSessionHandle, WorkdirSessionRouter, +}; #[cfg(test)] use worker::WorkerController; @@ -730,6 +773,29 @@ fn runtime_local_workdir_session( )) } +fn runtime_local_workdir_router( + attachments: &BTreeMap, + scope: manifest::SharedScope, +) -> Result, String> { + let router = Arc::new(WorkdirSessionRouter::new()); + for (alias, binding) in attachments { + router + .attach( + alias.clone(), + runtime_local_workdir_session( + &binding.working_directory.id, + binding.root(), + binding.cwd(), + scope.clone(), + binding.command_environment(), + binding.session_resources(), + ), + ) + .map_err(|error| format!("bind Workdir attachment `{alias}`: {error}"))?; + } + Ok(router) +} + fn bind_workspace_memory_settings( manifest: &mut manifest::WorkerManifest, request: &CreateWorkerRequest, @@ -845,15 +911,14 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { ) -> Result { let worker_name = Self::runtime_worker_name(&request); let profile = Self::runtime_profile(&request); - let has_local_filesystem = request.working_directory.is_some(); - let worker_root = request - .working_directory - .as_ref() + let has_local_filesystem = request.workdir_attachments.len() == 1; + let only_binding = (request.workdir_attachments.len() == 1) + .then(|| request.workdir_attachments.values().next()) + .flatten(); + let worker_root = only_binding .map(|binding| binding.root().to_path_buf()) .unwrap_or_else(|| self.profile_base_dir.clone()); - let filesystem_authority = request - .working_directory - .as_ref() + let filesystem_authority = only_binding .map(|binding| { WorkerFilesystemAuthority::local( binding.root().to_path_buf(), @@ -965,18 +1030,10 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { })?; let worker = prepared.worker_mut(); validate_worker_memory_settings(worker.manifest(), &request.request)?; - if let Some(binding) = request.working_directory.as_ref() { - worker.bind_workdir_session(Some(runtime_local_workdir_session( - &binding.working_directory.id, - binding.root(), - binding.cwd(), - worker.scope().clone(), - binding.command_environment(), - binding.session_resources(), - ))); - } else { - worker.bind_workdir_session(None); - } + worker.bind_workdir_sessions(runtime_local_workdir_router( + &request.workdir_attachments, + worker.scope().clone(), + )?); if let (Some(runtime_id), Some(workspace_id)) = (observation_runtime_id, observation_workspace_id.clone()) && observation_enabled @@ -1095,9 +1152,10 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { request: WorkerExecutionRestoreRequest, ) -> Result { let worker_name = Self::runtime_worker_name_for_ref(&request.worker_ref); - let filesystem_authority = request - .working_directory - .as_ref() + let only_binding = (request.workdir_attachments.len() == 1) + .then(|| request.workdir_attachments.values().next()) + .flatten(); + let filesystem_authority = only_binding .map(|binding| { WorkerFilesystemAuthority::local( binding.root().to_path_buf(), @@ -1201,18 +1259,10 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { }; validate_worker_memory_settings(worker.manifest(), &request.request)?; let flow_transition_enabled = worker.manifest().feature.flow.enabled; - if let Some(binding) = request.working_directory.as_ref() { - worker.bind_workdir_session(Some(runtime_local_workdir_session( - &binding.working_directory.id, - binding.root(), - binding.cwd(), - worker.scope().clone(), - binding.command_environment(), - binding.session_resources(), - ))); - } else { - worker.bind_workdir_session(None); - } + worker.bind_workdir_sessions(runtime_local_workdir_router( + &request.workdir_attachments, + worker.scope().clone(), + )?); if let (Some(runtime_id), Some(workspace_id)) = (observation_runtime_id, observation_workspace_id.clone()) && observation_enabled @@ -1738,7 +1788,7 @@ where handle: WorkerHandle, shutdown: Arc>>, controller_task: tokio::task::JoinHandle<()>, - working_directory: Option, + workdir_attachments: BTreeMap, workspace_client: Option>, ) -> WorkerExecutionSpawnResult { #[cfg(feature = "ws-server")] @@ -1818,8 +1868,13 @@ where Err(_) => { let retained_state = worker_state.read().ok().map(|state| state.clone()); - let retained_working_directory = - working_directory.as_ref().map(|binding| binding.status()); + let retained_workdir_attachments = workdir_attachments + .iter() + .map(|(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus { + alias: alias.clone(), + working_directory: binding.status(), + }) + .collect(); let retained_handle = self .retain_uncertain_unconnected_controller( &worker_ref, @@ -1834,7 +1889,7 @@ where result, handle: retained_handle, worker_state: retained_state, - working_directory: retained_working_directory, + workdir_attachments: retained_workdir_attachments, } } } @@ -1878,7 +1933,7 @@ where result, handle: Some(existing_handle), worker_state: existing_worker_state, - working_directory: None, + workdir_attachments: Vec::new(), }, } } else { @@ -1904,7 +1959,15 @@ where WorkerExecutionSpawnResult::Connected { handle: WorkerExecutionHandle::new(worker_ref, self.backend_id()), worker_state: connected_worker_state, - working_directory: working_directory.map(|binding| binding.status()), + workdir_attachments: workdir_attachments + .into_iter() + .map( + |(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus { + alias, + working_directory: binding.status(), + }, + ) + .collect(), } } } @@ -2118,67 +2181,134 @@ where } let mut request = request; - let mut rollback_working_directory = None; - let working_directory = match ( - request.request.working_directory_request.as_ref(), - request.request.working_directory.as_ref(), - ) { - (Some(_), Some(_)) => { + let Some(materializer) = self.working_directory_materializer.as_ref().or_else(|| { + (request.request.workdir_attachment_requests.is_empty() + && request.request.workdir_attachments.is_empty()) + .then_some(()) + .and(None) + }) else { + if request.request.workdir_attachment_requests.is_empty() + && request.request.workdir_attachments.is_empty() + { + // No provider is needed for a Workdir-less Worker. + let factory = self.factory.clone(); + let bridge_context = request.context.clone(); + let worker_ref = request.worker_ref.clone(); + let spawn_result = self + .run_cancellable_on_adapter_runtime(self.spawn_restore_timeout, async move { + factory.spawn_controller(request).await + }); + let controller = match spawn_result { + Ok(controller) => controller, + Err(message) => { + return WorkerExecutionSpawnResult::Errored( + WorkerExecutionResult::errored( + WorkerExecutionOperation::Spawn, + message, + ), + ); + } + }; + return self.connect_handle( + WorkerExecutionOperation::Spawn, + worker_ref, + bridge_context, + controller.handle, + controller.shutdown, + controller.controller_task, + BTreeMap::new(), + Some(controller.workspace_client), + ); + } + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( + WorkerExecutionOperation::Spawn, + "Workdir attachments were requested, but no materializer is configured for this Runtime backend", + )); + }; + let mut rollback_workdirs: Vec = Vec::new(); + for attachment in request.request.workdir_attachment_requests.clone() { + if request.workdir_attachments.contains_key(&attachment.alias) { + let message = format!("duplicate Workdir attachment alias `{}`", attachment.alias); + if let Err(result) = rollback_materialized_spawn_workdirs( + materializer.as_ref(), + &request.workdir_attachments, + &rollback_workdirs, + &message, + ) { + return result; + } return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( WorkerExecutionOperation::Spawn, - "worker spawn cannot specify both working_directory_request and working_directory", + message, )); } - (Some(working_directory_request), None) => { - let Some(materializer) = self.working_directory_materializer.as_ref() else { + match materializer.materialize(&request.worker_ref, &attachment.working_directory) { + Ok(binding) => { + rollback_workdirs.push(binding.working_directory.id.clone()); + request + .workdir_attachments + .insert(attachment.alias, binding); + } + Err(error) => { + let message = error.to_string(); + if let Err(result) = rollback_materialized_spawn_workdirs( + materializer.as_ref(), + &request.workdir_attachments, + &rollback_workdirs, + &message, + ) { + return result; + } return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( WorkerExecutionOperation::Spawn, - "working directory materialization requested, but no materializer is configured for this runtime backend", + message, )); - }; - match materializer.materialize(&request.worker_ref, working_directory_request) { - Ok(binding) => { - request.working_directory = Some(binding.clone()); - rollback_working_directory = Some(binding.clone()); - Some(binding) - } - Err(error) => { - return WorkerExecutionSpawnResult::Rejected( - WorkerExecutionResult::rejected( - WorkerExecutionOperation::Spawn, - error.to_string(), - ), - ); - } } } - (None, Some(working_directory)) => { - let Some(materializer) = self.working_directory_materializer.as_ref() else { - return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( - WorkerExecutionOperation::Spawn, - "working directory working_directory requested, but no materializer is configured for this runtime backend", - )); - }; - match materializer.bind_working_directory( - &working_directory.working_directory_id, - working_directory.relative_cwd.as_deref(), + } + for attachment in request.request.workdir_attachments.clone() { + if request.workdir_attachments.contains_key(&attachment.alias) { + let message = format!("duplicate Workdir attachment alias `{}`", attachment.alias); + if let Err(result) = rollback_materialized_spawn_workdirs( + materializer.as_ref(), + &request.workdir_attachments, + &rollback_workdirs, + &message, ) { - Ok(binding) => { - request.working_directory = Some(binding.clone()); - Some(binding) - } - Err(error) => { - return WorkerExecutionSpawnResult::Rejected( - WorkerExecutionResult::rejected( - WorkerExecutionOperation::Spawn, - error.to_string(), - ), - ); + return result; + } + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( + WorkerExecutionOperation::Spawn, + message, + )); + } + match materializer.bind_working_directory( + &attachment.working_directory_id, + attachment.relative_cwd.as_deref(), + ) { + Ok(binding) => { + request + .workdir_attachments + .insert(attachment.alias, binding); + } + Err(error) => { + let message = error.to_string(); + if let Err(result) = rollback_materialized_spawn_workdirs( + materializer.as_ref(), + &request.workdir_attachments, + &rollback_workdirs, + &message, + ) { + return result; } + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( + WorkerExecutionOperation::Spawn, + message, + )); } } - (None, None) => None, - }; + } + let workdir_attachments = request.workdir_attachments.clone(); let factory = self.factory.clone(); let bridge_context = request.context.clone(); @@ -2191,11 +2321,13 @@ where let controller = match spawn_result { Ok(controller) => controller, Err(message) => { - if let (Some(materializer), Some(binding)) = ( - self.working_directory_materializer.as_ref(), - rollback_working_directory.as_ref(), + if let Err(result) = rollback_materialized_spawn_workdirs( + materializer.as_ref(), + &workdir_attachments, + &rollback_workdirs, + &message, ) { - let _ = materializer.cleanup_working_directory(&binding.working_directory.id); + return result; } return WorkerExecutionSpawnResult::Errored(WorkerExecutionResult::errored( WorkerExecutionOperation::Spawn, @@ -2211,7 +2343,7 @@ where controller.handle, controller.shutdown, controller.controller_task, - working_directory, + workdir_attachments, Some(controller.workspace_client), ) } @@ -2232,46 +2364,56 @@ where } drop(workers); - if let Some(previous) = request.previous_working_directory.as_ref() { + if !request.previous_workdir_attachments.is_empty() { let materializer = self .working_directory_materializer .as_ref() .ok_or_else(|| { WorkerExecutionResult::rejected( WorkerExecutionOperation::Restore, - "Persisted Worker Workdir binding cannot be restored by this Runtime", + "Persisted Worker Workdir attachments cannot be restored by this Runtime", ) })?; - let current = materializer - .working_directory_status(&previous.summary.working_directory_id) - .map_err(|message| { - WorkerExecutionResult::rejected( + for previous in &request.previous_workdir_attachments { + let current = materializer + .working_directory_status( + &previous.working_directory.summary.working_directory_id, + ) + .map_err(|message| { + WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + format!( + "Persisted Worker Workdir attachment `{}` is unavailable: {message}", + previous.alias + ), + ) + })?; + if current.summary.status != WorkingDirectoryStatusKind::Active { + return Err(WorkerExecutionResult::rejected( WorkerExecutionOperation::Restore, - format!("Persisted Worker Workdir is unavailable: {message}"), - ) - })?; - if current.summary.status != WorkingDirectoryStatusKind::Active { - return Err(WorkerExecutionResult::rejected( - WorkerExecutionOperation::Restore, - "Persisted Worker Workdir is not active", - )); + format!( + "Persisted Worker Workdir attachment `{}` is not active", + previous.alias + ), + )); + } } } - if request.previous_working_directory.is_none() - && request.request.working_directory_request.is_some() + if request.previous_workdir_attachments.is_empty() + && !request.request.workdir_attachment_requests.is_empty() { return Err(WorkerExecutionResult::rejected( WorkerExecutionOperation::Restore, - "Persisted Worker Workdir allocation is unavailable", + "Persisted Worker Workdir allocations are unavailable", )); } - if request.previous_working_directory.is_none() - && request.request.working_directory.is_some() + if request.previous_workdir_attachments.is_empty() + && !request.request.workdir_attachments.is_empty() && self.working_directory_materializer.is_none() { return Err(WorkerExecutionResult::rejected( WorkerExecutionOperation::Restore, - "Persisted Worker Workdir claim cannot be restored by this Runtime", + "Persisted Worker Workdir claims cannot be restored by this Runtime", )); } if let Some(workspace_api) = request.request.workspace_api.as_ref() { @@ -2321,58 +2463,40 @@ where } drop(workers); - let working_directory = match request.previous_working_directory.clone() { - Some(status) => { - let Some(materializer) = self.working_directory_materializer.as_ref() else { - return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( - WorkerExecutionOperation::Restore, - "persisted worker has a working directory binding, but no materializer is configured for this runtime backend", - )); - }; - let relative_cwd = request - .request - .working_directory - .as_ref() - .and_then(|working_directory| working_directory.relative_cwd.as_deref()); - match materializer - .bind_working_directory(&status.summary.working_directory_id, relative_cwd) - { - Ok(binding) => { - request.working_directory = Some(binding.clone()); - Some(binding) - } - Err(error) => { - return WorkerExecutionSpawnResult::Rejected( - WorkerExecutionResult::rejected( - WorkerExecutionOperation::Restore, - error.to_string(), - ), - ); - } - } - } - None if request.request.working_directory_request.is_some() => { + let mut workdir_attachments = BTreeMap::new(); + if !request.previous_workdir_attachments.is_empty() { + let Some(materializer) = self.working_directory_materializer.as_ref() else { return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( WorkerExecutionOperation::Restore, - "persisted worker requested a working directory, but no persisted working directory binding is available to restore", + "persisted Worker has Workdir attachments, but no materializer is configured for this Runtime backend", )); - } - None if request.request.working_directory.is_some() => { - let Some(materializer) = self.working_directory_materializer.as_ref() else { - return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( - WorkerExecutionOperation::Restore, - "persisted worker has a working directory claim, but no materializer is configured for this runtime backend", - )); - }; - let working_directory = - request.request.working_directory.as_ref().expect("checked"); + }; + for status in request.previous_workdir_attachments.clone() { + let relative_cwd = request + .request + .workdir_attachments + .iter() + .find(|claim| claim.alias == status.alias) + .and_then(|claim| claim.relative_cwd.as_deref()); match materializer.bind_working_directory( - &working_directory.working_directory_id, - working_directory.relative_cwd.as_deref(), + &status.working_directory.summary.working_directory_id, + relative_cwd, ) { Ok(binding) => { - request.working_directory = Some(binding.clone()); - Some(binding) + if workdir_attachments + .insert(status.alias.clone(), binding) + .is_some() + { + return WorkerExecutionSpawnResult::Rejected( + WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + format!( + "duplicate persisted Workdir attachment alias `{}`", + status.alias + ), + ), + ); + } } Err(error) => { return WorkerExecutionSpawnResult::Rejected( @@ -2384,8 +2508,43 @@ where } } } - None => None, - }; + } else if !request.request.workdir_attachments.is_empty() { + let Some(materializer) = self.working_directory_materializer.as_ref() else { + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + "persisted Worker has Workdir claims, but no materializer is configured for this Runtime backend", + )); + }; + for claim in request.request.workdir_attachments.clone() { + match materializer.bind_working_directory( + &claim.working_directory_id, + claim.relative_cwd.as_deref(), + ) { + Ok(binding) => { + if workdir_attachments + .insert(claim.alias.clone(), binding) + .is_some() + { + return WorkerExecutionSpawnResult::Rejected( + WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + format!("duplicate Workdir attachment alias `{}`", claim.alias), + ), + ); + } + } + Err(error) => { + return WorkerExecutionSpawnResult::Rejected( + WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + error.to_string(), + ), + ); + } + } + } + } + request.workdir_attachments = workdir_attachments.clone(); let factory = self.factory.clone(); let bridge_context = request.context.clone(); @@ -2405,7 +2564,7 @@ where ), handle: None, worker_state: None, - working_directory: None, + workdir_attachments: Vec::new(), }; } }; @@ -2417,7 +2576,7 @@ where controller.handle, controller.shutdown, controller.controller_task, - working_directory, + workdir_attachments, Some(controller.workspace_client), ) } @@ -2739,8 +2898,8 @@ mod tests { use crate::Runtime as EmbeddedRuntime; use crate::catalog::{ ConfigBundleRef, CreateWorkerRequest, MaterializerKind, ProfileSelector, - RepositorySelector, WorkingDirectoryClaim, WorkingDirectoryRepository, - WorkingDirectoryRequest, WorkspaceApiRef, + RepositorySelector, WorkingDirectoryAttachmentClaim, WorkingDirectoryAttachmentRequest, + WorkingDirectoryRepository, WorkingDirectoryRequest, WorkspaceApiRef, }; use crate::execution::WorkerExecutionContext; use crate::identity::WorkerId; @@ -3279,18 +3438,17 @@ mod tests { FsStore::new(&self.store_dir).map_err(|err| err.to_string())?, FsWorkerStore::new(&self.worker_metadata_dir).map_err(|err| err.to_string())?, ); - let filesystem_authority = request - .working_directory - .as_ref() + let only_workdir = (request.workdir_attachments.len() == 1) + .then(|| request.workdir_attachments.values().next()) + .flatten(); + let filesystem_authority = only_workdir .map(|binding| { let cwd = binding.cwd().to_path_buf(); self.observed_cwds.lock().unwrap().push(cwd.clone()); WorkerFilesystemAuthority::local(binding.root().to_path_buf(), cwd) }) .unwrap_or(WorkerFilesystemAuthority::None); - let scope_root = request - .working_directory - .as_ref() + let scope_root = only_workdir .map(|binding| binding.root().to_path_buf()) .unwrap_or_else(|| self.cwd.clone()); let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request( @@ -3348,7 +3506,7 @@ mod tests { request: request.request, workspace_scope: request.workspace_scope, context: request.context, - working_directory: request.working_directory, + workdir_attachments: request.workdir_attachments, config_bundle: request.config_bundle, }; self.spawn_controller(request).await @@ -3513,8 +3671,8 @@ mod tests { digest: bundle.metadata.digest, }), initial_input: None, - working_directory_request: None, - working_directory: None, + workdir_attachment_requests: Vec::new(), + workdir_attachments: Vec::new(), worker_observation_enabled: false, worker_observation_grants: Vec::new(), workspace_api: None, @@ -3595,6 +3753,7 @@ mod tests { source_fingerprint: "sha256:test".to_string(), selector: Some(RepositorySelector::from("HEAD")), }, + display_name: None, materializer: MaterializerKind::RuntimeGitClone, backend_workdir_id: None, materialization: None, @@ -3698,7 +3857,7 @@ mod tests { request: create_request("1"), workspace_scope: None, context: test_execution_context(worker_ref), - working_directory: None, + workdir_attachments: BTreeMap::new(), config_bundle: None, }; @@ -3830,8 +3989,8 @@ mod tests { "server-main", )), context: test_execution_context(worker_ref), - previous_working_directory: None, - working_directory: None, + previous_workdir_attachments: Vec::new(), + workdir_attachments: BTreeMap::new(), config_bundle: None, }) .await @@ -3920,8 +4079,8 @@ mod tests { request: create_request("embedded restore"), workspace_scope: None, context: test_execution_context(worker_ref), - previous_working_directory: None, - working_directory: None, + previous_workdir_attachments: Vec::new(), + workdir_attachments: BTreeMap::new(), config_bundle: None, }) .await @@ -3988,7 +4147,7 @@ mod tests { &[ "WorkerBootstrap::new(", ".prepare()", - "worker.bind_workdir_session(", + "worker.bind_workdir_sessions(", "worker.bind_worker_observation_provider(", "install_runtime_flow_transition_feature()", "prepared.start()", @@ -3998,7 +4157,7 @@ mod tests { restore, &[ "Worker::restore_from_worker_metadata_with_context(", - "worker.bind_workdir_session(", + "worker.bind_workdir_sessions(", "worker.bind_worker_observation_provider(", "install_runtime_flow_transition_feature()", "PreparedWorker::new(", @@ -4354,7 +4513,10 @@ mod tests { .unwrap(); runtime.store_config_bundle(test_bundle()).unwrap(); let mut request = create_request("chat"); - request.working_directory_request = Some(working_directory_request(repo.path())); + request.workdir_attachment_requests = vec![WorkingDirectoryAttachmentRequest { + alias: workdir::WorkdirAttachmentAlias::new("checkout").unwrap(), + working_directory: working_directory_request(repo.path()), + }]; let detail = runtime.create_worker(request).unwrap(); runtime @@ -4376,7 +4538,7 @@ mod tests { ); } - assert!(detail.working_directory.is_some()); + assert_eq!(detail.workdir_attachments.len(), 1); let cwds = observed_cwds.lock().unwrap(); assert_eq!(cwds.len(), 1); let cwd = &cwds[0]; @@ -4387,6 +4549,71 @@ mod tests { observed_workspace_clients.lock().unwrap().as_slice(), &[("unavailable".to_string(), None, false)] ); + + let workdir_id = detail.workdir_attachments[0] + .working_directory + .summary + .working_directory_id + .clone(); + let second = runtime.create_worker(create_request("second")).unwrap(); + let conflict = runtime + .replace_worker_workdir_attachments( + &second.worker_ref, + vec![WorkingDirectoryAttachmentClaim { + alias: workdir::WorkdirAttachmentAlias::new("conflict").unwrap(), + working_directory_id: workdir_id.clone(), + relative_cwd: None, + }], + ) + .unwrap_err(); + assert!(matches!( + conflict, + crate::error::RuntimeError::InvalidRequest(_) + )); + let inactive = runtime + .create_working_directory(working_directory_request(repo.path())) + .unwrap(); + let inactive_id = inactive.summary.working_directory_id; + runtime.cleanup_working_directory(&inactive_id).unwrap(); + let inactive_error = runtime + .replace_worker_workdir_attachments( + &second.worker_ref, + vec![WorkingDirectoryAttachmentClaim { + alias: workdir::WorkdirAttachmentAlias::new("inactive").unwrap(), + working_directory_id: inactive_id, + relative_cwd: None, + }], + ) + .unwrap_err(); + assert!(matches!( + inactive_error, + crate::error::RuntimeError::InvalidRequest(_) + | crate::error::RuntimeError::WorkingDirectory(_) + )); + + let detached = runtime + .replace_worker_workdir_attachments(&detail.worker_ref, Vec::new()) + .unwrap(); + assert!(detached.workdir_attachments.is_empty()); + runtime.stop_worker(&detail.worker_ref, None).unwrap(); + let restored_without_workdir = runtime.restore_worker(&detail.worker_ref).unwrap(); + assert!(restored_without_workdir.workdir_attachments.is_empty()); + + let alias = workdir::WorkdirAttachmentAlias::new("restored").unwrap(); + runtime + .replace_worker_workdir_attachments( + &detail.worker_ref, + vec![WorkingDirectoryAttachmentClaim { + alias: alias.clone(), + working_directory_id: workdir_id, + relative_cwd: None, + }], + ) + .unwrap(); + runtime.stop_worker(&detail.worker_ref, None).unwrap(); + let stopped = runtime.worker_detail(&detail.worker_ref).unwrap(); + assert_eq!(stopped.workdir_attachments.len(), 1); + assert_eq!(stopped.workdir_attachments[0].alias, alias); } #[test] @@ -4620,12 +4847,13 @@ mod tests { .unwrap(); runtime.store_config_bundle(test_bundle()).unwrap(); let mut request = create_request("chat"); - request.working_directory_request = Some(working_directory_request(repo.path())); + request.workdir_attachment_requests = vec![WorkingDirectoryAttachmentRequest { + alias: workdir::WorkdirAttachmentAlias::new("checkout").unwrap(), + working_directory: working_directory_request(repo.path()), + }]; let detail = runtime.create_worker(request).unwrap(); - let workdir_id = detail + let workdir_id = detail.workdir_attachments[0] .working_directory - .as_ref() - .unwrap() .summary .working_directory_id .clone(); @@ -4642,7 +4870,6 @@ mod tests { crate::catalog::WorkingDirectoryStatusKind::Active ); assert_eq!(status.summary.cleanliness.as_deref(), Some("clean")); - assert_eq!(status.summary.primary_worker_id, None); } #[test] @@ -4663,10 +4890,11 @@ mod tests { let clone_root = materialized_clone_root(runtime_base.path(), &workdir_id); assert!(clone_root.join("README.md").exists()); let mut request = create_request("chat"); - request.working_directory = Some(WorkingDirectoryClaim { + request.workdir_attachments = vec![WorkingDirectoryAttachmentClaim { + alias: workdir::WorkdirAttachmentAlias::new("workdir").unwrap(), working_directory_id: workdir_id.clone(), relative_cwd: None, - }); + }]; let error = runtime.create_worker(request).unwrap_err(); @@ -4691,7 +4919,17 @@ mod tests { .unwrap(); runtime.store_config_bundle(test_bundle()).unwrap(); let mut request = create_request("chat"); - request.working_directory_request = Some(working_directory_request(repo.path())); + let second_repo = create_clean_repo(); + request.workdir_attachment_requests = vec![ + WorkingDirectoryAttachmentRequest { + alias: workdir::WorkdirAttachmentAlias::new("checkout").unwrap(), + working_directory: working_directory_request(repo.path()), + }, + WorkingDirectoryAttachmentRequest { + alias: workdir::WorkdirAttachmentAlias::new("docs").unwrap(), + working_directory: working_directory_request(second_repo.path()), + }, + ]; let error = runtime.create_worker(request).unwrap_err(); @@ -4708,4 +4946,41 @@ mod tests { assert_eq!(remaining_workdirs, 0); assert!(!working_directories_root.join(".repository-cache").exists()); } + + #[test] + fn second_materialization_failure_rolls_back_first_attachment() { + let runtime_base = tempfile::tempdir().unwrap(); + let repo = create_clean_repo(); + let missing_repo = runtime_base.path().join("missing-repository"); + let backend = WorkerRuntimeExecutionBackend::new(FailingFactory) + .unwrap() + .with_working_directory_materializer(RuntimeGitMaterializer::new(runtime_base.path())); + let runtime = + EmbeddedRuntime::with_execution_backend(RuntimeOptions::default(), Arc::new(backend)) + .unwrap(); + runtime.store_config_bundle(test_bundle()).unwrap(); + let mut request = create_request("chat"); + request.workdir_attachment_requests = vec![ + WorkingDirectoryAttachmentRequest { + alias: workdir::WorkdirAttachmentAlias::new("checkout").unwrap(), + working_directory: working_directory_request(repo.path()), + }, + WorkingDirectoryAttachmentRequest { + alias: workdir::WorkdirAttachmentAlias::new("docs").unwrap(), + working_directory: working_directory_request(&missing_repo), + }, + ]; + + let error = runtime.create_worker(request).unwrap_err(); + assert!(format!("{error:?}").contains("repository")); + let remaining_workdirs = fs::read_dir(runtime_base.path()) + .map(|entries| { + entries + .flatten() + .filter(|entry| !entry.file_name().to_string_lossy().starts_with('.')) + .count() + }) + .unwrap_or(0); + assert_eq!(remaining_workdirs, 0); + } } diff --git a/crates/worker-runtime/src/working_directory.rs b/crates/worker-runtime/src/working_directory.rs index 79cb50190..ff44d8b74 100644 --- a/crates/worker-runtime/src/working_directory.rs +++ b/crates/worker-runtime/src/working_directory.rs @@ -57,6 +57,8 @@ pub struct WorkingDirectoryEvidence { #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct WorkingDirectory { pub id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub display_name: Option, pub repository_id: String, pub materializer_kind: MaterializerKind, pub evidence: WorkingDirectoryEvidence, @@ -68,6 +70,7 @@ impl WorkingDirectory { pub fn status_summary(&self) -> WorkingDirectorySummary { WorkingDirectorySummary { working_directory_id: self.id.clone(), + display_name: self.display_name.clone(), repository_id: self.repository_id.clone(), creation_selector: self.evidence.requested_selector.clone(), creation_ref: Some(self.evidence.resolved_commit.clone()), @@ -80,7 +83,6 @@ impl WorkingDirectory { cleanup_target: Some(self.cleanup_target.clone()), status: self.status.clone(), cleanliness: None, - primary_worker_id: None, occupied_by: None, } } @@ -717,6 +719,7 @@ impl RuntimeGitMaterializer { WorkingDirectoryStatus { summary: WorkingDirectorySummary { working_directory_id: working_directory_id.to_string(), + display_name: None, repository_id: "unknown".to_string(), creation_selector: None, creation_ref: None, @@ -733,7 +736,6 @@ impl RuntimeGitMaterializer { }), status: WorkingDirectoryStatusKind::Corrupted, cleanliness: Some("unknown".to_string()), - primary_worker_id: None, occupied_by: None, }, } @@ -1072,6 +1074,7 @@ impl RuntimeGitMaterializer { let context = request.materialization.as_ref(); let working_directory = WorkingDirectory { id: working_directory_id.clone(), + display_name: request.display_name.clone(), repository_id: request.repository.id.clone(), materializer_kind: MaterializerKind::RuntimeGitClone, evidence: WorkingDirectoryEvidence { @@ -1211,6 +1214,7 @@ impl WorkingDirectoryMaterializer for RuntimeGitMaterializer { let working_request = WorkingDirectoryRequest { repository: request.repository.clone(), + display_name: None, materializer: MaterializerKind::RuntimeGitClone, backend_workdir_id: None, materialization: request.materialization.clone(), @@ -3316,6 +3320,7 @@ mod tests { source_fingerprint: "sha256:test".to_string(), selector: Some(RepositorySelector::from("HEAD")), }, + display_name: None, materializer: MaterializerKind::RuntimeGitClone, backend_workdir_id: None, materialization: None, diff --git a/crates/worker/src/bootstrap.rs b/crates/worker/src/bootstrap.rs index 234e53e69..587ed11c9 100644 --- a/crates/worker/src/bootstrap.rs +++ b/crates/worker/src/bootstrap.rs @@ -145,7 +145,7 @@ where .map_err(WorkerBootstrapError::Worker)?; if let Some(workdir_session) = self.workdir_session { - worker.bind_workdir_session(Some(workdir_session)); + worker.bind_single_workdir_session(Some(workdir_session)); } Ok(PreparedWorker::new(worker, self.layout, self.transport)) } @@ -168,7 +168,7 @@ where .map_err(WorkerBootstrapError::Worker)?; if let Some(workdir_session) = self.workdir_session { - worker.bind_workdir_session(Some(workdir_session)); + worker.bind_single_workdir_session(Some(workdir_session)); } Ok(PreparedWorker::new(worker, self.layout, self.transport)) } @@ -222,7 +222,7 @@ where C: LlmClient + Clone + 'static, St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, { - let cleanup_session = worker.workdir_session().cloned(); + let cleanup_sessions = worker.workdir_sessions(); let controller = match layout { WorkerBootstrapLayout::Direct { runtime_base, @@ -257,10 +257,7 @@ where controller_task, }), Err(source) => { - let cleanup_failed = match cleanup_session { - Some(session) => session.close().await.is_err(), - None => false, - }; + let cleanup_failed = cleanup_sessions.close_all().await.is_err(); Err(WorkerBootstrapError::Controller { source, cleanup_failed, diff --git a/crates/worker/src/compact/metrics_tracker.rs b/crates/worker/src/compact/metrics_tracker.rs index c0e81033c..dcea2131a 100644 --- a/crates/worker/src/compact/metrics_tracker.rs +++ b/crates/worker/src/compact/metrics_tracker.rs @@ -2,9 +2,8 @@ //! Engine callbacks (which run synchronously and cannot themselves //! perform `async` store writes). //! -//! Worker drains this buffer in `persist_turn` and writes each metric via -//! `session_metrics::record_metric`, alongside the regular `LlmUsage` -//! entries. +//! Worker drains this buffer at the next successful request accounting +//! boundary, with terminal `persist_turn` as a fallback. use std::sync::Mutex; @@ -21,12 +20,12 @@ impl MetricsTracker { } } - /// Queue a metric for the next `persist_turn` flush. + /// Queue a metric for the next request or terminal accounting flush. pub(crate) fn push(&self, metric: Metric) { self.pending.lock().unwrap().push(metric); } - /// Drain all queued metrics. Called by Worker after a run completes. + /// Drain all queued metrics at an accounting boundary. pub(crate) fn drain(&self) -> Vec { std::mem::take(&mut *self.pending.lock().unwrap()) } diff --git a/crates/worker/src/compact/usage_tracker.rs b/crates/worker/src/compact/usage_tracker.rs index 3310ca338..b0d068f5c 100644 --- a/crates/worker/src/compact/usage_tracker.rs +++ b/crates/worker/src/compact/usage_tracker.rs @@ -7,17 +7,24 @@ //! - **`on_usage` callback** (sync closure): receives the aggregated final //! `UsageEvent` for that request after the stream completes. //! -//! Pairing the two yields one `UsageRecord` per LLM call. Worker drains them -//! in `persist_turn` and writes them as `LogEntry::LlmUsage` entries. +//! Pairing the two yields one `UsageRecord` per LLM call. Worker commits +//! successful requests from the assistant-turn boundary and uses terminal +//! `persist_turn` only for records that have not already crossed that boundary. //! //! Multiple LLM calls per Worker run (tool loop) are supported: each call //! produces its own `(history_len, UsageEvent)` pair, and the records are //! buffered in chronological order. +use std::collections::VecDeque; use std::sync::Mutex; use agen::UsageRecord; use agen::timeline::event::UsageEvent; +use session_metrics::Metric; +use session_store::{LogEntry, StoreError, segment_log}; + +use crate::compact::state::CompactState; +use crate::compact::telemetry::correlated_post_request_metric; /// The metric emitted after the next measured provider request. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -41,11 +48,15 @@ pub(crate) struct PostRequestLink { pub(crate) metric: PostRequestMetric, } -/// One drained measurement and its causal metric links. +/// One measured request and its causal metric links. #[derive(Debug, Clone)] pub(crate) struct RecordedUsage { pub(crate) record: UsageRecord, pub(crate) post_requests: Vec, + /// True only when the provider stream completed normally and reported an + /// input occupancy. Partial/error/cancelled streams remain billable usage + /// records, but cannot satisfy a post-compaction rearm. + pub(crate) completed_with_occupancy: bool, } #[derive(Debug, Clone, Default, PartialEq, Eq)] @@ -56,29 +67,52 @@ pub(crate) struct UsageSnapshot { pub(crate) output_tokens: u64, } +#[derive(Debug)] +struct PendingRequest { + id: u64, + history_len: usize, + usage_recorded: bool, +} + +#[derive(Debug)] +struct PendingRecordedUsage { + request_id: u64, + has_input_occupancy: bool, + usage: RecordedUsage, +} + +#[derive(Debug, Default)] +struct UsageTrackerState { + next_request_id: u64, + pending_request: Option, + pending_correlations: Vec, + pending_records: VecDeque, +} + /// Shared between the pre-request hook, the `on_usage` callback, and Worker. pub(crate) struct UsageTracker { - /// `history.len()` captured at the most recent `pre_llm_request`. - /// Cleared when paired with an incoming `on_usage` event. - pending_history_len: Mutex>, - /// Optional causal link consumed by the next measured request. - pending_correlations: Mutex>, - /// Records accumulated during the current run; drained by Worker. - pending_records: Mutex>, + state: Mutex, } impl UsageTracker { pub(crate) fn new() -> Self { Self { - pending_history_len: Mutex::new(None), - pending_correlations: Mutex::new(Vec::new()), - pending_records: Mutex::new(Vec::new()), + state: Mutex::new(UsageTrackerState::default()), } } /// Called from a `pre_llm_request` hook with the current history length. + /// Replacing an older pending request classifies that request as incomplete; + /// any partial usage already recorded for it remains available for billing. pub(crate) fn note_request(&self, history_len: usize) { - *self.pending_history_len.lock().unwrap() = Some(history_len); + let mut state = self.state.lock().unwrap(); + let id = state.next_request_id; + state.next_request_id = state.next_request_id.wrapping_add(1); + state.pending_request = Some(PendingRequest { + id, + history_len, + usage_recorded: false, + }); } /// Pair a prune event with the next provider request. @@ -92,25 +126,31 @@ impl UsageTracker { } fn note_post_request(&self, id: String, metric: PostRequestMetric) { - let mut pending = self.pending_correlations.lock().unwrap(); - pending.retain(|link| link.metric != metric); - pending.push(PostRequestLink { + let mut state = self.state.lock().unwrap(); + state + .pending_correlations + .retain(|link| link.metric != metric); + state.pending_correlations.push(PostRequestLink { correlation_id: id, metric, }); } /// Called from the `on_usage` callback with the aggregated final - /// UsageEvent. If a `history_len` was previously stashed via - /// `note_request`, builds a `RecordedUsage` and pushes it onto the - /// buffer. If not (e.g. test code that fires Usage outside a request), - /// drops the event. + /// UsageEvent. The record remains associated with the current request until + /// [`Self::request_completed`] proves that the provider stream completed + /// normally. Error/cancel paths still retain the measurement for billing. pub(crate) fn record_usage(&self, event: &UsageEvent) { - let history_len = match self.pending_history_len.lock().unwrap().take() { - Some(n) => n, - None => return, + let mut state = self.state.lock().unwrap(); + let Some(pending) = state.pending_request.as_mut() else { + return; }; - let post_requests = std::mem::take(&mut *self.pending_correlations.lock().unwrap()); + if pending.usage_recorded { + return; + } + pending.usage_recorded = true; + let request_id = pending.id; + let history_len = pending.history_len; // UsageEvent.input_tokens は scheme 層で「占有量（プロンプト全長）」に // 正規化済みである前提（Anthropic は cache_read + cache_creation を // 加算して emit する）。 @@ -118,35 +158,166 @@ impl UsageTracker { let cache_read = event.cache_read_input_tokens.unwrap_or(0); let cache_write = event.cache_creation_input_tokens.unwrap_or(0); let output = event.output_tokens.unwrap_or(0); - self.pending_records.lock().unwrap().push(RecordedUsage { - record: UsageRecord { - history_len, - input_total_tokens: input_total, - cache_read_tokens: cache_read, - cache_write_tokens: cache_write, - output_tokens: output, + state.pending_records.push_back(PendingRecordedUsage { + request_id, + has_input_occupancy: event.input_tokens.is_some(), + usage: RecordedUsage { + record: UsageRecord { + history_len, + input_total_tokens: input_total, + cache_read_tokens: cache_read, + cache_write_tokens: cache_write, + output_tokens: output, + }, + post_requests: Vec::new(), + completed_with_occupancy: false, }, - post_requests, }); } - /// Return a clone of the accumulated `UsageRecord`s without clearing them. + /// Mark the current request as normally completed. A post-request causal + /// link is consumed only when that exact request also reported a valid + /// input occupancy; callbacks from partial/error/cancelled streams cannot + /// rearm the compaction guard. + pub(crate) fn request_completed(&self) { + let mut state = self.state.lock().unwrap(); + let Some(pending) = state.pending_request.take() else { + return; + }; + let Some(recorded_index) = state + .pending_records + .iter() + .rposition(|recorded| recorded.request_id == pending.id) + else { + return; + }; + let has_occupancy = state.pending_records[recorded_index].has_input_occupancy; + if !has_occupancy { + return; + } + let post_requests = std::mem::take(&mut state.pending_correlations); + let recorded = &mut state.pending_records[recorded_index].usage; + recorded.completed_with_occupancy = true; + recorded.post_requests = post_requests; + } + + /// Return a clone of accumulated measurements without clearing them. /// Used by request-time circuit breakers that need the same occupancy /// projection as Worker persistence while the run is still active. pub(crate) fn records(&self) -> Vec { - self.pending_records + self.state .lock() .unwrap() + .pending_records .iter() - .map(|r| r.record.clone()) + .map(|r| r.usage.record.clone()) .collect() } - /// Drain accumulated records. Called by Worker after a run completes, - /// before persisting the turn. - pub(crate) fn drain(&self) -> Vec { - std::mem::take(&mut *self.pending_records.lock().unwrap()) + /// Number of measurements that have not crossed the durable commit + /// boundary. Session rewrites must not activate a replacement Segment while + /// this is non-zero, because their history coordinates belong to the current + /// Segment. + pub(crate) fn pending_record_count(&self) -> usize { + self.state.lock().unwrap().pending_records.len() } + + /// Remove the oldest measurement for one durable commit attempt. + pub(crate) fn take_next(&self) -> Option { + self.state + .lock() + .unwrap() + .pending_records + .pop_front() + .map(|recorded| recorded.usage) + } + + /// Restore a measurement after its durable commit failed. + pub(crate) fn restore_front(&self, usage: RecordedUsage) { + self.state + .lock() + .unwrap() + .pending_records + .push_front(PendingRecordedUsage { + // Persist ordering is the only remaining identity requirement. + request_id: u64::MAX, + has_input_occupancy: usage.completed_with_occupancy, + usage, + }); + } + + /// Drain accumulated records at a terminal Worker boundary. + pub(crate) fn drain(&self) -> Vec { + let mut state = self.state.lock().unwrap(); + state.pending_request = None; + state + .pending_records + .drain(..) + .map(|recorded| recorded.usage) + .collect() + } +} + +/// Commit all currently measured usage records in request order. +/// +/// A record is removed from the tracker only after its `LlmUsage` append +/// succeeds. This is the single accounting boundary shared by the in-run +/// interceptor and terminal `persist_turn`, preventing both loss on append +/// failure and duplicate terminal writes. Correlated metrics are handed to the +/// same commit boundary so they stay on the usage record's Segment; the return +/// value is only useful as test/diagnostic evidence and must not be re-emitted. +pub(crate) fn persist_pending_usage( + tracker: &UsageTracker, + usage_history: &Mutex>, + compact_state: Option<&CompactState>, + mut commit: impl FnMut(LogEntry, &mut dyn FnMut(), &[Metric]) -> Result<(), StoreError>, +) -> Result, StoreError> { + let mut metrics = Vec::new(); + while let Some(recorded) = tracker.take_next() { + let RecordedUsage { + record, + post_requests, + completed_with_occupancy, + } = recorded; + let post_metrics = post_requests + .iter() + .map(|link| correlated_post_request_metric(link.metric, &link.correlation_id, &record)) + .collect::>(); + let usage_entry = LogEntry::LlmUsage { + ts: segment_log::now_millis(), + history_len: record.history_len, + input_total_tokens: record.input_total_tokens, + cache_read_tokens: record.cache_read_tokens, + cache_write_tokens: record.cache_write_tokens, + output_tokens: record.output_tokens, + }; + let mut publish_committed_usage = || { + usage_history + .lock() + .expect("usage_history poisoned") + .push(record.clone()); + + if completed_with_occupancy + && post_requests + .iter() + .any(|link| link.metric == PostRequestMetric::Compaction) + { + if let Some(state) = compact_state { + state.post_compact_request_committed(); + } + } + }; + if let Err(error) = commit(usage_entry, &mut publish_committed_usage, &post_metrics) { + tracker.restore_front(RecordedUsage { + record, + post_requests, + completed_with_occupancy, + }); + return Err(error); + } + metrics.extend(post_metrics); + } + Ok(metrics) } #[cfg(test)] @@ -230,6 +401,7 @@ mod tests { tracker.note_correlation_id("prune-id".into()); tracker.note_request(5); tracker.record_usage(&make_event(100, 10, 2, 20)); + tracker.request_completed(); let records = tracker.drain(); assert_eq!(records[0].post_requests.len(), 2); @@ -248,9 +420,11 @@ mod tests { tracker.note_correlation_id("abc".into()); tracker.note_request(5); tracker.record_usage(&make_event(100, 0, 0, 20)); + tracker.request_completed(); // Next request without a fresh stash → no correlation_id. tracker.note_request(10); tracker.record_usage(&make_event(200, 50, 0, 30)); + tracker.request_completed(); let records = tracker.drain(); assert_eq!(records.len(), 2); @@ -259,4 +433,130 @@ mod tests { assert_eq!(records[0].post_requests[0].metric, PostRequestMetric::Prune); assert!(records[1].post_requests.is_empty()); } + + #[test] + fn partial_usage_does_not_consume_post_compact_link() { + let tracker = UsageTracker::new(); + tracker.note_compaction_correlation_id("compact-id".into()); + tracker.note_request(5); + tracker.record_usage(&make_event(90, 0, 0, 2)); + + // Starting another request proves the prior stream did not complete. + tracker.note_request(6); + tracker.record_usage(&make_event(100, 0, 0, 3)); + tracker.request_completed(); + + let records = tracker.drain(); + assert_eq!(records.len(), 2); + assert!(!records[0].completed_with_occupancy); + assert!(records[0].post_requests.is_empty()); + assert!(records[1].completed_with_occupancy); + assert_eq!(records[1].post_requests.len(), 1); + assert_eq!(records[1].post_requests[0].correlation_id, "compact-id"); + } + + #[test] + fn completed_request_without_input_occupancy_does_not_consume_link() { + let tracker = UsageTracker::new(); + tracker.note_compaction_correlation_id("compact-id".into()); + tracker.note_request(5); + let mut missing_input = make_event(0, 0, 0, 2); + missing_input.input_tokens = None; + tracker.record_usage(&missing_input); + tracker.request_completed(); + + tracker.note_request(6); + tracker.record_usage(&make_event(100, 0, 0, 3)); + tracker.request_completed(); + + let records = tracker.drain(); + assert!(!records[0].completed_with_occupancy); + assert!(records[0].post_requests.is_empty()); + assert!(records[1].completed_with_occupancy); + assert_eq!(records[1].post_requests[0].correlation_id, "compact-id"); + } + + #[test] + fn completed_request_without_usage_does_not_complete_previous_partial_record() { + let tracker = UsageTracker::new(); + tracker.note_compaction_correlation_id("compact-id".into()); + tracker.note_request(5); + tracker.record_usage(&make_event(90, 0, 0, 2)); + + tracker.note_request(6); + tracker.request_completed(); + + let records = tracker.drain(); + assert_eq!(records.len(), 1); + assert!(!records[0].completed_with_occupancy); + assert!(records[0].post_requests.is_empty()); + } + + #[test] + fn failed_commit_can_restore_record_without_reordering() { + let tracker = UsageTracker::new(); + tracker.note_request(1); + tracker.record_usage(&make_event(10, 0, 0, 1)); + tracker.request_completed(); + tracker.note_request(2); + tracker.record_usage(&make_event(20, 0, 0, 2)); + tracker.request_completed(); + + let first = tracker.take_next().expect("first record"); + tracker.restore_front(first); + let records = tracker.drain(); + assert_eq!(records.len(), 2); + assert_eq!(records[0].record.input_total_tokens, 10); + assert_eq!(records[1].record.input_total_tokens, 20); + } + + #[test] + fn compaction_rearms_only_after_usage_append_succeeds() { + use crate::compact::state::{AutomaticCompactDecision, CompactionOutcome}; + + let tracker = UsageTracker::new(); + let state = CompactState::new(None, Some(10), 0); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Start(_) + )); + assert!(state.complete_automatic(CompactionOutcome::Succeeded)); + + tracker.note_compaction_correlation_id("compact-id".into()); + tracker.note_request(1); + tracker.record_usage(&make_event(5, 0, 0, 1)); + tracker.request_completed(); + let usage_history = Mutex::new(Vec::new()); + + let error = persist_pending_usage(&tracker, &usage_history, Some(&state), |_, _, _| { + Err(StoreError::Io(std::io::Error::other("synthetic failure"))) + }) + .expect_err("usage append must fail"); + assert!(error.to_string().contains("synthetic failure")); + assert!(usage_history.lock().unwrap().is_empty()); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Block(_) + )); + + let mut entries = Vec::new(); + let metrics = persist_pending_usage( + &tracker, + &usage_history, + Some(&state), + |entry, after_commit, _post_metrics| { + entries.push(entry); + after_commit(); + Ok(()) + }, + ) + .expect("retry should commit"); + assert_eq!(entries.len(), 1); + assert_eq!(usage_history.lock().unwrap().len(), 1); + assert_eq!(metrics.len(), 1); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Start(_) + )); + } } diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 2ea198113..babae8e21 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -617,7 +617,7 @@ impl WorkerController { C: LlmClient + Clone + 'static, St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, { - let session = worker.workdir_session().cloned(); + let sessions = worker.workdir_sessions(); let result = Self::spawn_initialized( worker, runtime_base, @@ -628,10 +628,9 @@ impl WorkerController { ) .await; if result.is_err() - && let Some(session) = session - && let Err(error) = session.close().await + && let Err(error) = sessions.close_all().await { - tracing::warn!(%error, "Workdir session close after controller startup failure failed"); + tracing::warn!(%error, "Workdir sessions close after controller startup failure failed"); } result } @@ -1144,7 +1143,7 @@ pub(crate) async fn register_worker_tools( runtime_base: PathBuf, spawned_registry: Arc, parent_method_tx: Option>, - inherited_workdir_tool_broker: Option, + inherited_workdir_tool_broker: Option, ) -> std::io::Result> where C: LlmClient + Clone + 'static, @@ -1153,26 +1152,13 @@ where // Worker-immutable snapshots taken before the mutable worker borrow // below so the worker borrow doesn't conflict with reads on `worker`. let feature_config = worker.manifest().feature.clone(); - let mut workdir_tool_broker = inherited_workdir_tool_broker; - if feature_config.manage_workdir.enabled && worker.workdir_session().is_none() { - let workspace_client = worker.workspace_client_handle(); - let broker = workdir::WorkdirToolBroker::new( - crate::feature::builtin::manage_workdir::WorkspaceAttachedWorkdirSession::handle( - workspace_client, - ), - ); - worker.bind_workdir_session(Some(broker.tool_session())); - workdir_tool_broker = Some(broker); - } else if workdir_tool_broker.is_none() - && let Some(existing) = worker.workdir_session().cloned() - { - let broker = workdir::WorkdirToolBroker::new(existing); - worker.bind_workdir_session(Some(broker.tool_session())); - workdir_tool_broker = Some(broker); - } - let worker_workdir = workdir_tool_broker - .as_ref() - .map(workdir::WorkdirToolBroker::tool_session); + let workdir_sessions = worker.workdir_sessions(); + let restrict_workdir_tools_to_current_capabilities = inherited_workdir_tool_broker.is_some(); + let workdir_tool_broker = Some( + inherited_workdir_tool_broker + .unwrap_or_else(|| workdir::WorkdirToolBrokerRouter::new(workdir_sessions.clone())), + ); + let workdir_for_view = workdir_sessions.only_session(); let local_filesystem = worker.local_working_directory().cloned(); let local_workspace_root = local_filesystem.as_ref().map(|local| local.root.clone()); let task_feature = worker.task_feature(); @@ -1207,28 +1193,29 @@ where paste_store, paste_session_id, )); - // Resolve the existing Worker–Workdir binding into the domain provider. - // Tools only consume the provider handle; they do not own its root, cwd, - // scope, or lifecycle. No-workdir Workers expose no local tools. - let (workdir_for_view, tracker) = if let Some(workdir) = worker_workdir { - let tracker = tools::Tracker::new(); + // Keep the alias-routed schemas installed for the Worker lifetime. The + // attachment set may be empty or mutate after registration, so target + // resolution and provider capability checks happen per invocation. + let tracker = tools::Tracker::new(); + let workdir_tools = if restrict_workdir_tools_to_current_capabilities { + tools::routed_builtin_tools_for_current_capabilities( + workdir_sessions.clone(), + tracker.clone(), + bash_output_dir.clone(), + ) + } else { + tools::routed_builtin_tools( + workdir_sessions.clone(), + tracker.clone(), + bash_output_dir.clone(), + ) + }; + worker.engine_mut().register_tools(workdir_tools); + if feature_config.image.enabled && model_supports_image_attachments(&spawner_manifest.model) { worker .engine_mut() - .register_tools(tools::core_builtin_tools( - workdir.clone(), - tracker.clone(), - bash_output_dir.clone(), - )); - if feature_config.image.enabled && model_supports_image_attachments(&spawner_manifest.model) - { - worker - .engine_mut() - .register_tool(tools::view_image_tool(workdir.clone())); - } - (Some(workdir), Some(tracker)) - } else { - (None, None) - }; + .register_tool(tools::routed_view_image_tool(workdir_sessions)); + } if feature_config.web.enabled { worker .engine_mut() @@ -1344,6 +1331,7 @@ where feature_registry.add_module( crate::feature::builtin::manage_workdir::ManageWorkdirFeature::with_child_lifecycle( workspace_client, + worker.workdir_sessions(), Arc::new(move || { let child_registry = shutdown_registry.clone(); Box::pin(async move { child_registry.shutdown_internal().await }) @@ -1485,9 +1473,7 @@ where if let Some((resident_summary, system_prompt_override)) = memory_prompt_contribution { worker.install_system_prompt_contribution(resident_summary, system_prompt_override); } - if let Some(tracker) = tracker { - worker.attach_tracker(tracker); - } + worker.attach_tracker(tracker); Ok(workdir_for_view) } @@ -2610,14 +2596,13 @@ async fn controller_loop( } } - if let Some(session) = worker.workdir_session() { - loop { - match session.close().await { - Ok(()) => break, - Err(error) => { - tracing::warn!(%error, "Workdir session close failed; retrying"); - tokio::time::sleep(std::time::Duration::from_millis(100)).await; - } + let sessions = worker.workdir_sessions(); + while !sessions.is_empty() { + match sessions.close_all().await { + Ok(()) => break, + Err(error) => { + tracing::warn!(%error, "Workdir session set close failed; retrying"); + tokio::time::sleep(std::time::Duration::from_millis(100)).await; } } } diff --git a/crates/worker/src/feature/builtin/manage_workdir.rs b/crates/worker/src/feature/builtin/manage_workdir.rs index 371718d9d..2df835b34 100644 --- a/crates/worker/src/feature/builtin/manage_workdir.rs +++ b/crates/worker/src/feature/builtin/manage_workdir.rs @@ -52,9 +52,9 @@ const DETACH_TOOL: &str = "WorkdirDetach"; const DELETE_TOOL: &str = "WorkdirDelete"; const LIST_DESCRIPTION: &str = "List persistent Workdirs in the current Workspace through Backend Workspace API authority. The result contains safe summaries and diagnostics, never host paths or Runtime connection details."; -const CREATE_DESCRIPTION: &str = "Materialize a persistent Workdir on a selected Runtime from a Workspace repository and optional selector. This does not change this Worker's attachment; use WorkdirAttach explicitly after creation."; -const ATTACH_DESCRIPTION: &str = "Attach this Worker to one existing Workdir. The Backend enforces one active Workdir per Worker and one active Worker per Workdir, then opens an ephemeral operation session."; -const DETACH_DESCRIPTION: &str = "Detach this Worker from its active Workdir and release Workdir occupancy. Any ephemeral operation session is closed."; +const CREATE_DESCRIPTION: &str = "Materialize a persistent Workdir with an optional human-facing display name on a selected Runtime from a Workspace repository and optional selector. The display name is independent of attachment aliases. This does not change this Worker's attachment; use WorkdirAttach explicitly after creation."; +const ATTACH_DESCRIPTION: &str = "Attach one existing Workdir to this Worker under a stable Worker-local alias. The Backend preserves one active Worker per Workdir while allowing this Worker to own multiple aliases."; +const DETACH_DESCRIPTION: &str = "Detach one Workdir attachment by its Worker-local alias. Active operations, commands, or Internal SubWorkers cause safe rejection rather than implicit authority loss."; pub(crate) type BeforeWorkdirRelease = Arc Pin> + Send>> + Send + Sync>; pub(crate) type AfterWorkdirAttach = Arc; @@ -64,6 +64,7 @@ const DELETE_DESCRIPTION: &str = "Request removal of one persistent Workdir by i #[derive(Clone)] pub struct ManageWorkdirFeature { client: Arc, + session_router: Arc, before_workdir_release: Option, after_workdir_attach: Option, } @@ -82,6 +83,7 @@ impl ManageWorkdirFeature { pub fn new(client: Arc) -> Self { Self { client, + session_router: Arc::new(workdir::WorkdirSessionRouter::new()), before_workdir_release: None, after_workdir_attach: None, } @@ -89,11 +91,13 @@ impl ManageWorkdirFeature { pub(crate) fn with_child_lifecycle( client: Arc, + session_router: Arc, before_workdir_release: BeforeWorkdirRelease, after_workdir_attach: AfterWorkdirAttach, ) -> Self { Self { client, + session_router, before_workdir_release: Some(before_workdir_release), after_workdir_attach: Some(after_workdir_attach), } @@ -116,10 +120,12 @@ impl FeatureModule for ManageWorkdirFeature { } fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> { - let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone()).with_child_lifecycle( - self.before_workdir_release.clone(), - self.after_workdir_attach.clone(), - ); + let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone()) + .with_session_router(self.session_router.clone()) + .with_child_lifecycle( + self.before_workdir_release.clone(), + self.after_workdir_attach.clone(), + ); for (name, definition) in [ ( LIST_TOOL, @@ -183,6 +189,8 @@ impl FeatureModule for ManageWorkdirFeature { #[derive(Clone)] struct WorkspaceHttpWorkdirBackend { client: Arc, + session_router: Arc, + attachment_mutation_lock: Arc>, before_workdir_release: Option, after_workdir_attach: Option, } @@ -205,14 +213,27 @@ impl std::fmt::Debug for WorkspaceHttpWorkdirBackend { #[derive(Debug)] pub struct WorkspaceAttachedWorkdirSession { client: Arc, + target_workdir: String, workdir: Workdir, } impl WorkspaceAttachedWorkdirSession { - pub fn handle(client: Arc) -> WorkdirSessionHandle { + pub fn handle( + client: Arc, + target_workdir: impl Into, + ) -> WorkdirSessionHandle { + Self::handle_for_workdir(client, target_workdir, "workspace-attachment") + } + + pub fn handle_for_workdir( + client: Arc, + target_workdir: impl Into, + workdir_id: impl Into, + ) -> WorkdirSessionHandle { Arc::new(Self { client, - workdir: Workdir::new("workspace-attachment"), + target_workdir: target_workdir.into(), + workdir: Workdir::new(workdir_id), }) } @@ -229,13 +250,15 @@ impl WorkspaceAttachedWorkdirSession { "/api/w/{}/workers/self/workdir-session/operations", encode_path_segment(workspace_id) ), - serde_json::to_string(&WorkspaceWorkdirSessionOperationRequest { operation }).map_err( - |error| { - WorkdirError::Transport(format!( - "failed to encode Workspace Workdir operation: {error}" - )) - }, - )?, + serde_json::to_string(&WorkspaceWorkdirSessionOperationRequest { + target_workdir: self.target_workdir.clone(), + operation, + }) + .map_err(|error| { + WorkdirError::Transport(format!( + "failed to encode Workspace Workdir operation: {error}" + )) + })?, ); let response = self .client @@ -399,11 +422,18 @@ impl WorkspaceHttpWorkdirBackend { fn new(client: Arc) -> Self { Self { client, + session_router: Arc::new(workdir::WorkdirSessionRouter::new()), + attachment_mutation_lock: Arc::new(tokio::sync::Mutex::new(())), before_workdir_release: None, after_workdir_attach: None, } } + fn with_session_router(mut self, session_router: Arc) -> Self { + self.session_router = session_router; + self + } + fn with_child_lifecycle( mut self, before_workdir_release: Option, @@ -449,10 +479,12 @@ impl WorkspaceHttpWorkdirBackend { .transpose()?; let repository_key = validate_identity(&input.repository_key, CREATE_TOOL, "repository_key")?; + let display_name = validate_optional_display_name(input.display_name)?; let selector = validate_optional_selector(input.selector)?; let workspace_id = encode_path_segment(self.workspace_id()?); let request = WorkdirCreateRequest { runtime_id: runtime_id.map(str::to_string), + display_name, repository_key: repository_key.to_string(), selector, operation_id: Some(operation_id), @@ -472,32 +504,71 @@ 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) + let alias = validate_identity(&input.alias, ATTACH_TOOL, "alias")?; + workdir::WorkdirAttachmentAlias::new(alias) + .map_err(|error| ToolError::InvalidArgument(error.to_string()))?; + let working_directory_id = validate_identity( + &input.working_directory_id, + ATTACH_TOOL, + "working_directory_id", + )?; + let alias_key = workdir::WorkdirAttachmentAlias::new(alias) + .map_err(|error| ToolError::InvalidArgument(error.to_string()))?; + if self.session_router.session(&alias_key).is_some() { + return Err(ToolError::InvalidArgument(format!( + "attachment alias `{alias}` already exists" + ))); + } + let response = self.attach_response(alias, working_directory_id)?; + if let Err(error) = self.session_router.attach( + alias_key.clone(), + WorkspaceAttachedWorkdirSession::handle_for_workdir( + self.client.clone(), + alias, + working_directory_id, + ), + ) { + let _ = self.detach(alias); + return Err(ToolError::ExecutionFailed(format!( + "attach Workdir session route: {error}" + ))); + } + workdir_output( + format!("Attached Workdir {working_directory_id} as {alias}"), + &response, + ) } - fn attach_response(&self, workdir_id: &str) -> Result { + fn attach_response( + &self, + alias: &str, + working_directory_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"), + format!("/api/w/{workspace_id}/workers/self/workdir-attachments"), serde_json::to_string(&WorkdirAttachRequest { - workdir_id: workdir_id.to_string(), + alias: alias.to_string(), + working_directory_id: working_directory_id.to_string(), }) .map_err(decode_error)?, )) } - fn detach(&self) -> Result { + fn detach(&self, alias: &str) -> Result { let workspace_id = encode_path_segment(self.workspace_id()?); + let alias_path = encode_path_segment(alias); let response = self.execute_json::(WorkspaceRequest { method: WorkspaceRequestMethod::Delete, - path: format!("/api/w/{workspace_id}/workers/self/workdir-attachment"), + path: format!("/api/w/{workspace_id}/workers/self/workdir-attachments/{alias_path}"), body: None, })?; workdir_output( - format!("Detached from Workdir {}", response.workdir_id), + format!( + "Detached Workdir {} ({})", + response.working_directory_id, response.alias + ), &response, ) } @@ -590,6 +661,7 @@ impl Tool for WorkspaceHttpWorkdirTool { ctx.call_id.to_string(), ), WorkdirOperation::Attach => { + let _mutation_guard = self.backend.attachment_mutation_lock.lock().await; let result = self .backend .attach(parse_input::(input_json)?); @@ -601,15 +673,41 @@ impl Tool for WorkspaceHttpWorkdirTool { result } WorkdirOperation::Detach => { - let _input = parse_input::(input_json)?; - if let Some(before_release) = &self.backend.before_workdir_release { - before_release().await.map_err(|error| { + let _mutation_guard = self.backend.attachment_mutation_lock.lock().await; + let input = parse_input::(input_json)?; + let alias = validate_identity(&input.alias, DETACH_TOOL, "alias")?; + let alias_key = workdir::WorkdirAttachmentAlias::new(alias) + .map_err(|error| ToolError::InvalidArgument(error.to_string()))?; + self.backend + .session_router + .begin_detach(&alias_key) + .await + .map_err(|error| ToolError::ExecutionFailed(error.to_string()))?; + if let Some(before_release) = &self.backend.before_workdir_release + && let Err(error) = before_release().await + { + self.backend.session_router.cancel_detach(&alias_key); + return Err(ToolError::ExecutionFailed(format!( + "stop Internal SubWorkers before Workdir detach: {error}" + ))); + } + let result = match self.backend.detach(alias) { + Ok(result) => result, + Err(error) => { + self.backend.session_router.cancel_detach(&alias_key); + return Err(error); + } + }; + self.backend + .session_router + .finish_detach(&alias_key) + .await + .map_err(|error| { ToolError::ExecutionFailed(format!( - "stop Internal SubWorkers before Workdir detach: {error}" + "detach Workdir session route after Backend release: {error}" )) })?; - } - self.backend.detach() + Ok(result) } WorkdirOperation::Delete => self .backend @@ -672,6 +770,25 @@ fn validate_delete_reason(reason: &str) -> Result<&str, ToolError> { Ok(reason) } +fn validate_optional_display_name( + display_name: Option, +) -> Result, ToolError> { + let Some(display_name) = display_name else { + return Ok(None); + }; + let display_name = display_name.trim(); + if display_name.is_empty() + || display_name.len() > 80 + || display_name.chars().any(char::is_control) + { + return Err(ToolError::InvalidArgument( + "WorkdirCreate display_name must be non-empty, contain no control characters, and be at most 80 bytes" + .to_string(), + )); + } + Ok(Some(display_name.to_string())) +} + fn validate_optional_selector(selector: Option) -> Result, ToolError> { let Some(selector) = selector else { return Ok(None); @@ -725,6 +842,7 @@ fn create_schema() -> serde_json::Value { "required": ["repository_key"], "properties": { "runtime_id": {"type": ["string", "null"], "minLength": 1}, + "display_name": {"type": ["string", "null"], "minLength": 1, "maxLength": 80}, "repository_key": {"type": "string", "minLength": 1}, "selector": {"type": ["string", "null"], "minLength": 1} } @@ -735,15 +853,23 @@ fn attach_schema() -> serde_json::Value { json!({ "type": "object", "additionalProperties": false, - "required": ["workdir_id"], + "required": ["alias", "working_directory_id"], "properties": { - "workdir_id": {"type": "string", "minLength": 1} + "alias": {"type": "string", "minLength": 1, "maxLength": 64}, + "working_directory_id": {"type": "string", "minLength": 1} } }) } fn detach_schema() -> serde_json::Value { - list_schema() + json!({ + "type": "object", + "additionalProperties": false, + "required": ["alias"], + "properties": { + "alias": {"type": "string", "minLength": 1, "maxLength": 64} + } + }) } fn delete_schema() -> serde_json::Value { @@ -767,6 +893,8 @@ struct WorkdirListInput {} struct WorkdirCreateInput { #[serde(default)] runtime_id: Option, + #[serde(default)] + display_name: Option, repository_key: String, #[serde(default)] selector: Option, @@ -775,22 +903,27 @@ struct WorkdirCreateInput { #[derive(Debug, Deserialize)] #[serde(deny_unknown_fields)] struct WorkdirAttachInput { - workdir_id: String, + alias: String, + working_directory_id: String, } #[derive(Debug, Deserialize)] #[serde(deny_unknown_fields)] -struct WorkdirDetachInput {} +struct WorkdirDetachInput { + alias: String, +} #[derive(Debug, Serialize)] struct WorkdirAttachRequest { - workdir_id: String, + alias: String, + working_directory_id: String, } #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)] struct WorkdirAttachmentResponse { workspace_id: String, - workdir_id: String, + alias: String, + working_directory_id: String, attached: bool, } @@ -1021,7 +1154,10 @@ mod tests { ); assert!(create["properties"].get("path").is_none()); assert!(create["properties"].get("session_id").is_none()); - assert_eq!(attach_schema()["required"], json!(["workdir_id"])); + assert_eq!( + attach_schema()["required"], + json!(["alias", "working_directory_id"]) + ); assert!(attach_schema()["properties"].get("session_id").is_none()); assert_eq!( delete_schema()["required"], @@ -1057,12 +1193,14 @@ mod tests { })), response(json!({ "workspace_id": "workspace/test", - "workdir_id": "wd-created", + "alias": "checkout", + "working_directory_id": "wd-created", "attached": true })), response(json!({ "workspace_id": "workspace/test", - "workdir_id": "wd-created", + "alias": "checkout", + "working_directory_id": "wd-created", "attached": false })), response(json!({ @@ -1089,6 +1227,7 @@ mod tests { .create( WorkdirCreateInput { runtime_id: Some("runtime/one".to_string()), + display_name: Some("Review checkout".to_string()), repository_key: "main".to_string(), selector: Some("refs/heads/topic".to_string()), }, @@ -1106,10 +1245,11 @@ mod tests { assert_eq!(client.requests().len(), 2); backend .attach(WorkdirAttachInput { - workdir_id: "wd-created".to_string(), + alias: "checkout".to_string(), + working_directory_id: "wd-created".to_string(), }) .unwrap(); - backend.detach().unwrap(); + backend.detach("checkout").unwrap(); backend .delete(WorkdirDeleteInput { working_directory_id: "wd-created".to_string(), @@ -1132,16 +1272,17 @@ mod tests { serde_json::from_str(requests[1].body.as_deref().unwrap()).unwrap(); assert_eq!(body["repository_key"], "main"); assert_eq!(body["runtime_id"], "runtime/one"); + assert_eq!(body["display_name"], "Review checkout"); assert_eq!(body["operation_id"], "call-create-1"); assert_eq!(body["selector"], "refs/heads/topic"); assert_eq!( requests[2].path, - "/api/w/workspace%2Ftest/workers/self/workdir-attachment" + "/api/w/workspace%2Ftest/workers/self/workdir-attachments" ); assert_eq!(requests[2].method, WorkspaceRequestMethod::Post); assert_eq!( requests[3].path, - "/api/w/workspace%2Ftest/workers/self/workdir-attachment" + "/api/w/workspace%2Ftest/workers/self/workdir-attachments/checkout" ); assert_eq!(requests[3].method, WorkspaceRequestMethod::Delete); assert_eq!( @@ -1160,7 +1301,7 @@ mod tests { "operation": "stat", "result": {"path": "visible.txt", "kind": "file", "size": 8} }))])); - let session = WorkspaceAttachedWorkdirSession::handle(client.clone()); + let session = WorkspaceAttachedWorkdirSession::handle(client.clone(), "workdir"); let result = session .stat(StatRequest { path: workdir::WorkdirPath::new("visible.txt").unwrap(), @@ -1176,6 +1317,7 @@ mod tests { ); let body: serde_json::Value = serde_json::from_str(requests[0].body.as_deref().unwrap()).unwrap(); + assert_eq!(body["target_workdir"], "workdir"); assert_eq!(body["operation"]["operation"], "stat"); assert!(body.get("expected_session_fence").is_none()); assert!(body.get("runtime_id").is_none()); @@ -1188,7 +1330,7 @@ mod tests { "operation": "command_start", "result": "command-1" }))])); - let session = WorkspaceAttachedWorkdirSession::handle(client.clone()); + let session = WorkspaceAttachedWorkdirSession::handle(client.clone(), "workdir"); let handle = session .start_command(CommandRequest { @@ -1226,7 +1368,7 @@ mod tests { workdir::http::WorkdirTransportErrorCode::InvalidRequest, "Workdir operation request is invalid", )])); - let session = WorkspaceAttachedWorkdirSession::handle(client); + let session = WorkspaceAttachedWorkdirSession::handle(client, "workdir"); let error = session .glob(workdir::GlobRequest { @@ -1246,7 +1388,7 @@ mod tests { status: 502, body: "secret token and /host/private/path".to_string(), }])); - let session = WorkspaceAttachedWorkdirSession::handle(client); + let session = WorkspaceAttachedWorkdirSession::handle(client, "workdir"); let error = session .stat(StatRequest { @@ -1273,6 +1415,7 @@ mod tests { ])); let broker = workdir::WorkdirToolBroker::new(WorkspaceAttachedWorkdirSession::handle( client.clone(), + "workdir", )); let scoped = broker .scope(workdir::WorkdirToolScope { @@ -1323,6 +1466,7 @@ mod tests { .create( WorkdirCreateInput { runtime_id: None, + display_name: None, repository_key: "main".to_string(), selector: None, }, @@ -1348,6 +1492,7 @@ mod tests { .create( WorkdirCreateInput { runtime_id: Some(" ".to_string()), + display_name: None, repository_key: "main".to_string(), selector: None, }, @@ -1363,9 +1508,17 @@ mod tests { async fn detach_stops_internal_subworkers_before_backend_release() { let client = Arc::new(RecordingWorkspaceClient::new(vec![response(json!({ "workspace_id": "workspace/test", - "workdir_id": "wd-attached", + "alias": "checkout", + "working_directory_id": "wd-attached", "attached": false }))])); + let router = Arc::new(workdir::WorkdirSessionRouter::new()); + router + .attach( + workdir::WorkdirAttachmentAlias::new("checkout").unwrap(), + WorkspaceAttachedWorkdirSession::handle(client.clone(), "workdir"), + ) + .unwrap(); let cleanup_calls = Arc::new(AtomicUsize::new(0)); let cleanup_calls_for_guard = cleanup_calls.clone(); let before_release: BeforeWorkdirRelease = Arc::new(move || { @@ -1377,11 +1530,12 @@ mod tests { }); let tool = WorkspaceHttpWorkdirTool { backend: WorkspaceHttpWorkdirBackend::new(client.clone()) + .with_session_router(router) .with_child_lifecycle(Some(before_release), None), operation: WorkdirOperation::Detach, }; - tool.execute("{}", ToolExecutionContext::default()) + tool.execute(r#"{"alias":"checkout"}"#, ToolExecutionContext::default()) .await .unwrap(); @@ -1389,23 +1543,31 @@ mod tests { assert_eq!(client.requests().len(), 1); assert_eq!( client.requests()[0].path, - "/api/w/workspace%2Ftest/workers/self/workdir-attachment" + "/api/w/workspace%2Ftest/workers/self/workdir-attachments/checkout" ); } #[tokio::test] async fn detach_does_not_release_backend_when_child_cleanup_fails() { let client = Arc::new(RecordingWorkspaceClient::new(Vec::new())); + let router = Arc::new(workdir::WorkdirSessionRouter::new()); + router + .attach( + workdir::WorkdirAttachmentAlias::new("checkout").unwrap(), + WorkspaceAttachedWorkdirSession::handle(client.clone(), "workdir"), + ) + .unwrap(); let before_release: BeforeWorkdirRelease = Arc::new(|| Box::pin(async { Err(std::io::Error::other("child cleanup failed")) })); let tool = WorkspaceHttpWorkdirTool { backend: WorkspaceHttpWorkdirBackend::new(client.clone()) + .with_session_router(router) .with_child_lifecycle(Some(before_release), None), operation: WorkdirOperation::Detach, }; let error = tool - .execute("{}", ToolExecutionContext::default()) + .execute(r#"{"alias":"checkout"}"#, ToolExecutionContext::default()) .await .unwrap_err(); @@ -1413,11 +1575,77 @@ mod tests { assert!(client.requests().is_empty()); } + #[tokio::test] + async fn detach_reconciles_naturally_completed_workspace_command() { + let client = Arc::new(RecordingWorkspaceClient::new(vec![ + response(json!({ + "workspace_id": "workspace/test", + "alias": "checkout", + "working_directory_id": "wd-attached", + "attached": true + })), + response(json!({ + "operation": "command_start", + "result": "command-1" + })), + response(json!({ + "operation": "command_status", + "result": "completed" + })), + response(json!({ + "workspace_id": "workspace/test", + "alias": "checkout", + "working_directory_id": "wd-attached", + "attached": false + })), + ])); + let backend = WorkspaceHttpWorkdirBackend::new(client.clone()); + backend + .attach(WorkdirAttachInput { + alias: "checkout".to_string(), + working_directory_id: "wd-attached".to_string(), + }) + .unwrap(); + let alias = workdir::WorkdirAttachmentAlias::new("checkout").unwrap(); + backend + .session_router + .session(&alias) + .unwrap() + .start_command(CommandRequest { + command: "printf done".to_string(), + timeout_secs: 5, + output_limit: 1024, + cwd: None, + spill_dir: None, + tool_call_id: None, + }) + .await + .unwrap(); + let tool = WorkspaceHttpWorkdirTool { + backend, + operation: WorkdirOperation::Detach, + }; + + tool.execute(r#"{"alias":"checkout"}"#, ToolExecutionContext::default()) + .await + .unwrap(); + + let requests = client.requests(); + assert_eq!(requests.len(), 4); + assert!( + requests[2] + .body + .as_deref() + .is_some_and(|body| body.contains("command_status")) + ); + } + #[tokio::test] async fn successful_attach_reopens_internal_subworker_admission() { let client = Arc::new(RecordingWorkspaceClient::new(vec![response(json!({ "workspace_id": "workspace/test", - "workdir_id": "wd-attached", + "alias": "checkout", + "working_directory_id": "wd-attached", "attached": true }))])); let reopen_calls = Arc::new(AtomicUsize::new(0)); @@ -1432,7 +1660,7 @@ mod tests { }; tool.execute( - r#"{"workdir_id":"wd-attached"}"#, + r#"{"alias":"checkout","working_directory_id":"wd-attached"}"#, ToolExecutionContext::default(), ) .await diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index d51858331..0b50042ce 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -165,7 +165,7 @@ where })?; worker.disable_manifest_lifecycle_features(); if let Some(session) = inherited_workdir_session { - worker.bind_workdir_session(Some(session)); + worker.bind_single_workdir_session(Some(session)); } let install_report = worker.install_features(features); @@ -608,7 +608,7 @@ pub(crate) async fn spawn_internal_worker_session( })?; worker.disable_manifest_lifecycle_features(); if let Some(session) = inherited_workdir_session { - worker.bind_workdir_session(Some(session)); + worker.bind_single_workdir_session(Some(session)); } let install_report = worker.install_features(features); let installed_tools = install_report.installed_tool_names(); @@ -701,7 +701,7 @@ pub(crate) fn prepare_internal_worker_from_spec( })?; worker.disable_manifest_lifecycle_features(); if let Some(session) = inherited_workdir_session { - worker.bind_workdir_session(Some(session)); + worker.bind_single_workdir_session(Some(session)); } let install_report = worker.install_features(features); let installed_tools = install_report.installed_tool_names(); @@ -772,7 +772,7 @@ pub(crate) async fn prepare_internal_worker_session( visibility: InternalWorkerVisibility, child_registry: Option>, on_turn_end: Option>, - command_event_broker: Option, + command_event_broker: Option, ) -> Result { let (event_tx, _event_rx) = broadcast::channel(256); let sink = worker.sink(); @@ -780,9 +780,14 @@ pub(crate) async fn prepare_internal_worker_session( let alerter = Alerter::new(event_tx.clone()); let in_flight = InFlightEvents::new(event_tx.clone()); if let Some(broker) = command_event_broker.as_ref() { - wire_workdir_command_events(&broker.tool_session(), &in_flight); + let sessions = broker.sessions(); + for alias in sessions.aliases() { + if let Ok(selected) = sessions.resolve(Some(alias.as_str())) { + wire_workdir_command_events(&selected.session, &in_flight); + } + } } else if let Some(session) = worker.workdir_session() { - wire_workdir_command_events(session, &in_flight); + wire_workdir_command_events(&session, &in_flight); } let actor_in_flight = in_flight.clone(); worker.attach_alerter(alerter.clone()); diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index 773c2298a..802d0c7d7 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -23,11 +23,12 @@ use agen::interceptor::{ use agen::tool::ToolOutput; use arc_swap::ArcSwap; use async_trait::async_trait; -use tracing::info; +use tracing::{info, warn}; +use crate::compact::metrics_tracker::MetricsTracker; use crate::compact::state::{AutomaticCompactDecision, CompactState}; -use crate::compact::usage_tracker::UsageTracker; -use session_store::SystemItem; +use crate::compact::usage_tracker::{UsageTracker, persist_pending_usage}; +use session_store::{SystemItem, segment_log}; use crate::hook::{ HookEventKind, HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookPromptAction, @@ -81,6 +82,10 @@ pub(crate) struct WorkerInterceptor { /// persisted into `usage_history`. Subsequent tool-loop LLM calls must /// see these records during pre-request safety accounting. usage_tracker: Option>, + /// Metrics produced synchronously while projecting the request (notably + /// prune.fire/skip). Flushed before the matching request usage so metric + /// correlation order remains stable at the in-run commit boundary. + metrics_tracker: Option>, /// Pending-notification buffer drained into `worker.history` /// via [`Self::pending_history_appends`] just before the next LLM /// request. The Engine `extend`s these into its persistent history @@ -152,6 +157,7 @@ impl WorkerInterceptor { compact_state, usage_history, usage_tracker: None, + metrics_tracker: None, pending_notifies, pending_attachments, prompts, @@ -168,6 +174,11 @@ impl WorkerInterceptor { self } + pub(crate) fn with_metrics_tracker(mut self, metrics_tracker: Arc) -> Self { + self.metrics_tracker = Some(metrics_tracker); + self + } + pub(crate) fn with_prompt_workspace_id(mut self, workspace_id: Option) -> Self { self.prompt_workspace_id = workspace_id; self @@ -272,6 +283,56 @@ impl WorkerInterceptor { } } } + fn try_persist_metric(writer: &Arc, metric: &session_metrics::Metric) { + let payload = serde_json::to_value(metric).expect("Metric is Serialize"); + if writer + .commit_log_entry(session_store::LogEntry::Extension { + ts: segment_log::now_millis(), + domain: session_metrics::DOMAIN.into(), + payload, + }) + .is_err() + { + warn!(name = %metric.name, "failed to record session metric; dropping"); + } + } + + fn persist_completed_request_usage(&self) -> InterceptorResult<()> { + let Some(tracker) = self.usage_tracker.as_ref() else { + return Ok(()); + }; + tracker.request_completed(); + + let (Some(writer), Some(usage_history)) = + (self.log_writer.as_ref(), self.usage_history.as_ref()) + else { + // Low-level Worker tests may omit the type-erased writer. Their + // terminal persist_turn remains the durable fallback. + return Ok(()); + }; + if let Some(metrics_tracker) = self.metrics_tracker.as_ref() { + for metric in metrics_tracker.drain() { + Self::try_persist_metric(writer, &metric); + } + } + + persist_pending_usage( + tracker, + usage_history, + self.compact_state.as_deref(), + |entry, after_commit, post_metrics| { + writer.commit_usage_boundary(entry, after_commit, post_metrics) + }, + ) + .map_err(|error| { + InterceptorError::new( + InterceptorErrorCategory::Dependency, + format!("usage persistence failed: {error}"), + ) + })?; + Ok(()) + } + fn attach_prompt_provenance(&self, items: &mut [SystemItem]) { let prompts = self.prompts.load(); let projection = prompts.projection(); @@ -596,6 +657,11 @@ impl Interceptor for WorkerInterceptor { &self, context: AssistantTurnEndContext<'_, SessionHistoryMetadata>, ) -> InterceptorResult { + // `on_assistant_turn_end` is reached only after a normal provider + // stream completed and its usage callback was flushed. Commit that + // exact request before tools can extend the same logical run. + self.persist_completed_request_usage()?; + let history = context.history; let final_text_preview = history .iter() @@ -752,6 +818,25 @@ mod tests { } } + struct FailingUsageCommitter { + attempts: Arc, + } + + impl SystemItemCommitter for FailingUsageCommitter { + fn commit_log_entry( + &self, + entry: session_store::LogEntry, + ) -> Result<(), session_store::StoreError> { + if matches!(entry, session_store::LogEntry::LlmUsage { .. }) { + self.attempts.fetch_add(1, Ordering::SeqCst); + return Err(session_store::StoreError::Io(std::io::Error::other( + "synthetic usage failure", + ))); + } + Ok(()) + } + } + struct AppendingPreRequestHook { saw_handle: Arc, } @@ -955,6 +1040,68 @@ mod tests { ); } + #[tokio::test] + async fn completed_request_usage_failure_stops_before_rearm() { + use crate::compact::state::{AutomaticCompactDecision, CompactionOutcome}; + + let registry = Arc::new(HookRegistryBuilder::new().build()); + let state = Arc::new(CompactState::new(None, Some(10), 0)); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Start(_) + )); + assert!(state.complete_automatic(CompactionOutcome::Succeeded)); + + let usage_history = Arc::new(Mutex::new(Vec::new())); + let tracker = Arc::new(UsageTracker::new()); + tracker.note_compaction_correlation_id("compact-id".into()); + tracker.note_request(1); + tracker.record_usage(&agen::event::UsageEvent { + input_tokens: Some(5), + output_tokens: Some(1), + total_tokens: Some(6), + cache_read_input_tokens: Some(0), + cache_creation_input_tokens: Some(0), + }); + let attempts = Arc::new(AtomicUsize::new(0)); + let interceptor = WorkerInterceptor::new( + registry, + Some(Arc::clone(&state)), + Some(Arc::clone(&usage_history)), + NotifyBuffer::new(), + Arc::new(Mutex::new(Vec::new())), + test_prompts(), + Some(Arc::new(FailingUsageCommitter { + attempts: Arc::clone(&attempts), + })), + ) + .with_usage_tracker(Arc::clone(&tracker)); + + let error = interceptor + .on_assistant_turn_end(AssistantTurnEndContext { + invocation: Default::default(), + assistant_entries: &[], + history: &[], + tool_calls: &[], + }) + .await + .expect_err("usage storage failure must stop the request boundary"); + + assert_eq!(error.category(), InterceptorErrorCategory::Dependency); + assert!(error.diagnostic().contains("usage persistence failed")); + assert_eq!(attempts.load(Ordering::SeqCst), 1); + assert_eq!( + tracker.records().len(), + 1, + "failed record must remain pending" + ); + assert!(usage_history.lock().unwrap().is_empty()); + assert!(matches!( + state.evaluate_request(11), + AutomaticCompactDecision::Block(_) + )); + } + #[tokio::test] async fn pre_llm_request_counts_in_flight_usage_records() { let registry = Arc::new(HookRegistryBuilder::new().build()); diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 7c2606f74..8cac5d7e5 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -27,7 +27,7 @@ use session_store::{ }; use tokio::sync::{Notify, broadcast}; use tracing::warn; -use workdir::WorkdirScopeLease; +use workdir::WorkdirScopeLeaseSet; use crate::internal_worker::{InternalWorkerSessionHandle, InternalWorkerVisibility}; use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord}; @@ -70,7 +70,7 @@ pub(crate) struct SubWorkerStopSummary { pub(crate) struct InternalSpawnedWorkerRecord { pub worker_name: String, pub scope_delegated: Vec, - pub workdir_tool_scope: Arc, + pub workdir_tool_scope: Arc, #[cfg(test)] pub installed_tools: Arc<[String]>, pub session: InternalWorkerSessionHandle, @@ -89,7 +89,7 @@ impl InternalSpawnedWorkerRecord { pub(crate) fn new( worker_name: String, scope_delegated: Vec, - workdir_tool_scope: WorkdirScopeLease, + workdir_tool_scope: WorkdirScopeLeaseSet, #[cfg(test)] installed_tools: Vec, session: InternalWorkerSessionHandle, child_registry: Arc, @@ -1179,7 +1179,11 @@ mod tests { InternalSpawnedWorkerRecord::new( name.into(), Vec::new(), - delegation, + workdir::WorkdirScopeLeaseSet::from_single( + workdir::WorkdirAttachmentAlias::new("workdir").unwrap(), + delegation, + ) + .unwrap(), Vec::new(), session, registry(), diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index 60efb8c68..18dfa2a79 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -5,6 +5,7 @@ //! in-process Internal Worker session actor. No Runtime Worker record, OS process, PID, Unix socket, //! or machine-wide child allocation is created. +use std::collections::BTreeMap; use std::path::{Path, PathBuf}; use std::sync::Arc; @@ -22,7 +23,7 @@ use manifest::{ use serde::Deserialize; use tokio::sync::mpsc; use workdir::{ - WorkdirToolBroker, WorkdirToolScope, WorkdirToolScopePermission, WorkdirToolScopeRule, + WorkdirToolBrokerRouter, WorkdirToolScope, WorkdirToolScopePermission, WorkdirToolScopeRule, }; use crate::PromptCatalogSource; @@ -81,6 +82,10 @@ struct ReviewerHandoffInput { #[derive(Debug, Deserialize, schemars::JsonSchema)] struct ScopeRuleInput { + /// Worker-local alias of the parent Workdir attachment. May be omitted only + /// when the parent has exactly one attachment. + #[serde(default)] + target_workdir: Option, /// Logical Workdir-relative target such as `.` or `src`. Absolute host /// paths and parent traversal are rejected. target: String, @@ -310,7 +315,7 @@ pub struct SubWorkerSpawnTool { /// Directory the spawned SubWorker's tools should use when the LLM did not /// override it. Defaults to the spawner's cwd. /// Parent-owned broker for scoped Workdir tool execution. - workdir_tool_broker: Option, + workdir_tool_broker: Option, /// Parent-owned in-memory registry shared by the five SubWorker tools. registry: Arc, /// Spawner's resolved Manifest. `profile = "inherit"` derives the @@ -337,7 +342,7 @@ impl SubWorkerSpawnTool { runtime_base: PathBuf, bash_output_dir: PathBuf, workspace_root: PathBuf, - workdir_tool_broker: Option, + workdir_tool_broker: Option, registry: Arc, spawner_manifest: WorkerManifest, prompt_loader: PromptCatalogSource, @@ -428,14 +433,19 @@ impl Tool for SubWorkerSpawnTool { )) })?; let workdir_tool_broker = require_workdir_tool_broker(self.workdir_tool_broker.as_ref())?; - let tool_scope = workdir_tool_scope(input.cwd.as_deref(), workdir_rules, input.command)?; + let tool_scopes = routed_workdir_tool_scopes( + workdir_tool_broker, + input.cwd.as_deref(), + workdir_rules, + input.command, + )?; let workdir_scope = workdir_tool_broker - .scope(tool_scope) + .scope(tool_scopes) .await .map_err(|error| { ToolError::InvalidArgument(format!("scope parent-owned Workdir tools: {error}")) })?; - let child_workdir_tool_broker = workdir_scope.broker(); + let child_workdir_tool_broker = workdir_scope.broker_router(); let spawn_selector = parse_spawn_profile_selector(input.profile.as_deref()).map_err(|msg| { @@ -537,6 +547,7 @@ impl Tool for SubWorkerSpawnTool { )) })?; let child_scope = child.scope().clone(); + child.bind_workdir_sessions(workdir_scope.session_router()); let child_registry = SpawnedWorkerRegistry::new_internal(input.name.clone(), child_scope); register_worker_tools( &mut child, @@ -717,22 +728,27 @@ fn logical_workdir_path(value: &str, field: &str) -> Result { }) } -fn parse_workdir_scope(rules: &[ScopeRuleInput]) -> Result, ToolError> { +fn parse_workdir_scope( + rules: &[ScopeRuleInput], +) -> Result, WorkdirToolScopeRule)>, ToolError> { if rules.is_empty() { return Err(ToolError::InvalidArgument("scope must not be empty".into())); } rules .iter() .map(|rule| { - Ok(WorkdirToolScopeRule { - target: logical_workdir_path(&rule.target, "scope.target")?, - permission: match rule.permission { - PermissionInput::Read => WorkdirToolScopePermission::Read, - PermissionInput::Write => WorkdirToolScopePermission::Write, + Ok(( + rule.target_workdir.clone(), + WorkdirToolScopeRule { + target: logical_workdir_path(&rule.target, "scope.target")?, + permission: match rule.permission { + PermissionInput::Read => WorkdirToolScopePermission::Read, + PermissionInput::Write => WorkdirToolScopePermission::Write, + }, + recursive: rule.recursive, + symlink_policy: rule.symlink_policy.into(), }, - recursive: rule.recursive, - symlink_policy: rule.symlink_policy.into(), - }) + )) }) .collect() } @@ -749,9 +765,28 @@ fn workdir_tool_scope( }) } +fn routed_workdir_tool_scopes( + broker: &WorkdirToolBrokerRouter, + cwd: Option<&str>, + rules: Vec<(Option, WorkdirToolScopeRule)>, + command: bool, +) -> Result, ToolError> { + let mut grouped = BTreeMap::>::new(); + for (target_workdir, rule) in rules { + let alias = broker + .resolve_alias(target_workdir.as_deref()) + .map_err(|error| ToolError::InvalidArgument(error.to_string()))?; + grouped.entry(alias).or_default().push(rule); + } + grouped + .into_iter() + .map(|(alias, rules)| workdir_tool_scope(cwd, rules, command).map(|scope| (alias, scope))) + .collect() +} + fn require_workdir_tool_broker( - broker: Option<&WorkdirToolBroker>, -) -> Result<&WorkdirToolBroker, ToolError> { + broker: Option<&WorkdirToolBrokerRouter>, +) -> Result<&WorkdirToolBrokerRouter, ToolError> { broker.ok_or_else(|| { ToolError::InvalidArgument( "SubWorkerSpawn requires parent-owned Workdir tools; attach a Workdir before granting filesystem access" @@ -992,7 +1027,7 @@ pub(crate) fn sub_worker_spawn_tool( runtime_base: PathBuf, bash_output_dir: PathBuf, workspace_root: PathBuf, - workdir_tool_broker: Option, + workdir_tool_broker: Option, registry: Arc, spawner_manifest: WorkerManifest, prompts: Arc>, @@ -1018,7 +1053,7 @@ fn sub_worker_spawn_tool_impl( runtime_base: PathBuf, bash_output_dir: PathBuf, workspace_root: PathBuf, - workdir_tool_broker: Option, + workdir_tool_broker: Option, registry: Arc, spawner_manifest: WorkerManifest, prompts: Arc>, @@ -1096,12 +1131,14 @@ mod tests { fn workdir_scope_uses_logical_relative_paths() { let rules = parse_workdir_scope(&[ ScopeRuleInput { + target_workdir: None, target: ".".to_string(), permission: PermissionInput::Read, recursive: true, symlink_policy: Default::default(), }, ScopeRuleInput { + target_workdir: Some("checkout".to_string()), target: "src".to_string(), permission: PermissionInput::Write, recursive: false, @@ -1109,12 +1146,14 @@ mod tests { }, ]) .unwrap(); - assert_eq!(rules[0].target.as_str(), ""); - assert_eq!(rules[1].target.as_str(), "src"); - assert_eq!(rules[0].symlink_policy, SymlinkPolicy::Resolved); - assert_eq!(rules[1].symlink_policy, SymlinkPolicy::Logical); + assert_eq!(rules[0].1.target.as_str(), ""); + assert_eq!(rules[1].0.as_deref(), Some("checkout")); + assert_eq!(rules[1].1.target.as_str(), "src"); + assert_eq!(rules[0].1.symlink_policy, SymlinkPolicy::Resolved); + assert_eq!(rules[1].1.symlink_policy, SymlinkPolicy::Logical); for target in ["/host/path", "../escape"] { let error = parse_workdir_scope(&[ScopeRuleInput { + target_workdir: None, target: target.to_string(), permission: PermissionInput::Read, recursive: true, @@ -1278,15 +1317,17 @@ enabled = false let fail_requests = Arc::new(AtomicBool::new(false)); let prompt_loader = PromptCatalogSource::builtins_only(); let (parent_method_tx, mut parent_method_rx) = mpsc::channel(8); - let workdir_tool_broker = workdir::WorkdirToolBroker::new(Arc::new( - workdir::LocalWorkdirSession::materialized_bound( + let workdir_tool_broker = workdir::WorkdirToolBrokerRouter::from_single( + workdir::WorkdirAttachmentAlias::new("workdir").unwrap(), + Arc::new(workdir::LocalWorkdirSession::materialized_bound( workdir::Workdir::new("test-workdir"), workspace_root.clone(), workspace_root.clone(), spawner_scope.clone(), workdir::WorkdirSessionCapabilities::ALL, - ), - )); + )), + ) + .unwrap(); let tool = SubWorkerSpawnTool::new( "parent".into(), workspace_context, @@ -1528,9 +1569,11 @@ enabled = false Arc::new(AvailableWorkspaceClient), ); let remote_client = Arc::new(StrictRemoteWorkdirWorkspaceClient::default()); - let workdir_tool_broker = workdir::WorkdirToolBroker::new( - WorkspaceAttachedWorkdirSession::handle(remote_client.clone()), - ); + let workdir_tool_broker = workdir::WorkdirToolBrokerRouter::from_single( + workdir::WorkdirAttachmentAlias::new("workdir").unwrap(), + WorkspaceAttachedWorkdirSession::handle(remote_client.clone(), "workdir"), + ) + .unwrap(); let calls = Arc::new(AtomicUsize::new(0)); let (parent_method_tx, _parent_method_rx) = mpsc::channel(8); let tool = SubWorkerSpawnTool::new( @@ -1607,6 +1650,7 @@ enabled = false assert!(properties.contains_key("command"), "schema: {schema}"); let schema_text = serde_json::to_string(&schema).unwrap(); assert!(schema_text.contains("symlink_policy"), "schema: {schema}"); + assert!(schema_text.contains("target_workdir"), "schema: {schema}"); assert!(schema_text.contains("logical"), "schema: {schema}"); let required = schema .get("required") diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 255933ab2..3ec8a9a0a 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -45,9 +45,9 @@ use crate::compact::state::{ }; use crate::compact::telemetry::{ CompactAttempt, CompactFailureCategory, CompactMode, CompactSuccessStats, - CompactThresholdPolicy, correlated_post_request_metric, new_compact_metric_correlation_id, + CompactThresholdPolicy, new_compact_metric_correlation_id, }; -use crate::compact::usage_tracker::UsageTracker; +use crate::compact::usage_tracker::{UsageTracker, persist_pending_usage}; use crate::feature::background::{BackgroundTaskRewriteGuard, FeatureBackgroundTaskRegistry}; use crate::feature::builtin::memory::WorkspaceMemoryBackendError; use crate::feature::builtin::{TaskFeature, WorkerObservationProvider}; @@ -374,7 +374,8 @@ use protocol::{ use tokio::net::UnixStream; use tokio::sync::broadcast; use workdir::{ - LocalWorkdirSession, ReadOnlyWorkdirSession, WorkdirSessionCapabilities, WorkdirSessionHandle, + LocalWorkdirSession, ReadOnlyWorkdirSession, WorkdirAttachmentAlias, + WorkdirSessionCapabilities, WorkdirSessionHandle, WorkdirSessionRouter, }; const RESTORE_RECONCILIATION_REACHABILITY_TIMEOUT: Duration = Duration::from_millis(500); @@ -1946,6 +1947,28 @@ impl PendingSubmissionHandle { pub trait SystemItemCommitter: Send + Sync { fn commit_log_entry(&self, entry: LogEntry) -> Result<(), StoreError>; + /// Commit one authoritative usage record while holding the same append + /// barrier across its in-memory publication, guard rearm, and correlated + /// post-request metrics. Metric appends remain best-effort. + fn commit_usage_boundary( + &self, + entry: LogEntry, + after_commit: &mut dyn FnMut(), + post_metrics: &[session_metrics::Metric], + ) -> Result<(), StoreError> { + self.commit_log_entry(entry)?; + after_commit(); + for metric in post_metrics { + let payload = serde_json::to_value(metric).expect("Metric is Serialize"); + let _ = self.commit_log_entry(LogEntry::Extension { + ts: segment_log::now_millis(), + domain: session_metrics::DOMAIN.into(), + payload, + }); + } + Ok(()) + } + fn commit_system_item_with_extensions( &self, item: SystemItem, @@ -1977,6 +2000,30 @@ where fn commit_log_entry(&self, entry: LogEntry) -> Result<(), StoreError> { self.append_entry(entry) } + + fn commit_usage_boundary( + &self, + entry: LogEntry, + after_commit: &mut dyn FnMut(), + post_metrics: &[session_metrics::Metric], + ) -> Result<(), StoreError> { + let _append_guard = self + .state + .append_lock + .lock() + .expect("segment append lock poisoned"); + self.append_entry_locked(entry)?; + after_commit(); + for metric in post_metrics { + let payload = serde_json::to_value(metric).expect("Metric is Serialize"); + let _ = self.append_entry_locked(LogEntry::Extension { + ts: segment_log::now_millis(), + domain: session_metrics::DOMAIN.into(), + payload, + }); + } + Ok(()) + } } struct SessionFlowRuntimeStateCommitter { @@ -2128,9 +2175,9 @@ pub struct Worker { /// Explicit local filesystem authority, or `None` for Workers with no /// local cwd and no filesystem/Bash tool surface. filesystem_authority: WorkerFilesystemAuthority, - /// Live WorkdirSession provider derived once from the Worker–Workdir binding. - /// Local tools, file views, and compaction workers clone this handle. - workdir_session: Option, + /// Alias-keyed live Workdir sessions. Provider transport remains hidden + /// behind each handle and no attachment is designated primary. + workdir_sessions: Arc, /// Path-free workspace identity/client context injected by Runtime/host. /// This never grants local filesystem authority. workspace_context: WorkerWorkspaceContext, @@ -2154,16 +2201,16 @@ pub struct Worker { /// Shared compaction state (present when threshold is configured). compact_state: Option>, /// Per-LLM-request Usage tracker. Always present after construction. - /// Captures `(history_len, UsageEvent)` pairs during a run; drained - /// in `persist_turn` and persisted as `LogEntry::LlmUsage` entries. + /// Captures `(history_len, UsageEvent)` pairs until their authoritative + /// request-boundary or terminal `LlmUsage` append succeeds. usage_tracker: Arc, /// Sync-side buffer for `Metric` values queued from inside Engine - /// callbacks (currently the prune observer). Drained in `persist_turn` - /// and written via `session_metrics::record_metric` alongside - /// `LogEntry::LlmUsage`. Always present after construction. + /// callbacks (currently the prune observer). Drained at the same request + /// boundary as usage, with terminal `persist_turn` as the fallback. metrics_tracker: Arc, /// Cumulative Usage measurement timeline, one entry per LLM call. - /// Restored from session log on `restore`, appended on each persist. + /// Restored from session log on `restore`, appended after each durable + /// request-boundary usage commit. /// Read by token-accounting APIs (`Worker::total_tokens`, etc.). /// /// Wrapped in `Arc` so that callbacks injected into the @@ -2464,7 +2511,7 @@ impl Worker { let prompts = Arc::new(ArcSwap::from(PromptCatalog::builtins_only()?)); DelegationScope::from_config(&manifest.delegation_scope).map_err(WorkerError::Scope)?; let scope = SharedScope::new(scope); - let workdir_session = workdir_session_from_authority(&filesystem_authority, &scope); + let workdir_sessions = workdir_sessions_from_authority(&filesystem_authority, &scope); let mut worker = Self { manifest, engine: Some(worker), @@ -2475,7 +2522,7 @@ impl Worker { worker_metadata_segment_cas: None, segment_state: SegmentState::new(session_id, segment_id, 0), filesystem_authority, - workdir_session, + workdir_sessions, workspace_context, flow_runtime_state: Arc::new(Mutex::new(None)), flow_feature_enabled: false, @@ -2652,15 +2699,36 @@ impl Worker { self.filesystem_authority.as_local() } - pub fn workdir_session(&self) -> Option<&WorkdirSessionHandle> { - self.workdir_session.as_ref() + /// Compatibility projection for tool schemas that do not yet carry a + /// target alias. It is available only when the attachment set has exactly + /// one member, so multiple attachments never acquire an implicit primary. + pub fn workdir_session(&self) -> Option { + self.workdir_sessions.only_session() } - /// Replace the constructor fallback with the provider binding resolved by - /// the owning Runtime. Runtime calls this before the Worker controller is - /// spawned, so tools only ever observe the Runtime-bound handle. - pub fn bind_workdir_session(&mut self, workdir_session: Option) { - self.workdir_session = workdir_session; + pub fn workdir_sessions(&self) -> Arc { + self.workdir_sessions.clone() + } + + /// Replace the constructor fallback with the complete alias-keyed provider + /// set resolved by the owning Runtime before controller startup. + pub fn bind_workdir_sessions(&mut self, workdir_sessions: Arc) { + self.workdir_sessions = workdir_sessions; + } + + /// Bind one session for direct/Internal Worker construction. Runtime Worker + /// creation uses [`Self::bind_workdir_sessions`] with the complete set. + pub fn bind_single_workdir_session(&mut self, workdir_session: Option) { + let router = Arc::new(WorkdirSessionRouter::new()); + if let Some(session) = workdir_session { + router + .attach( + WorkdirAttachmentAlias::new("workdir").expect("static alias is valid"), + session, + ) + .expect("fresh Workdir router accepts its only attachment"); + } + self.workdir_sessions = router; } /// Path-free workspace identity, if Runtime/host associated this Worker @@ -2928,12 +2996,7 @@ impl Worker { location.session_id, location.segment_id, ); - let read_only_tools = self - .workdir_session - .clone() - .map(|source| Arc::new(ReadOnlyWorkdirSession::new(source)) as WorkdirSessionHandle) - .map(tools::read_only_builtin_tools) - .unwrap_or_default(); + let read_only_tools = tools::read_only_routed_builtin_tools(self.workdir_sessions()); let client = self .engine .as_ref() @@ -3309,8 +3372,8 @@ impl Worker { /// Handle to the synchronous `MetricsTracker` buffer. /// /// Engine callbacks (e.g. the prune observer) clone this `Arc` and - /// `.push(metric)` into it; Worker drains it in `persist_turn` and - /// writes each metric via `session_metrics::record_metric`. + /// `.push(metric)` into it; Worker drains it at the next request accounting + /// boundary, with terminal persistence as fallback. pub(crate) fn metrics_tracker_handle( &self, ) -> Arc { @@ -3546,6 +3609,7 @@ impl Worker { self.pending_committed_history.clone(), ) .with_usage_tracker(self.usage_tracker.clone()) + .with_metrics_tracker(self.metrics_tracker.clone()) .with_prompt_workspace_id( self.workspace_context .workspace_id() @@ -4061,7 +4125,7 @@ impl Worker { /// unresolved placeholder stays in the flattened user message so the LLM /// still sees the intent. async fn resolve_file_refs(&self, segments: &[Segment]) -> Vec { - let Some(workdir) = self.workdir_session.clone() else { + let Some(workdir) = self.workdir_session() else { for seg in segments { if let Segment::FileRef { path } = seg { self.alert( @@ -4431,6 +4495,15 @@ impl Worker { self.handle_worker_result(result, history_before).await } + fn ensure_no_pending_usage_for_segment_replacement(&self) -> Result<(), WorkerError> { + let pending_records = self.usage_tracker.pending_record_count(); + if pending_records == 0 { + Ok(()) + } else { + Err(WorkerError::PendingUsagePersistence { pending_records }) + } + } + /// Ensure the session exists and the writer's tally still matches /// the on-disk entry count. /// @@ -4472,6 +4545,7 @@ impl Worker { if store_count == entries_written { return Ok(()); } + self.ensure_no_pending_usage_for_segment_replacement()?; // Auto-fork within the same Session: mint a fresh Segment and // switch to it. The source segment is left immutable (no terminal // marker is written back); the fork relationship is recorded @@ -5008,45 +5082,38 @@ impl Worker { self.try_record_metric(&metric); } - // Persist any LLM Usage measurements collected during this run. - // One LogEntry::LlmUsage per LLM call (the tool loop may have run - // many calls within a single Worker::run). Each is also appended to - // the in-memory `usage_history` so token-accounting APIs see it - // before the next run. Records carrying a `correlation_id` (set - // by an upstream observer such as the prune projection) also get - // a paired `prune.post_request` metric so cache_read/write can be - // joined back to the originating event. - let usage_records = self.usage_tracker.drain(); - for recorded in usage_records { - let crate::compact::usage_tracker::RecordedUsage { - record, - post_requests, - } = recorded; - let committed_post_compact_request = post_requests.iter().any(|link| { - link.metric == crate::compact::usage_tracker::PostRequestMetric::Compaction - }); - self.commit_entry(LogEntry::LlmUsage { - ts: segment_log::now_millis(), - history_len: record.history_len, - input_total_tokens: record.input_total_tokens, - cache_read_tokens: record.cache_read_tokens, - cache_write_tokens: record.cache_write_tokens, - output_tokens: record.output_tokens, - })?; - if committed_post_compact_request { - if let Some(state) = &self.compact_state { - state.post_compact_request_committed(); - } - } - for link in post_requests { - let metric = - correlated_post_request_metric(link.metric, &link.correlation_id, &record); + // Persist any request measurements that did not already cross the + // in-run request boundary (for example partial usage from an errored or + // cancelled stream, or low-level Workers without an attached writer). + // `persist_pending_usage` removes each record only after its authoritative + // LlmUsage append succeeds, so the in-run and terminal paths cannot + // double-write or silently lose a failed append. + let usage_tracker = self.usage_tracker.clone(); + let usage_history = self.usage_history.clone(); + let compact_state = self.compact_state.clone(); + if let Some(writer) = self.log_writer.clone() { + let _ = persist_pending_usage( + &usage_tracker, + &usage_history, + compact_state.as_deref(), + |entry, after_commit, post_metrics| { + writer.commit_usage_boundary(entry, after_commit, post_metrics) + }, + )?; + } else { + let post_request_metrics = persist_pending_usage( + &usage_tracker, + &usage_history, + compact_state.as_deref(), + |entry, after_commit, _post_metrics| { + self.commit_entry(entry)?; + after_commit(); + Ok(()) + }, + )?; + for metric in post_request_metrics { self.try_record_metric(&metric); } - self.usage_history - .lock() - .expect("usage_history poisoned") - .push(record); } let interrupted = matches!( @@ -5106,6 +5173,7 @@ impl Worker { mut cancel: Option>, trigger: CompactionTrigger, ) -> Result { + self.ensure_no_pending_usage_for_segment_replacement()?; self.release_pending_compaction_service().await?; let _rewrite_guard = self .prepare_session_rewrite(SessionRewriteKind::Compact) @@ -5380,11 +5448,14 @@ impl Worker { // Default references: the N most-recently-touched files in the // session, surfaced so the compact worker can inspect them and // decide which (if any) the next session needs. - let default_refs: Vec = self - .tracker - .as_ref() - .map(|t| t.recent_files(manifest::defaults::COMPACT_DEFAULT_REFERENCE_COUNT)) - .unwrap_or_default(); + let default_refs: Vec = if self.workdir_sessions.len() == 1 { + self.tracker + .as_ref() + .map(|t| t.recent_files(manifest::defaults::COMPACT_DEFAULT_REFERENCE_COUNT)) + .unwrap_or_default() + } else { + Vec::new() + }; // Input text fed to the compact worker. Includes the default // references, current TaskStore snapshot, current TaskStore snapshot, and the (pruned) conversation text. @@ -5428,7 +5499,7 @@ impl Worker { // Build a normal parent-owned Internal Worker over a pinned immutable // capture. Only SessionExplore and compaction output tools are installed. - let workdir = self.workdir_session.clone(); + let workdir = self.workdir_session(); let read_only_workdir = workdir.clone().map(|session| { Arc::new(ReadOnlyWorkdirSession::new(session)) as workdir::WorkdirSessionHandle }); @@ -6139,7 +6210,8 @@ where let worker_metadata_writer = Some(worker_metadata_writer_for_store(&store)); let worker_metadata_segment_cas = Some(worker_metadata_segment_cas_for_store(&store)); let scope = SharedScope::new(common.scope); - let workdir_session = workdir_session_from_authority(&common.filesystem_authority, &scope); + let workdir_sessions = + workdir_sessions_from_authority(&common.filesystem_authority, &scope); let mut worker = Self { manifest, @@ -6151,7 +6223,7 @@ where worker_metadata_segment_cas, segment_state: SegmentState::new(session_id, segment_id, 0), filesystem_authority: common.filesystem_authority, - workdir_session, + workdir_sessions, workspace_context: common.workspace_context, flow_runtime_state: Arc::new(Mutex::new(None)), flow_feature_enabled: false, @@ -6228,7 +6300,8 @@ where apply_worker_manifest(&mut engine, &manifest.engine); engine.set_cache_key(Some(segment_id.to_string())); let scope = SharedScope::new(common.scope); - let workdir_session = workdir_session_from_authority(&common.filesystem_authority, &scope); + let workdir_sessions = + workdir_sessions_from_authority(&common.filesystem_authority, &scope); let mut worker = Self { manifest, engine: Some(engine), @@ -6239,7 +6312,7 @@ where worker_metadata_segment_cas: None, segment_state: SegmentState::new(session_id, segment_id, 0), filesystem_authority: common.filesystem_authority, - workdir_session, + workdir_sessions, workspace_context: common.workspace_context, flow_runtime_state: Arc::new(Mutex::new(None)), flow_feature_enabled: false, @@ -6350,7 +6423,8 @@ where let worker_metadata_writer = Some(worker_metadata_writer_for_store(&store)); let worker_metadata_segment_cas = Some(worker_metadata_segment_cas_for_store(&store)); let scope = SharedScope::new(common.scope); - let workdir_session = workdir_session_from_authority(&common.filesystem_authority, &scope); + let workdir_sessions = + workdir_sessions_from_authority(&common.filesystem_authority, &scope); let mut worker = Self { manifest, @@ -6362,7 +6436,7 @@ where worker_metadata_segment_cas, segment_state: SegmentState::new(session_id, segment_id, 0), filesystem_authority: common.filesystem_authority, - workdir_session, + workdir_sessions, workspace_context: common.workspace_context, flow_runtime_state: Arc::new(Mutex::new(None)), flow_feature_enabled: false, @@ -6726,7 +6800,8 @@ where let worker_metadata_writer = Some(worker_metadata_writer_for_store(&store)); let worker_metadata_segment_cas = Some(worker_metadata_segment_cas_for_store(&store)); let scope = SharedScope::new(common.scope); - let workdir_session = workdir_session_from_authority(&common.filesystem_authority, &scope); + let workdir_sessions = + workdir_sessions_from_authority(&common.filesystem_authority, &scope); let mut worker = Self { manifest, @@ -6738,7 +6813,7 @@ where worker_metadata_segment_cas, segment_state: SegmentState::new(session_id, segment_id, state.entries_count), filesystem_authority: common.filesystem_authority, - workdir_session, + workdir_sessions, workspace_context: common.workspace_context, flow_runtime_state: Arc::new(Mutex::new(restored_flow_runtime_state( &state.extensions, @@ -7413,7 +7488,9 @@ fn compact_failure_category(error: &WorkerError) -> CompactFailureCategory { | WorkerError::SegmentActivationIncomplete { .. } => { CompactFailureCategory::ActiveSegmentCommit } - WorkerError::Store(_) => CompactFailureCategory::Storage, + WorkerError::Store(_) | WorkerError::PendingUsagePersistence { .. } => { + CompactFailureCategory::Storage + } WorkerError::InvalidState(_) | WorkerError::Engine(_) => { CompactFailureCategory::InternalWorker } @@ -7477,6 +7554,11 @@ pub enum WorkerError { #[error(transparent)] Provider(#[from] crate::model_client::ProviderError), + #[error( + "cannot replace the active Segment while {pending_records} usage record(s) await durable persistence" + )] + PendingUsagePersistence { pending_records: usize }, + #[error("active Segment changed before compaction commit")] CompactActiveSegmentChanged, @@ -7580,18 +7662,26 @@ pub enum WorkerError { }, } -fn workdir_session_from_authority( +fn workdir_sessions_from_authority( authority: &WorkerFilesystemAuthority, scope: &SharedScope, -) -> Option { - authority.as_local().map(|local| { - Arc::new(LocalWorkdirSession::materialized( +) -> Arc { + let router = Arc::new(WorkdirSessionRouter::new()); + if let Some(local) = authority.as_local() { + let session = Arc::new(LocalWorkdirSession::materialized( local.root.clone(), local.cwd.clone(), scope.clone(), WorkdirSessionCapabilities::ALL, - )) as WorkdirSessionHandle - }) + )) as WorkdirSessionHandle; + router + .attach( + WorkdirAttachmentAlias::new("workdir").expect("static alias is valid"), + session, + ) + .expect("fresh Workdir router accepts its only attachment"); + } + router } /// Bundle of resources that every high-level Worker constructor needs: @@ -8534,6 +8624,117 @@ mod build_summary_prompt_tests { } } + #[tokio::test] + async fn failed_usage_append_fences_segment_rotation_until_retry_commits_once() { + let dir = tempfile::tempdir().unwrap(); + let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap(); + let mut worker = Worker::new( + minimal_manifest(), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), + store.clone(), + WorkerWorkspaceContext::no_workspace(), + WorkerFilesystemAuthority::None, + Scope::empty(), + ) + .await + .unwrap(); + worker.ensure_segment_head().await.unwrap(); + let session_id = worker.session_id(); + let source_segment = worker.segment_id(); + + worker + .usage_tracker + .note_compaction_correlation_id("compact-id".into()); + worker.usage_tracker.note_request(1); + worker.usage_tracker.record_usage(&agen::event::UsageEvent { + input_tokens: Some(10), + output_tokens: Some(1), + total_tokens: Some(11), + cache_read_input_tokens: Some(0), + cache_creation_input_tokens: Some(0), + }); + worker.usage_tracker.request_completed(); + + let append_error = persist_pending_usage( + &worker.usage_tracker, + &worker.usage_history, + worker.compact_state.as_deref(), + |_, _, _| Err(StoreError::Io(std::io::Error::other("synthetic failure"))), + ) + .expect_err("first usage append must fail"); + assert!(append_error.to_string().contains("synthetic failure")); + assert_eq!(worker.usage_tracker.pending_record_count(), 1); + + // Simulate another writer extending the active Segment. The normal run + // prelude must not auto-fork while the failed record still belongs to + // this Segment. + store + .append( + session_id, + source_segment, + &LogEntry::TurnEnd { + ts: segment_log::now_millis(), + turn_count: 99, + }, + ) + .unwrap(); + let run_error = worker + .run_text("must stop before provider ownership") + .await + .expect_err("head drift must not fork past pending usage"); + assert!(matches!( + run_error, + WorkerError::PendingUsagePersistence { pending_records: 1 } + )); + assert_eq!(worker.segment_id(), source_segment); + assert_eq!( + store.list_segments(session_id).unwrap(), + vec![source_segment] + ); + + let error = worker + .compact(0) + .await + .expect_err("pending usage must fence Segment replacement"); + assert!(matches!( + error, + WorkerError::PendingUsagePersistence { pending_records: 1 } + )); + assert_eq!(worker.segment_id(), source_segment); + assert_eq!( + store.list_segments(session_id).unwrap(), + vec![source_segment] + ); + + let metrics = persist_pending_usage( + &worker.usage_tracker, + &worker.usage_history, + worker.compact_state.as_deref(), + |entry, after_commit, _| { + store.append(session_id, source_segment, &entry)?; + after_commit(); + Ok(()) + }, + ) + .expect("retry should commit on the originating Segment"); + assert_eq!(worker.usage_tracker.pending_record_count(), 0); + assert_eq!(worker.usage_history.lock().unwrap().len(), 1); + assert_eq!(metrics.len(), 1, "correlation must be emitted once"); + assert_eq!(metrics[0].name, "compact.post_request"); + + let usage_entries = store + .read_all(session_id, source_segment) + .unwrap() + .into_iter() + .filter(|entry| matches!(entry, LogEntry::LlmUsage { .. })) + .count(); + assert_eq!(usage_entries, 1, "retry must not duplicate accounting"); + assert_eq!( + store.list_segments(session_id).unwrap(), + vec![source_segment] + ); + } + #[tokio::test] async fn fresh_run_clears_interrupted_budget_before_pre_run_compaction() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/worker/tests/compact_events_test.rs b/crates/worker/tests/compact_events_test.rs index c6c37367c..7e086533f 100644 --- a/crates/worker/tests/compact_events_test.rs +++ b/crates/worker/tests/compact_events_test.rs @@ -7,15 +7,16 @@ //! (driven by `compact_request_threshold` → `PreRequestAction::Yield`) use std::pin::Pin; -use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; use agen::Engine; use agen::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent, UsageEvent}; use agen::llm_client::types::Item; use agen::llm_client::{ClientError, LlmClient, Request}; +use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use async_trait::async_trait; -use futures::Stream; +use futures::{Stream, StreamExt}; use protocol::{Event, Method, RunResult}; use session_store::{ CombinedStore, FsStore, FsWorkerStore, LogEntry, Store, WorkerMetadata, WorkerMetadataStore, @@ -125,10 +126,13 @@ fn annotated(item: Item) -> session_store::LoggedHistoryEntry { } } +type CallObserver = Arc; + #[derive(Clone)] struct MockClient { responses: Arc>>, call_count: Arc, + observer: Arc>>, } impl MockClient { @@ -136,8 +140,13 @@ impl MockClient { Self { responses: Arc::new(responses), call_count: Arc::new(AtomicUsize::new(0)), + observer: Arc::new(Mutex::new(None)), } } + + fn set_observer(&self, observer: impl Fn(usize) + Send + Sync + 'static) { + *self.observer.lock().unwrap() = Some(Arc::new(observer)); + } } #[async_trait] @@ -152,6 +161,9 @@ impl LlmClient for MockClient { ) -> Result> + Send>>, ClientError> { let count = self.call_count.fetch_add(1, Ordering::SeqCst); + if let Some(observer) = self.observer.lock().unwrap().clone() { + observer(count); + } if count >= self.responses.len() { return Err(ClientError::Config("mock client exhausted".into())); } @@ -161,6 +173,90 @@ impl LlmClient for MockClient { } } +#[derive(Clone)] +struct PauseAfterUsageClient { + responses: Arc>>, + call_count: Arc, + hanging_call: usize, +} + +impl PauseAfterUsageClient { + fn new(responses: Vec>, hanging_call: usize) -> Self { + Self { + responses: Arc::new(responses), + call_count: Arc::new(AtomicUsize::new(0)), + hanging_call, + } + } +} + +#[async_trait] +impl LlmClient for PauseAfterUsageClient { + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + async fn stream( + &self, + _request: Request, + ) -> Result> + Send>>, ClientError> + { + let count = self.call_count.fetch_add(1, Ordering::SeqCst); + let events = self + .responses + .get(count) + .cloned() + .ok_or_else(|| ClientError::Config("pause mock client exhausted".into()))?; + let finite = futures::stream::iter(events.into_iter().map(Ok)); + if count == self.hanging_call { + Ok(Box::pin(finite.chain(futures::stream::pending()))) + } else { + Ok(Box::pin(finite)) + } + } +} + +#[derive(Clone)] +struct GrowTool; + +#[async_trait] +impl Tool for GrowTool { + async fn execute( + &self, + _input: &str, + _ctx: agen::tool::ToolExecutionContext, + ) -> Result { + Ok(ToolOutput { + summary: "grew context".into(), + content: Some("x".repeat(400_000)), + attachments: Vec::new(), + }) + } +} + +fn grow_tool_definition() -> ToolDefinition { + Arc::new(|| { + ( + ToolMeta::new("grow") + .description("test tool that grows request context") + .input_schema(serde_json::json!({"type": "object"})), + Arc::new(GrowTool) as Arc, + ) + }) +} + +fn grow_tool_use_events(call_id: &str, input_tokens: u64) -> Vec { + vec![ + LlmEvent::tool_use_start(0, call_id, "grow"), + LlmEvent::tool_input_delta(0, "{}"), + LlmEvent::tool_use_stop(0), + LlmEvent::usage(input_tokens, 1), + LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] +} + #[derive(Clone)] struct BlockingCompactClient { calls: Arc, @@ -365,12 +461,13 @@ where let scope = worker::Scope::writable(&pwd).unwrap(); std::mem::forget(pwd_tmp); - let worker = + let mut engine = Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client); + engine.register_tool(grow_tool_definition()); let observed_store = store.clone(); let mut worker = Worker::new( manifest, - worker, + engine, store, worker::WorkerWorkspaceContext::local_filesystem(None), worker::WorkerFilesystemAuthority::local(pwd.clone(), pwd.clone()), @@ -379,6 +476,7 @@ where .await .unwrap(); worker.enable_worker_metadata_write_through().unwrap(); + worker.attach_log_writer(Arc::new(worker.log_writer_handle())); (worker, observed_store) } @@ -1057,6 +1155,289 @@ async fn request_threshold_compact_publishes_runtime_progress() { assert_eq!(post.metric.correlation_id.as_deref(), Some(correlation_id)); } +#[tokio::test] +async fn completed_post_compact_usage_rearms_within_one_tool_loop() { + let client = MockClient::new(vec![ + text_events_with_usage("seed", 100_000), + write_summary_tool_use_events("compact-1", "first summary"), + single_text_events("done"), + grow_tool_use_events("grow-1", 1_000), + write_summary_tool_use_events("compact-2", "second summary"), + single_text_events("done"), + text_events_with_usage("finished", 800), + ]); + let call_count = Arc::clone(&client.call_count); + let observed_client = client.clone(); + let manifest = MID_TURN_MANIFEST_TOML.replace( + "compact_request_threshold = 100", + "compact_threshold = 80000\ncompact_request_threshold = 50000", + ); + let (mut worker, store) = make_worker_with_manifest_and_store(&manifest, client).await; + let session_id = worker.session_id(); + let durable_at_next_request = Arc::new(AtomicBool::new(false)); + let observed_barrier = Arc::clone(&durable_at_next_request); + let observed_store = store.clone(); + observed_client.set_observer(move |call_index| { + if call_index != 4 { + return; + } + let usage_is_durable = observed_store + .list_segments(session_id) + .unwrap() + .into_iter() + .flat_map(|segment_id| observed_store.read_all(session_id, segment_id).unwrap()) + .any(|entry| { + matches!( + entry, + LogEntry::LlmUsage { + input_total_tokens: 1_000, + .. + } + ) + }); + assert!( + usage_is_durable, + "the first post-compact usage must be durable before the next provider request" + ); + observed_barrier.store(true, Ordering::SeqCst); + }); + + worker.run_text("seed input").await.unwrap(); + worker + .run_text("continue through two compactions") + .await + .expect("a successful post-compact request must rearm before the tool-loop threshold"); + + assert_eq!( + call_count.load(Ordering::SeqCst), + 7, + "both compactor cycles and both normal post-compact requests must run" + ); + assert!(durable_at_next_request.load(Ordering::SeqCst)); + + let metrics = session_metrics::read_session_metrics(&store, session_id).unwrap(); + let compact_starts = metrics + .iter() + .filter(|record| record.metric.name == "compact.start") + .collect::>(); + assert_eq!(compact_starts.len(), 2); + assert!(compact_starts.iter().any(|record| { + record.metric.dimensions.get("trigger").map(String::as_str) == Some("pre_run") + })); + assert!(compact_starts.iter().any(|record| { + record.metric.dimensions.get("trigger").map(String::as_str) == Some("request_threshold") + })); + let compact_posts = metrics + .iter() + .filter(|record| record.metric.name == "compact.post_request") + .collect::>(); + assert_eq!(compact_posts.len(), 2); + assert_ne!( + compact_posts[0].metric.correlation_id, + compact_posts[1].metric.correlation_id + ); + + let usage = store + .list_segments(session_id) + .unwrap() + .into_iter() + .flat_map(|segment_id| store.read_all(session_id, segment_id).unwrap()) + .filter_map(|entry| match entry { + LogEntry::LlmUsage { + input_total_tokens, .. + } => Some(input_total_tokens), + _ => None, + }) + .collect::>(); + assert_eq!(usage.iter().filter(|tokens| **tokens == 1_000).count(), 1); + assert_eq!(usage.iter().filter(|tokens| **tokens == 800).count(), 1); + + let restored = session_store::restore(&store, session_id, worker.segment_id()).unwrap(); + assert_eq!(restored.usage_history.len(), 1); + assert_eq!(restored.usage_history[0].input_total_tokens, 800); +} + +#[tokio::test] +async fn post_compact_pause_bills_partial_without_rearming_or_consuming_correlation() { + let client = PauseAfterUsageClient::new( + vec![ + text_events_with_usage("seed", 100_000), + write_summary_tool_use_events("manual-compact", "pause test summary"), + single_text_events("done"), + vec![ + LlmEvent::text_block_start(0), + LlmEvent::text_delta(0, "partial"), + LlmEvent::usage(90, 2), + ], + grow_tool_use_events("grow-after-resume", 80), + write_summary_tool_use_events("automatic-compact", "rearmed summary"), + single_text_events("done"), + text_events_with_usage("finished", 50), + ], + 3, + ); + let call_count = Arc::clone(&client.call_count); + let manifest = MID_TURN_MANIFEST_TOML.replace( + "compact_request_threshold = 100", + "compact_threshold = 80000\ncompact_request_threshold = 50000", + ); + let (worker, store) = make_worker_with_manifest_and_store(&manifest, client).await; + let session_id = worker.session_id(); + + let runtime_tmp = tempfile::tempdir().unwrap(); + let bash_output_dir = runtime_tmp.path().join("bash-output"); + let (handle, shutdown_receiver) = + WorkerController::spawn(worker, runtime_tmp.path(), &bash_output_dir) + .await + .unwrap(); + let mut rx = handle.subscribe(); + handle + .send(Method::submit_text( + protocol::new_submission_request_id(), + "seed history", + )) + .await + .unwrap(); + loop { + if matches!( + tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) + .await + .expect("timeout waiting for seed") + .expect("event"), + Event::RunEnd { + result: RunResult::Finished + } + ) { + break; + } + } + handle + .send(Method::submit_text( + protocol::new_submission_request_id(), + "pause the first post-compact request", + )) + .await + .unwrap(); + + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while call_count.load(Ordering::SeqCst) < 4 { + tokio::task::yield_now().await; + } + }) + .await + .expect("post-compact provider request should start"); + handle + .send(Method::Pause { + command: protocol::WorkerCommandEnvelope::new(1), + }) + .await + .unwrap(); + loop { + let event = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) + .await + .expect("timeout waiting for pause") + .expect("event"); + if matches!( + event, + Event::RunEnd { + result: RunResult::Paused + } + ) { + break; + } + } + + let paused_usage = store + .list_segments(session_id) + .unwrap() + .into_iter() + .flat_map(|segment_id| store.read_all(session_id, segment_id).unwrap()) + .filter_map(|entry| match entry { + LogEntry::LlmUsage { + input_total_tokens, .. + } => Some(input_total_tokens), + _ => None, + }) + .collect::>(); + assert_eq!( + paused_usage.iter().filter(|tokens| **tokens == 90).count(), + 1 + ); + let paused_metrics = session_metrics::read_session_metrics(&store, session_id).unwrap(); + assert!( + paused_metrics + .iter() + .all(|record| record.metric.name != "compact.post_request"), + "partial Pause usage must not consume compact correlation or rearm" + ); + + handle + .send(Method::Resume { + command: protocol::WorkerCommandEnvelope::new(2), + }) + .await + .unwrap(); + loop { + if matches!( + tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) + .await + .expect("timeout waiting for resume completion") + .expect("event"), + Event::RunEnd { + result: RunResult::Finished + } + ) { + break; + } + } + + assert_eq!( + call_count.load(Ordering::SeqCst), + 8, + "the completed resumed request must rearm and permit request-threshold compaction" + ); + let metrics = session_metrics::read_session_metrics(&store, session_id).unwrap(); + let compact_starts = metrics + .iter() + .filter(|record| record.metric.name == "compact.start") + .collect::>(); + assert_eq!(compact_starts.len(), 2); + let pre_run_correlation = compact_starts + .iter() + .find(|record| { + record.metric.dimensions.get("trigger").map(String::as_str) == Some("pre_run") + }) + .and_then(|record| record.metric.correlation_id.as_deref()) + .expect("pre-run automatic compact correlation"); + let pre_run_posts = metrics + .iter() + .filter(|record| { + record.metric.name == "compact.post_request" + && record.metric.correlation_id.as_deref() == Some(pre_run_correlation) + }) + .collect::>(); + assert_eq!(pre_run_posts.len(), 1); + assert_eq!( + pre_run_posts[0].metric.dimensions["input_total_tokens"], + "80" + ); + assert!(metrics.iter().all(|record| { + record.metric.name != "compact.post_request" + || record.metric.dimensions["input_total_tokens"] != "90" + })); + + handle + .send(Method::Shutdown { + command: protocol::WorkerCommandEnvelope::new(3), + }) + .await + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(2), shutdown_receiver) + .await + .expect("controller shutdown timeout") + .expect("shutdown confirmation"); +} + #[tokio::test] async fn compacted_context_above_request_threshold_fails_before_provider_request() { let client = MockClient::new(vec![ diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index 2193945fa..0f996122b 100644 --- a/crates/worker/tests/controller_test.rs +++ b/crates/worker/tests/controller_test.rs @@ -3,7 +3,9 @@ use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use agen::Engine; -use agen::llm_client::event::{ErrorEvent, Event as LlmEvent, ResponseStatus, StatusEvent}; +use agen::llm_client::event::{ + ErrorEvent, Event as LlmEvent, ResponseStatus, StatusEvent, UsageEvent, +}; use agen::llm_client::types::Item; use agen::llm_client::{ClientError, LlmClient, Request}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; @@ -349,7 +351,7 @@ async fn shutdown_closes_bound_workdir_session() { }) .await .unwrap(); - worker.bind_workdir_session(Some(Arc::clone(&session))); + worker.bind_single_workdir_session(Some(Arc::clone(&session))); let runtime_base = tempfile::tempdir().unwrap(); let bash_output_dir = runtime_base.path().join("bash-output"); @@ -384,7 +386,7 @@ async fn controller_projects_workdir_command_events_and_snapshot_state() { worker.scope().clone(), WorkdirSessionCapabilities::ALL, )); - worker.bind_workdir_session(Some(Arc::clone(&session))); + worker.bind_single_workdir_session(Some(Arc::clone(&session))); let handle = spawn_controller(worker).await; let mut events = handle.subscribe(); @@ -495,7 +497,7 @@ async fn controller_refreshes_command_snapshot_after_high_output_provider_lag() worker.scope().clone(), WorkdirSessionCapabilities::ALL, )); - worker.bind_workdir_session(Some(Arc::clone(&session))); + worker.bind_single_workdir_session(Some(Arc::clone(&session))); let handle = spawn_controller(worker).await; // Local command telemetry uses 8 KiB chunks and a 256-event channel. One @@ -572,7 +574,7 @@ async fn controller_startup_failure_closes_bound_workdir_session() { worker.scope().clone(), WorkdirSessionCapabilities::ALL, )); - worker.bind_workdir_session(Some(Arc::clone(&session))); + worker.bind_single_workdir_session(Some(Arc::clone(&session))); let runtime_base = tempfile::tempdir().unwrap(); let invalid_runtime_base = runtime_base.path().join("not-a-directory"); std::fs::write(&invalid_runtime_base, "file").unwrap(); @@ -2580,18 +2582,32 @@ async fn drain_until bool>( /// the interrupted turn's history consistency. #[tokio::test] async fn pause_then_resume_preserves_notifications_and_history_consistency() { - // Response 1: hang after opening a text block (no stop / completed), - // so the Engine is parked inside the stream read and `cancel_rx` - // races it cleanly on Method::Pause. + // Response 1: report billable partial usage, then hang after opening a text + // block (no stop / completed). Pause must flush that usage without treating + // the provider request as normally completed. let hang = MockResponse::Hang(vec![ LlmEvent::text_block_start(0), LlmEvent::text_delta(0, "partial..."), + LlmEvent::Usage(UsageEvent { + input_tokens: Some(90), + output_tokens: Some(2), + total_tokens: Some(92), + cache_read_input_tokens: Some(0), + cache_creation_input_tokens: Some(0), + }), ]); // Response 2: a clean assistant reply delivered on Resume. let ok = MockResponse::Complete(vec![ LlmEvent::text_block_start(0), LlmEvent::text_delta(0, "resumed output"), LlmEvent::text_block_stop(0, None), + LlmEvent::Usage(UsageEvent { + input_tokens: Some(100), + output_tokens: Some(3), + total_tokens: Some(103), + cache_read_input_tokens: Some(0), + cache_creation_input_tokens: Some(0), + }), LlmEvent::Status(StatusEvent { status: ResponseStatus::Completed, }), @@ -2645,6 +2661,23 @@ async fn pause_then_resume_preserves_notifications_and_history_consistency() { tokio::time::sleep(std::time::Duration::from_millis(50)).await; assert_eq!(handle.shared_state.catalog_status(), WorkerStatus::Paused); + let (paused_entries, _) = handle.sink.subscribe_with_snapshot(); + let paused_usage = paused_entries + .iter() + .filter_map(|entry| match entry { + LogEntry::LlmUsage { + input_total_tokens, + output_tokens, + .. + } => Some((*input_total_tokens, *output_tokens)), + _ => None, + }) + .collect::>(); + assert_eq!( + paused_usage, + vec![(90, 2)], + "Pause must durably account for usage already observed on the interrupted request" + ); handle .send(Method::Notify { @@ -2729,6 +2762,23 @@ async fn pause_then_resume_preserves_notifications_and_history_consistency() { }) .unwrap_or_default(); assert_eq!(assistant_text, "resumed output"); + let (final_entries, _) = handle.sink.subscribe_with_snapshot(); + let final_usage = final_entries + .iter() + .filter_map(|entry| match entry { + LogEntry::LlmUsage { + input_total_tokens, + output_tokens, + .. + } => Some((*input_total_tokens, *output_tokens)), + _ => None, + }) + .collect::>(); + assert_eq!( + final_usage, + vec![(90, 2), (100, 3)], + "the paused request must remain a distinct billing record rather than merging into resume" + ); let has_tool_call = history.iter().any(|i| i.is_tool_call()); assert!(!has_tool_call, "no orphan tool_call in history"); } diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index fd9c19172..7dba7d0b3 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -32,8 +32,9 @@ use worker_runtime::catalog::{ ConfigBundleRef, CreateWorkerRequest, ProfileSelector, ProfileSourceArchiveSource, RepositoryRefObservation, RepositoryRefObservationRequest, WorkerDetail as EmbeddedWorkerDetail, WorkerStatus as EmbeddedWorkerStatus, - WorkingDirectoryClaim, WorkingDirectoryRepositoryAccessRequest, WorkingDirectoryRequest, - WorkingDirectoryStatus, WorkingDirectorySummary, WorkspaceApiRef, + WorkingDirectoryAttachmentClaim, WorkingDirectoryAttachmentRequest, + WorkingDirectoryAttachmentStatus, WorkingDirectoryRepositoryAccessRequest, + WorkingDirectoryRequest, WorkingDirectoryStatus, WorkspaceApiRef, }; use worker_runtime::config_bundle::{ConfigBundle, ConfigBundleAvailability, ConfigBundleSummary}; #[cfg(test)] @@ -270,8 +271,8 @@ pub struct WorkerSummary { pub retention_state: String, pub implementation: WorkerImplementationSummary, pub capabilities: WorkerCapabilitySummary, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub working_directory: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub workdir_attachments: Vec, pub diagnostics: Vec, } @@ -336,7 +337,7 @@ impl From for server_api::RuntimeSummary { pub(crate) fn workspace_worker_summary( summary: WorkerSummary, resource_key: String, - working_directory: Option, + workdir_attachments: Vec, ) -> server_api::WorkerSummary { server_api::WorkerSummary { runtime_id: summary.worker.runtime_id, @@ -366,7 +367,7 @@ pub(crate) fn workspace_worker_summary( can_stop: summary.capabilities.can_stop, can_spawn_followup: summary.capabilities.can_spawn_followup, }, - working_directory, + workdir_attachments, diagnostics: summary.diagnostics.into_iter().map(Into::into).collect(), } } @@ -483,6 +484,13 @@ pub struct WorkerCreateBinding { pub create_fingerprint: String, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct WorkerSpawnWorkingDirectoryAttachmentRequest { + pub alias: workdir::WorkdirAttachmentAlias, + pub working_directory: WorkerSpawnWorkingDirectoryRequest, +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[serde(deny_unknown_fields)] pub struct WorkerSpawnRequest { @@ -495,15 +503,13 @@ pub struct WorkerSpawnRequest { pub ticket_assignment: Option, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub initial_submit: Vec, - /// Optional safe working-directory creation request. The Workspace server resolves - /// this into a runtime-internal `WorkingDirectoryRequest` from configured - /// repositories before calling a host. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub working_directory_request: Option, + /// Safe working-directory materialization requests, keyed by stable Worker-local aliases. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub workdir_attachment_requests: Vec, #[serde(skip, default)] - pub resolved_working_directory_request: Option, + pub resolved_workdir_attachment_requests: Vec, #[serde(skip, default)] - pub resolved_working_directory: Option, + pub resolved_workdir_attachments: Vec, #[serde(skip, default)] pub resolved_config_bundle: Option, #[serde(skip, default)] @@ -905,6 +911,24 @@ pub trait WorkspaceWorkerRuntime: Send + Sync { } } + fn replace_worker_workdir_attachments( + &self, + worker_id: &str, + _attachments: Vec, + ) -> WorkerWorkspaceApiResult { + WorkerWorkspaceApiResult { + state: WorkerOperationState::Unsupported, + worker: None, + diagnostics: vec![diagnostic( + "worker_workdir_attachments_replace_unsupported", + DiagnosticSeverity::Info, + format!( + "runtime does not support replacing Workdir attachments for worker `{worker_id}`" + ), + )], + } + } + fn create_working_directory( &self, _request: WorkingDirectoryRequest, @@ -1489,6 +1513,19 @@ impl RuntimeRegistry { Ok(runtime.replace_worker_workspace_api(worker_id, workspace_api)) } + pub fn replace_worker_workdir_attachments( + &self, + worker: &RuntimeWorkerRef, + attachments: Vec, + ) -> Result { + let runtime_id = worker.runtime_id.as_str(); + let worker_id = worker.worker_id.as_str(); + validate_backend_identifier("runtime_id", runtime_id)?; + validate_backend_identifier("worker_id", worker_id)?; + let runtime = self.runtime(runtime_id)?; + Ok(runtime.replace_worker_workdir_attachments(worker_id, attachments)) + } + pub fn observe_workspace_prompt_projection( &self, projection: worker::WorkspacePromptProjection, @@ -2176,7 +2213,7 @@ impl EmbeddedWorkerRuntime { && self.can_stop_embedded_worker(summary.status), can_spawn_followup: false, }, - working_directory: summary.working_directory.map(|status| status.summary), + workdir_attachments: summary.workdir_attachments, diagnostics: embedded_worker_projection_diagnostics( summary.execution_metadata_available, ), @@ -2220,7 +2257,7 @@ impl EmbeddedWorkerRuntime { && self.can_stop_embedded_worker(detail.status), can_spawn_followup: false, }, - working_directory: detail.working_directory.map(|status| status.summary), + workdir_attachments: detail.workdir_attachments, diagnostics: embedded_worker_projection_diagnostics( detail.execution_metadata_available, ), @@ -2429,6 +2466,39 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime { } } + fn replace_worker_workdir_attachments( + &self, + worker_id: &str, + attachments: Vec, + ) -> WorkerWorkspaceApiResult { + let Some(worker_ref) = self.worker_ref(worker_id) else { + return WorkerWorkspaceApiResult { + state: WorkerOperationState::Rejected, + worker: None, + diagnostics: vec![diagnostic( + "embedded_worker_id_invalid", + DiagnosticSeverity::Warning, + "Worker id was empty and cannot change Workdir attachments".to_string(), + )], + }; + }; + match self + .runtime + .replace_worker_workdir_attachments(&worker_ref, attachments) + { + Ok(detail) => WorkerWorkspaceApiResult { + state: WorkerOperationState::Accepted, + worker: Some(self.map_worker_detail(detail)), + diagnostics: Vec::new(), + }, + Err(err) => WorkerWorkspaceApiResult { + state: WorkerOperationState::Rejected, + worker: None, + diagnostics: vec![embedded_runtime_diagnostic(&err)], + }, + } + } + fn create_working_directory( &self, _request: WorkingDirectoryRequest, @@ -2491,8 +2561,8 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime { request: WorkerSpawnRequest, ) -> WorkerSpawnResult { let mut diagnostics = Vec::new(); - if request.resolved_working_directory_request.is_some() - || request.resolved_working_directory.is_some() + if !request.resolved_workdir_attachment_requests.is_empty() + || !request.resolved_workdir_attachments.is_empty() { diagnostics.push(embedded_workdir_unsupported_diagnostic()); return WorkerSpawnResult { @@ -2573,8 +2643,8 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime { config_bundle, profile_source, initial_input: initial_worker_input(&request.initial_submit), - working_directory_request: request.resolved_working_directory_request.clone(), - working_directory: request.resolved_working_directory.clone(), + workdir_attachment_requests: request.resolved_workdir_attachment_requests.clone(), + workdir_attachments: request.resolved_workdir_attachments.clone(), worker_observation_enabled: request.resolved_worker_observation_enabled, worker_observation_grants: request.resolved_worker_observation_grants.clone(), workspace_api: Some(workspace_api), @@ -3982,7 +4052,7 @@ impl RemoteWorkerRuntime { && runtime_worker_can_stop(true, summary.status), can_spawn_followup: false, }, - working_directory: summary.working_directory.map(|status| status.summary), + workdir_attachments: summary.workdir_attachments, diagnostics: remote_worker_projection_diagnostics(summary.execution_metadata_available), } } @@ -4024,7 +4094,7 @@ impl RemoteWorkerRuntime { && runtime_worker_can_stop(true, detail.status), can_spawn_followup: false, }, - working_directory: detail.working_directory.map(|status| status.summary), + workdir_attachments: detail.workdir_attachments, diagnostics: remote_worker_projection_diagnostics(detail.execution_metadata_available), } } @@ -4388,6 +4458,51 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime { } } + fn replace_worker_workdir_attachments( + &self, + worker_id: &str, + attachments: Vec, + ) -> WorkerWorkspaceApiResult { + let workdir_attachments = match runtime_contract_convert(attachments) { + Ok(attachments) => attachments, + Err(diagnostic) => { + return WorkerWorkspaceApiResult { + state: WorkerOperationState::Rejected, + worker: None, + diagnostics: vec![diagnostic], + }; + } + }; + let request = runtime_api::WorkerWorkdirAttachmentsRequest { + workdir_attachments, + }; + let worker_id = worker_id.to_string(); + match self + .run_runtime_api( + self.request_timeout, + MAX_REMOTE_RUNTIME_RESPONSE_BYTES, + move |client| async move { + client + .replace_worker_workdir_attachments(worker_id, request) + .await + }, + ) + .and_then(|value| { + runtime_contract_convert::<_, worker_runtime::catalog::WorkerDetail>(value.worker) + }) { + Ok(response) => WorkerWorkspaceApiResult { + state: WorkerOperationState::Accepted, + worker: Some(self.map_worker_detail(response)), + diagnostics: Vec::new(), + }, + Err(diagnostic) => WorkerWorkspaceApiResult { + state: WorkerOperationState::Rejected, + worker: None, + diagnostics: vec![diagnostic], + }, + } + } + fn create_working_directory( &self, request: WorkingDirectoryRequest, @@ -4567,8 +4682,8 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime { config_bundle, profile_source, initial_input: initial_worker_input(&request.initial_submit), - working_directory_request: request.resolved_working_directory_request.clone(), - working_directory: request.resolved_working_directory.clone(), + workdir_attachment_requests: request.resolved_workdir_attachment_requests.clone(), + workdir_attachments: request.resolved_workdir_attachments.clone(), worker_observation_enabled: request.resolved_worker_observation_enabled, worker_observation_grants: request.resolved_worker_observation_grants.clone(), workspace_api: Some(workspace_api), @@ -5724,7 +5839,7 @@ pub fn placeholder_worker(host_id: impl Into) -> WorkerSummary { can_stop: false, can_spawn_followup: false, }, - working_directory: None, + workdir_attachments: Vec::new(), diagnostics: vec![diagnostic( "runtime_capability_unsupported", DiagnosticSeverity::Info, @@ -6265,10 +6380,16 @@ mod tests { worker_state: protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, - working_directory: request - .working_directory - .as_ref() - .map(|binding| binding.status()), + workdir_attachments: request + .workdir_attachments + .iter() + .map(|(alias, binding)| { + worker_runtime::catalog::WorkingDirectoryAttachmentStatus { + alias: alias.clone(), + working_directory: binding.status(), + } + }) + .collect(), } } @@ -6360,7 +6481,7 @@ mod tests { can_stop: false, can_spawn_followup: false, }, - working_directory: None, + workdir_attachments: Vec::new(), diagnostics: Vec::new(), }], observed_prompt_revisions: Arc::new(Mutex::new(Vec::new())), @@ -6650,9 +6771,9 @@ mod tests { profile: ProfileSelector::Builtin("builtin:coder".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -6890,9 +7011,9 @@ mod tests { profile: ProfileSelector::Builtin("builtin:coder".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -6990,9 +7111,9 @@ mod tests { profile: ProfileSelector::Builtin("builtin:coder".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -7029,9 +7150,9 @@ mod tests { profile: ProfileSelector::Builtin("builtin:companion".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), diff --git a/crates/workspace-server/src/latest_schema.sql b/crates/workspace-server/src/latest_schema.sql index 7cd9de68c..2b0a096a8 100644 --- a/crates/workspace-server/src/latest_schema.sql +++ b/crates/workspace-server/src/latest_schema.sql @@ -667,6 +667,7 @@ CREATE TABLE workdir_create_credential_revision_retentions ( CREATE TABLE "workdir_registry" ( workspace_id TEXT NOT NULL, workdir_id TEXT NOT NULL, + display_name TEXT, runtime_id TEXT NOT NULL, repository_id TEXT NOT NULL, creation_selector TEXT, @@ -859,10 +860,10 @@ CREATE TABLE "worker_workdir_links" ( runtime_id TEXT NOT NULL, worker_id TEXT NOT NULL, workdir_id TEXT NOT NULL, - role TEXT NOT NULL, + alias TEXT NOT NULL, linked_at TEXT NOT NULL, unlinked_at TEXT, - PRIMARY KEY (workspace_id, worker_id, workdir_id, role), + PRIMARY KEY (workspace_id, worker_id, workdir_id, alias), FOREIGN KEY (workspace_id, worker_id) REFERENCES "worker_registry"(workspace_id, worker_id) ON DELETE CASCADE, FOREIGN KEY (workspace_id, workdir_id) @@ -1093,8 +1094,8 @@ CREATE INDEX worker_removal_operations_worker_idx ON worker_removal_operations(w CREATE UNIQUE INDEX worker_workdir_links_active_workdir_unique ON worker_workdir_links(workspace_id, workdir_id) WHERE unlinked_at IS NULL; -CREATE UNIQUE INDEX worker_workdir_links_active_worker_unique - ON worker_workdir_links(workspace_id, worker_id) +CREATE UNIQUE INDEX worker_workdir_links_active_alias_unique + ON worker_workdir_links(workspace_id, worker_id, alias) WHERE unlinked_at IS NULL; CREATE INDEX worker_workdir_links_workdir ON worker_workdir_links(workspace_id, workdir_id); diff --git a/crates/workspace-server/src/retention.rs b/crates/workspace-server/src/retention.rs index 6aa4925c6..7b4cc6869 100644 --- a/crates/workspace-server/src/retention.rs +++ b/crates/workspace-server/src/retention.rs @@ -2,7 +2,10 @@ //! Runtime receives only resolved dispositions and stable ids; policy authority //! never comes from prompts, profiles, or model input. -use crate::{Error as StoreError, store::SqliteWorkspaceStore}; +use crate::{ + Error as StoreError, + store::{SqliteWorkspaceStore, WorkerRemovalCommit, commit_worker_catalog_removal}, +}; use chrono::Utc; use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params}; use serde::{Deserialize, Serialize}; @@ -417,13 +420,19 @@ impl SqliteWorkspaceStore { operation_id: &str, fp: &str, result: &WorkerRetentionExecutionResult, - ) -> Result { + ) -> Result { self.with_conn_mut(|conn|{ let tx=conn.transaction_with_behavior(TransactionBehavior::Immediate)?; let mut plan=load_plan_op(&tx,operation_id)?.ok_or_else(||StoreError::InvalidInput("operation missing".into()))?; if plan.workspace_id!=workspace_id{return Err(StoreError::InvalidInput("cross-workspace".into()));} if plan.input_fingerprint!=fp||result.input_fingerprint!=fp||result.operation_id!=operation_id{return Err(StoreError::InvalidInput(format!("fingerprint:{operation_id}")));} - if plan.state==WorkerRemovalPlanState::Succeeded{tx.commit()?;return Ok(plan);} + if plan.state==WorkerRemovalPlanState::Succeeded{ + // Repair projection fences left by the pre-atomic implementation and + // return a replayable change for a broadcast missed after commit. + let catalog=commit_worker_catalog_removal(&tx,workspace_id,&plan.worker)?; + tx.commit()?; + return Ok(WorkerRemovalCommit { plan, catalog }); + } if plan.state != WorkerRemovalPlanState::Executing { return Err(StoreError::InvalidInput(format!("stale:{}:plan state {} is not committable", plan.plan_id, state_s(plan.state)))); } @@ -460,12 +469,14 @@ impl SqliteWorkspaceStore { if plan.metadata_disposition==MetadataDisposition::Tombstone{ tx.execute("INSERT OR IGNORE INTO worker_tombstones(workspace_id,runtime_id,worker_id,display_name,profile,worker_created_at,removed_at,archive_id,policy_id,policy_revision,operation_id) VALUES(?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11)",params![workspace_id,plan.worker.runtime_id,plan.worker.worker_id,worker.display_name,worker.profile,worker.created_at,now,plan.archive_id,plan.policy_id,plan.policy_revision,operation_id])?; } + tx.execute("UPDATE worker_workdir_links SET unlinked_at=?4 WHERE workspace_id=?1 AND runtime_id=?2 AND worker_id=?3 AND unlinked_at IS NULL",params![workspace_id,plan.worker.runtime_id,plan.worker.worker_id,now])?; let deleted=tx.execute("DELETE FROM worker_registry WHERE workspace_id=?1 AND runtime_id=?2 AND worker_id=?3 AND updated_at=?4",params![workspace_id,plan.worker.runtime_id,plan.worker.worker_id,plan.worker_revision])?; tx.execute("UPDATE worker_create_reservations SET state='removed',updated_at=?4 WHERE workspace_id=?1 AND runtime_id=?2 AND worker_id=?3 AND state='created'",params![workspace_id,plan.worker.runtime_id,plan.worker.worker_id,now])?; if deleted!=1{return Err(StoreError::InvalidInput(format!("stale:{}:removal fence changed",plan.plan_id)));} + let catalog=commit_worker_catalog_removal(&tx,workspace_id,&plan.worker)?; tx.execute("UPDATE worker_removal_operations SET state='succeeded',failure_category=NULL,updated_at=?1 WHERE operation_id=?2",params![now,operation_id])?; tx.execute("INSERT OR IGNORE INTO worker_retention_audit_events(event_id,operation_id,workspace_id,event_kind,detail,created_at) VALUES(?1,?2,?3,'worker_removed',?4,?5)",params![stable("wre",operation_id),operation_id,workspace_id,format!("runtime_id={} worker_id={} session={} metadata={} diagnostics={}",plan.worker.runtime_id,plan.worker.worker_id,sess(plan.session_disposition),meta(plan.metadata_disposition),diag(plan.diagnostics_disposition)),now])?; - tx.commit()?;plan.state=WorkerRemovalPlanState::Succeeded;plan.updated_at=now;Ok(plan) + tx.commit()?;plan.state=WorkerRemovalPlanState::Succeeded;plan.updated_at=now;Ok(WorkerRemovalCommit { plan, catalog }) }).map_err(map_error) } @@ -948,7 +959,9 @@ fn parse_state(v: &str) -> rusqlite::Result { #[cfg(test)] mod tests { use super::*; - use crate::store::{ControlPlaneStore, TicketCoderAssignmentRecord, WorkerRegistryRecord}; + use crate::store::{ + ControlPlaneStore, TicketCoderAssignmentRecord, WorkerCatalogChange, WorkerRegistryRecord, + }; use worker_runtime::identity::WorkerId; fn worker_id() -> WorkerId { WorkerId::from_legacy_u64(1) @@ -1203,6 +1216,7 @@ mod tests { assert_eq!( s.commit_worker_removal("w", &p.operation_id, &p.input_fingerprint, &result) .unwrap() + .plan .state, WorkerRemovalPlanState::Succeeded ); @@ -1219,6 +1233,7 @@ mod tests { assert_eq!( s.commit_worker_removal("w", &p.operation_id, &p.input_fingerprint, &result) .unwrap() + .plan .state, WorkerRemovalPlanState::Succeeded ); @@ -1351,6 +1366,98 @@ mod tests { .unwrap(); assert!(s.worker_tombstone("w", &p.worker).unwrap().is_none()); } + #[test] + fn terminal_catalog_failure_rolls_back_removal_projection_and_operation() { + let store = setup(); + let purge = WorkerRetentionPolicyUpdate { + policy_id: "purge".into(), + session_disposition: SessionDisposition::Purge, + metadata_disposition: MetadataDisposition::Tombstone, + archive_retention: ArchiveRetention::Forever, + diagnostics_disposition: DiagnosticsDisposition::Purge, + diagnostics_retention_seconds: None, + }; + store + .update_worker_retention_policy("w", 1, &purge) + .unwrap(); + let plan = store.plan_worker_removal(&req(), &inv()).unwrap(); + store + .prepare_worker_removal_execution("w", &plan.plan_id, &plan.input_fingerprint) + .unwrap(); + let result = WorkerRetentionExecutionResult { + operation_id: plan.operation_id.clone(), + input_fingerprint: plan.input_fingerprint.clone(), + expected_worker_revision: plan.worker_revision.clone(), + worker_id: worker_id(), + session_disposition: SessionDisposition::Purge, + diagnostics_disposition: DiagnosticsDisposition::Purge, + archive: None, + source_removed: true, + diagnostics_retained: false, + }; + store + .with_conn(|conn| { + conn.execute_batch( + "CREATE TRIGGER fail_worker_projection_removal + BEFORE INSERT ON worker_registry_projection_removals + BEGIN SELECT RAISE(ABORT, 'injected projection failure'); END;", + )?; + Ok(()) + }) + .unwrap(); + + assert!( + store + .commit_worker_removal("w", &plan.operation_id, &plan.input_fingerprint, &result,) + .is_err() + ); + let (registry, removal, state, tombstone): (i64, i64, String, i64) = store + .with_conn(|conn| { + Ok(( + conn.query_row( + "SELECT COUNT(*) FROM worker_registry WHERE workspace_id='w' AND runtime_id='r' AND worker_id=?1", + [worker_id().to_string()], + |row| row.get(0), + )?, + conn.query_row( + "SELECT COUNT(*) FROM worker_registry_projection_removals WHERE workspace_id='w' AND runtime_id='r' AND worker_id=?1", + [worker_id().to_string()], + |row| row.get(0), + )?, + conn.query_row( + "SELECT state FROM worker_removal_operations WHERE operation_id=?1", + [&plan.operation_id], + |row| row.get(0), + )?, + conn.query_row( + "SELECT COUNT(*) FROM worker_tombstones WHERE workspace_id='w' AND runtime_id='r' AND worker_id=?1", + [worker_id().to_string()], + |row| row.get(0), + )?, + )) + }) + .unwrap(); + assert_eq!(registry, 1, "catalog deletion must roll back"); + assert_eq!(removal, 0, "removal fence must not partially commit"); + assert_eq!(state, "executing", "operation must remain recoverable"); + assert_eq!(tombstone, 0, "retention metadata must roll back"); + + store + .with_conn(|conn| { + conn.execute_batch("DROP TRIGGER fail_worker_projection_removal;")?; + Ok(()) + }) + .unwrap(); + let committed = store + .commit_worker_removal("w", &plan.operation_id, &plan.input_fingerprint, &result) + .unwrap(); + assert!(matches!( + committed.catalog.changes.as_slice(), + [WorkerCatalogChange::Removed(worker)] if worker == &plan.worker + )); + assert_eq!(committed.plan.state, WorkerRemovalPlanState::Succeeded); + } + #[test] fn commit_requires_executing_state_and_exact_runtime_result() { let store = setup(); @@ -1597,6 +1704,26 @@ mod tests { .unwrap() .is_none() ); + s.with_conn(|conn| { + conn.execute( + "DELETE FROM worker_registry_projection_removals WHERE workspace_id='w' AND runtime_id='r' AND worker_id=?1", + [worker_id().to_string()], + )?; + Ok(()) + }) + .unwrap(); + let repaired = s + .recover_succeeded_worker_removal_catalog("w", &request.worker) + .unwrap(); + assert!(matches!( + repaired.changes.as_slice(), + [WorkerCatalogChange::Removed(worker)] if worker == &request.worker + )); + assert!( + s.get_worker_registry("w", &request.worker) + .unwrap() + .is_none() + ); let recovered = s .recover_worker_removal_execution("w", &request.worker) .unwrap() diff --git a/crates/workspace-server/src/runtime_subscription_tests.rs b/crates/workspace-server/src/runtime_subscription_tests.rs index 9d618c199..07529ec24 100644 --- a/crates/workspace-server/src/runtime_subscription_tests.rs +++ b/crates/workspace-server/src/runtime_subscription_tests.rs @@ -22,7 +22,7 @@ impl WorkerExecutionBackend for TestExecutionBackend { protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, - None, + Vec::new(), ) } @@ -64,8 +64,8 @@ fn create_request(name: &str) -> CreateWorkerRequest { .unwrap(), }, initial_input: None, - working_directory_request: None, - working_directory: None, + workdir_attachment_requests: Vec::new(), + workdir_attachments: Vec::new(), worker_observation_enabled: false, worker_observation_grants: Vec::new(), workspace_api: None, diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 48804f051..69bec5a00 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -195,7 +195,9 @@ use worker_runtime::catalog::{ ConfigBundleRef, ProfileSelector, RepositoryMaterializationContext, RepositoryRefObservation, RepositoryRefObservationRequest, RepositorySelector as RuntimeRepositorySelector, RepositorySshCredentialCandidate, RepositorySshMaterializationAccess, SensitiveString, - WorkingDirectoryClaim, WorkingDirectoryRepository, WorkingDirectoryRequest, WorkspaceApiRef, + WorkingDirectoryAttachmentClaim, WorkingDirectoryAttachmentRequest, + WorkingDirectoryAttachmentStatus, WorkingDirectoryRepository, WorkingDirectoryRequest, + WorkspaceApiRef, }; use worker_runtime::config_bundle::ConfigBundle; use worker_runtime::http_server::MAX_WORKER_FILE_UPLOAD_BYTES; @@ -381,6 +383,7 @@ static EMBEDDED_RUNTIME_REQUEST_IDENTITY: std::sync::LazyLock< #[derive(Clone)] struct WorkdirCommandSession { + attachment_alias: String, source: WorkdirSessionHandle, provider_handle: CommandHandle, } @@ -404,10 +407,6 @@ struct WorkdirSessionRegistry { } impl WorkdirSessionRegistry { - fn attachment(&self, worker: &RuntimeWorkerRef) -> Option { - self.attachments.get(worker).cloned() - } - fn insert_attachment( &mut self, worker: RuntimeWorkerRef, @@ -423,6 +422,7 @@ impl WorkdirSessionRegistry { fn register_command( &mut self, worker: RuntimeWorkerRef, + attachment_alias: String, source: WorkdirSessionHandle, provider_handle: CommandHandle, ) -> CommandHandle { @@ -438,6 +438,7 @@ impl WorkdirSessionRegistry { self.commands.insert( (worker, external_handle.clone()), WorkdirCommandSession { + attachment_alias, source, provider_handle, }, @@ -448,13 +449,42 @@ impl WorkdirSessionRegistry { fn command( &self, worker: &RuntimeWorkerRef, + attachment_alias: &str, external_handle: &CommandHandle, ) -> Option { self.commands .get(&(worker.clone(), external_handle.clone())) + .filter(|session| session.attachment_alias == attachment_alias) .cloned() } + fn take_worker_commands_for_alias( + &mut self, + worker: &RuntimeWorkerRef, + attachment_alias: &str, + ) -> Vec { + let command_handles = self + .commands + .iter() + .filter(|((owner, _), session)| { + owner == worker && session.attachment_alias == attachment_alias + }) + .map(|((_, handle), _)| handle.clone()) + .collect::>(); + command_handles + .into_iter() + .filter_map(|external_handle| { + self.commands + .remove(&(worker.clone(), external_handle.clone())) + .map(|session| RegisteredWorkdirSession::Command { + worker: worker.clone(), + external_handle, + session, + }) + }) + .collect() + } + fn take_worker(&mut self, worker: &RuntimeWorkerRef) -> Vec { let mut sessions = Vec::new(); if let Some(session) = self.attachments.remove(worker) { @@ -500,14 +530,10 @@ impl WorkdirSessionRegistry { } } -async fn close_worker_workdir_sessions( +async fn close_registered_workdir_sessions( registry: &Arc>, - worker: &RuntimeWorkerRef, + registered: Vec, ) -> std::result::Result<(), String> { - let registered = registry - .lock() - .map_err(|_| "Workdir session registry was poisoned".to_string())? - .take_worker(worker); let mut failed = Vec::new(); let mut errors = Vec::new(); for item in registered { @@ -535,6 +561,29 @@ async fn close_worker_workdir_sessions( } } +async fn close_worker_workdir_sessions( + registry: &Arc>, + worker: &RuntimeWorkerRef, +) -> std::result::Result<(), String> { + let registered = registry + .lock() + .map_err(|_| "Workdir session registry was poisoned".to_string())? + .take_worker(worker); + close_registered_workdir_sessions(registry, registered).await +} + +async fn close_worker_workdir_command_sessions_for_alias( + registry: &Arc>, + worker: &RuntimeWorkerRef, + attachment_alias: &str, +) -> std::result::Result<(), String> { + let registered = registry + .lock() + .map_err(|_| "Workdir session registry was poisoned".to_string())? + .take_worker_commands_for_alias(worker, attachment_alias); + close_registered_workdir_sessions(registry, registered).await +} + const ATTACHMENT_UPLOAD_GRANT_TTL_MS: u64 = 5 * 60 * 1_000; #[derive(Clone)] @@ -659,7 +708,7 @@ impl From<&WorkspaceApi> for ServerAuthApi { } #[derive(Clone)] -struct WorkspaceWorkerRemoveExecutor { +struct WorkerRemovalService { workspace_id: String, store: Arc, runtime: Weak, @@ -667,10 +716,10 @@ struct WorkspaceWorkerRemoveExecutor { workdir_session_locks: Arc>>>>, worker_remove_locks: Arc>>>>, worker_control_locks: Arc>>>>, - worker_projection: std::sync::Weak, + worker_projection: Arc, } -impl WorkspaceWorkerRemoveExecutor { +impl WorkerRemovalService { fn new(api: &WorkspaceApi) -> Self { Self { workspace_id: api.config.workspace_id.clone(), @@ -680,10 +729,123 @@ impl WorkspaceWorkerRemoveExecutor { workdir_session_locks: api.workdir_session_locks.clone(), worker_remove_locks: api.worker_remove_locks.clone(), worker_control_locks: api.worker_control_locks.clone(), - worker_projection: Arc::downgrade(&api.worker_projection), + worker_projection: api.worker_projection.clone(), } } + /// Terminal catalog deletion for rollback-style compensation. Normal cleanup + /// and WorkerRemove use the retention commit below; compensation preserves + /// its pre-create semantics while sharing the atomic catalog contract. + fn commit_catalog_removal(&self, target: &RuntimeWorkerRef) -> crate::Result<()> { + let commit = self + .store + .delete_worker_registry(&self.workspace_id, target)?; + self.worker_projection.publish_commit(commit) + } + + async fn execute_cleanup_removal(&self, candidate: &CleanupWorkerCandidate) -> ApiResult<()> { + let target = RuntimeWorkerRef::new( + candidate.runtime_id.clone(), + candidate.runtime_worker_id.clone(), + ); + let runtime = self.runtime.upgrade().ok_or_else(|| { + Error::Store("Workspace Runtime registry is unavailable during cleanup".to_string()) + })?; + let response = self + .execute_target_removal( + runtime.as_ref(), + &target, + &format!("Browser cleanup: {}", candidate.reason), + None, + ) + .await + .map_err(|message| { + cleanup_api_error( + candidate.runtime_id.as_str(), + "workspace_cleanup_worker_removal_failed", + message.as_str(), + ) + })?; + if response.status != StatusCode::OK.as_u16() { + return Err(cleanup_api_error( + candidate.runtime_id.as_str(), + "workspace_cleanup_worker_removal_rejected", + response.body.as_str(), + )); + } + Ok(()) + } + + fn delete_spawn_compensation_runtime_worker( + &self, + worker_ref: &RuntimeWorkerRef, + ) -> (bool, Vec) { + let mut diagnostics = Vec::new(); + let Some(runtime) = self.runtime.upgrade() else { + diagnostics.push(spawn_compensation_diagnostic( + "worker_spawn_compensation_runtime_delete_failed", + "Workspace Runtime registry is unavailable during spawn compensation".to_string(), + )); + return (false, diagnostics); + }; + let lifecycle_request = WorkerLifecycleRequest { + reason: Some("Backend spawn finalize failed; compensating Runtime Worker".to_string()), + ticket_assignment: None, + }; + let cancellation = runtime.cancel_worker(worker_ref, lifecycle_request.clone()); + let stop = runtime.stop_worker(worker_ref, lifecycle_request); + let stop_accepted = stop + .as_ref() + .is_ok_and(|result| result.state == WorkerOperationState::Accepted); + let termination_detail = (!stop_accepted).then(|| { + let cancellation = lifecycle_failure_detail("cancel", &cancellation); + let stop = lifecycle_failure_detail("stop", &stop); + format!("{cancellation}; {stop}") + }); + + let runtime_deleted = match runtime.delete_worker(worker_ref) { + Ok(result) if result.state == WorkerOperationState::Accepted && result.deleted => true, + Ok(result) => { + let mut message = format!( + "Runtime did not delete Worker {}:{}: state={:?}, deleted={}", + worker_ref.runtime_id, worker_ref.worker_id, result.state, result.deleted + ); + if let Some(detail) = termination_detail.as_deref() { + message.push_str(&format!("; cancellation: {detail}")); + } + if !result.diagnostics.is_empty() { + message.push_str(&format!( + "; delete diagnostics: {}", + runtime_diagnostics_message(&result.diagnostics) + )); + } + diagnostics.push(spawn_compensation_diagnostic( + "worker_spawn_compensation_runtime_delete_failed", + message, + )); + false + } + Err(RuntimeRegistryError::UnknownWorker { .. }) => true, + Err(error) => { + let mut message = format!( + "Failed to delete Runtime Worker {}:{}: {}", + worker_ref.runtime_id, + worker_ref.worker_id, + error.message() + ); + if let Some(detail) = termination_detail.as_deref() { + message.push_str(&format!("; cancellation: {detail}")); + } + diagnostics.push(spawn_compensation_diagnostic( + "worker_spawn_compensation_runtime_delete_failed", + message, + )); + false + } + }; + (runtime_deleted, diagnostics) + } + async fn resume_worker_retention( &self, runtime: &RuntimeRegistry, @@ -707,10 +869,10 @@ impl WorkspaceWorkerRemoveExecutor { &prepared.plan.input_fingerprint, &result, ) { - Ok(_) => { - if let Some(worker_projection) = self.worker_projection.upgrade() { - let _ = worker_projection.publish_removed(target.clone()); - } + Ok(commit) => { + self.worker_projection + .publish_commit(commit.catalog) + .map_err(|error| error.to_string())?; Ok(worker_remove_success_response(target)) } Err(error) => Ok(worker_retention_error_response(error)), @@ -834,7 +996,14 @@ impl WorkspaceWorkerRemoveExecutor { .map_err(|_| "Worker removal recovery authority is unavailable".to_string())?; if let Some(prepared) = prepared { if prepared.plan.state == crate::retention::WorkerRemovalPlanState::Succeeded { - return Ok(worker_remove_success_response(&target)); + let commit = self + .store + .recover_succeeded_worker_removal_catalog(&self.workspace_id, target) + .map_err(|error| error.to_string())?; + self.worker_projection + .publish_commit(commit) + .map_err(|error| error.to_string())?; + return Ok(worker_remove_success_response(target)); } let prepared = if matches!( prepared.plan.state, @@ -1023,10 +1192,10 @@ impl WorkspaceWorkerRemoveExecutor { &plan.input_fingerprint, &retention_result, ) { - Ok(_) => { - if let Some(worker_projection) = self.worker_projection.upgrade() { - let _ = worker_projection.publish_removed(plan.worker.clone()); - } + Ok(commit) => { + self.worker_projection + .publish_commit(commit.catalog) + .map_err(|error| error.to_string())?; } Err(error) => { let _ = self.store.fail_worker_removal( @@ -1042,7 +1211,7 @@ impl WorkspaceWorkerRemoveExecutor { } } -impl crate::worker_source::VerifiedWorkerRemoveExecutor for WorkspaceWorkerRemoveExecutor { +impl crate::worker_source::VerifiedWorkerRemoveExecutor for WorkerRemovalService { fn execute( &self, source: crate::worker_source::VerifiedWorkerMutationSource, @@ -1252,7 +1421,7 @@ impl WorkspaceServerApi { { let worker_key = worker.display_name.clone(); let target = worker.worker; - let response = WorkspaceWorkerRemoveExecutor::new(&api) + let response = WorkerRemovalService::new(&api) .execute_target_removal( api.runtime.as_ref(), &target, @@ -2136,31 +2305,35 @@ async fn seed_test_registered_workspace( fn take_new_workdir_repository_access( request: &mut WorkerSpawnRequest, -) -> ApiResult> { - let Some(working_directory) = request.resolved_working_directory_request.as_mut() else { - return Ok(None); - }; - let Some(materialization) = working_directory.materialization.as_mut() else { - return Ok(None); - }; - if materialization.ssh.is_none() { - return Ok(None); +) -> ApiResult> { + let mut accesses = Vec::new(); + for attachment in &mut request.resolved_workdir_attachment_requests { + let working_directory = &mut attachment.working_directory; + let Some(materialization) = working_directory.materialization.as_mut() else { + continue; + }; + if materialization.ssh.is_none() { + continue; + } + let working_directory_id = + working_directory + .backend_workdir_id + .clone() + .ok_or_else(|| { + Error::Config( + "repository access authorization requires a Backend WorkingDirectory id" + .to_string(), + ) + })?; + accesses.push( + worker_runtime::catalog::WorkingDirectoryRepositoryAccessRequest { + working_directory_id, + materialization: materialization.clone(), + }, + ); + materialization.ssh = None; } - let working_directory_id = working_directory - .backend_workdir_id - .clone() - .ok_or_else(|| { - Error::Config( - "repository access authorization requires a Backend WorkingDirectory id" - .to_string(), - ) - })?; - let access = worker_runtime::catalog::WorkingDirectoryRepositoryAccessRequest { - working_directory_id, - materialization: materialization.clone(), - }; - materialization.ssh = None; - Ok(Some(access)) + Ok(accesses) } fn embedded_runtime_request_audience(config: &ServerConfig) -> crate::Result { @@ -2473,7 +2646,7 @@ impl WorkspaceApi { }; if let Some(dispatcher) = worker_remove_dispatcher { dispatcher - .install_executor(Arc::new(WorkspaceWorkerRemoveExecutor::new(&api))) + .install_executor(Arc::new(WorkerRemovalService::new(&api))) .map_err(|message| Error::Config(message.to_string()))?; } recover_workdir_removals(&api)?; @@ -2552,33 +2725,39 @@ impl WorkspaceApi { self.validate_worker_spawn_repository_scope(&request)?; let workspace_api = self.workspace_api_ref(runtime_id); request.resolved_workspace_api = Some(workspace_api.clone()); - if let Some(access) = take_new_workdir_repository_access(&mut request)? { + for access in take_new_workdir_repository_access(&mut request)? { self.runtime .authorize_working_directory_repository_access(runtime_id, access) .map_err(RuntimeRegistryError::into_error)?; } - if let Some(working_directory) = request.resolved_working_directory.as_ref() - && let Some(access) = repository_access_request_for_workdir( + for attachment in &request.resolved_workdir_attachments { + if let Some(access) = repository_access_request_for_workdir( self, runtime_id, - &working_directory.working_directory_id, + &attachment.working_directory_id, &format!("worker-spawn:{}", WorkerId::now_v7()), - )? - { - self.runtime - .authorize_working_directory_repository_access(runtime_id, access) - .map_err(RuntimeRegistryError::into_error)?; + )? { + self.runtime + .authorize_working_directory_repository_access(runtime_id, access) + .map_err(RuntimeRegistryError::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(), - ) - }); + let spawned_workdir_ids = request + .resolved_workdir_attachment_requests + .iter() + .filter_map(|attachment| attachment.working_directory.backend_workdir_id.clone()) + .collect::>(); + let attachment_reservations = request + .resolved_workdir_attachments + .iter() + .map(|attachment| { + ( + attachment.alias.to_string(), + attachment.working_directory_id.clone(), + Uuid::new_v4().to_string(), + ) + }) + .collect::>(); let request_fingerprint = worker_spawn_create_fingerprint(&request) .map_err(|message| Error::Config(message.to_string()))?; let current_memory_settings = self @@ -2618,42 +2797,42 @@ impl WorkspaceApi { }; let compensation_context = WorkerSpawnCompensationContext { assignment: None, - prepared_workdir_id: attachment_reservation - .as_ref() - .map(|(workdir_id, _)| workdir_id.as_str()), - cleanup_spawned_workdir: false, + spawned_workdir_ids: &spawned_workdir_ids, }; - if let Some((workdir_id, reservation_id)) = attachment_reservation.as_ref() - && let Err(error) = self.store.reserve_worker_workdir_attachment( + let mut reserved_attachments: Vec<(String, String, String)> = Vec::new(); + for (alias, workdir_id, reservation_id) in &attachment_reservations { + if let Err(error) = self.store.reserve_worker_workdir_attachment( &self.config.workspace_id, workdir_id, reservation_id, &now_registry_timestamp(), - ) - { - let mut diagnostics = Vec::new(); - if let Err(cleanup_error) = self.config_store.fail_worker_create_reservation( - &self.config.workspace_id, - runtime_id, - worker_id, - &reservation_fingerprint, ) { - diagnostics.push(spawn_compensation_diagnostic( - "worker_spawn_compensation_create_reservation_release_failed", - format!( - "Failed to terminalize Worker create reservation {} and release its resource key: {}", - worker_id, - sanitize_backend_error(&cleanup_error.to_string()) - ), - )); + let mut diagnostics = + release_worker_workdir_attachment_reservations(self, &reserved_attachments); + if let Err(cleanup_error) = self.config_store.fail_worker_create_reservation( + &self.config.workspace_id, + runtime_id, + worker_id, + &reservation_fingerprint, + ) { + diagnostics.push(spawn_compensation_diagnostic( + "worker_spawn_compensation_create_reservation_release_failed", + format!( + "Failed to terminalize Worker create reservation {} and release its resource key: {}", + worker_id, + sanitize_backend_error(&cleanup_error.to_string()) + ), + )); + } + write_workspace_worker_create_failure( + runtime_id, + worker_id, + "workdir_attachment_reserve", + &diagnostics, + ); + return Err(ApiError::with_diagnostics(error, diagnostics)); } - write_workspace_worker_create_failure( - runtime_id, - worker_id, - "workdir_attachment_reserve", - &diagnostics, - ); - return Err(ApiError::with_diagnostics(error, diagnostics)); + reserved_attachments.push((alias.clone(), workdir_id.clone(), reservation_id.clone())); } let mut result = match self .runtime @@ -2669,11 +2848,7 @@ impl WorkspaceApi { "runtime_spawn_transport", None, &compensation_context, - attachment_reservation - .as_ref() - .map(|(workdir_id, reservation_id)| { - (workdir_id.as_str(), reservation_id.as_str()) - }), + &reserved_attachments, ); return Err(ApiError::with_diagnostics(error.into_error(), diagnostics)); } @@ -2689,11 +2864,7 @@ impl WorkspaceApi { "runtime_spawn_rejected", None, &compensation_context, - attachment_reservation - .as_ref() - .map(|(workdir_id, reservation_id)| { - (workdir_id.as_str(), reservation_id.as_str()) - }), + &reserved_attachments, )); return Ok(result); }; @@ -2707,11 +2878,7 @@ impl WorkspaceApi { "runtime_worker_identity", Some(worker), &compensation_context, - attachment_reservation - .as_ref() - .map(|(workdir_id, reservation_id)| { - (workdir_id.as_str(), reservation_id.as_str()) - }), + &reserved_attachments, ); return Err(ApiError::with_diagnostics( Error::RuntimeOperationFailed { @@ -2734,11 +2901,7 @@ impl WorkspaceApi { "runtime_spawn_rejected", Some(worker), &compensation_context, - attachment_reservation - .as_ref() - .map(|(workdir_id, reservation_id)| { - (workdir_id.as_str(), reservation_id.as_str()) - }), + &reserved_attachments, ); result.diagnostics.extend(diagnostics); return Ok(result); @@ -2757,11 +2920,7 @@ impl WorkspaceApi { "workspace_api_transport", Some(worker), &compensation_context, - attachment_reservation - .as_ref() - .map(|(workdir_id, reservation_id)| { - (workdir_id.as_str(), reservation_id.as_str()) - }), + &reserved_attachments, ); return Err(ApiError::with_diagnostics(error.into_error(), diagnostics)); } @@ -2775,11 +2934,7 @@ impl WorkspaceApi { "workspace_api_rejected", Some(worker), &compensation_context, - attachment_reservation - .as_ref() - .map(|(workdir_id, reservation_id)| { - (workdir_id.as_str(), reservation_id.as_str()) - }), + &reserved_attachments, ); return Err(ApiError::with_diagnostics( Error::RuntimeOperationFailed { @@ -2796,58 +2951,53 @@ impl WorkspaceApi { diagnostics, )); } - if let Some((workdir_id, reservation_id)) = attachment_reservation.as_ref() { - if let Err(error) = parse_runtime_worker_id_for_registry(&worker.worker.worker_id) { - let diagnostics = compensate_failed_workspace_worker_create( - self, - runtime_id, - worker_id, - &reservation_fingerprint, - "worker_registry_identity", - Some(worker), - &compensation_context, - Some((workdir_id.as_str(), reservation_id.as_str())), - ); - return Err(api_error_with_additional_diagnostics(error, diagnostics)); - } - let compensation_context = WorkerSpawnCompensationContext { - assignment: None, - prepared_workdir_id: Some(workdir_id.as_str()), - cleanup_spawned_workdir: false, - }; - let registry_result = record_worker_summary( + if let Err(error) = parse_runtime_worker_id_for_registry(&worker.worker.worker_id) { + let diagnostics = compensate_failed_workspace_worker_create( self, - worker, - worker.label.as_str(), - worker.profile.clone(), - WorkerRegistryDisplayNamePolicy::PreserveExisting, - ) - .map(|_| ()); - if let Err(error) = finalize_worker_spawn_stage( - self, - worker, + runtime_id, + worker_id, + &reservation_fingerprint, + "worker_registry_identity", + Some(worker), &compensation_context, - WorkerSpawnFinalizeStage::WorkerRegistry, - registry_result, - ) { - let diagnostics = compensate_failed_workspace_worker_create( - self, - runtime_id, - worker_id, - &reservation_fingerprint, - "worker_registry_finalize", - None, - &compensation_context, - Some((workdir_id.as_str(), reservation_id.as_str())), - ); - return Err(api_error_with_additional_diagnostics(error, diagnostics)); - } + &reserved_attachments, + ); + return Err(api_error_with_additional_diagnostics(error, diagnostics)); + } + let registry_result = record_worker_summary( + self, + worker, + worker.label.as_str(), + worker.profile.clone(), + WorkerRegistryDisplayNamePolicy::PreserveExisting, + ) + .map(|_| ()); + if let Err(error) = finalize_worker_spawn_stage( + self, + worker, + &compensation_context, + WorkerSpawnFinalizeStage::WorkerRegistry, + registry_result, + ) { + let diagnostics = compensate_failed_workspace_worker_create( + self, + runtime_id, + worker_id, + &reservation_fingerprint, + "worker_registry_finalize", + None, + &compensation_context, + &reserved_attachments, + ); + return Err(api_error_with_additional_diagnostics(error, diagnostics)); + } + for (alias, workdir_id, reservation_id) in &reserved_attachments { let attachment = WorkerWorkdirLinkRecord { workspace_id: self.config.workspace_id.clone(), worker: worker_ref.clone(), workdir_id: workdir_id.clone(), - role: "attachment".to_string(), + alias: alias.clone(), linked_at: now_registry_timestamp(), unlinked_at: None, }; @@ -2870,20 +3020,11 @@ impl WorkspaceApi { "workdir_attachment_finalize", None, &compensation_context, - Some((workdir_id.as_str(), reservation_id.as_str())), + &reserved_attachments, ); return Err(api_error_with_additional_diagnostics(error, diagnostics)); } } - if attachment_reservation.is_none() { - record_worker_summary( - self, - worker, - worker.label.as_str(), - worker.profile.clone(), - WorkerRegistryDisplayNamePolicy::PreserveExisting, - )?; - } if let Err(error) = self .config_store .complete_worker_create_reservation(&self.config.workspace_id, worker_id) @@ -2896,7 +3037,7 @@ impl WorkspaceApi { "create_reservation_complete", Some(worker), &compensation_context, - None, + &reserved_attachments, ); return Err(ApiError::with_diagnostics( Error::Config(error.to_string()), @@ -2910,11 +3051,11 @@ impl WorkspaceApi { &self, worker: &RuntimeWorkerRef, ) -> ApiResult { - if let Some(link) = self + for link in self .store .list_worker_workdir_links(&self.config.workspace_id, worker)? .into_iter() - .find(|link| link.unlinked_at.is_none()) + .filter(|link| link.unlinked_at.is_none()) { let workdir_runtime_id = registered_workdir_runtime_id(self, &link.workdir_id)?; if let Some(access) = repository_access_request_for_workdir( @@ -2994,40 +3135,40 @@ impl WorkspaceApi { &self, request: &WorkerSpawnRequest, ) -> ApiResult<()> { - let (selected_repository_id, selected_ref_selector) = - if let Some(working_directory) = request.resolved_working_directory_request.as_ref() { - let repository_id = working_directory.repository.id.as_str(); - self.require_workspace_repository(repository_id)?; - ( - Some(repository_id.to_string()), - working_directory - .repository - .selector - .as_deref() - .map(str::to_owned), - ) - } else if let Some(claim) = request.resolved_working_directory.as_ref() { - let workdir = self - .store - .get_workdir_registry(&self.config.workspace_id, &claim.working_directory_id)? - .ok_or_else(|| { - ApiError::from(Error::Config(format!( - "unknown working directory `{}` in this Workspace", - claim.working_directory_id - ))) - })?; - self.require_workspace_repository(&workdir.repository_id)?; - (Some(workdir.repository_id), workdir.creation_selector) - } else { - (None, None) - }; + let mut selected_repositories = Vec::new(); + for attachment in &request.resolved_workdir_attachment_requests { + let working_directory = &attachment.working_directory; + let repository_id = working_directory.repository.id.as_str(); + self.require_workspace_repository(repository_id)?; + selected_repositories.push(( + repository_id.to_string(), + working_directory + .repository + .selector + .as_deref() + .map(str::to_owned), + )); + } + for claim in &request.resolved_workdir_attachments { + let workdir = self + .store + .get_workdir_registry(&self.config.workspace_id, &claim.working_directory_id)? + .ok_or_else(|| { + ApiError::from(Error::Config(format!( + "unknown working directory `{}` in this Workspace", + claim.working_directory_id + ))) + })?; + self.require_workspace_repository(&workdir.repository_id)?; + selected_repositories.push((workdir.repository_id, workdir.creation_selector)); + } if let WorkerSpawnIntent::TicketRole { ticket_id, .. } = &request.intent { let ticket = self.authority.ticket(ticket_id)?; // Workdir-less Ticket Workers cannot execute repository implementation. // Preserve that control-plane launch while still validating any persisted // target (including its Workspace ownership) when one exists. - if selected_repository_id.is_none() && ticket.repository_key.is_none() { + if selected_repositories.is_empty() && ticket.repository_key.is_none() { return Ok(()); } let repository_key = ticket.repository_key.as_deref().ok_or_else(|| { @@ -3050,16 +3191,18 @@ impl WorkspaceApi { "Ticket implementation target is no longer resolvable: {error:?}" ))) })?; - if selected_repository_id.as_deref() != Some(repository_id) { - return Err(ApiError::from(Error::Config(format!( - "Ticket `{ticket_id}` targets Repository `{repository_key}`, but the Worker launch resolves a different Repository" - )))); - } - if selected_ref_selector.as_deref() != Some(ref_selector) { - return Err(ApiError::from(Error::Config(format!( - "Ticket `{ticket_id}` targets selector `{ref_selector}`, but the Worker launch resolves `{}`", - selected_ref_selector.as_deref().unwrap_or("none") - )))); + for (selected_repository_id, selected_ref_selector) in &selected_repositories { + if selected_repository_id != repository_id { + return Err(ApiError::from(Error::Config(format!( + "Ticket `{ticket_id}` targets Repository `{repository_key}`, but a Worker attachment resolves a different Repository" + )))); + } + if selected_ref_selector.as_deref() != Some(ref_selector) { + return Err(ApiError::from(Error::Config(format!( + "Ticket `{ticket_id}` targets selector `{ref_selector}`, but a Worker attachment resolves `{}`", + selected_ref_selector.as_deref().unwrap_or("none") + )))); + } } } Ok(()) @@ -3614,9 +3757,12 @@ fn build_inner_router(api: WorkspaceApi) -> Router { 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), + "/api/w/{workspace_id}/workers/self/workdir-attachments", + post(scoped_attach_current_worker_workdir), + ) + .route( + "/api/w/{workspace_id}/workers/self/workdir-attachments/{alias}", + delete(scoped_detach_current_worker_workdir), ) .route( "/api/w/{workspace_id}/workers/self/workdir-session/operations", @@ -4221,13 +4367,15 @@ struct PutFlowRequest { #[derive(Debug, Deserialize)] #[serde(deny_unknown_fields)] struct AttachCurrentWorkerWorkdirRequest { - workdir_id: String, + alias: String, + working_directory_id: String, } #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)] struct CurrentWorkerWorkdirAttachmentResponse { workspace_id: String, - workdir_id: String, + alias: String, + working_directory_id: String, attached: bool, } @@ -6766,34 +6914,46 @@ fn require_assigned_workdir_source( .runtime .worker(&assignment.worker) .map_err(|error| error.into_error())?; - let attached_workdir = worker.working_directory.ok_or_else(|| { - Error::MergeRequest(merge_request::MergeRequestError::Conflict( - "merge_request_source_workdir_missing: current Coder has no attached Workdir".into(), - )) - })?; - let workdir = api - .runtime - .working_directory( - &assignment.worker.runtime_id, - &attached_workdir.working_directory_id, - ) - .map_err(|error| error.into_error())? - .working_directory - .ok_or_else(|| { + if worker.workdir_attachments.is_empty() { + return Err( Error::MergeRequest(merge_request::MergeRequestError::Conflict( - "merge_request_source_workdir_unavailable: current Coder Workdir could not be observed" + "merge_request_source_workdir_missing: current Coder has no attached Workdir" .into(), )) - })? - .summary; - validate_assigned_workdir_source(&workdir, repository_id, selector, revision_ref).map_err( - |diagnostic| Error::RuntimeOperationFailed { - runtime_id: assignment.worker.runtime_id.clone(), - code: diagnostic.code, - message: diagnostic.message, - }, - )?; - Ok(()) + .into(), + ); + } + let mut last_diagnostic = None; + for attachment in &worker.workdir_attachments { + let Some(workdir) = api + .runtime + .working_directory( + &assignment.worker.runtime_id, + &attachment.working_directory.summary.working_directory_id, + ) + .map_err(|error| error.into_error())? + .working_directory + .map(|status| status.summary) + else { + continue; + }; + match validate_assigned_workdir_source(&workdir, repository_id, selector, revision_ref) { + Ok(()) => return Ok(()), + Err(diagnostic) => last_diagnostic = Some(diagnostic), + } + } + let diagnostic = last_diagnostic.unwrap_or_else(|| { + worker_runtime::working_directory::WorkingDirectoryDiagnostic { + code: "merge_request_source_workdir_unavailable".to_string(), + message: "No attached Coder Workdir could be observed".to_string(), + } + }); + Err(Error::RuntimeOperationFailed { + runtime_id: assignment.worker.runtime_id.clone(), + code: diagnostic.code, + message: diagnostic.message, + } + .into()) } fn validate_assigned_workdir_source( @@ -8303,56 +8463,6 @@ fn current_worker_identity( authenticate_worker_mutation_source(api, workspace_id, headers) } -fn current_worker_active_attachment( - api: &WorkspaceApi, - worker: &RuntimeWorkerRef, -) -> ApiResult { - if let Some(link) = api - .store - .list_worker_workdir_links(&api.config.workspace_id, worker)? - .into_iter() - .next() - { - return Ok(link); - } - - if api - .store - .worker_workdir_link_history_exists(&api.config.workspace_id, worker)? - { - return Err(Error::WorkdirAttachmentConflict(format!( - "Worker {}:{} has no active Workdir attachment", - worker.runtime_id, worker.worker_id - )) - .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 observed_worker = api - .runtime - .worker(worker) - .map_err(|error| error.into_error())?; - if observed_worker.working_directory.is_some() { - sync_worker_observation(api, &observed_worker)?; - if let Some(link) = api - .store - .list_worker_workdir_links(&api.config.workspace_id, worker)? - .into_iter() - .next() - { - return Ok(link); - } - } - Err(Error::WorkdirAttachmentConflict(format!( - "Worker {}:{} has no active Workdir attachment", - worker.runtime_id, worker.worker_id - )) - .into()) -} - fn current_worker_session_lock( api: &WorkspaceApi, worker: &RuntimeWorkerRef, @@ -8415,13 +8525,11 @@ async fn open_current_worker_workdir_session_locked( api.runtime .authorize_working_directory_repository_access(&workdir.runtime_id, access) .map_err(RuntimeRegistryError::into_error)?; - } else if let Some(session) = api - .workdir_sessions - .lock() - .expect("Workdir session registry lock poisoned") - .attachment(worker) - { - return Ok(session); + } else { + // The registry is an operation-session cache, not routing authority. + // Close any prior alias session before opening the explicitly selected + // attachment so multiple aliases can never share one provider namespace. + close_current_worker_attachment_session_locked(api, worker).await?; } let owner_worker_id = runtime_local_owner_worker_id(worker, &workdir.runtime_id); let session = api @@ -8472,6 +8580,20 @@ async fn close_current_worker_attachment_session_locked( Ok(()) } +async fn close_current_worker_command_sessions_for_alias_locked( + api: &WorkspaceApi, + worker: &RuntimeWorkerRef, + attachment_alias: &str, +) -> Result<()> { + close_worker_workdir_command_sessions_for_alias(&api.workdir_sessions, worker, attachment_alias) + .await + .map_err(|message| Error::RuntimeOperationFailed { + runtime_id: worker.runtime_id.clone(), + code: "workdir_session_close_failed".to_string(), + message, + }) +} + async fn close_current_worker_session_locked( api: &WorkspaceApi, worker: &RuntimeWorkerRef, @@ -8485,6 +8607,60 @@ async fn close_current_worker_session_locked( }) } +fn sync_runtime_worker_workdir_attachments( + api: &WorkspaceApi, + worker: &RuntimeWorkerRef, +) -> ApiResult<()> { + let attachments = api + .store + .list_worker_workdir_links(&api.config.workspace_id, worker)? + .into_iter() + .filter(|link| link.unlinked_at.is_none()) + .map(|link| { + Ok(WorkingDirectoryAttachmentClaim { + alias: workdir::WorkdirAttachmentAlias::new(link.alias) + .map_err(|error| Error::InvalidInput(error.to_string()))?, + working_directory_id: link.workdir_id, + relative_cwd: None, + }) + }) + .collect::>>()?; + let result = api + .runtime + .replace_worker_workdir_attachments(worker, attachments) + .map_err(RuntimeRegistryError::into_error)?; + if result.state != WorkerOperationState::Accepted { + return Err(Error::RuntimeOperationFailed { + runtime_id: worker.runtime_id.clone(), + code: "worker_workdir_attachments_replace_failed".to_string(), + message: result + .diagnostics + .first() + .map(|diagnostic| diagnostic.message.clone()) + .unwrap_or_else(|| "Runtime rejected Workdir attachment persistence".to_string()), + } + .into()); + } + Ok(()) +} + +async fn refresh_current_worker_session_locked( + api: &WorkspaceApi, + worker: &RuntimeWorkerRef, +) -> Result<()> { + close_current_worker_attachment_session_locked(api, worker).await?; + let remaining = api + .store + .list_worker_workdir_links(&api.config.workspace_id, worker)? + .into_iter() + .filter(|link| link.unlinked_at.is_none()) + .collect::>(); + if remaining.len() == 1 { + open_current_worker_workdir_session_locked(api, worker, &remaining[0]).await?; + } + Ok(()) +} + async fn scoped_attach_current_worker_workdir( State(api): State, AxumPath(path): AxumPath, @@ -8493,75 +8669,182 @@ async fn scoped_attach_current_worker_workdir( ) -> ApiResult> { validate_workspace_scope(&api, &path.workspace_id)?; let worker = current_worker_identity(&api, &path.workspace_id, &headers)?; - let workdir_id = request.workdir_id.trim(); + let alias = workdir::WorkdirAttachmentAlias::new(request.alias) + .map_err(|error| Error::InvalidInput(error.to_string()))?; + let workdir_id = request.working_directory_id.trim(); if workdir_id.is_empty() || workdir_id.chars().any(char::is_control) { - return Err(Error::InvalidRecordId(request.workdir_id).into()); + return Err(Error::InvalidRecordId(request.working_directory_id).into()); } - if !api + let workdir = api .store .list_workdir_registry(&api.config.workspace_id, 10_000)? - .iter() - .any(|workdir| workdir.workdir_id == workdir_id) - { - return Err(Error::RuntimeOperationFailed { + .into_iter() + .find(|workdir| workdir.workdir_id == workdir_id) + .ok_or_else(|| Error::RuntimeOperationFailed { runtime_id: worker.runtime_id.clone(), code: "working_directory_not_found".to_string(), message: format!("unknown Workdir `{workdir_id}`"), - } + })?; + let status = runtime_workdir_summary_from_record(&workdir).status; + if status != WorkingDirectoryStatusKind::Active { + return Err(Error::WorkdirAttachmentConflict(format!( + "Workdir `{workdir_id}` is not active ({status})" + )) .into()); } let session_lock = current_worker_session_lock(&api, &worker); let _session_guard = session_lock.lock().await; + if let Some(existing) = api + .store + .list_worker_workdir_links(&api.config.workspace_id, &worker)? + .into_iter() + .find(|link| link.unlinked_at.is_none() && link.alias == alias.as_str()) + { + if existing.workdir_id != workdir_id { + return Err(Error::WorkdirAttachmentConflict(format!( + "Worker {}:{} attachment alias `{alias}` is already bound to Workdir {}", + worker.runtime_id, worker.worker_id, existing.workdir_id + )) + .into()); + } + sync_runtime_worker_workdir_attachments(&api, &worker)?; + api.worker_projection + .refresh(&worker) + .map_err(ApiError::from)?; + return Ok(Json(CurrentWorkerWorkdirAttachmentResponse { + workspace_id: api.config.workspace_id.clone(), + alias: alias.to_string(), + working_directory_id: workdir_id.to_string(), + attached: true, + })); + } + close_current_worker_attachment_session_locked(&api, &worker).await?; let link = api.store.attach_worker_workdir(&WorkerWorkdirLinkRecord { workspace_id: api.config.workspace_id.clone(), worker: worker.clone(), workdir_id: workdir_id.to_string(), - role: "attachment".to_string(), + alias: alias.to_string(), linked_at: Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true), unlinked_at: None, })?; - if let Err(error) = open_current_worker_workdir_session_locked(&api, &worker, &link).await { + if let Err(error) = sync_runtime_worker_workdir_attachments(&api, &worker) { let _ = api.store.detach_worker_workdir( &api.config.workspace_id, &worker, Some(workdir_id), &Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true), ); + return Err(error); + } + let active_count = api + .store + .list_worker_workdir_links(&api.config.workspace_id, &worker)? + .into_iter() + .filter(|link| link.unlinked_at.is_none()) + .count(); + if active_count == 1 + && let Err(error) = open_current_worker_workdir_session_locked(&api, &worker, &link).await + { + let _ = api.store.detach_worker_workdir( + &api.config.workspace_id, + &worker, + Some(workdir_id), + &Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true), + ); + let _ = sync_runtime_worker_workdir_attachments(&api, &worker); return Err(error.into()); } - api.worker_projection - .publish_catalog_change(&worker) - .map_err(ApiError::from)?; + if let Err(error) = api + .worker_projection + .refresh(&worker) + .map_err(ApiError::from) + { + api.store.detach_worker_workdir( + &api.config.workspace_id, + &worker, + Some(workdir_id), + &Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true), + )?; + sync_runtime_worker_workdir_attachments(&api, &worker)?; + refresh_current_worker_session_locked(&api, &worker).await?; + return Err(error); + } Ok(Json(CurrentWorkerWorkdirAttachmentResponse { workspace_id: api.config.workspace_id.clone(), - workdir_id: workdir_id.to_string(), + alias: alias.to_string(), + working_directory_id: workdir_id.to_string(), attached: true, })) } async fn scoped_detach_current_worker_workdir( State(api): State, - AxumPath(path): AxumPath, + AxumPath((workspace_id, alias)): AxumPath<(String, String)>, headers: HeaderMap, ) -> ApiResult> { - validate_workspace_scope(&api, &path.workspace_id)?; - let worker = current_worker_identity(&api, &path.workspace_id, &headers)?; + validate_workspace_scope(&api, &workspace_id)?; + let alias = workdir::WorkdirAttachmentAlias::new(alias) + .map_err(|error| Error::InvalidInput(error.to_string()))?; + let worker = current_worker_identity(&api, &workspace_id, &headers)?; let session_lock = current_worker_session_lock(&api, &worker); let _session_guard = session_lock.lock().await; - let link = current_worker_active_attachment(&api, &worker)?; - close_current_worker_session_locked(&api, &worker).await?; + let link = api + .store + .list_worker_workdir_links(&api.config.workspace_id, &worker)? + .into_iter() + .find(|link| link.unlinked_at.is_none() && link.alias == alias.as_str()) + .ok_or_else(|| { + Error::WorkdirAttachmentConflict(format!( + "Worker {}:{} has no Workdir attachment `{alias}`", + worker.runtime_id, worker.worker_id + )) + })?; + close_current_worker_attachment_session_locked(&api, &worker).await?; + close_current_worker_command_sessions_for_alias_locked(&api, &worker, alias.as_str()).await?; api.store.detach_worker_workdir( &api.config.workspace_id, &worker, Some(&link.workdir_id), &Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true), )?; - api.worker_projection - .publish_catalog_change(&worker) - .map_err(ApiError::from)?; + if let Err(error) = sync_runtime_worker_workdir_attachments(&api, &worker) { + let mut restored = link.clone(); + restored.unlinked_at = None; + api.store.attach_worker_workdir(&restored)?; + return Err(error); + } + let remaining = api + .store + .list_worker_workdir_links(&api.config.workspace_id, &worker)? + .into_iter() + .filter(|link| link.unlinked_at.is_none()) + .collect::>(); + if remaining.len() == 1 + && let Err(error) = + open_current_worker_workdir_session_locked(&api, &worker, &remaining[0]).await + { + let mut restored = link.clone(); + restored.unlinked_at = None; + api.store.attach_worker_workdir(&restored)?; + let _ = sync_runtime_worker_workdir_attachments(&api, &worker); + return Err(error.into()); + } + if let Err(error) = api + .worker_projection + .refresh(&worker) + .map_err(ApiError::from) + { + let mut restored = link.clone(); + restored.unlinked_at = None; + api.store.attach_worker_workdir(&restored)?; + sync_runtime_worker_workdir_attachments(&api, &worker)?; + refresh_current_worker_session_locked(&api, &worker).await?; + return Err(error); + } Ok(Json(CurrentWorkerWorkdirAttachmentResponse { workspace_id: api.config.workspace_id.clone(), - workdir_id: link.workdir_id, + alias: alias.to_string(), + working_directory_id: link.workdir_id, attached: false, })) } @@ -8569,8 +8852,43 @@ async fn scoped_detach_current_worker_workdir( fn validated_current_worker_attachment( api: &WorkspaceApi, worker: &RuntimeWorkerRef, + target_workdir: &str, ) -> ApiResult { - current_worker_active_attachment(api, worker) + let find = || -> Result> { + Ok(api + .store + .list_worker_workdir_links(&api.config.workspace_id, worker)? + .into_iter() + .find(|link| link.unlinked_at.is_none() && link.alias == target_workdir)) + }; + if let Some(link) = find()? { + return Ok(link); + } + + // A Runtime can start the Worker immediately after reserving its local + // bindings, before the outer spawn handler projects them into the Backend. + // Import that exact alias set once; never substitute another attachment. + if !api + .store + .worker_workdir_link_history_exists(&api.config.workspace_id, worker)? + { + let observed_worker = api + .runtime + .worker(worker) + .map_err(|error| error.into_error())?; + if !observed_worker.workdir_attachments.is_empty() { + sync_worker_observation(api, &observed_worker)?; + if let Some(link) = find()? { + return Ok(link); + } + } + } + + Err(Error::WorkdirAttachmentConflict(format!( + "Worker {}:{} has no Workdir attachment alias `{target_workdir}`", + worker.runtime_id, worker.worker_id + )) + .into()) } #[derive(Debug)] @@ -8625,11 +8943,12 @@ async fn scoped_execute_current_worker_workdir_operation( ) -> std::result::Result, WorkdirOperationApiError> { validate_workspace_scope(&api, &path.workspace_id)?; let worker = current_worker_identity(&api, &path.workspace_id, &headers)?; + let target_workdir = request.target_workdir; let result = match request.operation { WorkdirSessionOperation::CommandStart(command) => { let session_lock = current_worker_session_lock(&api, &worker); let _session_guard = session_lock.lock().await; - let link = validated_current_worker_attachment(&api, &worker)?; + let link = validated_current_worker_attachment(&api, &worker, &target_workdir)?; let source = open_current_worker_workdir_session_locked(&api, &worker, &link).await?; let provider_handle = source .start_command(command) @@ -8650,12 +8969,17 @@ async fn scoped_execute_current_worker_workdir_operation( .workdir_sessions .lock() .expect("Workdir session registry lock poisoned") - .register_command(worker.clone(), registered_source, provider_handle); + .register_command( + worker.clone(), + target_workdir.clone(), + registered_source, + provider_handle, + ); WorkdirSessionOperationResult::CommandStart(external_handle) } WorkdirSessionOperation::CommandStatus(external_handle) => { let (session, provider_handle) = - current_worker_command_session(&api, &worker, &external_handle)?; + current_worker_command_session(&api, &worker, &target_workdir, &external_handle)?; session .command_status(provider_handle) .await @@ -8664,7 +8988,7 @@ async fn scoped_execute_current_worker_workdir_operation( } WorkdirSessionOperation::CommandOutput(mut output) => { let (session, provider_handle) = - current_worker_command_session(&api, &worker, &output.handle)?; + current_worker_command_session(&api, &worker, &target_workdir, &output.handle)?; output.handle = provider_handle; session .command_output(output) @@ -8674,7 +8998,7 @@ async fn scoped_execute_current_worker_workdir_operation( } WorkdirSessionOperation::CommandCancel(external_handle) => { let (session, provider_handle) = - current_worker_command_session(&api, &worker, &external_handle)?; + current_worker_command_session(&api, &worker, &target_workdir, &external_handle)?; session .cancel_command(provider_handle) .await @@ -8692,7 +9016,7 @@ async fn scoped_execute_current_worker_workdir_operation( | WorkdirSessionOperation::Grep(_)) => { let session_lock = current_worker_session_lock(&api, &worker); let _session_guard = session_lock.lock().await; - let link = validated_current_worker_attachment(&api, &worker)?; + let link = validated_current_worker_attachment(&api, &worker, &target_workdir)?; let source = open_current_worker_workdir_session_locked(&api, &worker, &link).await?; execute_workdir_session_operation(&source, operation) .await @@ -8705,14 +9029,15 @@ async fn scoped_execute_current_worker_workdir_operation( fn current_worker_command_session( api: &WorkspaceApi, worker: &RuntimeWorkerRef, + target_workdir: &str, external_handle: &CommandHandle, ) -> std::result::Result<(WorkdirSessionHandle, CommandHandle), WorkdirOperationApiError> { - let _link = validated_current_worker_attachment(api, worker)?; + let _link = validated_current_worker_attachment(api, worker, target_workdir)?; let command = api .workdir_sessions .lock() .expect("Workdir session registry lock poisoned") - .command(worker, external_handle) + .command(worker, target_workdir, external_handle) .ok_or_else(|| { WorkdirOperationApiError::Provider(current_worker_workdir_operation_error( worker, @@ -9168,9 +9493,9 @@ fn start_memory_staging_consolidation( profile: profile_selector, ticket_assignment: None, initial_submit, - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -10008,7 +10333,7 @@ async fn scoped_worker_remove_source_boundary( .await { Ok(source) => { - let executor = WorkspaceWorkerRemoveExecutor::new(&api); + let executor = WorkerRemovalService::new(&api); match executor .execute_async( source, @@ -10707,9 +11032,9 @@ async fn scoped_start_workspace_orchestrator( profile: ProfileSelector::Builtin("builtin:orchestrator".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: true, resolved_worker_observation_grants: Vec::new(), @@ -10770,7 +11095,16 @@ fn worker_launch_worker_summary(worker: WorkerSummary) -> WorkerLaunchWorkerSumm can_stop: worker.capabilities.can_stop, can_spawn_followup: worker.capabilities.can_spawn_followup, }, - working_directory: worker.working_directory.map(Into::into), + workdir_attachments: worker + .workdir_attachments + .into_iter() + .map( + |attachment| server_api::RuntimeWorkerWorkdirAttachmentSummary { + alias: attachment.alias.to_string(), + working_directory: attachment.working_directory.summary.into(), + }, + ) + .collect(), diagnostics: worker .diagnostics .into_iter() @@ -11062,9 +11396,24 @@ async fn create_workspace_working_directory( api: &WorkspaceApi, workspace_id: &str, route_runtime_id: Option<&str>, - request: BrowserWorkingDirectoryCreateRequest, + mut request: BrowserWorkingDirectoryCreateRequest, ) -> ApiResult<(StatusCode, Json)> { validate_workspace_scope(api, workspace_id)?; + request.display_name = request + .display_name + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|value| { + if value.len() > 80 || value.chars().any(char::is_control) { + return Err(Error::InvalidInput( + "display_name must contain no control characters and be at most 80 bytes" + .to_string(), + )); + } + Ok(value.to_string()) + }) + .transpose()?; if let (Some(route_runtime_id), Some(request_runtime_id)) = (route_runtime_id, request.runtime_id.as_deref()) && route_runtime_id != request_runtime_id @@ -11134,6 +11483,7 @@ async fn create_workspace_working_directory( &request.repository_key, selector.as_deref(), requested_runtime_id.as_deref(), + working_directory_request.display_name.as_deref(), &working_directory_request.repository.source_fingerprint, working_directory_request.repository.source_revision, ); @@ -11903,19 +12253,22 @@ async fn set_worker_retention( } } let retention_state = if pinned { "pinned" } else { "normal" }; - let changed = api.store.update_worker_retention( + let commit = api.store.update_worker_retention( &api.config.workspace_id, &worker_ref, retention_state, now_registry_timestamp().as_str(), )?; - if !changed { + if commit.changes.is_empty() { return Err(cleanup_api_error( runtime_id.as_str(), "workspace_worker_retention_unknown_worker", "Worker is not known to the Backend registry", )); } + api.worker_projection + .publish_commit(commit) + .map_err(ApiError::from)?; Ok(Json(WorkerRetentionResponse { workspace_id: api.config.workspace_id, worker_ref, @@ -11978,6 +12331,11 @@ fn build_runtime_cleanup_plan( &record.worker, )?; let is_running = live_running_worker_ids.contains(&record.worker); + let is_internal = api + .runtime + .worker(&record.worker) + .ok() + .is_some_and(|worker| worker.singleton_key.is_some()); let pinned = record.retention_state == "pinned"; let blocking_reason = if let Some(assignment) = current_assignment { Some(format!( @@ -11987,6 +12345,8 @@ fn build_runtime_cleanup_plan( )) } else if pinned { Some("worker is pinned".to_string()) + } else if is_internal { + Some("internal service Worker cannot be deleted".to_string()) } else if is_running { Some("worker is running".to_string()) } else { @@ -12196,13 +12556,9 @@ async fn execute_runtime_cleanup( )); } parse_runtime_worker_id_for_registry(&candidate.runtime_worker_id)?; - let session_lock = current_worker_session_lock(api, &worker); - let _session_guard = session_lock.lock().await; - close_current_worker_session_locked(api, &worker).await?; - cleanup_runtime_worker_for_execution(api, runtime_id, candidate)?; - api.store - .delete_worker_registry(&api.config.workspace_id, &worker)?; - drop(_session_guard); + WorkerRemovalService::new(api) + .execute_cleanup_removal(candidate) + .await?; api.workdir_session_locks .lock() .expect("Workdir session lock registry poisoned") @@ -12297,49 +12653,6 @@ async fn execute_runtime_cleanup( }) } -fn cleanup_runtime_worker_for_execution( - api: &WorkspaceApi, - runtime_id: &str, - candidate: &CleanupWorkerCandidate, -) -> ApiResult<()> { - let worker = RuntimeWorkerRef::new(runtime_id, &candidate.runtime_worker_id); - match api.runtime.stop_worker( - &worker, - WorkerLifecycleRequest { - reason: Some("cleanup worker before deletion".to_string()), - ticket_assignment: None, - }, - ) { - Ok(result) if result.state == WorkerOperationState::Accepted => {} - Ok(result) => { - return Err(ApiError::with_diagnostics( - Error::RuntimeOperationFailed { - runtime_id: runtime_id.to_string(), - code: "workspace_cleanup_worker_runtime_stop_rejected".to_string(), - message: "Runtime did not stop selected Worker before deletion".to_string(), - }, - result.diagnostics, - )); - } - Err(RuntimeRegistryError::UnknownWorker { .. }) => return Ok(()), - Err(error) => return Err(error.into_error().into()), - } - - match api.runtime.delete_worker(&worker) { - Ok(result) if result.deleted && result.state == WorkerOperationState::Accepted => Ok(()), - Ok(result) => Err(ApiError::with_diagnostics( - Error::RuntimeOperationFailed { - runtime_id: runtime_id.to_string(), - code: "workspace_cleanup_worker_runtime_delete_rejected".to_string(), - message: "Runtime did not delete selected Worker after stopping it".to_string(), - }, - result.diagnostics, - )), - Err(RuntimeRegistryError::UnknownWorker { .. }) => Ok(()), - Err(error) => Err(error.into_error().into()), - } -} - fn cleanup_api_error(runtime_id: &str, code: &str, message: &str) -> ApiError { Error::RuntimeOperationFailed { runtime_id: runtime_id.to_string(), @@ -14797,6 +15110,7 @@ fn working_directory_request_from_repository( }) .or_else(|| Some(RuntimeRepositorySelector::from("HEAD"))), }, + display_name: None, materializer: MaterializerKind::RuntimeGitClone, backend_workdir_id: None, materialization: None, @@ -14939,7 +15253,7 @@ async fn create_workspace_worker_inner( profile, ticket_assignment, initial_submit, - working_directory, + workdir_attachments, control_operation_id: _, } = request; let config_state = api @@ -14999,17 +15313,23 @@ async fn create_workspace_worker_inner( &initial_submit, ticket_assignment.is_some(), )?; - let selected_working_directory_id = working_directory - .as_ref() - .map(|selection| selection.working_directory_id.clone()); - let resolved_working_directory = working_directory.map(|selection| WorkingDirectoryClaim { - working_directory_id: selection.working_directory_id, - relative_cwd: selection.relative_cwd, - }); - validate_working_directory_claim_for_browser(resolved_working_directory.as_ref())?; - if resolved_working_directory.is_none() { - reject_no_workdir_for_non_embedded_runtime(&runtime_id)?; - } + let resolved_workdir_attachments = workdir_attachments + .into_iter() + .map(|selection| { + let alias = workdir::WorkdirAttachmentAlias::new(selection.alias).map_err(|_| { + settings_bad_request( + "invalid_workdir_attachment_alias", + "attachment alias must be a valid Worker-local routing key", + ) + })?; + Ok(WorkingDirectoryAttachmentClaim { + alias, + working_directory_id: selection.working_directory_id, + relative_cwd: selection.relative_cwd, + }) + }) + .collect::>>()?; + validate_working_directory_claims_for_browser(&resolved_workdir_attachments)?; let (intent, acceptance, ticket_assignment) = browser_worker_spawn_policy(ticket_assignment, &initial_submit)?; let request = WorkerSpawnRequest { @@ -15019,9 +15339,9 @@ async fn create_workspace_worker_inner( profile: profile_selector, ticket_assignment, initial_submit, - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments, resolved_config_bundle, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -15060,7 +15380,6 @@ async fn create_workspace_worker_inner( &api, worker, display_name, - selected_working_directory_id.as_deref(), Vec::new(), Some(assignment), )?)); @@ -15081,7 +15400,6 @@ async fn create_workspace_worker_inner( &api, runtime_id, display_name, - selected_working_directory_id, result, assignment.as_ref(), )?)) @@ -15091,7 +15409,6 @@ fn record_browser_worker_spawn( api: &WorkspaceApi, requested_runtime_id: String, display_name: String, - selected_working_directory_id: Option, result: WorkerSpawnResult, assignment: Option<&WorkerTicketAssignmentRequest>, ) -> ApiResult { @@ -15106,21 +15423,13 @@ fn record_browser_worker_spawn( code: "workspace_worker_create_missing_summary".to_string(), message: "Runtime completed worker creation without returning a Worker summary".to_string(), })?; - browser_worker_response_from_summary( - api, - worker, - display_name, - selected_working_directory_id.as_deref(), - result.diagnostics, - assignment, - ) + browser_worker_response_from_summary(api, worker, display_name, result.diagnostics, assignment) } fn browser_worker_response_from_summary( api: &WorkspaceApi, worker: WorkerSummary, display_name: String, - selected_working_directory_id: Option<&str>, diagnostics: Vec, assignment: Option<&WorkerTicketAssignmentRequest>, ) -> ApiResult { @@ -15134,21 +15443,20 @@ fn browser_worker_response_from_summary( Ok(record) => record, Err(error) => return Err(error.into()), }; + let compensation = WorkerSpawnCompensationContext { + assignment, + spawned_workdir_ids: &[], + }; if let Some(assignment) = assignment { if let Err(error) = api.store.bind_ticket_assignment_operation_worker( &api.config.workspace_id, &assignment.operation_id, &worker.worker.worker_id, ) { - let context = WorkerSpawnCompensationContext { - assignment: Some(assignment), - prepared_workdir_id: selected_working_directory_id, - cleanup_spawned_workdir: false, - }; return finalize_worker_spawn_stage( api, &worker, - &context, + &compensation, WorkerSpawnFinalizeStage::TicketAssignmentBind, Err(error.into()), ); @@ -15159,70 +15467,21 @@ fn browser_worker_response_from_summary( &worker.worker.runtime_id, &worker.worker.worker_id, ) { - let context = WorkerSpawnCompensationContext { - assignment: Some(assignment), - prepared_workdir_id: selected_working_directory_id, - cleanup_spawned_workdir: false, - }; return finalize_worker_spawn_stage( api, &worker, - &context, + &compensation, WorkerSpawnFinalizeStage::TicketAssignmentBind, Err(error.into()), ); } } - if let Some(working_directory) = worker.working_directory.as_ref() { - let workdir_record = - workdir_record_from_summary(api, worker.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, - None, - )?; - } - if let Some(workdir_id) = selected_working_directory_id { - if api - .store - .get_workdir_registry(&api.config.workspace_id, workdir_id)? - .is_none() - { - if let Ok(result) = api - .runtime - .working_directory(worker.worker.runtime_id.as_str(), workdir_id) - .map_err(|err| err.into_error()) - { - if let Some(status) = result.working_directory { - let record = workdir_record_from_summary( - api, - worker.worker.runtime_id.as_str(), - &status.summary, - ); - api.store.upsert_workdir_registry(&record)?; - } - } - } - if api - .store - .get_workdir_registry(&api.config.workspace_id, workdir_id)? - .is_some() - { - link_worker_to_workdir(api, &worker_record, workdir_id, None)?; - } - } + finalize_spawned_worker_workdir_attachments(api, &worker, &worker_record, &compensation)?; if let Some(assignment) = assignment { - let context = WorkerSpawnCompensationContext { - assignment: Some(assignment), - prepared_workdir_id: selected_working_directory_id, - cleanup_spawned_workdir: false, - }; finalize_worker_spawn_stage( api, &worker, - &context, + &compensation, WorkerSpawnFinalizeStage::TicketStateAccept, accept_queued_ticket_after_worker_spawn(api, assignment).map_err(ApiError::from), )?; @@ -15253,6 +15512,70 @@ fn browser_worker_response_from_summary( }) } +fn finalize_spawned_worker_workdir_attachments( + api: &WorkspaceApi, + worker: &WorkerSummary, + worker_record: &WorkerRegistryRecord, + compensation: &WorkerSpawnCompensationContext<'_>, +) -> ApiResult<()> { + let existing_links = finalize_worker_spawn_stage( + api, + worker, + compensation, + WorkerSpawnFinalizeStage::WorkdirAttachment, + api.store + .list_worker_workdir_links(&api.config.workspace_id, &worker_record.worker) + .map_err(ApiError::from), + )?; + for attachment in &worker.workdir_attachments { + let summary = &attachment.working_directory.summary; + let mut workdir_record = + workdir_record_from_summary(api, worker.worker.runtime_id.as_str(), summary); + let previous_record = finalize_worker_spawn_stage( + api, + worker, + compensation, + WorkerSpawnFinalizeStage::WorkdirRegistry, + api.store + .get_workdir_registry(&api.config.workspace_id, &summary.working_directory_id) + .map_err(ApiError::from), + )?; + preserve_workdir_identity_for_corrupted_summary( + &mut workdir_record, + previous_record.as_ref(), + ); + finalize_worker_spawn_stage( + api, + worker, + compensation, + WorkerSpawnFinalizeStage::WorkdirRegistry, + api.store + .upsert_workdir_registry(&workdir_record) + .map_err(ApiError::from), + )?; + if !existing_links.iter().any(|link| { + link.unlinked_at.is_none() + && link.alias == attachment.alias.as_str() + && link.workdir_id == summary.working_directory_id + }) { + finalize_worker_spawn_stage( + api, + worker, + compensation, + WorkerSpawnFinalizeStage::WorkdirAttachment, + link_worker_to_workdir( + api, + worker_record, + attachment.alias.as_str(), + &summary.working_directory_id, + None, + ), + )?; + } + } + Ok(()) +} + async fn scoped_post_internal_runtime_resource_fetch( State(api): State, AxumPath(workspace_id): AxumPath, @@ -15505,26 +15828,6 @@ fn reject_workdir_for_embedded_runtime(runtime_id: &str, has_workdir: bool) -> A )) } -fn reject_no_workdir_for_non_embedded_runtime(runtime_id: &str) -> ApiResult<()> { - if runtime_id == EMBEDDED_WORKER_RUNTIME_ID { - return Ok(()); - } - Err(ApiError::with_diagnostics( - Error::RuntimeOperationFailed { - runtime_id: runtime_id.to_string(), - code: "workspace_worker_workdir_required".to_string(), - message: "Only the embedded Runtime can launch a Worker without a working directory" - .to_string(), - }, - vec![RuntimeDiagnostic { - code: "workspace_worker_workdir_required".to_string(), - severity: DiagnosticSeverity::Error, - message: "Select a working directory for this Runtime, or choose the embedded Runtime for a conversation-only Worker." - .to_string(), - }], - )) -} - async fn list_runtime_workers( State(api): State, AxumPath(runtime_id): AxumPath, @@ -15584,8 +15887,7 @@ impl WorkerSpawnFinalizeStage { struct WorkerSpawnCompensationContext<'a> { assignment: Option<&'a crate::hosts::WorkerTicketAssignmentRequest>, - prepared_workdir_id: Option<&'a str>, - cleanup_spawned_workdir: bool, + spawned_workdir_ids: &'a [String], } fn finalize_worker_spawn_stage( @@ -15674,6 +15976,29 @@ fn api_error_with_additional_diagnostics( error } +fn release_worker_workdir_attachment_reservations( + api: &WorkspaceApi, + reservations: &[(String, String, String)], +) -> Vec { + let mut diagnostics = Vec::new(); + for (_, workdir_id, reservation_id) in reservations { + if let Err(error) = api.store.release_worker_workdir_attachment_reservation( + &api.config.workspace_id, + workdir_id, + reservation_id, + ) { + diagnostics.push(spawn_compensation_diagnostic( + "worker_spawn_compensation_attachment_reservation_release_failed", + format!( + "Failed to release Workdir `{workdir_id}` attachment reservation `{reservation_id}`: {}", + sanitize_backend_error(&error.to_string()) + ), + )); + } + } + diagnostics +} + fn compensate_failed_workspace_worker_create( api: &WorkspaceApi, runtime_id: &str, @@ -15682,46 +16007,46 @@ fn compensate_failed_workspace_worker_create( failure_phase: &str, worker: Option<&WorkerSummary>, context: &WorkerSpawnCompensationContext<'_>, - attachment_reservation: Option<(&str, &str)>, + attachment_reservations: &[(String, String, String)], ) -> Vec { let mut diagnostics = Vec::new(); let reserved_worker_ref = RuntimeWorkerRef::new(runtime_id.to_string(), reservation_worker_id.to_string()); let mut reserved_worker_absent = false; + let mut backend_compensated = false; if let Some(worker) = worker { let (returned_worker_absent, delete_diagnostics) = - delete_runtime_worker_for_spawn_compensation(api, &worker.worker); + WorkerRemovalService::new(api).delete_spawn_compensation_runtime_worker(&worker.worker); diagnostics.extend(delete_diagnostics); if returned_worker_absent { diagnostics.extend(finalize_spawn_compensation_after_worker_delete( - api, worker, context, + api, + &worker.worker, + context, )); + backend_compensated = true; } if worker.worker == reserved_worker_ref { reserved_worker_absent = returned_worker_absent; } } if !reserved_worker_absent { - let (worker_absent, delete_diagnostics) = - delete_runtime_worker_for_spawn_compensation(api, &reserved_worker_ref); + let (worker_absent, delete_diagnostics) = WorkerRemovalService::new(api) + .delete_spawn_compensation_runtime_worker(&reserved_worker_ref); reserved_worker_absent = worker_absent; diagnostics.extend(delete_diagnostics); } - if let Some((workdir_id, reservation_id)) = attachment_reservation - && let Err(error) = api.store.release_worker_workdir_attachment_reservation( - &api.config.workspace_id, - workdir_id, - reservation_id, - ) - { - diagnostics.push(spawn_compensation_diagnostic( - "worker_spawn_compensation_attachment_reservation_release_failed", - format!( - "Failed to release Workdir `{workdir_id}` attachment reservation `{reservation_id}`: {}", - sanitize_backend_error(&error.to_string()) - ), + if reserved_worker_absent && !backend_compensated { + diagnostics.extend(finalize_spawn_compensation_after_worker_delete( + api, + &reserved_worker_ref, + context, )); } + diagnostics.extend(release_worker_workdir_attachment_reservations( + api, + attachment_reservations, + )); if reserved_worker_absent { if let Err(error) = api.config_store.fail_worker_create_reservation( &api.config.workspace_id, @@ -15762,84 +16087,21 @@ fn compensate_failed_worker_spawn( context: &WorkerSpawnCompensationContext<'_>, ) -> Vec { let (runtime_deleted, mut diagnostics) = - delete_runtime_worker_for_spawn_compensation(api, &worker.worker); + WorkerRemovalService::new(api).delete_spawn_compensation_runtime_worker(&worker.worker); if !runtime_deleted { return diagnostics; } diagnostics.extend(finalize_spawn_compensation_after_worker_delete( - api, worker, context, + api, + &worker.worker, + context, )); diagnostics } -fn delete_runtime_worker_for_spawn_compensation( - api: &WorkspaceApi, - worker_ref: &RuntimeWorkerRef, -) -> (bool, Vec) { - let mut diagnostics = Vec::new(); - let lifecycle_request = WorkerLifecycleRequest { - reason: Some("Backend spawn finalize failed; compensating Runtime Worker".to_string()), - ticket_assignment: None, - }; - let cancellation = api - .runtime - .cancel_worker(worker_ref, lifecycle_request.clone()); - let stop = api.runtime.stop_worker(worker_ref, lifecycle_request); - let stop_accepted = stop - .as_ref() - .is_ok_and(|result| result.state == WorkerOperationState::Accepted); - let termination_detail = (!stop_accepted).then(|| { - let cancellation = lifecycle_failure_detail("cancel", &cancellation); - let stop = lifecycle_failure_detail("stop", &stop); - format!("{cancellation}; {stop}") - }); - - let runtime_deleted = match api.runtime.delete_worker(worker_ref) { - Ok(result) if result.state == WorkerOperationState::Accepted && result.deleted => true, - Ok(result) => { - let mut message = format!( - "Runtime did not delete Worker {}:{}: state={:?}, deleted={}", - worker_ref.runtime_id, worker_ref.worker_id, result.state, result.deleted - ); - if let Some(detail) = termination_detail.as_deref() { - message.push_str(&format!("; cancellation: {detail}")); - } - if !result.diagnostics.is_empty() { - message.push_str(&format!( - "; delete diagnostics: {}", - runtime_diagnostics_message(&result.diagnostics) - )); - } - diagnostics.push(spawn_compensation_diagnostic( - "worker_spawn_compensation_runtime_delete_failed", - message, - )); - false - } - Err(RuntimeRegistryError::UnknownWorker { .. }) => true, - Err(error) => { - let mut message = format!( - "Failed to delete Runtime Worker {}:{}: {}", - worker_ref.runtime_id, - worker_ref.worker_id, - error.message() - ); - if let Some(detail) = termination_detail.as_deref() { - message.push_str(&format!("; cancellation: {detail}")); - } - diagnostics.push(spawn_compensation_diagnostic( - "worker_spawn_compensation_runtime_delete_failed", - message, - )); - false - } - }; - (runtime_deleted, diagnostics) -} - fn finalize_spawn_compensation_after_worker_delete( api: &WorkspaceApi, - worker: &WorkerSummary, + worker_ref: &RuntimeWorkerRef, context: &WorkerSpawnCompensationContext<'_>, ) -> Vec { let mut diagnostics = Vec::new(); @@ -15853,61 +16115,55 @@ fn finalize_spawn_compensation_after_worker_delete( format!( "Failed to roll back Ticket assignment operation `{}` for Worker {}:{}: {}", assignment.operation_id, - worker.worker.runtime_id, - worker.worker.worker_id, + worker_ref.runtime_id, + worker_ref.worker_id, sanitize_backend_error(&error.to_string()) ), )); } } - if let Err(error) = api - .store - .delete_worker_registry(&api.config.workspace_id, &worker.worker) - { + if let Err(error) = WorkerRemovalService::new(api).commit_catalog_removal(worker_ref) { diagnostics.push(spawn_compensation_diagnostic( "worker_spawn_compensation_registry_delete_failed", format!( "Failed to remove Backend Worker registry for {}:{}: {}", - worker.worker.runtime_id, - worker.worker.worker_id, + worker_ref.runtime_id, + worker_ref.worker_id, sanitize_backend_error(&error.to_string()) ), )); } - if context.cleanup_spawned_workdir { - if let Some(workdir_id) = context.prepared_workdir_id { - match execute_workdir_removal( - api, - workdir_id, - "backend:worker_spawn_compensation", - "remove Workdir created by rejected Worker spawn", - ) { - Ok(result) - if result.disposition == WorkingDirectoryRemovalDisposition::Removed => {} - Ok(result) => diagnostics.push(spawn_compensation_diagnostic( - "worker_spawn_compensation_workdir_cleanup_failed", - format!( - "Durable removal retained spawn-created Workdir `{workdir_id}` for Worker {}:{}: disposition={:?}, retryable={}", - worker.worker.runtime_id, - worker.worker.worker_id, - result.disposition, - result.retryable, - ), - )), - Err(error) => diagnostics.push(spawn_compensation_diagnostic( - "worker_spawn_compensation_workdir_cleanup_failed", - format!( - "Failed to reserve durable removal for spawn-created Workdir `{workdir_id}` for Worker {}:{}: {}", - worker.worker.runtime_id, - worker.worker.worker_id, - sanitize_backend_error(&error.to_string()) - ), - )), - } + for workdir_id in context.spawned_workdir_ids { + match execute_workdir_removal( + api, + workdir_id, + "backend:worker_spawn_compensation", + "remove Workdir created by rejected Worker spawn", + ) { + Ok(result) if result.disposition == WorkingDirectoryRemovalDisposition::Removed => {} + Ok(result) => diagnostics.push(spawn_compensation_diagnostic( + "worker_spawn_compensation_workdir_cleanup_failed", + format!( + "Durable removal retained spawn-created Workdir `{workdir_id}` for Worker {}:{}: disposition={:?}, retryable={}", + worker_ref.runtime_id, + worker_ref.worker_id, + result.disposition, + result.retryable, + ), + )), + Err(error) => diagnostics.push(spawn_compensation_diagnostic( + "worker_spawn_compensation_workdir_cleanup_failed", + format!( + "Failed to reserve durable removal for spawn-created Workdir `{workdir_id}` for Worker {}:{}: {}", + worker_ref.runtime_id, + worker_ref.worker_id, + sanitize_backend_error(&error.to_string()) + ), + )), } } - if let Ok(worker_id) = parse_runtime_worker_id_for_registry(&worker.worker.worker_id) + if let Ok(worker_id) = parse_runtime_worker_id_for_registry(&worker_ref.worker_id) && let Err(error) = api .config_store .release_removed_worker_resource_key(&api.config.workspace_id, worker_id) @@ -15916,7 +16172,7 @@ fn finalize_spawn_compensation_after_worker_delete( "worker_spawn_compensation_resource_key_release_failed", format!( "Failed to release removed Workspace Worker resource key {}: {}", - worker.worker.worker_id, + worker_ref.worker_id, sanitize_backend_error(&error.to_string()) ), )); @@ -15986,16 +16242,22 @@ async fn create_runtime_worker( validate_ticket_assignment_spawn(&api, &runtime_id, &request)?; reject_workdir_for_embedded_runtime( &runtime_id, - request.working_directory_request.is_some() || request.resolved_working_directory.is_some(), + !request.workdir_attachment_requests.is_empty() + || !request.resolved_workdir_attachments.is_empty(), )?; - if request.working_directory_request.is_none() && request.resolved_working_directory.is_none() { - reject_no_workdir_for_non_embedded_runtime(&runtime_id)?; - } - request.resolved_working_directory_request = request - .working_directory_request - .as_ref() - .map(|working_directory| configured_working_directory_request(&api, working_directory)) - .transpose()?; + request.resolved_workdir_attachment_requests = request + .workdir_attachment_requests + .iter() + .map(|attachment| { + Ok(WorkingDirectoryAttachmentRequest { + alias: attachment.alias.clone(), + working_directory: configured_working_directory_request( + &api, + &attachment.working_directory, + )?, + }) + }) + .collect::>>()?; let repository_operation_id = request .resolved_control_operation .as_ref() @@ -16006,11 +16268,9 @@ async fn create_runtime_worker( .as_ref() .map(|assignment| assignment.operation_id.clone()) }); - let prepared_workdir_id = if let Some(working_directory_request) = - request.resolved_working_directory_request.as_mut() - { + for attachment in &mut request.resolved_workdir_attachment_requests { let workdir_id = - upsert_pending_backend_workdir(&api, &runtime_id, working_directory_request)?; + upsert_pending_backend_workdir(&api, &runtime_id, &mut attachment.working_directory)?; let operation_id = repository_operation_id .clone() .unwrap_or_else(|| format!("worker-spawn-workdir:{workdir_id}")); @@ -16018,15 +16278,9 @@ async fn create_runtime_worker( &api, &runtime_id, &operation_id, - working_directory_request, + &mut attachment.working_directory, )?; - Some(workdir_id) - } else { - request - .resolved_working_directory - .as_ref() - .map(|claim| claim.working_directory_id.clone()) - }; + } let requested_worker_name = request.requested_worker_name.clone(); let spawn_idempotency = crate::hosts::worker_spawn_idempotency(&request).map_err(Error::Config)?; @@ -16043,13 +16297,16 @@ async fn create_runtime_worker( &Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), )?; } - let cleanup_spawned_workdir = request.resolved_working_directory_request.is_some(); + let created_workdir_ids = request + .resolved_workdir_attachment_requests + .iter() + .filter_map(|attachment| attachment.working_directory.backend_workdir_id.clone()) + .collect::>(); let result = api.spawn_workspace_worker(&runtime_id, request)?; if let Some(worker) = result.worker.as_ref() { let compensation = WorkerSpawnCompensationContext { assignment: lifecycle_assignment.as_ref(), - prepared_workdir_id: prepared_workdir_id.as_deref(), - cleanup_spawned_workdir, + spawned_workdir_ids: &created_workdir_ids, }; let display_name = requested_worker_name .as_deref() @@ -16104,38 +16361,18 @@ async fn create_runtime_worker( accept_queued_ticket_after_worker_spawn(&api, assignment).map_err(ApiError::from), )?; } - if worker.working_directory.is_none() { - if let Some(workdir_id) = prepared_workdir_id.as_deref() { - let workdir_exists = finalize_worker_spawn_stage( - &api, - worker, - &compensation, - WorkerSpawnFinalizeStage::WorkdirRegistry, - api.store - .get_workdir_registry(&api.config.workspace_id, workdir_id) - .map_err(ApiError::from), - )? - .is_some(); - if workdir_exists { - finalize_worker_spawn_stage( - &api, - worker, - &compensation, - WorkerSpawnFinalizeStage::WorkdirAttachment, - link_worker_to_workdir(&api, &record, workdir_id, None), - )?; - } + finalize_spawned_worker_workdir_attachments(&api, worker, &record, &compensation)?; + } else { + for workdir_id in &created_workdir_ids { + if let Some(mut record) = api + .store + .get_workdir_registry(&api.config.workspace_id, workdir_id)? + { + record.materialization_status = "failed".to_string(); + record.updated_at = now_registry_timestamp(); + api.store.upsert_workdir_registry(&record)?; } } - } else if let Some(workdir_id) = prepared_workdir_id.as_deref() { - if let Some(mut record) = api - .store - .get_workdir_registry(&api.config.workspace_id, workdir_id)? - { - record.materialization_status = "failed".to_string(); - record.updated_at = now_registry_timestamp(); - api.store.upsert_workdir_registry(&record)?; - } } Ok(Json(result)) } @@ -16760,19 +16997,25 @@ fn project_workspace_worker( summary.worker.worker_id )) })?; - let working_directory = summary - .working_directory - .as_ref() - .map(|working_directory| { - let record = - workdir_record_from_summary(api, &summary.worker.runtime_id, working_directory); - projected_workdir_summary_from_record(api, &record) + let workdir_attachments = summary + .workdir_attachments + .iter() + .map(|attachment| { + let record = workdir_record_from_summary( + api, + &summary.worker.runtime_id, + &attachment.working_directory.summary, + ); + Ok(server_api::WorkerWorkdirAttachmentSummary { + alias: attachment.alias.to_string(), + working_directory: projected_workdir_summary_from_record(api, &record)?.into(), + }) }) - .transpose()?; + .collect::>>()?; Ok(workspace_worker_summary( summary, resource_key, - working_directory.map(Into::into), + workdir_attachments, )) } @@ -17467,7 +17710,7 @@ fn available_working_directory_summaries( { continue; } - if summary.occupied_by.is_none() && summary.primary_worker_id.is_none() { + if summary.occupied_by.is_none() { available.push(summary); } } @@ -17554,7 +17797,10 @@ fn record_worker_summary( created_at: timestamp.clone(), updated_at: timestamp, }; - api.store.upsert_worker_registry(&record)?; + let catalog_commit = api.store.upsert_worker_registry(&record)?; + api.worker_projection + .publish_commit(catalog_commit) + .map_err(ApiError::from)?; if let Ok(worker_id) = protocol::subscription::SubscriptionWorkerId::new(worker.worker.worker_id.clone()) { @@ -17576,9 +17822,7 @@ fn record_worker_summary( workspace_id: worker.workspace.workspace_id.clone(), display_name: Some(worker.display_name.clone()), profile: worker.profile.clone(), - repository_id: None, - repository_key: None, - working_directory_id: None, + workdir_attachments: Vec::new(), }; api.worker_projection .seed_observation( @@ -17590,10 +17834,6 @@ fn record_worker_summary( .unwrap_or(record.updated_at.as_str()), ) .map_err(ApiError::from)?; - } else { - api.worker_projection - .publish_catalog_change(&worker_ref) - .map_err(ApiError::from)?; } Ok(api .store @@ -17628,7 +17868,7 @@ fn worker_summary_from_registry(record: &WorkerRegistryRecord) -> WorkerSummary kind: "backend_worker_registry".to_string(), display_hint: "Missing Worker".to_string(), }, - working_directory: None, + workdir_attachments: Vec::new(), diagnostics: vec![RuntimeDiagnostic { code: "backend_worker_missing".to_string(), severity: DiagnosticSeverity::Info, @@ -17691,21 +17931,30 @@ fn merge_worker_registry_projection( summary.profile = record.profile.clone(); summary.pinned = record.retention_state == "pinned"; summary.retention_state = record.retention_state.clone(); - summary.working_directory = links.iter().find_map(|link| { - workdirs - .iter() - .find(|workdir| workdir.workdir_id == link.workdir_id) - .map(|workdir| { - let mut workdir_summary = runtime_workdir_summary_from_record(workdir); - workdir_summary.occupied_by = Some(WorkingDirectoryOccupancy { - runtime_id: record.worker.runtime_id.clone(), - worker_id: record.worker.worker_id.clone(), - display_name: record.display_name.clone(), - linked_at: link.linked_at.clone(), - }); - workdir_summary - }) - }); + summary.workdir_attachments = links + .iter() + .filter_map(|link| { + workdirs + .iter() + .find(|workdir| workdir.workdir_id == link.workdir_id) + .map(|workdir| { + let mut workdir_summary = runtime_workdir_summary_from_record(workdir); + workdir_summary.occupied_by = Some(WorkingDirectoryOccupancy { + runtime_id: record.worker.runtime_id.clone(), + worker_id: record.worker.worker_id.clone(), + display_name: record.display_name.clone(), + linked_at: link.linked_at.clone(), + }); + WorkingDirectoryAttachmentStatus { + alias: workdir::WorkdirAttachmentAlias::new(link.alias.clone()) + .expect("persisted attachment aliases are validated on write"), + working_directory: worker_runtime::catalog::WorkingDirectoryStatus { + summary: workdir_summary, + }, + } + }) + }) + .collect(); summary } @@ -17720,11 +17969,27 @@ fn sync_worker_observation( worker.profile.clone(), WorkerRegistryDisplayNamePolicy::PreserveExisting, )?; - if let Some(working_directory) = worker.working_directory.as_ref() { + for attachment in &worker.workdir_attachments { + let working_directory = &attachment.working_directory.summary; let workdir_record = workdir_record_from_summary(api, worker.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, None)?; + let links = api + .store + .list_worker_workdir_links(&api.config.workspace_id, &record.worker)?; + if !links.iter().any(|link| { + link.unlinked_at.is_none() + && link.alias == attachment.alias.as_str() + && link.workdir_id == working_directory.working_directory_id + }) { + link_worker_to_workdir( + api, + &record, + attachment.alias.as_str(), + &working_directory.working_directory_id, + None, + )?; + } } Ok(record) } @@ -17743,6 +18008,7 @@ fn upsert_pending_backend_workdir( api.store.upsert_workdir_registry(&WorkdirRegistryRecord { workspace_id: api.config.workspace_id.clone(), workdir_id: workdir_id.clone(), + display_name: request.display_name.clone(), runtime_id: runtime_id.to_string(), repository_id: request.repository.id.clone(), creation_selector: request @@ -17936,6 +18202,7 @@ fn workdir_record_from_summary( WorkdirRegistryRecord { workspace_id: api.config.workspace_id.clone(), workdir_id: summary.working_directory_id.clone(), + display_name: summary.display_name.clone(), runtime_id: runtime_id.to_string(), repository_id: summary.repository_id.clone(), creation_selector: summary.creation_selector.clone(), @@ -18010,6 +18277,7 @@ fn runtime_workdir_summary_from_record( }; worker_runtime::catalog::WorkingDirectorySummary { working_directory_id: record.workdir_id.clone(), + display_name: record.display_name.clone(), repository_id: record.repository_id.clone(), creation_selector: record.creation_selector.clone(), creation_ref: record.creation_ref.clone(), @@ -18026,7 +18294,6 @@ fn runtime_workdir_summary_from_record( }), status, cleanliness: Some(record.cleanliness.clone()), - primary_worker_id: None, occupied_by: None, } } @@ -18045,6 +18312,7 @@ fn workdir_summary_from_record( }; WorkingDirectorySummary { working_directory_id: record.workdir_id.clone(), + display_name: record.display_name.clone(), repository_key: repository_key.to_string(), creation_selector: record.creation_selector.clone(), creation_ref: record.creation_ref.clone(), @@ -18061,7 +18329,6 @@ fn workdir_summary_from_record( }), status, cleanliness: Some(record.cleanliness.clone()), - primary_worker_id: None, occupied_by: None, } } @@ -18074,7 +18341,6 @@ fn apply_workdir_occupancy_projection( .store .list_workdir_worker_links(&api.config.workspace_id, &summary.working_directory_id)?; let Some(link) = links.first() else { - summary.primary_worker_id = None; summary.occupied_by = None; return Ok(()); }; @@ -18088,7 +18354,6 @@ fn apply_workdir_occupancy_projection( link.workdir_id, link.worker.runtime_id, link.worker.worker_id )) })?; - summary.primary_worker_id = None; summary.occupied_by = Some(WorkingDirectoryOccupancy { runtime_id: link.worker.runtime_id.clone(), worker_id: link.worker.worker_id.clone(), @@ -18114,6 +18379,7 @@ fn projected_workdir_summary_from_record( fn link_worker_to_workdir( api: &WorkspaceApi, worker_record: &WorkerRegistryRecord, + alias: &str, workdir_id: &str, reservation_id: Option<&str>, ) -> ApiResult<()> { @@ -18122,7 +18388,7 @@ fn link_worker_to_workdir( workspace_id: api.config.workspace_id.clone(), worker: worker_record.worker.clone(), workdir_id: workdir_id.to_string(), - role: "attachment".to_string(), + alias: alias.to_string(), linked_at: timestamp, unlinked_at: None, }; @@ -18133,37 +18399,50 @@ fn link_worker_to_workdir( api.store.attach_worker_workdir(&record)?; } api.worker_projection - .publish_catalog_change(&worker_record.worker) + .refresh(&worker_record.worker) .map_err(ApiError::from)?; Ok(()) } -fn validate_working_directory_claim_for_browser( - claim: Option<&WorkingDirectoryClaim>, +fn validate_working_directory_claims_for_browser( + claims: &[WorkingDirectoryAttachmentClaim], ) -> ApiResult<()> { - let Some(claim) = claim else { - return Ok(()); - }; - if let Some(relative_cwd) = claim.relative_cwd.as_deref() { - let path = Path::new(relative_cwd); - if path.is_absolute() - || path - .components() - .any(|component| !matches!(component, Component::CurDir | Component::Normal(_))) - { - return Err(ApiError::with_diagnostics( - Error::RuntimeOperationFailed { - runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), - code: "working_directory_relative_cwd_invalid".to_string(), - message: "working directory relative_cwd must stay inside the Runtime-owned working directory".to_string(), - }, - vec![RuntimeDiagnostic { - code: "working_directory_relative_cwd_invalid".to_string(), - severity: DiagnosticSeverity::Error, - message: "relative_cwd must be a relative path without parent traversal".to_string(), - }], + let mut aliases = HashSet::new(); + let mut workdir_ids = HashSet::new(); + for claim in claims { + if !aliases.insert(claim.alias.as_str()) { + return Err(settings_bad_request( + "duplicate_workdir_attachment_alias", + "attachment aliases must be unique within one Worker", )); } + if !workdir_ids.insert(claim.working_directory_id.as_str()) { + return Err(settings_bad_request( + "duplicate_workdir_attachment", + "a Workdir cannot be attached more than once to one Worker", + )); + } + if let Some(relative_cwd) = claim.relative_cwd.as_deref() { + let path = Path::new(relative_cwd); + if path.is_absolute() + || path + .components() + .any(|component| !matches!(component, Component::CurDir | Component::Normal(_))) + { + return Err(ApiError::with_diagnostics( + Error::RuntimeOperationFailed { + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + code: "working_directory_relative_cwd_invalid".to_string(), + message: "working directory relative_cwd must stay inside the Runtime-owned working directory".to_string(), + }, + vec![RuntimeDiagnostic { + code: "working_directory_relative_cwd_invalid".to_string(), + severity: DiagnosticSeverity::Error, + message: "relative_cwd must be a relative path without parent traversal".to_string(), + }], + )); + } + } } Ok(()) } @@ -18670,6 +18949,21 @@ fn working_directory_request_for_browser( request: BrowserWorkingDirectoryCreateRequest, ) -> ApiResult { let repository = api.require_configured_workspace_repository_by_key(&request.repository_key)?; + let display_name = request + .display_name + .as_deref() + .map(str::trim) + .filter(|name| !name.is_empty()) + .map(str::to_string); + if display_name + .as_deref() + .is_some_and(|name| name.len() > 120 || name.chars().any(char::is_control)) + { + return Err(settings_bad_request( + "working_directory_display_name_invalid", + "display_name must be at most 120 bytes and contain no control characters", + )); + } let selector = request .selector .or_else(|| repository.default_selector.clone()) @@ -18683,6 +18977,7 @@ fn working_directory_request_for_browser( source_fingerprint: repository.source_fingerprint.clone(), selector: selector.map(RuntimeRepositorySelector), }, + display_name, materializer: MaterializerKind::RuntimeGitClone, backend_workdir_id: None, materialization: None, @@ -19274,6 +19569,234 @@ mod tests { } } + #[derive(Default)] + struct WorkdirlessFixtureRuntime { + workers: Mutex>, + } + + impl WorkdirlessFixtureRuntime { + const RUNTIME_ID: &'static str = "workdirless-runtime"; + + fn workdir_status( + working_directory_id: &str, + ) -> worker_runtime::catalog::WorkingDirectoryStatus { + worker_runtime::catalog::WorkingDirectoryStatus { + summary: worker_runtime::catalog::WorkingDirectorySummary { + working_directory_id: working_directory_id.to_string(), + display_name: None, + repository_id: "repo-test".to_string(), + creation_selector: Some("HEAD".to_string()), + creation_ref: None, + creation_tree: None, + current_selector: Some("HEAD".to_string()), + current_ref: None, + current_tree: None, + observed_at_epoch_seconds: None, + materializer_kind: + workdir::workspace::WorkingDirectoryMaterializerKind::RuntimeGitClone, + cleanup_target: None, + status: WorkingDirectoryStatusKind::Active, + cleanliness: Some("clean".to_string()), + occupied_by: None, + }, + } + } + + fn worker_summary( + binding: WorkerCreateBinding, + request: &WorkerSpawnRequest, + ) -> WorkerSummary { + WorkerSummary { + worker: RuntimeWorkerRef::new(Self::RUNTIME_ID, binding.worker_id.to_string()), + host_id: "workdirless-host".to_string(), + display_name: request + .requested_worker_name + .clone() + .unwrap_or_else(|| "Workdirless Worker".to_string()), + label: request + .requested_worker_name + .clone() + .unwrap_or_else(|| "Workdirless Worker".to_string()), + profile: match &request.profile { + ProfileSelector::Builtin(profile) | ProfileSelector::Named(profile) => { + Some(profile.clone()) + } + }, + singleton_key: None, + tags: Vec::new(), + workspace: WorkerWorkspaceSummary { + visibility: "workspace".to_string(), + identity: TEST_WORKSPACE_ID.to_string(), + workspace_id: Some(TEST_WORKSPACE_ID.to_string()), + }, + state: "idle".to_string(), + worker_state: None, + last_seen_at: None, + pinned: false, + retention_state: "normal".to_string(), + implementation: WorkerImplementationSummary { + kind: "fixture".to_string(), + display_hint: "Workdirless fixture".to_string(), + }, + capabilities: WorkerCapabilitySummary { + can_stop: true, + can_spawn_followup: false, + }, + workdir_attachments: Vec::new(), + diagnostics: Vec::new(), + } + } + } + + impl crate::hosts::WorkspaceWorkerRuntime for WorkdirlessFixtureRuntime { + fn runtime_id(&self) -> &str { + Self::RUNTIME_ID + } + + fn runtime_summary(&self, _limit: usize) -> crate::hosts::RuntimeSummary { + crate::hosts::RuntimeSummary { + runtime_id: Self::RUNTIME_ID.to_string(), + label: "Workdirless Runtime".to_string(), + kind: "fixture".to_string(), + status: "active".to_string(), + source: crate::hosts::RuntimeSourceSummary::remote_http(), + host_ids: vec!["workdirless-host".to_string()], + worker_creation_available: true, + os: "linux".to_string(), + arch: "x86_64".to_string(), + diagnostics: Vec::new(), + } + } + + fn list_hosts(&self, _limit: usize) -> crate::hosts::RuntimeList { + crate::hosts::RuntimeList { + items: Vec::new(), + diagnostics: Vec::new(), + } + } + + fn list_workers(&self, limit: usize) -> crate::hosts::RuntimeList { + crate::hosts::RuntimeList { + items: self + .workers + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .iter() + .take(limit) + .cloned() + .collect(), + diagnostics: Vec::new(), + } + } + + fn worker(&self, worker_id: &str) -> crate::hosts::WorkerLookupResult { + crate::hosts::WorkerLookupResult { + worker: self + .workers + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .iter() + .find(|worker| worker.worker.worker_id == worker_id) + .cloned(), + diagnostics: Vec::new(), + } + } + + fn working_directory( + &self, + working_directory_id: &str, + ) -> crate::hosts::RuntimeWorkingDirectoryResult { + crate::hosts::RuntimeWorkingDirectoryResult { + state: WorkerOperationState::Accepted, + working_directory: Some(Self::workdir_status(working_directory_id)), + diagnostics: Vec::new(), + } + } + + fn cleanup_working_directory( + &self, + _working_directory_id: &str, + ) -> crate::hosts::RuntimeWorkingDirectoryResult { + crate::hosts::RuntimeWorkingDirectoryResult { + state: WorkerOperationState::Accepted, + working_directory: None, + diagnostics: Vec::new(), + } + } + + fn spawn_worker( + &self, + binding: WorkerCreateBinding, + request: WorkerSpawnRequest, + ) -> WorkerSpawnResult { + assert!(request.workdir_attachment_requests.is_empty()); + assert!(request.resolved_workdir_attachment_requests.is_empty()); + assert!(request.resolved_workdir_attachments.is_empty()); + let worker = Self::worker_summary(binding, &request); + self.workers + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(worker.clone()); + WorkerSpawnResult { + state: WorkerOperationState::Accepted, + worker: Some(worker), + acceptance_evidence: Vec::new(), + diagnostics: Vec::new(), + } + } + + fn stop_worker( + &self, + worker_id: &str, + _request: WorkerLifecycleRequest, + ) -> WorkerLifecycleResult { + WorkerLifecycleResult { + state: WorkerOperationState::Accepted, + worker: RuntimeWorkerRef::new(Self::RUNTIME_ID, worker_id), + diagnostics: Vec::new(), + } + } + + fn cancel_worker( + &self, + worker_id: &str, + _request: WorkerLifecycleRequest, + ) -> WorkerLifecycleResult { + WorkerLifecycleResult { + state: WorkerOperationState::Accepted, + worker: RuntimeWorkerRef::new(Self::RUNTIME_ID, worker_id), + diagnostics: Vec::new(), + } + } + + fn delete_worker(&self, worker_id: &str) -> crate::hosts::WorkerDeleteResult { + let mut workers = self + .workers + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let previous_len = workers.len(); + workers.retain(|worker| worker.worker.worker_id != worker_id); + crate::hosts::WorkerDeleteResult { + state: WorkerOperationState::Accepted, + worker: RuntimeWorkerRef::new(Self::RUNTIME_ID, worker_id), + deleted: workers.len() != previous_len, + diagnostics: Vec::new(), + } + } + + fn replace_worker_workspace_api( + &self, + worker_id: &str, + _workspace_api: WorkspaceApiRef, + ) -> crate::hosts::WorkerWorkspaceApiResult { + crate::hosts::WorkerWorkspaceApiResult { + state: WorkerOperationState::Accepted, + worker: self.worker(worker_id).worker, + diagnostics: Vec::new(), + } + } + } + #[tokio::test] async fn repository_conflict_rejections_use_the_typed_error_contract() { let response = repository_api_rejection( @@ -19502,6 +20025,7 @@ mod tests { fn assigned_workdir_must_be_clean_and_match_published_source() { let mut workdir = worker_runtime::catalog::WorkingDirectorySummary { working_directory_id: "workdir-1".to_string(), + display_name: Some("Ticket checkout".to_string()), repository_id: "repository-1".to_string(), creation_selector: Some("work/T-549".to_string()), creation_ref: Some("abc123".to_string()), @@ -19515,7 +20039,6 @@ mod tests { cleanup_target: None, status: worker_runtime::catalog::WorkingDirectoryStatusKind::Active, cleanliness: Some("clean".to_string()), - primary_worker_id: Some("worker-1".to_string()), occupied_by: None, }; assert!( @@ -19648,8 +20171,12 @@ mod tests { let mut registry = WorkdirSessionRegistry::default(); registry.insert_attachment(worker.clone(), source.clone()); let registered_source = registry.remove_attachment(&worker).unwrap(); - let external_handle = - registry.register_command(worker.clone(), registered_source, provider_handle.clone()); + let external_handle = registry.register_command( + worker.clone(), + "checkout".to_string(), + registered_source, + provider_handle.clone(), + ); assert_ne!(external_handle, provider_handle); let refreshed: WorkdirSessionHandle = Arc::new(workdir::LocalWorkdirSession::new( @@ -19664,7 +20191,14 @@ mod tests { .await .unwrap(); - let command = registry.command(&worker, &external_handle).unwrap(); + let command = registry + .command(&worker, "checkout", &external_handle) + .unwrap(); + assert!( + registry + .command(&worker, "other", &external_handle) + .is_none() + ); let output = command .source .command_output(workdir::CommandOutputRequest { @@ -19724,6 +20258,73 @@ mod tests { ); } + #[tokio::test] + async fn attachment_scoped_command_cleanup_preserves_other_aliases() { + let left_directory = tempfile::tempdir().unwrap(); + let right_directory = tempfile::tempdir().unwrap(); + let worker = RuntimeWorkerRef::new("runtime-command-alias", "worker-command-alias"); + let make_session = |path: &std::path::Path| -> WorkdirSessionHandle { + Arc::new(workdir::LocalWorkdirSession::new( + manifest::Scope::writable(path).unwrap(), + path.to_path_buf(), + )) + }; + let left = make_session(left_directory.path()); + let right = make_session(right_directory.path()); + let request = || workdir::CommandRequest { + command: "sleep 30".to_string(), + timeout_secs: 60, + output_limit: 4096, + cwd: None, + spill_dir: None, + tool_call_id: None, + }; + let left_provider = left.start_command(request()).await.unwrap(); + let right_provider = right.start_command(request()).await.unwrap(); + + let registry = Arc::new(Mutex::new(WorkdirSessionRegistry::default())); + let (left_external, right_external) = { + let mut sessions = registry.lock().unwrap(); + let left_external = + sessions.register_command(worker.clone(), "left".to_string(), left, left_provider); + let right_external = sessions.register_command( + worker.clone(), + "right".to_string(), + right, + right_provider, + ); + assert!(sessions.command(&worker, "right", &left_external).is_none()); + assert!(sessions.command(&worker, "left", &right_external).is_none()); + (left_external, right_external) + }; + + close_worker_workdir_command_sessions_for_alias(®istry, &worker, "right") + .await + .unwrap(); + + let left_command = { + let sessions = registry.lock().unwrap(); + assert!( + sessions + .command(&worker, "right", &right_external) + .is_none() + ); + sessions.command(&worker, "left", &left_external).unwrap() + }; + assert_eq!( + left_command + .source + .command_status(left_command.provider_handle.clone()) + .await + .unwrap(), + workdir::CommandStatus::Running + ); + + close_worker_workdir_sessions(®istry, &worker) + .await + .unwrap(); + } + fn seed_test_api_token(store: &dyn ControlPlaneStore, suffix: &str) -> String { let account_id = format!("account-{suffix}"); let user_id = format!("user-{suffix}"); @@ -20163,6 +20764,7 @@ mod tests { let workdir = WorkdirRegistryRecord { workspace_id: "workspace-1".to_string(), workdir_id: "0000019a00000000000".to_string(), + display_name: None, runtime_id: "embedded".to_string(), repository_id: "repo".to_string(), creation_selector: Some("develop".to_string()), @@ -20181,7 +20783,7 @@ mod tests { workspace_id: "workspace-1".to_string(), worker: worker.worker.clone(), workdir_id: workdir.workdir_id.clone(), - role: "attachment".to_string(), + alias: "attachment".to_string(), linked_at: "4".to_string(), unlinked_at: None, }; @@ -20189,7 +20791,7 @@ mod tests { let projected = merge_worker_registry_projection(None, &worker, vec![link], &[workdir]); assert_eq!(projected.state, "missing"); - let working_directory = projected.working_directory.as_ref().unwrap(); + let working_directory = &projected.workdir_attachments[0].working_directory.summary; assert_eq!( working_directory.status, WorkingDirectoryStatusKind::NotFound @@ -20204,7 +20806,6 @@ mod tests { let occupied_by = working_directory.occupied_by.as_ref().unwrap(); assert_eq!(occupied_by.runtime_id, "embedded"); assert_eq!(occupied_by.worker_id, "1"); - assert!(working_directory.primary_worker_id.is_none()); let occupancy = serde_json::to_value(occupied_by).unwrap(); assert_eq!(occupancy["runtime_id"], "embedded"); assert_eq!(occupancy["worker_id"], "1"); @@ -20431,6 +21032,7 @@ mod tests { source_fingerprint: source_fingerprint.clone(), selector: None, }, + display_name: None, materializer: Default::default(), backend_workdir_id: Some("workdir-a".to_string()), materialization: None, @@ -20594,12 +21196,15 @@ mod tests { initial_submit: vec![Segment::Flow { selector: "builtin:coder-review".to_string(), }], - working_directory_request: None, - resolved_working_directory_request: Some(working_directory_request_from_repository( - &foreign_repository, - Some("HEAD"), - )), - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: vec![WorkingDirectoryAttachmentRequest { + alias: workdir::WorkdirAttachmentAlias::new("checkout").unwrap(), + working_directory: working_directory_request_from_repository( + &foreign_repository, + Some("HEAD"), + ), + }], + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -20633,10 +21238,11 @@ mod tests { }, ) .unwrap(); - let working_directory = repository_access_launch - .resolved_working_directory_request - .as_mut() - .unwrap(); + let working_directory = &mut repository_access_launch + .resolved_workdir_attachment_requests + .first_mut() + .unwrap() + .working_directory; working_directory.backend_workdir_id = Some("working-directory-1".to_string()); working_directory.materialization = Some(worker_runtime::catalog::RepositoryMaterializationContext { @@ -20670,9 +21276,8 @@ mod tests { ), }); - let access = take_new_workdir_repository_access(&mut repository_access_launch) - .unwrap() - .unwrap(); + let accesses = take_new_workdir_repository_access(&mut repository_access_launch).unwrap(); + let access = accesses.first().unwrap(); assert_eq!(access.working_directory_id, "working-directory-1"); assert!(access.materialization.ssh.is_some()); @@ -20683,9 +21288,9 @@ mod tests { assert!(!serialized_access.contains("known_hosts_entry")); assert!( repository_access_launch - .resolved_working_directory_request - .as_ref() - .and_then(|request| request.materialization.as_ref()) + .resolved_workdir_attachment_requests + .first() + .and_then(|attachment| attachment.working_directory.materialization.as_ref()) .and_then(|materialization| materialization.ssh.as_ref()) .is_none() ); @@ -20852,9 +21457,9 @@ mod tests { initial_submit: vec![Segment::Flow { selector: "builtin:coder-review".to_string(), }], - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -20905,7 +21510,7 @@ mod tests { initial_submit: vec![Segment::Flow { selector: "builtin:coder-review".to_string(), }], - working_directory: None, + workdir_attachments: Vec::new(), control_operation_id: None, }), ) @@ -20961,7 +21566,7 @@ mod tests { initial_submit: vec![Segment::Flow { selector: "builtin:coder-review".to_string(), }], - working_directory: None, + workdir_attachments: Vec::new(), control_operation_id: None, }), ) @@ -20994,7 +21599,7 @@ mod tests { profile: Some("builtin:coder".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory: None, + workdir_attachments: Vec::new(), control_operation_id: None, }), ) @@ -21029,7 +21634,7 @@ mod tests { profile: Some("builtin:coder".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory: None, + workdir_attachments: Vec::new(), control_operation_id: None, }), ) @@ -21240,7 +21845,7 @@ mod tests { profile: None, ticket_assignment: None, initial_submit: Vec::new(), - working_directory: None, + workdir_attachments: Vec::new(), control_operation_id: None, }), ) @@ -21271,7 +21876,7 @@ mod tests { profile: None, ticket_assignment: None, initial_submit: Vec::new(), - working_directory: None, + workdir_attachments: Vec::new(), control_operation_id: Some("control-spawn-retry".to_string()), }; @@ -21350,7 +21955,7 @@ mod tests { profile: Some("builtin:orchestrator".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory: None, + workdir_attachments: Vec::new(), control_operation_id: None, }), ) @@ -21367,7 +21972,7 @@ mod tests { profile: Some("builtin:orchestrator".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory: None, + workdir_attachments: Vec::new(), control_operation_id: None, }), ) @@ -21541,6 +22146,7 @@ mod tests { .upsert_workdir_registry(&WorkdirRegistryRecord { workspace_id: TEST_WORKSPACE_ID.to_string(), workdir_id: "managed".to_string(), + display_name: Some("Managed checkout".to_string()), runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), repository_id: "repo".to_string(), creation_selector: None, @@ -21560,6 +22166,7 @@ mod tests { .upsert_workdir_registry(&WorkdirRegistryRecord { workspace_id: TEST_WORKSPACE_ID.to_string(), workdir_id: "runtime-direct".to_string(), + display_name: None, runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), repository_id: "repo".to_string(), creation_selector: None, @@ -21595,7 +22202,7 @@ mod tests { workspace_id: TEST_WORKSPACE_ID.to_string(), worker: RuntimeWorkerRef::new(EMBEDDED_WORKER_RUNTIME_ID, "7"), workdir_id: "managed".to_string(), - role: "attachment".to_string(), + alias: "attachment".to_string(), linked_at: "3".to_string(), unlinked_at: None, }) @@ -21634,6 +22241,7 @@ mod tests { let workdir = WorkdirRegistryRecord { workspace_id: "workspace-1".to_string(), workdir_id: "runtime-direct".to_string(), + display_name: None, runtime_id: "embedded".to_string(), repository_id: "repo".to_string(), creation_selector: None, @@ -22217,23 +22825,14 @@ mod tests { ), ); } - let working_directory = match request.request.working_directory.as_ref() { - Some(claim) => match self.materializer.bind_working_directory( - &claim.working_directory_id, - claim.relative_cwd.as_deref(), - ) { - Ok(binding) => Some(binding.status()), - Err(diagnostic) => { - return worker_runtime::execution::WorkerExecutionSpawnResult::Rejected( - worker_runtime::execution::WorkerExecutionResult::rejected( - worker_runtime::execution::WorkerExecutionOperation::Spawn, - diagnostic.to_string(), - ), - ); - } - }, - None => None, - }; + let workdir_attachments = request + .workdir_attachments + .iter() + .map(|(alias, binding)| WorkingDirectoryAttachmentStatus { + alias: alias.clone(), + working_directory: binding.status(), + }) + .collect(); self.contexts .lock() .unwrap() @@ -22246,7 +22845,7 @@ mod tests { worker_state: protocol::WorkerStateSnapshot { ..protocol::WorkerStatus::Idle.into() }, - working_directory, + workdir_attachments, } } @@ -22507,6 +23106,7 @@ mod tests { let foreign = WorkdirRegistryRecord { workspace_id: "other-workspace".to_string(), workdir_id: "foreign-workdir".to_string(), + display_name: None, runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), repository_id: "foreign-repository".to_string(), creation_selector: None, @@ -24059,9 +24659,9 @@ mod tests { profile: ProfileSelector::Builtin(MEMORY_CONSOLIDATION_PROFILE.to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -24288,9 +24888,9 @@ mod tests { profile: ProfileSelector::Builtin("builtin:coder".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -24633,9 +25233,9 @@ mod tests { profile: ProfileSelector::Builtin("builtin:coder".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -24736,9 +25336,9 @@ mod tests { profile: ProfileSelector::Builtin("builtin:coder".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -25361,9 +25961,9 @@ mod tests { profile: ProfileSelector::Builtin("builtin:coder".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -25427,9 +26027,9 @@ mod tests { profile: ProfileSelector::Builtin("builtin:coder".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -25775,9 +26375,9 @@ mod tests { initial_submit: vec![Segment::Flow { selector: "builtin:coder-review".to_string(), }], - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -26031,6 +26631,9 @@ mod tests { async fn worker_spawn_finalize_failure_reports_stage_and_compensation_outcomes() { let dir = tempfile::tempdir().unwrap(); let api = test_api(dir.path()).await; + register_test_runtime(&api, WorkdirlessFixtureRuntime::RUNTIME_ID).await; + api.runtime + .register_or_replace(WorkdirlessFixtureRuntime::default()); let mut ticket_input = ticket::NewTicket::new("Compensation test Ticket"); ticket_input.workflow_state = Some(TicketWorkflowState::InProgress); let ticket_id = browser_ticket_backend(&api) @@ -26057,9 +26660,9 @@ mod tests { initial_submit: vec![Segment::Flow { selector: "builtin:coder-review".to_string(), }], - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -26071,13 +26674,13 @@ mod tests { State(api.clone()), AxumPath(ScopedRuntimePath { workspace_id: TEST_WORKSPACE_ID.to_string(), - runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + runtime_id: WorkdirlessFixtureRuntime::RUNTIME_ID.to_string(), }), Json(request), ) .await .unwrap(); - let worker = created.worker.unwrap(); + let mut worker = created.worker.unwrap(); assert!( api.store .get_worker_registry(TEST_WORKSPACE_ID, &worker.worker) @@ -26085,34 +26688,98 @@ mod tests { .is_some() ); + let spawned_workdir_ids = vec![ + "spawned-workdir-left".to_string(), + "spawned-workdir-right".to_string(), + ]; + for workdir_id in &spawned_workdir_ids { + seed_cleanup_workdir_for_runtime( + &api, + workdir_id, + WorkdirlessFixtureRuntime::RUNTIME_ID, + "present", + "clean", + ); + } + worker.workdir_attachments = spawned_workdir_ids + .iter() + .enumerate() + .map(|(index, workdir_id)| { + let mut working_directory = WorkdirlessFixtureRuntime::workdir_status(workdir_id); + if index == 1 { + working_directory.summary.display_name = + Some("inject-registry-upsert-failure".to_string()); + } + WorkingDirectoryAttachmentStatus { + alias: workdir::WorkdirAttachmentAlias::new(if index == 0 { + "left" + } else { + "right" + }) + .unwrap(), + working_directory, + } + }) + .collect(); + let worker_record = api + .store + .get_worker_registry(TEST_WORKSPACE_ID, &worker.worker) + .unwrap() + .unwrap(); + rusqlite::Connection::open(&api.config.database_path) + .unwrap() + .execute_batch( + r#" + CREATE TABLE spawn_finalize_attachment_audit(alias TEXT NOT NULL); + CREATE TRIGGER audit_spawn_finalize_attachment + AFTER INSERT ON worker_workdir_links + BEGIN + INSERT INTO spawn_finalize_attachment_audit(alias) VALUES (NEW.alias); + END; + CREATE TRIGGER fail_second_spawn_workdir_upsert + BEFORE UPDATE ON workdir_registry + WHEN NEW.display_name = 'inject-registry-upsert-failure' + BEGIN + SELECT RAISE(ABORT, 'injected second Workdir registry failure'); + END; + "#, + ) + .unwrap(); + let context = WorkerSpawnCompensationContext { assignment: Some(&assignment), - prepared_workdir_id: None, - cleanup_spawned_workdir: false, + spawned_workdir_ids: &spawned_workdir_ids, }; - let error = finalize_worker_spawn_stage::<()>( - &api, - &worker, - &context, - WorkerSpawnFinalizeStage::WorkdirAttachment, - Err(Error::RegistryInconsistency("injected finalize failure".to_string()).into()), - ) - .unwrap_err(); + let error = + finalize_spawned_worker_workdir_attachments(&api, &worker, &worker_record, &context) + .unwrap_err(); assert!(matches!( error.error, Error::RuntimeOperationFailed { ref code, .. } - if code == "worker_spawn_finalize_workdir_attachment_failed" + if code == "worker_spawn_finalize_workdir_registry_failed" )); assert!(error.diagnostics.iter().any(|diagnostic| { - diagnostic.code == "worker_spawn_finalize_workdir_attachment_failed" + diagnostic.code == "worker_spawn_finalize_workdir_registry_failed" && diagnostic.message.contains(&worker.worker.runtime_id) && diagnostic.message.contains(&worker.worker.worker_id) })); + let audit_connection = rusqlite::Connection::open(&api.config.database_path).unwrap(); + let mut audit_statement = audit_connection + .prepare("SELECT alias FROM spawn_finalize_attachment_audit ORDER BY rowid") + .unwrap(); + let finalized_aliases = audit_statement + .query_map([], |row| row.get::<_, String>(0)) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(finalized_aliases, ["left"]); assert!( error .diagnostics .iter() - .any(|diagnostic| diagnostic.code == "worker_spawn_compensated") + .any(|diagnostic| diagnostic.code == "worker_spawn_compensated"), + "{:?}", + error.diagnostics ); assert!(matches!( api.runtime.worker(&worker.worker), @@ -26137,12 +26804,59 @@ mod tests { .is_none() ); + for workdir_id in &spawned_workdir_ids { + assert!( + api.store + .get_workdir_registry(TEST_WORKSPACE_ID, workdir_id) + .unwrap() + .is_none(), + "spawn compensation must clean every created Workdir: {workdir_id}" + ); + } + + let missing_summary_workdir_ids = vec![ + "missing-summary-workdir-left".to_string(), + "missing-summary-workdir-right".to_string(), + ]; + for workdir_id in &missing_summary_workdir_ids { + seed_cleanup_workdir_for_runtime( + &api, + workdir_id, + WorkdirlessFixtureRuntime::RUNTIME_ID, + "present", + "clean", + ); + } + let missing_summary_context = WorkerSpawnCompensationContext { + assignment: None, + spawned_workdir_ids: &missing_summary_workdir_ids, + }; + let reservation_worker_id = WorkerId::now_v7(); + let _ = compensate_failed_workspace_worker_create( + &api, + EMBEDDED_WORKER_RUNTIME_ID, + reservation_worker_id, + "missing-summary-fingerprint", + "runtime_spawn_missing_summary", + None, + &missing_summary_context, + &[], + ); + for workdir_id in &missing_summary_workdir_ids { + assert!( + api.store + .get_workdir_registry(TEST_WORKSPACE_ID, workdir_id) + .unwrap() + .is_none(), + "missing-summary compensation must clean every created Workdir: {workdir_id}" + ); + } + let mut residual_worker = worker.clone(); residual_worker.worker.runtime_id = "missing-runtime".to_string(); let residual_context = WorkerSpawnCompensationContext { assignment: None, - prepared_workdir_id: None, - cleanup_spawned_workdir: false, + spawned_workdir_ids: &[], }; let residual_error = finalize_worker_spawn_stage::<()>( &api, @@ -26994,7 +27708,7 @@ mod tests { permission: worker_runtime::auth::WORKER_REMOVE_PERMISSION.to_string(), jti: "caller-guard-proof".to_string(), }; - let executor = WorkspaceWorkerRemoveExecutor::new(&api); + let executor = WorkerRemovalService::new(&api); let self_response = executor .execute_async( verified_source(), @@ -27021,9 +27735,9 @@ mod tests { ), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: Some(runtime_test_bundle()), resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -27079,9 +27793,9 @@ mod tests { ), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: Some(runtime_test_bundle()), resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -27125,8 +27839,9 @@ mod tests { let summary = api.runtime.worker(&target).unwrap(); sync_worker_observation(&api, &summary).unwrap(); seed_worker_control_grant(&api, &source, &target, "embedded-valid-proof"); + let mut subscriber = api.worker_projection.subscribe(); - let response = WorkspaceWorkerRemoveExecutor::new(&api) + let response = WorkerRemovalService::new(&api) .execute_async( crate::worker_source::VerifiedWorkerMutationSource { runtime_id: source.runtime_id, @@ -27157,6 +27872,91 @@ mod tests { .unwrap() .is_none() ); + let removal = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + let event = subscriber.recv().await.expect("projection event"); + if event.changes.iter().any(|change| { + matches!(change, crate::store::WorkerCatalogChange::Removed(worker) if worker == &target) + }) { + return event; + } + } + }) + .await + .expect("WorkerRemove catalog removal broadcast"); + assert!(removal.revision > 0); + } + + #[tokio::test] + async fn published_spawn_compensation_uses_catalog_removal_commit() { + let temp = tempfile::tempdir().unwrap(); + let api = test_api(temp.path()).await; + let spawned = api + .spawn_workspace_worker( + EMBEDDED_WORKER_RUNTIME_ID, + WorkerSpawnRequest { + intent: WorkerSpawnIntent::WorkspaceCompanion, + requested_worker_name: Some("compensation-target".to_string()), + acceptance: WorkerSpawnAcceptanceRequirement::RunAccepted { + expected_segments: 0, + }, + profile: worker_runtime::catalog::ProfileSelector::Builtin( + "builtin:companion".to_string(), + ), + ticket_assignment: None, + initial_submit: Vec::new(), + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), + resolved_config_bundle: Some(runtime_test_bundle()), + resolved_worker_observation_enabled: false, + resolved_worker_observation_grants: Vec::new(), + resolved_workspace_api: None, + resolved_memory_settings: None, + resolved_control_operation: None, + }, + ) + .unwrap(); + let worker = spawned.worker.unwrap().worker; + let summary = api.runtime.worker(&worker).unwrap(); + sync_worker_observation(&api, &summary).unwrap(); + let mut subscriber = api.worker_projection.subscribe(); + let diagnostics = compensate_failed_worker_spawn( + &api, + &summary, + &WorkerSpawnCompensationContext { + assignment: None, + spawned_workdir_ids: &[], + }, + ); + assert!(diagnostics.is_empty(), "{diagnostics:?}"); + assert!( + api.store + .get_worker_registry(TEST_WORKSPACE_ID, &worker) + .unwrap() + .is_none() + ); + let removal = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + let event = subscriber.recv().await.expect("projection event"); + if event.changes.iter().any(|change| { + matches!(change, crate::store::WorkerCatalogChange::Removed(removed) if removed == &worker) + }) { + return event; + } + } + }) + .await + .expect("spawn compensation catalog removal broadcast"); + assert!(removal.revision > 0); + sync_worker_observation(&api, &summary).unwrap(); + assert!( + api.store + .get_worker_registry(TEST_WORKSPACE_ID, &worker) + .unwrap() + .is_none(), + "a late pre-compensation observation must not clear the removal fence" + ); } #[tokio::test] @@ -27647,13 +28447,24 @@ mod tests { } fn seed_cleanup_workdir(api: &WorkspaceApi, workdir_id: &str, status: &str, cleanliness: &str) { + seed_cleanup_workdir_for_runtime(api, workdir_id, "runtime-test", status, cleanliness); + } + + fn seed_cleanup_workdir_for_runtime( + api: &WorkspaceApi, + workdir_id: &str, + runtime_id: &str, + status: &str, + cleanliness: &str, + ) { seed_test_repository(api, "repo-test"); let now = now_registry_timestamp(); api.store .upsert_workdir_registry(&WorkdirRegistryRecord { workspace_id: api.config.workspace_id.clone(), workdir_id: workdir_id.to_string(), - runtime_id: "runtime-test".to_string(), + display_name: None, + runtime_id: runtime_id.to_string(), repository_id: "repo-test".to_string(), creation_selector: Some("HEAD".to_string()), creation_ref: None, @@ -27759,7 +28570,7 @@ mod tests { workspace_id: api.config.workspace_id.clone(), worker: RuntimeWorkerRef::new("runtime-test", runtime_worker_id.to_string()), workdir_id: workdir_id.to_string(), - role: "attachment".to_string(), + alias: "attachment".to_string(), linked_at: now_registry_timestamp(), unlinked_at: None, }) @@ -27803,10 +28614,12 @@ mod tests { Request::builder() .method("POST") .uri(format!( - "/api/w/{TEST_WORKSPACE_ID}/workers/self/workdir-attachment" + "/api/w/{TEST_WORKSPACE_ID}/workers/self/workdir-attachments" )) .header("content-type", "application/json") - .body(Body::from(json!({"workdir_id": "wd"}).to_string())) + .body(Body::from( + json!({"alias": "checkout", "working_directory_id": "wd"}).to_string(), + )) .unwrap(), ) .await @@ -28534,6 +29347,273 @@ mod tests { ); } + #[tokio::test] + async fn cleanup_service_rejects_internal_worker_before_runtime_removal() { + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let api = test_api(workspace.path()).await; + let Json(orchestrator) = scoped_start_workspace_orchestrator( + State(api.clone()), + AxumPath(ScopedWorkspacePath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + }), + ) + .await + .unwrap(); + let worker = orchestrator.worker.unwrap(); + let worker = RuntimeWorkerRef::new(worker.runtime_id, worker.worker_id); + api.runtime + .stop_worker( + &worker, + WorkerLifecycleRequest { + reason: Some("prepare internal cleanup guard".to_string()), + ticket_assignment: None, + }, + ) + .unwrap(); + let plan = build_runtime_cleanup_plan(&api, worker.runtime_id.as_str()) + .unwrap_or_else(|error| panic!("cleanup plan: {}", error.error)); + let candidate = plan + .workers + .iter() + .find(|candidate| candidate.runtime_worker_id == worker.worker_id) + .expect("internal Worker cleanup candidate") + .clone(); + assert_eq!( + candidate.blocking_reason.as_deref(), + Some("internal service Worker cannot be deleted") + ); + let mut subscriber = api.worker_projection.subscribe(); + + let error = WorkerRemovalService::new(&api) + .execute_cleanup_removal(&candidate) + .await + .unwrap_err(); + assert!(matches!( + error.error, + Error::RuntimeOperationFailed { ref code, .. } + if code == "workspace_cleanup_worker_removal_rejected" + )); + assert!(api.runtime.worker(&worker).is_ok()); + assert!( + api.store + .get_worker_registry(TEST_WORKSPACE_ID, &worker) + .unwrap() + .is_some() + ); + assert!(matches!( + subscriber.try_recv(), + Err(tokio::sync::broadcast::error::TryRecvError::Empty) + )); + } + + #[tokio::test] + async fn cleanup_http_entry_publishes_atomic_catalog_removal_without_runtime_event() { + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let api = test_api(workspace.path()).await; + let spawned = api + .spawn_workspace_worker( + EMBEDDED_WORKER_RUNTIME_ID, + WorkerSpawnRequest { + intent: WorkerSpawnIntent::WorkspaceCompanion, + requested_worker_name: Some("cleanup-projection-target".to_string()), + acceptance: WorkerSpawnAcceptanceRequirement::RunAccepted { + expected_segments: 0, + }, + profile: worker_runtime::catalog::ProfileSelector::Builtin( + "builtin:companion".to_string(), + ), + ticket_assignment: None, + initial_submit: Vec::new(), + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), + resolved_config_bundle: Some(runtime_test_bundle()), + resolved_worker_observation_enabled: false, + resolved_worker_observation_grants: Vec::new(), + resolved_workspace_api: None, + resolved_memory_settings: None, + resolved_control_operation: None, + }, + ) + .unwrap(); + let worker = spawned.worker.unwrap().worker; + let stopped = api + .runtime + .stop_worker( + &worker, + WorkerLifecycleRequest { + reason: Some("prepare cleanup projection regression".to_string()), + ticket_assignment: None, + }, + ) + .unwrap(); + assert_eq!(stopped.state, WorkerOperationState::Accepted); + let worker_root = workspace + .path() + .join(".test-embedded-runtime-store/workers") + .join(&worker.worker_id); + fs::create_dir_all(worker_root.join("session/segments")).unwrap(); + fs::write( + worker_root.join("session/session.json"), + serde_json::to_vec_pretty(&json!({ + "schema_version": 1, + "session_id": "cleanup-removal-session" + })) + .unwrap(), + ) + .unwrap(); + fs::write( + worker_root.join("session/segments/segment-a.jsonl"), + b"cleanup retention evidence\n", + ) + .unwrap(); + let summary = api.runtime.worker(&worker).unwrap(); + sync_worker_observation(&api, &summary).unwrap(); + // Deliberately drop Runtime observation delivery: catalog removal itself + // must drive every connected Workspace subscription. + api.worker_projection.shutdown(); + let worker_id = worker.worker_id.clone(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = build_inner_router(api.clone()).layer(Extension(test_browser_request_actor())); + let server = tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + let subscription_url = format!("ws://{address}/api/w/{TEST_WORKSPACE_ID}/protocol/ws"); + let (mut acting_socket, _) = connect_async(subscription_url.as_str()).await.unwrap(); + let (mut observing_socket, _) = connect_async(subscription_url.as_str()).await.unwrap(); + let acting_snapshot = + subscribe_workspace_worker_snapshot(&mut acting_socket, "cleanup-acting").await; + let observing_snapshot = + subscribe_workspace_worker_snapshot(&mut observing_socket, "cleanup-observing").await; + assert!( + acting_snapshot + .iter() + .any(|item| item.worker_id.as_str() == worker_id) + ); + assert!( + observing_snapshot + .iter() + .any(|item| item.worker_id.as_str() == worker_id) + ); + let mut subscriber = api.worker_projection.subscribe(); + let plan = build_runtime_cleanup_plan(&api, EMBEDDED_WORKER_RUNTIME_ID) + .unwrap_or_else(|error| panic!("cleanup plan: {}", error.error)); + let candidate = plan + .workers + .iter() + .find(|candidate| candidate.worker_id == worker_id) + .unwrap(); + let request = ExecuteRuntimeCleanupRequest { + expected_plan_revision: plan.revision, + expected_plan_digest: plan.digest, + worker_target_ids: vec![candidate.target_id.clone()], + workdir_target_ids: Vec::new(), + confirm_dirty_discard_target_ids: Vec::new(), + }; + + let Json(response) = scoped_execute_runtime_cleanup( + State(api.clone()), + AxumPath(ScopedRuntimePath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + }), + Json(request), + ) + .await + .unwrap(); + assert_eq!(response.results[0].status, "deleted"); + + let event = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + let event = subscriber.recv().await.expect("projection event"); + if event.changes.iter().any(|change| { + matches!(change, crate::store::WorkerCatalogChange::Removed(removed) if removed == &worker) + }) { + return event; + } + } + }) + .await + .expect("catalog removal broadcast"); + assert!(matches!( + event.changes.as_slice(), + [crate::store::WorkerCatalogChange::Removed(removed)] if removed == &worker + )); + let acting_revision = tokio::time::timeout( + std::time::Duration::from_secs(1), + next_workspace_worker_removal(&mut acting_socket, worker_id.as_str()), + ) + .await + .expect("acting Workspace Worker subscription removal"); + let observing_revision = tokio::time::timeout( + std::time::Duration::from_secs(1), + next_workspace_worker_removal(&mut observing_socket, worker_id.as_str()), + ) + .await + .expect("independent Workspace Worker subscription removal"); + assert_eq!(acting_revision, event.revision); + assert_eq!(observing_revision, event.revision); + assert!( + api.store + .get_worker_registry(&api.config.workspace_id, &worker) + .unwrap() + .is_none() + ); + let recovered = api + .store + .recover_worker_removal_execution(&api.config.workspace_id, &worker) + .unwrap() + .expect("Browser cleanup uses durable retention removal"); + assert_eq!( + recovered.plan.state, + crate::retention::WorkerRemovalPlanState::Succeeded + ); + let (snapshot_revision, snapshot) = api + .store + .worker_registry_projection_snapshot(&api.config.workspace_id, 100) + .unwrap(); + assert_eq!(snapshot_revision, event.revision); + assert!( + snapshot + .iter() + .all(|record| record.registry.worker != worker) + ); + let rest_workers = workers_response(api.clone()).unwrap(); + assert!( + rest_workers + .items + .iter() + .all(|item| item.worker_id != worker_id) + ); + let (mut fresh_socket, _) = connect_async(subscription_url.as_str()).await.unwrap(); + let fresh_snapshot = + subscribe_workspace_worker_snapshot(&mut fresh_socket, "cleanup-fresh").await; + assert!( + fresh_snapshot + .iter() + .all(|item| item.worker_id.as_str() != worker_id) + ); + + // A retry replays the durable fence at the same revision, healing an + // in-process broadcast omitted after the first DB commit. + WorkerRemovalService::new(&api) + .commit_catalog_removal(&worker) + .unwrap(); + let replay = tokio::time::timeout(std::time::Duration::from_secs(1), subscriber.recv()) + .await + .expect("catalog removal replay") + .expect("projection replay"); + assert_eq!(replay.revision, event.revision); + assert!(matches!( + replay.changes.as_slice(), + [crate::store::WorkerCatalogChange::Removed(removed)] if removed == &worker + )); + server.abort(); + } + #[tokio::test] async fn cleanup_blocks_assigned_worker_before_runtime_deletion() { let workspace = tempfile::tempdir().unwrap(); @@ -29033,7 +30113,7 @@ mod tests { can_stop: true, can_spawn_followup: false, }, - working_directory: None, + workdir_attachments: Vec::new(), diagnostics: Vec::new(), } } @@ -29130,7 +30210,7 @@ mod tests { can_stop: true, can_spawn_followup: false, }, - working_directory: None, + workdir_attachments: Vec::new(), diagnostics: Vec::new(), }], None, @@ -29459,8 +30539,8 @@ mod tests { digest: bundle.metadata.digest, }), initial_input: None, - working_directory_request: None, - working_directory: None, + workdir_attachment_requests: Vec::new(), + workdir_attachments: Vec::new(), worker_observation_enabled: false, worker_observation_grants: Vec::new(), workspace_api: None, @@ -29522,6 +30602,7 @@ mod tests { &repository_id, Some("HEAD"), Some(EMBEDDED_WORKER_RUNTIME_ID), + None, &repository.source_fingerprint, repository.source_revision, ); @@ -29828,10 +30909,11 @@ mod tests { "display_name": "Coding Worker", "profile": "builtin:coder", "initial_submit": [], - "working_directory": { + "workdir_attachments": [{ + "alias": "workdir", "working_directory_id": working_directory_id, "relative_cwd": "../escape" - } + }] })), StatusCode::BAD_REQUEST, ) @@ -30603,35 +31685,66 @@ mod tests { } #[tokio::test] - async fn browser_worker_create_rejects_non_embedded_no_workdir() { + async fn runtime_and_browser_worker_create_allow_non_embedded_zero_attachments() { let dir = tempfile::tempdir().unwrap(); - let app = test_app(dir.path()).await; - let response = request_json( - app, - "POST", - "/api/workers", - Some(serde_json::json!({ - "runtime_id": "remote-runtime", - "display_name": "Remote Worker", - "profile": "builtin:companion", - "initial_submit": [] - })), - StatusCode::BAD_REQUEST, + let api = test_api(dir.path()).await; + register_test_runtime(&api, WorkdirlessFixtureRuntime::RUNTIME_ID).await; + api.runtime + .register_or_replace(WorkdirlessFixtureRuntime::default()); + + let Json(runtime_response) = create_runtime_worker( + State(api.clone()), + AxumPath(WorkdirlessFixtureRuntime::RUNTIME_ID.to_string()), + Json(WorkerSpawnRequest { + intent: WorkerSpawnIntent::WorkspaceCompanion, + requested_worker_name: Some("Runtime Worker".to_string()), + acceptance: WorkerSpawnAcceptanceRequirement::RunAccepted { + expected_segments: 0, + }, + profile: ProfileSelector::Builtin("builtin:companion".to_string()), + ticket_assignment: None, + initial_submit: Vec::new(), + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), + resolved_config_bundle: None, + resolved_worker_observation_enabled: false, + resolved_worker_observation_grants: Vec::new(), + resolved_workspace_api: None, + resolved_memory_settings: None, + resolved_control_operation: None, + }), ) - .await; + .await + .unwrap(); assert!( - response["message"] - .as_str() - .unwrap_or_default() - .contains("workspace_worker_workdir_required") - ); - assert!( - response["diagnostics"] - .as_array() + runtime_response + .worker + .as_ref() .unwrap() - .iter() - .any(|diagnostic| { diagnostic["code"] == "workspace_worker_workdir_required" }) + .workdir_attachments + .is_empty() ); + + let Json(response) = create_workspace_worker_inner( + api, + HeaderMap::new(), + CreateWorkspaceWorkerRequest { + runtime_id: WorkdirlessFixtureRuntime::RUNTIME_ID.to_string(), + display_name: "Remote Worker".to_string(), + profile: Some("builtin:companion".to_string()), + ticket_assignment: None, + initial_submit: Vec::new(), + workdir_attachments: Vec::new(), + control_operation_id: None, + }, + None, + ) + .await + .unwrap(); + + assert_eq!(response.runtime_id, WorkdirlessFixtureRuntime::RUNTIME_ID); + assert!(response.worker.workdir_attachments.is_empty()); } #[tokio::test] @@ -30695,10 +31808,13 @@ mod tests { "kind": "builtin", "value": "builtin:coder" }, - "working_directory_request": { - "repository_key": "test-repository", - "selector": "HEAD" - } + "workdir_attachment_requests": [{ + "alias": "workdir", + "working_directory": { + "repository_key": "test-repository", + "selector": "HEAD" + } + }] })), StatusCode::BAD_REQUEST, ) @@ -31397,9 +32513,9 @@ mod tests { profile: ProfileSelector::Builtin("builtin:coder".to_string()), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: None, resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -31682,6 +32798,7 @@ mod tests { .upsert_workdir_registry(&WorkdirRegistryRecord { workspace_id: TEST_WORKSPACE_ID.to_string(), workdir_id: workdir_id.to_string(), + display_name: None, runtime_id: "external-workdir-runtime".to_string(), repository_id, creation_selector: None, @@ -31702,7 +32819,7 @@ mod tests { workspace_id: TEST_WORKSPACE_ID.to_string(), worker: RuntimeWorkerRef::new("embedded-worker-runtime", &worker_id), workdir_id: workdir_id.to_string(), - role: "attachment".to_string(), + alias: "attachment".to_string(), linked_at: "2".to_string(), unlinked_at: None, }) @@ -31877,6 +32994,74 @@ mod tests { serde_json::from_str(&text).unwrap() } + type TestWebSocket = tokio_tungstenite::WebSocketStream< + tokio_tungstenite::MaybeTlsStream, + >; + + async fn subscribe_workspace_worker_snapshot( + socket: &mut TestWebSocket, + request_id: &str, + ) -> Vec { + let frame = protocol::subscription::SubscriptionFrame::new( + protocol::subscription::SubscriptionFramePayload::Request( + protocol::subscription::SubscriptionRequest::SubscribeEvents { + request_id: protocol::subscription::SubscriptionRequestId::new(request_id) + .unwrap(), + selector: protocol::subscription::EventSubscriptionSelector::WorkspaceWorkers, + }, + ), + ); + socket + .send(Message::Text(serde_json::to_string(&frame).unwrap().into())) + .await + .unwrap(); + loop { + let Message::Text(text) = socket.next().await.unwrap().unwrap() else { + continue; + }; + let frame: protocol::subscription::SubscriptionFrame = + serde_json::from_str(text.as_str()).unwrap(); + if let protocol::subscription::SubscriptionFramePayload::Response( + protocol::subscription::SubscriptionResponse::Subscribed { + selector: protocol::subscription::EventSubscriptionSelector::WorkspaceWorkers, + snapshot: protocol::subscription::SubscriptionSnapshot::Workers { workers }, + .. + }, + ) = frame.payload + { + return workers; + } + } + } + + async fn next_workspace_worker_removal( + socket: &mut TestWebSocket, + expected_worker_id: &str, + ) -> u64 { + loop { + let Message::Text(text) = socket.next().await.unwrap().unwrap() else { + continue; + }; + let frame: protocol::subscription::SubscriptionFrame = + serde_json::from_str(text.as_str()).unwrap(); + if let protocol::subscription::SubscriptionFramePayload::Event( + protocol::subscription::SubscriptionEvent::Event { + subject_revision, + payload: + protocol::subscription::SubscriptionEventPayload::WorkerRemoved { + worker_id, + .. + }, + .. + }, + ) = frame.payload + && worker_id.as_str() == expected_worker_id + { + return subject_revision; + } + } + } + async fn spawn_runtime_worker() -> ( worker_runtime::Runtime, worker_runtime::identity::WorkerRef, @@ -31976,9 +33161,9 @@ mod tests { ), ticket_assignment: None, initial_submit: Vec::new(), - working_directory_request: None, - resolved_working_directory_request: None, - resolved_working_directory: None, + workdir_attachment_requests: Vec::new(), + resolved_workdir_attachment_requests: Vec::new(), + resolved_workdir_attachments: Vec::new(), resolved_config_bundle: Some(runtime_test_bundle()), resolved_worker_observation_enabled: false, resolved_worker_observation_grants: Vec::new(), @@ -32000,14 +33185,14 @@ mod tests { .find(|projection| projection.registry.worker.worker_id == worker_id) .and_then(|projection| projection.observation.as_ref()) .expect("spawned Worker observation exists"); - assert!(observation.worker.repository_id.is_none()); - assert!(observation.worker.working_directory_id.is_none()); + assert!(observation.worker.workdir_attachments.is_empty()); let workdir_id = "001a0b415233500000c"; api.store .upsert_workdir_registry(&WorkdirRegistryRecord { workspace_id: TEST_WORKSPACE_ID.to_string(), workdir_id: workdir_id.to_string(), + display_name: None, runtime_id: "arcadia".to_string(), repository_id: test_repository_id(&api), creation_selector: Some("HEAD".to_string()), @@ -32028,7 +33213,7 @@ mod tests { workspace_id: TEST_WORKSPACE_ID.to_string(), worker: RuntimeWorkerRef::new(EMBEDDED_WORKER_RUNTIME_ID, &worker_id), workdir_id: workdir_id.to_string(), - role: "attachment".to_string(), + alias: "attachment".to_string(), linked_at: "1".to_string(), unlinked_at: None, }) @@ -32079,16 +33264,16 @@ mod tests { .find(|worker| worker.worker_id.as_str() == worker_id) .expect("spawned Worker is projected"); assert_eq!(projected_worker.resource_key.as_deref(), Some("W-1")); + assert_eq!(projected_worker.workdir_attachments.len(), 1); + let projected_attachment = &projected_worker.workdir_attachments[0]; + assert_eq!(projected_attachment.alias, "attachment"); assert_eq!( - projected_worker.repository_key.as_deref(), + projected_attachment.repository_key.as_deref(), Some("test-repository") ); assert_eq!( - projected_worker - .working_directory_id - .as_ref() - .map(|workdir_id| workdir_id.as_str()), - Some(workdir_id) + projected_attachment.working_directory_id.as_str(), + workdir_id ); projection_api @@ -32102,7 +33287,7 @@ mod tests { .unwrap(); projection_api .worker_projection - .publish_catalog_change(&RuntimeWorkerRef::new( + .refresh(&RuntimeWorkerRef::new( EMBEDDED_WORKER_RUNTIME_ID, &worker_id, )) @@ -32123,8 +33308,7 @@ mod tests { other => panic!("expected Worker upsert after Workdir detach, got {other:?}"), }; assert_eq!(detached_worker.worker_id.as_str(), worker_id); - assert!(detached_worker.repository_key.is_none()); - assert!(detached_worker.working_directory_id.is_none()); + assert!(detached_worker.workdir_attachments.is_empty()); let subscribe_protocol = protocol::subscription::SubscriptionFrame::new( protocol::subscription::SubscriptionFramePayload::Request( @@ -32612,6 +33796,7 @@ VALUES ('0192f0e8-4d84-7d6e-a000-000000000001', 'ticket', 3); items: vec![ WorkingDirectorySummary { working_directory_id: "wd-1".to_string(), + display_name: Some("Checkout".to_string()), repository_key: "main".to_string(), creation_selector: None, creation_ref: None, @@ -32624,7 +33809,6 @@ VALUES ('0192f0e8-4d84-7d6e-a000-000000000001', 'ticket', 3); cleanup_target: None, status: WorkingDirectoryStatusKind::Active, cleanliness: Some("clean".to_string()), - primary_worker_id: None, occupied_by: Some(WorkingDirectoryOccupancy { runtime_id: "arcadia".to_string(), worker_id: "worker-opaque-64".to_string(), diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index 1e0f0cf66..2fdc2f986 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -21,7 +21,7 @@ use crate::workspace_deletion::WorkspaceDeletionStore; use crate::{Error, Result}; const OLDEST_SCHEMA_VERSION: i64 = 50; -const LATEST_SCHEMA_VERSION: i64 = 62; +const LATEST_SCHEMA_VERSION: i64 = 63; const SCHEMA_BASELINE_NAME: &str = "workspace schema baseline"; const WORKSPACE_RUNTIME_BINDINGS_MIGRATION_NAME: &str = "workspace runtime bindings"; const RUNTIME_BINDING_AUDIT_MIGRATION_NAME: &str = "workspace Runtime binding revision and audit"; @@ -41,6 +41,8 @@ const RUNTIME_REMOVAL_OPERATION_MIGRATION_NAME: &str = "guarded Runtime removal const REMOVE_WORKER_RUN_GENERATION_MIGRATION_NAME: &str = "remove obsolete Worker run generation"; const WORKER_REGISTRY_PROJECTION_MIGRATION_NAME: &str = "durable Runtime-backed Worker registry projection"; +const MULTI_WORKDIR_ATTACHMENTS_MIGRATION_NAME: &str = + "multiple Worker Workdir attachments and Workdir display names"; const WORKER_REGISTRY_PROJECTION_SCHEMA: &str = r#" CREATE TABLE IF NOT EXISTS worker_registry_observations ( workspace_id TEXT NOT NULL, @@ -146,6 +148,11 @@ const MIGRATIONS: &[Migration] = &[ name: WORKER_REGISTRY_PROJECTION_MIGRATION_NAME, apply: migrate_worker_registry_projection_v61_to_v62, }, + Migration { + version: 63, + name: MULTI_WORKDIR_ATTACHMENTS_MIGRATION_NAME, + apply: migrate_multi_workdir_attachments_v62_to_v63, + }, ]; #[derive(Clone, Copy)] @@ -563,10 +570,27 @@ pub struct WorkerRegistryProjectionRecord { pub observation: Option, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum WorkerCatalogChange { + Upsert(WorkerRegistryProjectionRecord), + Removed(RuntimeWorkerRef), +} + +/// A catalog mutation and the exact projection payload committed at its revision. +/// +/// Callers may broadcast this value, but must never re-read catalog rows to infer +/// what happened: a later catalog mutation could otherwise be paired with an +/// older revision. #[derive(Debug, Clone, PartialEq, Eq)] pub struct WorkerRegistryProjectionCommit { pub revision: u64, - pub changed_workers: Vec, + pub changes: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WorkerRemovalCommit { + pub plan: crate::retention::WorkerRemovalPlan, + pub catalog: WorkerRegistryProjectionCommit, } /// Durable authority describing which Runtime Worker another Runtime Worker may @@ -764,6 +788,7 @@ pub struct WorkdirCreateOperationRecord { pub struct WorkdirRegistryRecord { pub workspace_id: String, pub workdir_id: String, + pub display_name: Option, pub runtime_id: String, pub repository_id: String, pub creation_selector: Option, @@ -784,7 +809,7 @@ pub struct WorkerWorkdirLinkRecord { pub workspace_id: String, pub worker: RuntimeWorkerRef, pub workdir_id: String, - pub role: String, + pub alias: String, pub linked_at: String, pub unlinked_at: Option, } @@ -1130,16 +1155,23 @@ pub trait ControlPlaneStore: Send + Sync + WorkspaceDeletionStore { "Worker retention authority is unavailable".to_string(), )) } + fn recover_succeeded_worker_removal_catalog( + &self, + workspace_id: &str, + worker: &RuntimeWorkerRef, + ) -> Result { + let _ = (workspace_id, worker); + Err(Error::Store( + "Worker removal catalog recovery authority is unavailable".to_string(), + )) + } fn commit_worker_removal( &self, workspace_id: &str, operation_id: &str, input_fingerprint: &str, result: &worker_runtime::retention::WorkerRetentionExecutionResult, - ) -> std::result::Result< - crate::retention::WorkerRemovalPlan, - crate::retention::WorkerRetentionError, - > { + ) -> std::result::Result { let _ = (workspace_id, operation_id, input_fingerprint, result); Err(crate::retention::WorkerRetentionError::Invalid( "Worker retention authority is unavailable".to_string(), @@ -1317,7 +1349,10 @@ pub trait ControlPlaneStore: Send + Sync + WorkspaceDeletionStore { consumed_at: &str, ) -> Result>; - fn upsert_worker_registry(&self, record: &WorkerRegistryRecord) -> Result<()>; + fn upsert_worker_registry( + &self, + record: &WorkerRegistryRecord, + ) -> Result; fn get_worker_registry( &self, workspace_id: &str, @@ -1369,12 +1404,7 @@ pub trait ControlPlaneStore: Send + Sync + WorkspaceDeletionStore { subject_revision: Option, observed_at: &str, ) -> Result; - fn publish_worker_registry_catalog_change( - &self, - workspace_id: &str, - worker: &RuntimeWorkerRef, - ) -> Result; - fn publish_worker_registry_removal( + fn commit_worker_registry_projection_refresh( &self, workspace_id: &str, worker: &RuntimeWorkerRef, @@ -1385,9 +1415,12 @@ pub trait ControlPlaneStore: Send + Sync + WorkspaceDeletionStore { worker: &RuntimeWorkerRef, retention_state: &str, updated_at: &str, - ) -> Result; - fn delete_worker_registry(&self, workspace_id: &str, worker: &RuntimeWorkerRef) - -> Result; + ) -> Result; + fn delete_worker_registry( + &self, + workspace_id: &str, + worker: &RuntimeWorkerRef, + ) -> Result; fn create_worker_control_grant( &self, @@ -4330,16 +4363,36 @@ impl ControlPlaneStore for SqliteWorkspaceStore { ) } + fn recover_succeeded_worker_removal_catalog( + &self, + workspace_id: &str, + worker: &RuntimeWorkerRef, + ) -> Result { + self.with_conn_mut(|conn| { + let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; + let succeeded: bool = tx.query_row( + "SELECT EXISTS(SELECT 1 FROM worker_removal_operations WHERE workspace_id=?1 AND runtime_id=?2 AND worker_id=?3 AND state='succeeded')", + params![workspace_id, worker.runtime_id, worker.worker_id], + |row| row.get(0), + )?; + if !succeeded { + return Err(Error::InvalidInput( + "Worker removal operation has not succeeded".to_string(), + )); + } + let commit = commit_worker_catalog_removal(&tx, workspace_id, worker)?; + tx.commit()?; + Ok(commit) + }) + } + fn commit_worker_removal( &self, workspace_id: &str, operation_id: &str, input_fingerprint: &str, result: &worker_runtime::retention::WorkerRetentionExecutionResult, - ) -> std::result::Result< - crate::retention::WorkerRemovalPlan, - crate::retention::WorkerRetentionError, - > { + ) -> std::result::Result { SqliteWorkspaceStore::commit_worker_removal( self, workspace_id, @@ -5527,15 +5580,21 @@ impl ControlPlaneStore for SqliteWorkspaceStore { }) } - fn upsert_worker_registry(&self, record: &WorkerRegistryRecord) -> Result<()> { - self.with_conn(|conn| { - let tx = conn.unchecked_transaction()?; + fn upsert_worker_registry( + &self, + record: &WorkerRegistryRecord, + ) -> Result { + self.with_conn_mut(|conn| { + let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; let removal_blocks_upsert: bool = tx.query_row( "SELECT EXISTS( SELECT 1 FROM worker_removal_operations WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3 AND state IN ('executing', 'failed', 'succeeded') + UNION ALL + SELECT 1 FROM worker_registry_projection_removals + WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3 )", params![ record.workspace_id, @@ -5545,9 +5604,12 @@ impl ControlPlaneStore for SqliteWorkspaceStore { |row| row.get(0), )?; if removal_blocks_upsert { - return Ok(()); + return Ok(WorkerRegistryProjectionCommit { + revision: current_worker_projection_revision(&tx, &record.workspace_id)?, + changes: Vec::new(), + }); } - tx.execute( + let changed = tx.execute( r#"INSERT INTO worker_registry ( workspace_id, runtime_id, worker_id, display_name, profile, retention_state, transcript_ref, session_ref, summary_ref, @@ -5596,12 +5658,20 @@ impl ControlPlaneStore for SqliteWorkspaceStore { &record.worker.worker_id, &record.created_at, )?; - tx.execute( - "DELETE FROM worker_registry_projection_removals WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3", - params![record.workspace_id, record.worker.runtime_id, record.worker.worker_id], + let changed_workers = if changed == 0 { + Vec::new() + } else { + vec![record.worker.clone()] + }; + let revision = + commit_worker_projection_changes(&tx, &record.workspace_id, &changed_workers)?; + let changes = worker_catalog_upsert_changes( + &tx, + &record.workspace_id, + changed_workers.as_slice(), )?; tx.commit()?; - Ok(()) + Ok(WorkerRegistryProjectionCommit { revision, changes }) }) } @@ -5715,7 +5785,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore { tx.commit()?; return Ok(WorkerRegistryProjectionCommit { revision, - changed_workers: Vec::new(), + changes: Vec::new(), }); } upsert_worker_projection_cursor( @@ -5775,11 +5845,10 @@ impl ControlPlaneStore for SqliteWorkspaceStore { } } let revision = commit_worker_projection_changes(&tx, workspace_id, &changed_workers)?; + let changes = + worker_catalog_upsert_changes(&tx, workspace_id, changed_workers.as_slice())?; tx.commit()?; - Ok(WorkerRegistryProjectionCommit { - revision, - changed_workers, - }) + Ok(WorkerRegistryProjectionCommit { revision, changes }) }) } @@ -5810,7 +5879,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore { tx.commit()?; return Ok(WorkerRegistryProjectionCommit { revision, - changed_workers: Vec::new(), + changes: Vec::new(), }); } let changed = upsert_worker_observation( @@ -5825,11 +5894,10 @@ impl ControlPlaneStore for SqliteWorkspaceStore { )?; let changed_workers = if changed { vec![identity] } else { Vec::new() }; let revision = commit_worker_projection_changes(&tx, workspace_id, &changed_workers)?; + let changes = + worker_catalog_upsert_changes(&tx, workspace_id, changed_workers.as_slice())?; tx.commit()?; - Ok(WorkerRegistryProjectionCommit { - revision, - changed_workers, - }) + Ok(WorkerRegistryProjectionCommit { revision, changes }) }) } @@ -5878,15 +5946,14 @@ impl ControlPlaneStore for SqliteWorkspaceStore { } } let revision = commit_worker_projection_changes(&tx, workspace_id, &changed_workers)?; + let changes = + worker_catalog_upsert_changes(&tx, workspace_id, changed_workers.as_slice())?; tx.commit()?; - Ok(WorkerRegistryProjectionCommit { - revision, - changed_workers, - }) + Ok(WorkerRegistryProjectionCommit { revision, changes }) }) } - fn publish_worker_registry_catalog_change( + fn commit_worker_registry_projection_refresh( &self, workspace_id: &str, worker: &RuntimeWorkerRef, @@ -5899,46 +5966,10 @@ impl ControlPlaneStore for SqliteWorkspaceStore { Vec::new() }; let revision = commit_worker_projection_changes(&tx, workspace_id, &changed_workers)?; + let changes = + worker_catalog_upsert_changes(&tx, workspace_id, changed_workers.as_slice())?; tx.commit()?; - Ok(WorkerRegistryProjectionCommit { - revision, - changed_workers, - }) - }) - } - - fn publish_worker_registry_removal( - &self, - workspace_id: &str, - worker: &RuntimeWorkerRef, - ) -> Result { - self.with_conn_mut(|conn| { - let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; - let published_revision = tx - .query_row( - "SELECT projection_revision FROM worker_registry_projection_removals WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3", - params![workspace_id, worker.runtime_id, worker.worker_id], - |row| row.get::<_, i64>(0), - ) - .optional()?; - if let Some(revision) = published_revision { - tx.commit()?; - return Ok(WorkerRegistryProjectionCommit { - revision: revision.max(0) as u64, - changed_workers: Vec::new(), - }); - } - let changed_workers = vec![worker.clone()]; - let revision = commit_worker_projection_changes(&tx, workspace_id, &changed_workers)?; - tx.execute( - "INSERT INTO worker_registry_projection_removals (workspace_id, runtime_id, worker_id, projection_revision) VALUES (?1, ?2, ?3, ?4)", - params![workspace_id, worker.runtime_id, worker.worker_id, revision as i64], - )?; - tx.commit()?; - Ok(WorkerRegistryProjectionCommit { - revision, - changed_workers, - }) + Ok(WorkerRegistryProjectionCommit { revision, changes }) }) } @@ -5948,9 +5979,10 @@ impl ControlPlaneStore for SqliteWorkspaceStore { worker: &RuntimeWorkerRef, retention_state: &str, updated_at: &str, - ) -> Result { - self.with_conn(|conn| { - let changed = conn.execute( + ) -> Result { + self.with_conn_mut(|conn| { + let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; + let changed = tx.execute( r#"UPDATE worker_registry SET retention_state = ?4, updated_at = ?5 WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3 @@ -5968,7 +6000,16 @@ impl ControlPlaneStore for SqliteWorkspaceStore { updated_at ], )?; - Ok(changed > 0) + let changed_workers = if changed == 0 { + Vec::new() + } else { + vec![worker.clone()] + }; + let revision = commit_worker_projection_changes(&tx, workspace_id, &changed_workers)?; + let changes = + worker_catalog_upsert_changes(&tx, workspace_id, changed_workers.as_slice())?; + tx.commit()?; + Ok(WorkerRegistryProjectionCommit { revision, changes }) }) } @@ -5976,9 +6017,9 @@ impl ControlPlaneStore for SqliteWorkspaceStore { &self, workspace_id: &str, worker: &RuntimeWorkerRef, - ) -> Result { - self.with_conn(|conn| { - let tx = conn.unchecked_transaction()?; + ) -> Result { + self.with_conn_mut(|conn| { + let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; tx.execute( r#"UPDATE worker_workdir_links SET unlinked_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') @@ -5989,7 +6030,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore { "DELETE FROM worker_registry WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3", params![workspace_id, worker.runtime_id, worker.worker_id], )?; - if changed != 0 { + let commit = if changed != 0 { tx.execute( "UPDATE worker_create_reservations SET state = 'removed', updated_at = ?4 @@ -6002,9 +6043,27 @@ impl ControlPlaneStore for SqliteWorkspaceStore { chrono::Utc::now().to_rfc3339(), ], )?; - } + commit_worker_catalog_removal(&tx, workspace_id, worker)? + } else if let Some(revision) = tx + .query_row( + "SELECT projection_revision FROM worker_registry_projection_removals WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3", + params![workspace_id, worker.runtime_id, worker.worker_id], + |row| row.get::<_, i64>(0), + ) + .optional()? + { + WorkerRegistryProjectionCommit { + revision: revision.max(0) as u64, + changes: vec![WorkerCatalogChange::Removed(worker.clone())], + } + } else { + WorkerRegistryProjectionCommit { + revision: current_worker_projection_revision(&tx, workspace_id)?, + changes: Vec::new(), + } + }; tx.commit()?; - Ok(changed > 0) + Ok(commit) }) } @@ -7275,12 +7334,13 @@ impl ControlPlaneStore for SqliteWorkspaceStore { self.with_conn(|conn| { conn.execute( r#"INSERT INTO workdir_registry ( - workspace_id, workdir_id, runtime_id, repository_id, + workspace_id, workdir_id, display_name, runtime_id, repository_id, creation_selector, creation_ref, creation_tree, current_selector, current_ref, current_tree, observed_at_epoch_seconds, materialization_status, cleanliness, created_at, updated_at - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15) + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16) ON CONFLICT(workspace_id, workdir_id) DO UPDATE SET + display_name = excluded.display_name, runtime_id = excluded.runtime_id, repository_id = excluded.repository_id, creation_selector = excluded.creation_selector, @@ -7296,6 +7356,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore { params![ record.workspace_id, record.workdir_id, + record.display_name, record.runtime_id, record.repository_id, record.creation_selector, @@ -7474,10 +7535,15 @@ impl ControlPlaneStore for SqliteWorkspaceStore { record: &WorkerWorkdirLinkRecord, reservation_id: &str, ) -> Result { - if record.role != "attachment" || record.unlinked_at.is_some() { + workdir::WorkdirAttachmentAlias::new(record.alias.clone()).map_err(|error| { + Error::WorkdirAttachmentConflict(format!( + "invalid Workdir attachment alias `{}`: {error}", + record.alias + )) + })?; + if record.unlinked_at.is_some() { return Err(Error::WorkdirAttachmentConflict( - "reserved attachment finalization requires an active canonical attachment" - .to_string(), + "reserved attachment finalization requires an active attachment".to_string(), )); } self.with_conn(|conn| { @@ -7501,9 +7567,9 @@ impl ControlPlaneStore for SqliteWorkspaceStore { } tx.execute( r#"INSERT INTO worker_workdir_links ( - workspace_id, runtime_id, worker_id, workdir_id, role, linked_at, unlinked_at + workspace_id, runtime_id, worker_id, workdir_id, alias, linked_at, unlinked_at ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, NULL) - ON CONFLICT(workspace_id, worker_id, workdir_id, role) DO UPDATE SET + ON CONFLICT(workspace_id, worker_id, workdir_id, alias) DO UPDATE SET linked_at = excluded.linked_at, unlinked_at = NULL"#, params![ @@ -7511,7 +7577,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore { record.worker.runtime_id, record.worker.worker_id, record.workdir_id, - record.role, + record.alias, record.linked_at, ], ) @@ -7539,12 +7605,12 @@ impl ControlPlaneStore for SqliteWorkspaceStore { &self, record: &WorkerWorkdirLinkRecord, ) -> Result { - if record.role != "attachment" { - return Err(Error::WorkdirAttachmentConflict(format!( - "unsupported Workdir attachment role `{}`", - record.role - ))); - } + workdir::WorkdirAttachmentAlias::new(record.alias.clone()).map_err(|error| { + Error::WorkdirAttachmentConflict(format!( + "invalid Workdir attachment alias `{}`: {error}", + record.alias + )) + })?; if record.unlinked_at.is_some() { return Err(Error::WorkdirAttachmentConflict( "a new attachment cannot already be unlinked".to_string(), @@ -7583,32 +7649,34 @@ impl ControlPlaneStore for SqliteWorkspaceStore { record.workdir_id ))); } - let active_for_worker = tx + let active_for_alias = tx .query_row( - r#"SELECT workspace_id, runtime_id, worker_id, workdir_id, role, linked_at, unlinked_at + r#"SELECT workspace_id, runtime_id, worker_id, workdir_id, alias, linked_at, unlinked_at FROM worker_workdir_links - WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3 AND unlinked_at IS NULL"#, + WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3 + AND alias = ?4 AND unlinked_at IS NULL"#, params![ record.workspace_id, record.worker.runtime_id, record.worker.worker_id, + record.alias, ], read_worker_workdir_link_record, ) .optional()?; - if let Some(active) = active_for_worker { + if let Some(active) = active_for_alias { if active.workdir_id == record.workdir_id { tx.commit()?; return Ok(active); } return Err(Error::WorkdirAttachmentConflict(format!( - "Worker {}:{} is already attached to Workdir {}", - record.worker.runtime_id, record.worker.worker_id, active.workdir_id + "Worker {}:{} attachment alias `{}` is already bound to Workdir {}", + record.worker.runtime_id, record.worker.worker_id, record.alias, active.workdir_id ))); } let active_for_workdir = tx .query_row( - r#"SELECT workspace_id, runtime_id, worker_id, workdir_id, role, linked_at, unlinked_at + r#"SELECT workspace_id, runtime_id, worker_id, workdir_id, alias, 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], @@ -7623,9 +7691,9 @@ impl ControlPlaneStore for SqliteWorkspaceStore { } let write = tx.execute( r#"INSERT INTO worker_workdir_links ( - workspace_id, runtime_id, worker_id, workdir_id, role, linked_at, unlinked_at + workspace_id, runtime_id, worker_id, workdir_id, alias, linked_at, unlinked_at ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, NULL) - ON CONFLICT(workspace_id, worker_id, workdir_id, role) DO UPDATE SET + ON CONFLICT(workspace_id, worker_id, workdir_id, alias) DO UPDATE SET linked_at = excluded.linked_at, unlinked_at = NULL"#, params![ @@ -7633,7 +7701,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore { record.worker.runtime_id, record.worker.worker_id, record.workdir_id, - record.role, + record.alias, record.linked_at, ], ); @@ -7665,32 +7733,32 @@ impl ControlPlaneStore for SqliteWorkspaceStore { )?; let active = tx .query_row( - r#"SELECT workspace_id, runtime_id, worker_id, workdir_id, role, linked_at, unlinked_at + r#"SELECT workspace_id, runtime_id, worker_id, workdir_id, alias, linked_at, unlinked_at FROM worker_workdir_links - WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3 AND unlinked_at IS NULL"#, - params![workspace_id, worker.runtime_id, worker.worker_id], + WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3 + AND (?4 IS NULL OR workdir_id = ?4) AND unlinked_at IS NULL"#, + params![workspace_id, worker.runtime_id, worker.worker_id, expected_workdir_id], read_worker_workdir_link_record, ) .optional()?; let Some(active) = active else { + if let Some(expected_workdir_id) = expected_workdir_id { + return Err(Error::WorkdirAttachmentConflict(format!( + "Worker {}:{} is not attached to Workdir {expected_workdir_id}", + worker.runtime_id, worker.worker_id + ))); + } 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 {}:{} is attached to Workdir {}, not {expected_workdir_id}", - worker.runtime_id, worker.worker_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 worker_id = ?3 AND unlinked_at IS NULL"#, - params![workspace_id, worker.runtime_id, worker.worker_id, unlinked_at], + WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3 + AND (?5 IS NULL OR workdir_id = ?5) AND unlinked_at IS NULL"#, + params![workspace_id, worker.runtime_id, worker.worker_id, unlinked_at, expected_workdir_id], )?; - if changed != 1 { + if changed == 0 || (expected_workdir_id.is_some() && changed != 1) { return Err(Error::WorkdirAttachmentConflict(format!( "Worker {}:{} attachment changed during detach", worker.runtime_id, worker.worker_id @@ -7729,7 +7797,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore { ) -> Result> { self.with_conn(|conn| { let mut stmt = conn.prepare( - r#"SELECT workspace_id, runtime_id, worker_id, workdir_id, role, linked_at, unlinked_at + r#"SELECT workspace_id, runtime_id, worker_id, workdir_id, alias, linked_at, unlinked_at FROM worker_workdir_links WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3 AND unlinked_at IS NULL ORDER BY linked_at DESC"#, @@ -7750,7 +7818,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore { ) -> Result> { self.with_conn(|conn| { let mut stmt = conn.prepare( - r#"SELECT workspace_id, runtime_id, worker_id, workdir_id, role, linked_at, unlinked_at + r#"SELECT workspace_id, runtime_id, worker_id, workdir_id, alias, linked_at, unlinked_at FROM worker_workdir_links WHERE workspace_id = ?1 AND workdir_id = ?2 AND unlinked_at IS NULL ORDER BY linked_at DESC"#, @@ -8445,7 +8513,7 @@ fn read_worker_workdir_link_record( workspace_id: row.get(0)?, worker: RuntimeWorkerRef::new(row.get::<_, String>(1)?, row.get::<_, String>(2)?), workdir_id: row.get(3)?, - role: row.get(4)?, + alias: row.get(4)?, linked_at: row.get(5)?, unlinked_at: row.get(6)?, }) @@ -8621,9 +8689,7 @@ fn read_worker_registry_projection( workspace_id: Some(workspace_id.to_string()), display_name: Some(registry.display_name.clone()), profile: registry.profile.clone(), - repository_id: None, - repository_key: None, - working_directory_id: None, + workdir_attachments: Vec::new(), } }; observed_worker.availability = availability; @@ -8820,6 +8886,60 @@ fn set_worker_observation_unavailable( Ok(true) } +fn worker_catalog_upsert_changes( + conn: &Connection, + workspace_id: &str, + changed_workers: &[RuntimeWorkerRef], +) -> Result> { + changed_workers + .iter() + .map(|worker| { + read_worker_registry_projection(conn, workspace_id, worker)?.map_or_else( + || { + Err(Error::Store(format!( + "committed Worker catalog row disappeared before projection payload was captured: {}:{}", + worker.runtime_id, worker.worker_id + ))) + }, + |record| Ok(WorkerCatalogChange::Upsert(record)), + ) + }) + .collect() +} + +/// Commits the durable removal fence in the same transaction as catalog deletion. +/// Existing fences are returned as replayable changes so a retry can heal a +/// process-local broadcast missed after the original DB commit. +pub(crate) fn commit_worker_catalog_removal( + conn: &Connection, + workspace_id: &str, + worker: &RuntimeWorkerRef, +) -> Result { + if let Some(revision) = conn + .query_row( + "SELECT projection_revision FROM worker_registry_projection_removals WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3", + params![workspace_id, worker.runtime_id, worker.worker_id], + |row| row.get::<_, i64>(0), + ) + .optional()? + { + return Ok(WorkerRegistryProjectionCommit { + revision: revision.max(0) as u64, + changes: vec![WorkerCatalogChange::Removed(worker.clone())], + }); + } + let revision = + commit_worker_projection_changes(conn, workspace_id, std::slice::from_ref(worker))?; + conn.execute( + "INSERT INTO worker_registry_projection_removals (workspace_id, runtime_id, worker_id, projection_revision) VALUES (?1, ?2, ?3, ?4)", + params![workspace_id, worker.runtime_id, worker.worker_id, revision as i64], + )?; + Ok(WorkerRegistryProjectionCommit { + revision, + changes: vec![WorkerCatalogChange::Removed(worker.clone())], + }) +} + fn commit_worker_projection_changes( conn: &Connection, workspace_id: &str, @@ -9236,7 +9356,7 @@ fn require_expected_ticket_assignment( fn workdir_registry_select_sql(where_clause: &str) -> String { format!( - "SELECT workspace_id, workdir_id, runtime_id, repository_id, \ + "SELECT workspace_id, workdir_id, display_name, runtime_id, repository_id, \ creation_selector, creation_ref, creation_tree, \ current_selector, current_ref, current_tree, observed_at_epoch_seconds, \ materialization_status, cleanliness, created_at, updated_at \ @@ -9250,19 +9370,20 @@ fn read_workdir_registry_record( Ok(WorkdirRegistryRecord { workspace_id: row.get(0)?, workdir_id: row.get(1)?, - runtime_id: row.get(2)?, - repository_id: row.get(3)?, - creation_selector: row.get(4)?, - creation_ref: row.get(5)?, - creation_tree: row.get(6)?, - current_selector: row.get(7)?, - current_ref: row.get(8)?, - current_tree: row.get(9)?, - observed_at_epoch_seconds: row.get::<_, Option>(10)?.map(|value| value as u64), - materialization_status: row.get(11)?, - cleanliness: row.get(12)?, - created_at: row.get(13)?, - updated_at: row.get(14)?, + display_name: row.get(2)?, + runtime_id: row.get(3)?, + repository_id: row.get(4)?, + creation_selector: row.get(5)?, + creation_ref: row.get(6)?, + creation_tree: row.get(7)?, + current_selector: row.get(8)?, + current_ref: row.get(9)?, + current_tree: row.get(10)?, + observed_at_epoch_seconds: row.get::<_, Option>(11)?.map(|value| value as u64), + materialization_status: row.get(12)?, + cleanliness: row.get(13)?, + created_at: row.get(14)?, + updated_at: row.get(15)?, }) } @@ -10500,10 +10621,31 @@ fn migrate_worker_registry_projection_v61_to_v62(conn: &Connection) -> Result<() tx.execute_batch(WORKER_REGISTRY_PROJECTION_SCHEMA)?; tx.execute( "INSERT INTO __yoi_schema_migrations (version, name) VALUES (?1, ?2)", - params![ - LATEST_SCHEMA_VERSION, - WORKER_REGISTRY_PROJECTION_MIGRATION_NAME - ], + params![62, WORKER_REGISTRY_PROJECTION_MIGRATION_NAME], + )?; + tx.commit()?; + Ok(()) +} + +fn migrate_multi_workdir_attachments_v62_to_v63(conn: &Connection) -> Result<()> { + let tx = rusqlite::Transaction::new_unchecked(conn, TransactionBehavior::Exclusive)?; + if !column_exists(&tx, "workdir_registry", "display_name")? { + tx.execute_batch("ALTER TABLE workdir_registry ADD COLUMN display_name TEXT;")?; + } + if column_exists(&tx, "worker_workdir_links", "role")? + && !column_exists(&tx, "worker_workdir_links", "alias")? + { + tx.execute_batch("ALTER TABLE worker_workdir_links RENAME COLUMN role TO alias;")?; + } + tx.execute_batch( + "DROP INDEX IF EXISTS worker_workdir_links_active_worker_unique; + CREATE UNIQUE INDEX IF NOT EXISTS worker_workdir_links_active_alias_unique + ON worker_workdir_links(workspace_id, worker_id, alias) + WHERE unlinked_at IS NULL;", + )?; + tx.execute( + "INSERT INTO __yoi_schema_migrations (version, name) VALUES (?1, ?2)", + params![63, MULTI_WORKDIR_ATTACHMENTS_MIGRATION_NAME], )?; tx.commit()?; Ok(()) @@ -11884,6 +12026,10 @@ mod tests { version: 62, name: WORKER_REGISTRY_PROJECTION_MIGRATION_NAME.to_string(), }, + WorkspaceSchemaMigrationStep { + version: 63, + name: MULTI_WORKDIR_ATTACHMENTS_MIGRATION_NAME.to_string(), + }, ] ); @@ -11935,6 +12081,7 @@ mod tests { REMOVE_WORKER_RUN_GENERATION_MIGRATION_NAME.to_string(), ), (62, WORKER_REGISTRY_PROJECTION_MIGRATION_NAME.to_string()), + (63, MULTI_WORKDIR_ATTACHMENTS_MIGRATION_NAME.to_string()), ] ); assert!(!table_exists(conn, "trusted_runtime_records")?); @@ -12005,7 +12152,7 @@ mod tests { .iter() .map(|migration| migration.version) .collect::>(), - vec![52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62] + vec![52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63] ); SqliteWorkspaceStore::migrate_database(&path).unwrap(); let conn = Connection::open(&path).unwrap(); @@ -12013,7 +12160,9 @@ mod tests { current_schema_version(&conn).unwrap(), LATEST_SCHEMA_VERSION ); - assert_eq!(workspace_schema_migration_history(&conn).unwrap().len(), 13); + assert_eq!(workspace_schema_migration_history(&conn).unwrap().len(), 14); + assert!(column_exists(&conn, "worker_workdir_links", "alias").unwrap()); + assert!(!column_exists(&conn, "worker_workdir_links", "role").unwrap()); } #[test] @@ -13499,9 +13648,11 @@ mod tests { .is_err() ); assert!( - store + !store .delete_worker_registry("workspace-a", &reserved_worker) .unwrap() + .changes + .is_empty() ); let removed_state: String = store .with_conn(|conn| { @@ -14202,9 +14353,11 @@ INSERT INTO worker_registry ( .is_none() ); assert!( - store + !store .delete_worker_registry("workspace-role", &worker.worker) .unwrap() + .changes + .is_empty() ); let removed_worker_assignment = TicketRoleAssignmentRecord { assignment_id: "contributor-removed-worker".to_string(), @@ -14722,6 +14875,7 @@ INSERT INTO worker_registry ( let workdir = WorkdirRegistryRecord { workspace_id: "local-dev".to_string(), workdir_id: "0000019a00000000001".to_string(), + display_name: Some("Checkout".to_string()), runtime_id: "embedded".to_string(), repository_id: "repo".to_string(), creation_selector: Some("develop".to_string()), @@ -14740,6 +14894,7 @@ INSERT INTO worker_registry ( let unmanaged_workdir = WorkdirRegistryRecord { workspace_id: "local-dev".to_string(), workdir_id: "runtime-direct".to_string(), + display_name: None, runtime_id: "embedded".to_string(), repository_id: "repo".to_string(), creation_selector: Some("feature".to_string()), @@ -14760,7 +14915,7 @@ INSERT INTO worker_registry ( workspace_id: "local-dev".to_string(), worker: worker.worker.clone(), workdir_id: workdir.workdir_id.clone(), - role: "attachment".to_string(), + alias: "checkout".to_string(), linked_at: "4".to_string(), unlinked_at: None, }; @@ -14825,6 +14980,30 @@ INSERT INTO worker_registry ( Err(Error::WorkdirAttachmentConflict(_)) )); + let invalid_alias = WorkerWorkdirLinkRecord { + alias: "invalid/alias".to_string(), + ..worker_conflict.clone() + }; + assert!(matches!( + store.attach_worker_workdir(&invalid_alias), + Err(Error::WorkdirAttachmentConflict(_)) + )); + let second_attachment = WorkerWorkdirLinkRecord { + alias: "docs".to_string(), + ..worker_conflict + }; + assert_eq!( + store.attach_worker_workdir(&second_attachment).unwrap(), + second_attachment + ); + assert_eq!( + store + .list_worker_workdir_links("local-dev", &worker.worker) + .unwrap() + .len(), + 2 + ); + let second_worker = WorkerRegistryRecord { worker: RuntimeWorkerRef::new("embedded", "2"), display_name: "Browser 2".to_string(), @@ -14919,9 +15098,7 @@ INSERT INTO worker_registry ( workspace_id: Some("local-dev".to_string()), display_name: Some("Known Worker".to_string()), profile: None, - repository_id: None, - repository_key: None, - working_directory_id: None, + workdir_attachments: Vec::new(), }; let orphan = SubscriptionWorker { worker_id: SubscriptionWorkerId::new("orphan").unwrap(), @@ -14937,7 +15114,10 @@ INSERT INTO worker_registry ( "2", ) .unwrap(); - assert_eq!(first.changed_workers, vec![catalog.worker.clone()]); + assert!(matches!( + first.changes.as_slice(), + [WorkerCatalogChange::Upsert(record)] if record.registry.worker == catalog.worker + )); let (_, records) = store .worker_registry_projection_snapshot("local-dev", 10) .unwrap(); @@ -14959,7 +15139,10 @@ INSERT INTO worker_registry ( let unavailable = store .reconcile_worker_registry_snapshot("local-dev", "embedded", 1, 11, &[], "3") .unwrap(); - assert_eq!(unavailable.changed_workers, vec![catalog.worker.clone()]); + assert!(matches!( + unavailable.changes.as_slice(), + [WorkerCatalogChange::Upsert(record)] if record.registry.worker == catalog.worker + )); let projection = store .worker_registry_projection("local-dev", &catalog.worker) .unwrap() @@ -14978,7 +15161,7 @@ INSERT INTO worker_registry ( let duplicate = store .reconcile_worker_registry_snapshot("local-dev", "embedded", 1, 11, &[], "4") .unwrap(); - assert!(duplicate.changed_workers.is_empty()); + assert!(duplicate.changes.is_empty()); let reconnected = store .reconcile_worker_registry_snapshot( @@ -14990,7 +15173,10 @@ INSERT INTO worker_registry ( "5", ) .unwrap(); - assert_eq!(reconnected.changed_workers, vec![catalog.worker.clone()]); + assert!(matches!( + reconnected.changes.as_slice(), + [WorkerCatalogChange::Upsert(record)] if record.registry.worker == catalog.worker + )); let projection = store .worker_registry_projection("local-dev", &catalog.worker) .unwrap() @@ -15003,7 +15189,7 @@ INSERT INTO worker_registry ( let duplicate_reconnect = store .reconcile_worker_registry_snapshot("local-dev", "embedded", 1, 11, &[observed], "6") .unwrap(); - assert!(duplicate_reconnect.changed_workers.is_empty()); + assert!(duplicate_reconnect.changes.is_empty()); } #[tokio::test] @@ -15037,20 +15223,44 @@ INSERT INTO worker_registry ( updated_at: "1".to_string(), }) .unwrap(); + let first_removal = store + .delete_worker_registry("local-dev", &worker_ref) + .unwrap(); + assert!(matches!( + first_removal.changes.as_slice(), + [WorkerCatalogChange::Removed(worker)] if worker == &worker_ref + )); + let duplicate_removal = store + .delete_worker_registry("local-dev", &worker_ref) + .unwrap(); + assert!(matches!( + duplicate_removal.changes.as_slice(), + [WorkerCatalogChange::Removed(worker)] if worker == &worker_ref + )); + assert_eq!(duplicate_removal.revision, first_removal.revision); + let stale_upsert = store + .upsert_worker_registry(&WorkerRegistryRecord { + workspace_id: "local-dev".to_string(), + worker: worker_ref.clone(), + display_name: "Stale Removed Worker".to_string(), + profile: None, + retention_state: "normal".to_string(), + transcript_ref: None, + session_ref: None, + summary_ref: None, + diagnostics_ref: None, + created_at: "1".to_string(), + updated_at: "stale-after-removal".to_string(), + }) + .unwrap(); + assert!(stale_upsert.changes.is_empty()); + assert_eq!(stale_upsert.revision, first_removal.revision); assert!( store - .delete_worker_registry("local-dev", &worker_ref) + .get_worker_registry("local-dev", &worker_ref) .unwrap() + .is_none() ); - let first_removal = store - .publish_worker_registry_removal("local-dev", &worker_ref) - .unwrap(); - assert_eq!(first_removal.changed_workers, vec![worker_ref.clone()]); - let duplicate_removal = store - .publish_worker_registry_removal("local-dev", &worker_ref) - .unwrap(); - assert!(duplicate_removal.changed_workers.is_empty()); - assert_eq!(duplicate_removal.revision, first_removal.revision); let event = SubscriptionWorker { worker_id: SubscriptionWorkerId::new("removed").unwrap(), runtime_id: Some("embedded".to_string()), @@ -15063,14 +15273,12 @@ INSERT INTO worker_registry ( workspace_id: Some("local-dev".to_string()), display_name: None, profile: None, - repository_id: None, - repository_key: None, - working_directory_id: None, + workdir_attachments: Vec::new(), }; let commit = store .apply_worker_registry_observation("local-dev", "embedded", 2, &event, "2") .unwrap(); - assert!(commit.changed_workers.is_empty()); + assert!(commit.changes.is_empty()); assert!( store .worker_registry_projection("local-dev", &worker_ref) @@ -15194,9 +15402,11 @@ INSERT INTO worker_registry ( .is_empty() ); assert!( - store + !store .delete_worker_registry("workspace-control", &subject_record.worker) .unwrap() + .changes + .is_empty() ); assert!( store diff --git a/crates/workspace-server/src/workdir_create_operations.rs b/crates/workspace-server/src/workdir_create_operations.rs index 46daeee99..89d842ec5 100644 --- a/crates/workspace-server/src/workdir_create_operations.rs +++ b/crates/workspace-server/src/workdir_create_operations.rs @@ -25,6 +25,7 @@ pub fn request_fingerprint( repository_id: &str, selector: Option<&str>, requested_runtime_id: Option<&str>, + display_name: Option<&str>, repository_source_fingerprint: &str, repository_source_revision: u64, ) -> String { @@ -33,6 +34,7 @@ pub fn request_fingerprint( Some(repository_id), selector, requested_runtime_id, + display_name, Some(repository_source_fingerprint), ] { match value { @@ -569,6 +571,7 @@ mod tests { "main", Some("develop"), None, + None, "sha256:test", 1, ), @@ -718,7 +721,7 @@ mod tests { ); let mut changed_input = record.clone(); changed_input.request_fingerprint = - request_fingerprint("main", Some("main"), None, "sha256:test", 1); + request_fingerprint("main", Some("main"), None, None, "sha256:test", 1); assert!( store .reserve_workdir_create_operation(&changed_input) diff --git a/crates/workspace-server/src/workdir_removal.rs b/crates/workspace-server/src/workdir_removal.rs index 0cb2d4d78..3fccda15c 100644 --- a/crates/workspace-server/src/workdir_removal.rs +++ b/crates/workspace-server/src/workdir_removal.rs @@ -716,7 +716,7 @@ fn load_workdir_record( workdir_id: &str, ) -> Result> { conn.query_row( - r#"SELECT workspace_id, workdir_id, runtime_id, repository_id, + r#"SELECT workspace_id, workdir_id, display_name, runtime_id, repository_id, creation_selector, creation_ref, creation_tree, current_selector, current_ref, current_tree, observed_at_epoch_seconds, materialization_status, cleanliness, created_at, updated_at @@ -726,19 +726,20 @@ fn load_workdir_record( Ok(WorkdirRegistryRecord { workspace_id: row.get(0)?, workdir_id: row.get(1)?, - runtime_id: row.get(2)?, - repository_id: row.get(3)?, - creation_selector: row.get(4)?, - creation_ref: row.get(5)?, - creation_tree: row.get(6)?, - current_selector: row.get(7)?, - current_ref: row.get(8)?, - current_tree: row.get(9)?, - observed_at_epoch_seconds: row.get::<_, Option>(10)?.map(|value| value as u64), - materialization_status: row.get(11)?, - cleanliness: row.get(12)?, - created_at: row.get(13)?, - updated_at: row.get(14)?, + display_name: row.get(2)?, + runtime_id: row.get(3)?, + repository_id: row.get(4)?, + creation_selector: row.get(5)?, + creation_ref: row.get(6)?, + creation_tree: row.get(7)?, + current_selector: row.get(8)?, + current_ref: row.get(9)?, + current_tree: row.get(10)?, + observed_at_epoch_seconds: row.get::<_, Option>(11)?.map(|value| value as u64), + materialization_status: row.get(12)?, + cleanliness: row.get(13)?, + created_at: row.get(14)?, + updated_at: row.get(15)?, }) }, ) @@ -934,6 +935,7 @@ mod tests { let workdir = WorkdirRegistryRecord { workspace_id: "workspace-a".to_string(), workdir_id: "workdir-a".to_string(), + display_name: None, runtime_id: "runtime-a".to_string(), repository_id: "repository-a".to_string(), creation_selector: Some("refs/heads/develop".to_string()), diff --git a/crates/workspace-server/src/worker_projection.rs b/crates/workspace-server/src/worker_projection.rs index 36b01a001..22095f8c0 100644 --- a/crates/workspace-server/src/worker_projection.rs +++ b/crates/workspace-server/src/worker_projection.rs @@ -11,18 +11,12 @@ use tokio::task::JoinHandle; use crate::runtime_subscription::{ BrokerSubscriptionEvent, RuntimeSubscriptionBroker, RuntimeSubscriptionBrokerError, }; -use crate::store::{ - ControlPlaneStore, WorkerRegistryProjectionCommit, WorkerRegistryProjectionRecord, -}; +use crate::store::{ControlPlaneStore, WorkerCatalogChange, WorkerRegistryProjectionCommit}; use worker_runtime::identity::RuntimeWorkerRef; const PROJECTION_EVENT_CAPACITY: usize = 256; -#[derive(Debug, Clone)] -pub enum WorkerProjectionChange { - Upsert(WorkerRegistryProjectionRecord), - Removed(RuntimeWorkerRef), -} +pub type WorkerProjectionChange = WorkerCatalogChange; #[derive(Debug, Clone)] pub struct WorkerProjectionEvent { @@ -144,27 +138,13 @@ impl WorkerProjectionService { self.publish_commit(commit) } - pub fn publish_catalog_change(&self, worker: &RuntimeWorkerRef) -> crate::Result<()> { + pub fn refresh(&self, worker: &RuntimeWorkerRef) -> crate::Result<()> { let commit = self .store - .publish_worker_registry_catalog_change(self.workspace_id.as_str(), worker)?; + .commit_worker_registry_projection_refresh(self.workspace_id.as_str(), worker)?; self.publish_commit(commit) } - pub fn publish_removed(&self, worker: RuntimeWorkerRef) -> crate::Result<()> { - let commit = self - .store - .publish_worker_registry_removal(self.workspace_id.as_str(), &worker)?; - if commit.changed_workers.is_empty() { - return Ok(()); - } - let _ = self.events.send(WorkerProjectionEvent { - revision: commit.revision, - changes: vec![WorkerProjectionChange::Removed(worker)], - }); - Ok(()) - } - async fn apply_broker_event( &self, runtime_id: &str, @@ -236,38 +216,21 @@ impl WorkerProjectionService { self.publish_commit(commit) } - fn publish_commit(&self, commit: WorkerRegistryProjectionCommit) -> crate::Result<()> { - if commit.changed_workers.is_empty() { + /// Broadcast an already committed catalog payload. Persistence always happens + /// before this call; a process crash therefore closes subscribers, whose next + /// subscription starts from the durable snapshot. Retrying a removal returns + /// the same durable change and can heal an in-process missed broadcast. + pub(crate) fn publish_commit( + &self, + commit: WorkerRegistryProjectionCommit, + ) -> crate::Result<()> { + if commit.changes.is_empty() { return Ok(()); } - let changes = commit - .changed_workers - .iter() - .filter_map(|worker| { - match self - .store - .worker_registry_projection(self.workspace_id.as_str(), worker) - { - Ok(Some(record)) => Some(WorkerProjectionChange::Upsert(record)), - Ok(None) => None, - Err(error) => { - tracing::warn!( - runtime_id = %worker.runtime_id, - worker_id = %worker.worker_id, - error = %error, - "failed to load committed Worker projection" - ); - None - } - } - }) - .collect::>(); - if !changes.is_empty() { - let _ = self.events.send(WorkerProjectionEvent { - revision: commit.revision, - changes, - }); - } + let _ = self.events.send(WorkerProjectionEvent { + revision: commit.revision, + changes: commit.changes, + }); Ok(()) } } diff --git a/crates/workspace-server/src/workspace_deletion.rs b/crates/workspace-server/src/workspace_deletion.rs index 3dbfa46f7..e4c5f5b54 100644 --- a/crates/workspace-server/src/workspace_deletion.rs +++ b/crates/workspace-server/src/workspace_deletion.rs @@ -1198,7 +1198,7 @@ mod tests { ); assert!( - store + !store .delete_worker_registry( &workspace_id, &worker_runtime::identity::RuntimeWorkerRef { @@ -1207,6 +1207,8 @@ mod tests { }, ) .expect("remove worker registry") + .changes + .is_empty() ); let state: String = store .with_conn(|conn| { diff --git a/crates/workspace-server/src/workspace_subscription.rs b/crates/workspace-server/src/workspace_subscription.rs index e036d0ed3..dac5df2a2 100644 --- a/crates/workspace-server/src/workspace_subscription.rs +++ b/crates/workspace-server/src/workspace_subscription.rs @@ -413,9 +413,7 @@ fn project_registry_worker( workspace_id: Some(record.registry.workspace_id.clone()), display_name: Some(record.registry.display_name.clone()), profile: record.registry.profile.clone(), - repository_id: None, - repository_key: None, - working_directory_id: None, + workdir_attachments: Vec::new(), } }; worker.runtime_id = Some(runtime_id); @@ -440,20 +438,30 @@ fn project_working_directory( else { return false; }; - if let Some(link) = links.first() { - let Ok(Some(workdir)) = api - .store - .get_workdir_registry(&api.config.workspace_id, &link.workdir_id) - else { - return false; - }; - let Ok(working_directory_id) = - protocol::subscription::SubscriptionWorkdirId::new(workdir.workdir_id) - else { - return false; - }; - worker.repository_id = Some(workdir.repository_id); - worker.working_directory_id = Some(working_directory_id); + if !links.is_empty() { + let mut attachments = Vec::with_capacity(links.len()); + for link in links { + let Ok(Some(workdir)) = api + .store + .get_workdir_registry(&api.config.workspace_id, &link.workdir_id) + else { + return false; + }; + let Ok(working_directory_id) = + protocol::subscription::SubscriptionWorkdirId::new(workdir.workdir_id) + else { + return false; + }; + attachments.push( + protocol::subscription::SubscriptionWorkerWorkdirAttachment { + alias: link.alias, + repository_id: Some(workdir.repository_id), + repository_key: None, + working_directory_id, + }, + ); + } + worker.workdir_attachments = attachments; } else { let Ok(has_link_history) = api .store @@ -462,26 +470,28 @@ fn project_working_directory( return false; }; if has_link_history { - worker.repository_id = None; - worker.repository_key = None; - worker.working_directory_id = None; + worker.workdir_attachments.clear(); } } - project_repository_key(api, worker) + project_repository_keys(api, worker) } -fn project_repository_key(api: &WorkspaceApi, worker: &mut SubscriptionWorker) -> bool { - let Some(repository_id) = worker.repository_id.take() else { - worker.repository_key = None; - return true; - }; - let Ok(Some(repository)) = api - .store - .get_repository(&api.config.workspace_id, &repository_id) - else { - return false; - }; - worker.repository_key = Some(repository.repository_key); +fn project_repository_keys(api: &WorkspaceApi, worker: &mut SubscriptionWorker) -> bool { + for attachment in &mut worker.workdir_attachments { + let Some(repository_id) = attachment.repository_id.take() else { + if attachment.repository_key.is_none() { + return false; + } + continue; + }; + let Ok(Some(repository)) = api + .store + .get_repository(&api.config.workspace_id, &repository_id) + else { + return false; + }; + attachment.repository_key = Some(repository.repository_key); + } true } diff --git a/web/workspace/src/lib/generated/protocol.ts b/web/workspace/src/lib/generated/protocol.ts index c96d3b430..200a59665 100644 --- a/web/workspace/src/lib/generated/protocol.ts +++ b/web/workspace/src/lib/generated/protocol.ts @@ -236,6 +236,12 @@ export type SubscriptionWorkerAvailability = "observed" | "unavailable"; export type EventSubscriptionSelector = { "topic": "runtime_workers" } | { "topic": "worker_lifecycle", worker_ids: SubscriptionWorkerIds, } | { "topic": "worker_protocol", worker_id: SubscriptionWorkerId, runtime_id?: string | null, } | { "topic": "workspace_workers" } | { "topic": "workspace_workdirs" }; +export type SubscriptionWorkerWorkdirAttachment = { +/** + * Stable Worker-local routing alias. + */ +alias: string, repository_key?: string | null, working_directory_id: SubscriptionWorkdirId, }; + export type SubscriptionWorker = { worker_id: SubscriptionWorkerId, /** * Set by the Workspace Server when projecting a Runtime-owned Worker to clients. @@ -264,14 +270,9 @@ worker_state?: WorkerStateSnapshot | null, /** * Runtime catalog lifecycle compatibility projection; not foreground-state authority. */ -state: SubscriptionWorkerState, has_running_internal_workers: boolean, workspace_id?: string | null, display_name?: string | null, profile?: string | null, -/** - * Workspace-facing Repository key. Runtime producers leave this unset and - * Workspace Server projections replace `repository_id` with this field. - */ -repository_key?: string | null, working_directory_id?: SubscriptionWorkdirId | null, }; +state: SubscriptionWorkerState, has_running_internal_workers: boolean, workspace_id?: string | null, display_name?: string | null, profile?: string | null, workdir_attachments?: Array, }; -export type WorkspaceSubscriptionWorkdir = { working_directory_id: SubscriptionWorkdirId, repository_key: string, state: string, primary_worker_id?: SubscriptionWorkerId | null, }; +export type WorkspaceSubscriptionWorkdir = { working_directory_id: SubscriptionWorkdirId, repository_key: string, state: string, }; export type SubscriptionSnapshot = { "topic": "workers", "data": { workers: Array, } } | { "topic": "worker_protocol", "data": { worker_id: SubscriptionWorkerId, events: Array, } } | { "topic": "workspace_workdirs", "data": { workdirs: Array, } }; diff --git a/web/workspace/src/lib/generated/workdir-api.ts b/web/workspace/src/lib/generated/workdir-api.ts index 7673d592d..b4701eca3 100644 --- a/web/workspace/src/lib/generated/workdir-api.ts +++ b/web/workspace/src/lib/generated/workdir-api.ts @@ -13,9 +13,17 @@ export type WorkingDirectoryCleanupTarget = { kind: string, working_directory_id export type WorkingDirectoryOccupancy = { runtime_id: string, worker_id: string, display_name: string, linked_at: string, }; -export type WorkingDirectorySummary = { working_directory_id: string, repository_key: string, creation_selector?: string | null, creation_ref?: string | null, creation_tree?: string | null, current_selector?: string | null, current_ref?: string | null, current_tree?: string | null, observed_at_epoch_seconds?: number | null, materializer_kind: WorkingDirectoryMaterializerKind, cleanup_target?: WorkingDirectoryCleanupTarget | null, status: WorkingDirectoryStatusKind, cleanliness?: string | null, primary_worker_id?: string | null, occupied_by?: WorkingDirectoryOccupancy | null, }; +export type WorkingDirectorySummary = { working_directory_id: string, +/** + * Optional human-facing label, never a routing key. + */ +display_name?: string | null, repository_key: string, creation_selector?: string | null, creation_ref?: string | null, creation_tree?: string | null, current_selector?: string | null, current_ref?: string | null, current_tree?: string | null, observed_at_epoch_seconds?: number | null, materializer_kind: WorkingDirectoryMaterializerKind, cleanup_target?: WorkingDirectoryCleanupTarget | null, status: WorkingDirectoryStatusKind, cleanliness?: string | null, occupied_by?: WorkingDirectoryOccupancy | null, }; -export type WorkingDirectoryCreateRequest = { runtime_id?: string | null, repository_key: string, selector?: string | null, operation_id?: string | null, }; +export type WorkingDirectoryCreateRequest = { runtime_id?: string | null, +/** + * Optional human-facing label, independent of every Worker attachment alias. + */ +display_name?: string | null, repository_key: string, selector?: string | null, operation_id?: string | null, }; export type WorkingDirectoryListResponse = { workspace_id: string, items: Array, diagnostics: Array, }; diff --git a/web/workspace/src/lib/generated/worker-launch-api.ts b/web/workspace/src/lib/generated/worker-launch-api.ts index 7ed93f50d..022413353 100644 --- a/web/workspace/src/lib/generated/worker-launch-api.ts +++ b/web/workspace/src/lib/generated/worker-launch-api.ts @@ -15,11 +15,19 @@ export type WorkingDirectoryCleanupTarget = { kind: string, working_directory_id export type RuntimeWorkingDirectoryCleanupTarget = { kind: string, working_directory_id: string, repository_id: string, }; -export type RuntimeWorkingDirectorySummary = { working_directory_id: string, repository_id: string, creation_selector?: string | null, creation_ref?: string | null, creation_tree?: string | null, current_selector?: string | null, current_ref?: string | null, current_tree?: string | null, observed_at_epoch_seconds?: number | null, materializer_kind: WorkingDirectoryMaterializerKind, cleanup_target?: RuntimeWorkingDirectoryCleanupTarget | null, status: WorkingDirectoryStatusKind, cleanliness?: string | null, primary_worker_id?: string | null, occupied_by?: WorkingDirectoryOccupancy | null, }; +export type RuntimeWorkingDirectorySummary = { working_directory_id: string, +/** + * Optional human-facing label, never a routing key. + */ +display_name?: string | null, repository_id: string, creation_selector?: string | null, creation_ref?: string | null, creation_tree?: string | null, current_selector?: string | null, current_ref?: string | null, current_tree?: string | null, observed_at_epoch_seconds?: number | null, materializer_kind: WorkingDirectoryMaterializerKind, cleanup_target?: RuntimeWorkingDirectoryCleanupTarget | null, status: WorkingDirectoryStatusKind, cleanliness?: string | null, occupied_by?: WorkingDirectoryOccupancy | null, }; export type WorkingDirectoryOccupancy = { runtime_id: string, worker_id: string, display_name: string, linked_at: string, }; -export type WorkingDirectorySummary = { working_directory_id: string, repository_key: string, creation_selector?: string | null, creation_ref?: string | null, creation_tree?: string | null, current_selector?: string | null, current_ref?: string | null, current_tree?: string | null, observed_at_epoch_seconds?: number | null, materializer_kind: WorkingDirectoryMaterializerKind, cleanup_target?: WorkingDirectoryCleanupTarget | null, status: WorkingDirectoryStatusKind, cleanliness?: string | null, primary_worker_id?: string | null, occupied_by?: WorkingDirectoryOccupancy | null, }; +export type WorkingDirectorySummary = { working_directory_id: string, +/** + * Optional human-facing label, never a routing key. + */ +display_name?: string | null, repository_key: string, creation_selector?: string | null, creation_ref?: string | null, creation_tree?: string | null, current_selector?: string | null, current_ref?: string | null, current_tree?: string | null, observed_at_epoch_seconds?: number | null, materializer_kind: WorkingDirectoryMaterializerKind, cleanup_target?: WorkingDirectoryCleanupTarget | null, status: WorkingDirectoryStatusKind, cleanliness?: string | null, occupied_by?: WorkingDirectoryOccupancy | null, }; export type WorkerWorkspaceSummary = { visibility: string, identity: string, workspace_id?: string | null, }; @@ -27,7 +35,9 @@ export type WorkerImplementationSummary = { kind: string, display_hint: string, export type WorkerCapabilitySummary = { can_stop: boolean, can_spawn_followup: boolean, }; -export type WorkerLaunchWorkerSummary = { runtime_id: string, worker_id: string, host_id: string, display_name: string, label: string, profile: string | null, singleton_key: string | null, tags: Array, workspace: WorkerWorkspaceSummary, state: string, last_seen_at: string | null, pinned: boolean, retention_state: string, implementation: WorkerImplementationSummary, capabilities: WorkerCapabilitySummary, working_directory?: RuntimeWorkingDirectorySummary | null, diagnostics: Array, }; +export type RuntimeWorkerWorkdirAttachmentSummary = { alias: string, working_directory: RuntimeWorkingDirectorySummary, }; + +export type WorkerLaunchWorkerSummary = { runtime_id: string, worker_id: string, host_id: string, display_name: string, label: string, profile: string | null, singleton_key: string | null, tags: Array, workspace: WorkerWorkspaceSummary, state: string, last_seen_at: string | null, pinned: boolean, retention_state: string, implementation: WorkerImplementationSummary, capabilities: WorkerCapabilitySummary, workdir_attachments?: Array, diagnostics: Array, }; export type WorkerLaunchRuntimeOption = { runtime_id: string, display_name: string, built_in: boolean, worker_creation_available: boolean, working_directory_required: boolean, status: string, diagnostics: Array, }; @@ -37,11 +47,15 @@ export type WorkingDirectoryRepositoryOption = { repository_key: string, default export type WorkerLaunchOptionsResponse = { workspace_id: string, runtimes: Array, default_profile: string | null, profiles: Array, repositories: Array, working_directories: Array, diagnostics: Array, }; -export type BrowserWorkerWorkingDirectorySelection = { working_directory_id: string, relative_cwd: string | null, }; +export type BrowserWorkerWorkingDirectorySelection = { +/** + * Stable Worker-local routing alias. + */ +alias: string, working_directory_id: string, relative_cwd: string | null, }; export type CreateWorkspaceWorkerTicketAssignmentRequest = { ticket_id: string, operation_id: string, }; -export type CreateWorkspaceWorkerRequest = { runtime_id: string, display_name: string, profile: string | null, ticket_assignment: CreateWorkspaceWorkerTicketAssignmentRequest | null, initial_submit: Array, working_directory: BrowserWorkerWorkingDirectorySelection | null, +export type CreateWorkspaceWorkerRequest = { runtime_id: string, display_name: string, profile: string | null, ticket_assignment: CreateWorkspaceWorkerTicketAssignmentRequest | null, initial_submit: Array, workdir_attachments: Array, /** * Backend idempotency key used only for authenticated Worker-owned spawn/control. */ diff --git a/web/workspace/src/lib/workspace/api/workdirs.ts b/web/workspace/src/lib/workspace/api/workdirs.ts index f7494f70b..5cd6dbaa3 100644 --- a/web/workspace/src/lib/workspace/api/workdirs.ts +++ b/web/workspace/src/lib/workspace/api/workdirs.ts @@ -11,6 +11,7 @@ import type { const SUMMARY_KEYS = new Set([ "working_directory_id", + "display_name", "repository_key", "creation_selector", "creation_ref", @@ -23,11 +24,11 @@ const SUMMARY_KEYS = new Set([ "cleanup_target", "status", "cleanliness", - "primary_worker_id", "occupied_by", ]); const CREATE_REQUEST_KEYS = new Set([ "runtime_id", + "display_name", "repository_key", "selector", "operation_id", @@ -84,6 +85,7 @@ export function validateWorkingDirectoryCreateRequest( repository_key: stringField(record, "repository_key"), }; assignOptionalString(request, record, "runtime_id"); + assignOptionalString(request, record, "display_name"); assignOptionalString(request, record, "selector"); assignOptionalString(request, record, "operation_id"); return request; @@ -124,6 +126,7 @@ export function parseWorkingDirectorySummary( "unknown", ]), }; + assignOptionalString(summary, record, "display_name"); assignOptionalString(summary, record, "creation_selector"); assignOptionalString(summary, record, "creation_ref"); assignOptionalString(summary, record, "creation_tree"); @@ -131,7 +134,6 @@ export function parseWorkingDirectorySummary( assignOptionalString(summary, record, "current_ref"); assignOptionalString(summary, record, "current_tree"); assignOptionalString(summary, record, "cleanliness"); - assignOptionalString(summary, record, "primary_worker_id"); if (record.observed_at_epoch_seconds !== undefined) { const observedAt = record.observed_at_epoch_seconds; if (observedAt === null) { diff --git a/web/workspace/src/lib/workspace/api/workers.test.ts b/web/workspace/src/lib/workspace/api/workers.test.ts index f80f632cb..de930e42d 100644 --- a/web/workspace/src/lib/workspace/api/workers.test.ts +++ b/web/workspace/src/lib/workspace/api/workers.test.ts @@ -139,10 +139,11 @@ Deno.test("Worker create request parser requires the complete shared request", ( { kind: "text", content: "Implement T-565." }, { kind: "flow", selector: "builtin:coder-review" }, ], - working_directory: { + workdir_attachments: [{ + alias: "source", working_directory_id: "workdir-a", relative_cwd: null, - }, + }], control_operation_id: null, }; diff --git a/web/workspace/src/lib/workspace/api/workers.ts b/web/workspace/src/lib/workspace/api/workers.ts index 15e5ee522..f085c4727 100644 --- a/web/workspace/src/lib/workspace/api/workers.ts +++ b/web/workspace/src/lib/workspace/api/workers.ts @@ -8,6 +8,7 @@ import type { DiagnosticSeverity, RuntimeWorkingDirectoryCleanupTarget, RuntimeWorkingDirectorySummary, + RuntimeWorkerWorkdirAttachmentSummary, WorkerCapabilitySummary, WorkerImplementationSummary, WorkerLaunchOptionsResponse, @@ -261,6 +262,7 @@ function runtimeWorkingDirectory( item, [ "working_directory_id", + "display_name", "repository_id", "creation_selector", "creation_ref", @@ -273,7 +275,6 @@ function runtimeWorkingDirectory( "cleanup_target", "status", "cleanliness", - "primary_worker_id", "occupied_by", ], label, @@ -318,6 +319,11 @@ function runtimeWorkingDirectory( item.working_directory_id, `${label}.working_directory_id`, ), + display_name: optional( + item.display_name, + `${label}.display_name`, + string, + ), repository_id: string(item.repository_id, `${label}.repository_id`), creation_selector: optional( item.creation_selector, @@ -350,15 +356,25 @@ function runtimeWorkingDirectory( ), status: status as RuntimeWorkingDirectorySummary["status"], cleanliness: optional(item.cleanliness, `${label}.cleanliness`, string), - primary_worker_id: optional( - item.primary_worker_id, - `${label}.primary_worker_id`, - string, - ), occupied_by: occupied, }; } +function runtimeWorkdirAttachment( + value: unknown, + label: string, +): RuntimeWorkerWorkdirAttachmentSummary { + const item = record(value, label); + exact(item, ["alias", "working_directory"], label); + return { + alias: string(item.alias, `${label}.alias`), + working_directory: runtimeWorkingDirectory( + item.working_directory, + `${label}.working_directory`, + ), + }; +} + function workerSummary( value: unknown, label: string, @@ -382,7 +398,7 @@ function workerSummary( "retention_state", "implementation", "capabilities", - "working_directory", + "workdir_attachments", "diagnostics", ], label, @@ -406,11 +422,13 @@ function workerSummary( `${label}.implementation`, ), capabilities: capabilitySummary(item.capabilities, `${label}.capabilities`), - working_directory: optional( - item.working_directory, - `${label}.working_directory`, - runtimeWorkingDirectory, - ), + workdir_attachments: item.workdir_attachments === undefined + ? undefined + : array( + item.workdir_attachments, + `${label}.workdir_attachments`, + runtimeWorkdirAttachment, + ), diagnostics: array(item.diagnostics, `${label}.diagnostics`, diagnostic), }; } @@ -464,8 +482,9 @@ function workingDirectorySelection( label: string, ): BrowserWorkerWorkingDirectorySelection { const item = record(value, label); - exact(item, ["working_directory_id", "relative_cwd"], label); + exact(item, ["alias", "working_directory_id", "relative_cwd"], label); return { + alias: string(item.alias, `${label}.alias`), working_directory_id: string( item.working_directory_id, `${label}.working_directory_id`, @@ -638,7 +657,7 @@ export function parseCreateWorkspaceWorkerRequest( "profile", "ticket_assignment", "initial_submit", - "working_directory", + "workdir_attachments", "control_operation_id", ], "Worker create request", @@ -651,9 +670,11 @@ export function parseCreateWorkspaceWorkerRequest( ? null : ticketAssignment(item.ticket_assignment, "ticket_assignment"), initial_submit: array(item.initial_submit, "initial_submit", segment), - working_directory: item.working_directory === null - ? null - : workingDirectorySelection(item.working_directory, "working_directory"), + workdir_attachments: array( + item.workdir_attachments, + "workdir_attachments", + workingDirectorySelection, + ), control_operation_id: nullableString( item.control_operation_id, "control_operation_id", diff --git a/web/workspace/src/lib/workspace/console/worker-console.ui.test.ts b/web/workspace/src/lib/workspace/console/worker-console.ui.test.ts index c4f21edaa..774df4ab1 100644 --- a/web/workspace/src/lib/workspace/console/worker-console.ui.test.ts +++ b/web/workspace/src/lib/workspace/console/worker-console.ui.test.ts @@ -256,8 +256,10 @@ Deno.test("workspace Worker list lives on the dedicated Workers page", async () workersNav.includes("worker.display_name || worker.label") && workersNav.includes("worker-status-dot") && workersNav.includes("worker-status-spinner") && - workersNav.includes("worker.repository_key ?? '—'") && - workersNav.includes("worker.working_directory_id ?? '—'") && + workersNav.includes("workdirMeta(worker)") && + workersNav.includes("attachment.alias") && + workersNav.includes("attachment.repository_key ?? '—'") && + workersNav.includes("attachment.working_directory_id") && !workersNav.includes('aria-disabled="true"'), "Workers sidebar should link to the Worker list page and show state indicators with repository/workdir metadata", ); diff --git a/web/workspace/src/lib/workspace/sidebar/WorkersNavSection.svelte b/web/workspace/src/lib/workspace/sidebar/WorkersNavSection.svelte index e527f0ede..3260def2e 100644 --- a/web/workspace/src/lib/workspace/sidebar/WorkersNavSection.svelte +++ b/web/workspace/src/lib/workspace/sidebar/WorkersNavSection.svelte @@ -37,6 +37,14 @@ Math.max(0, workers.length - COLLAPSED_WORKER_COUNT), ); + function workdirMeta(worker: SidebarWorker): string { + const attachments = worker.workdir_attachments; + if (attachments.length === 0) return '—'; + return attachments + .map((attachment) => `${attachment.alias}: ${attachment.repository_key ?? '—'}・${attachment.working_directory_id}`) + .join(', '); + } + function workerKey(worker: SidebarWorker): string { return `${worker.runtime_id}:${worker.worker_id}`; } @@ -190,7 +198,7 @@ {label} - {worker.repository_key ?? '—'}・{worker.working_directory_id ?? '—'} + {workdirMeta(worker)} + + +
+ {#each workdirAttachments as attachment, index} +
+ + + + +
+ {/each} +
+ + {#if attachmentValidationError} +

{attachmentValidationError}

+ {/if} {/if} {#if isNewWorkingDirectorySelected}
-

New workdir

+

New Workdir for “{workdirAttachments[newWorkingDirectoryAttachmentIndex]?.alias || `attachment ${newWorkingDirectoryAttachmentIndex + 1}`}”

+ - {/if}
diff --git a/web/workspace/tests/workdir-api.test.ts b/web/workspace/tests/workdir-api.test.ts index 3b4b06a12..e93ea2c3d 100644 --- a/web/workspace/tests/workdir-api.test.ts +++ b/web/workspace/tests/workdir-api.test.ts @@ -10,6 +10,7 @@ import { const summary = { working_directory_id: "workdir-1", + display_name: "Checkout", repository_key: "main", materializer_kind: "runtime_git_clone", status: "active", @@ -27,6 +28,9 @@ Deno.test("Workdir REST validation accepts the generated list and create contrac items: [summary], diagnostics: [], }); + if (list.items[0]?.display_name !== "Checkout") { + throw new Error("display name was not preserved"); + } if (list.items[0]?.occupied_by?.runtime_id !== "arcadia") { throw new Error("occupancy subject was not preserved"); } @@ -60,6 +64,7 @@ Deno.test("Workdir REST validation rejects stale response JSON", () => { Deno.test("Workdir REST validation enforces create operation fields", () => { const request = validateWorkingDirectoryCreateRequest({ runtime_id: "arcadia", + display_name: "Review checkout", repository_key: "main", selector: "develop", operation_id: "operation-1", @@ -67,6 +72,9 @@ Deno.test("Workdir REST validation enforces create operation fields", () => { if (request.operation_id !== "operation-1") { throw new Error("operation id was not preserved"); } + if (request.display_name !== "Review checkout") { + throw new Error("display name was not preserved"); + } for ( const invalid of [