diff --git a/crates/workspace-server/src/retention.rs b/crates/workspace-server/src/retention.rs index f4e06eb6..90b89935 100644 --- a/crates/workspace-server/src/retention.rs +++ b/crates/workspace-server/src/retention.rs @@ -106,12 +106,15 @@ pub struct WorkerRemovalPlan { pub reason: String, pub created_at: String, pub updated_at: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub failure_category: Option, } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct PreparedWorkerRemoval { pub plan: WorkerRemovalPlan, pub runtime_request: WorkerRetentionExecutionRequest, + pub prior_failure_category: Option, } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -342,7 +345,7 @@ impl SqliteWorkspaceStore { return Err(stale_error(&plan,"current assignment added")); } let now=Utc::now().to_rfc3339(); - tx.execute("UPDATE worker_removal_operations SET state='executing',failure_category=NULL,updated_at=?1 WHERE operation_id=?2",params![now,plan.operation_id])?; + tx.execute("UPDATE worker_removal_operations SET state='executing',updated_at=?1 WHERE operation_id=?2",params![now,plan.operation_id])?; plan.state=WorkerRemovalPlanState::Executing;plan.updated_at=now;tx.commit()?;Ok(plan) }).map_err(map_error) } @@ -369,6 +372,7 @@ impl SqliteWorkspaceStore { ) })?; let removed_at = plan.created_at.clone(); + let prior_failure_category = plan.failure_category.clone(); Ok(PreparedWorkerRemoval { runtime_request: WorkerRetentionExecutionRequest { operation_id: plan.operation_id.clone(), @@ -389,6 +393,7 @@ impl SqliteWorkspaceStore { diagnostics_disposition: plan.diagnostics_disposition, }, plan, + prior_failure_category, }) } @@ -429,6 +434,7 @@ impl SqliteWorkspaceStore { let Some(plan) = plan else { return Ok(None); }; + let prior_failure_category = plan.failure_category.clone(); let worker_number = plan.worker.worker_id.parse::().map_err(|_| { WorkerRetentionError::Invalid( "Runtime Worker id is not a canonical unsigned integer".to_string(), @@ -468,6 +474,7 @@ impl SqliteWorkspaceStore { diagnostics_disposition: plan.diagnostics_disposition, }, plan, + prior_failure_category, })) } @@ -758,7 +765,7 @@ fn load_plan_q(c: &Connection, key: &str, id: &str) -> crate::Result crate::Result = s + .with_conn(|conn| { + conn.query_row( + "SELECT failure_category FROM worker_removal_operations WHERE operation_id=?1", + params![p.operation_id], + |row| row.get(0), + ) + .map_err(StoreError::from) + }) + .unwrap(); + assert_eq!( + persisted.as_deref(), + Some("workdir_attachment_release_failed") + ); + } + #[test] fn prepared_execution_is_derived_from_pinned_plan_generation() { let s = setup(); @@ -1532,6 +1577,36 @@ mod tests { ); } + #[test] + fn recovery_preserves_attachment_failure_stage_for_retry_ordering() { + let s = setup(); + let request = req(); + let plan = s.plan_worker_removal(&request, &inv()).unwrap(); + s.prepare_worker_removal_execution("w", &plan.plan_id, &plan.input_fingerprint) + .unwrap(); + s.fail_worker_removal( + "w", + &plan.operation_id, + &plan.input_fingerprint, + "workdir_attachment_release_failed", + ) + .unwrap(); + let recovered = s + .recover_worker_removal_execution( + "w", + &request.worker, + &request.expected_worker_revision, + &request.reason, + ) + .unwrap() + .unwrap(); + assert_eq!( + recovered.prior_failure_category.as_deref(), + Some("workdir_attachment_release_failed") + ); + assert_eq!(recovered.plan.state, WorkerRemovalPlanState::Failed); + } + #[test] fn old_schema_upgrade_seeds_existing_workspace() { let temp = tempfile::tempdir().unwrap(); diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index f65acf93..5abd8a82 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -396,6 +396,11 @@ impl WorkspaceWorkerRemoveExecutor { if prepared.plan.state == crate::retention::WorkerRemovalPlanState::Succeeded { return Ok(worker_remove_success_response(&target)); } + let must_close_session = + prepared.prior_failure_category.as_deref() == Some("workdir_session_close_failed"); + let must_release_attachment = must_close_session + || prepared.prior_failure_category.as_deref() + == Some("workdir_attachment_release_failed"); let prepared = if prepared.plan.state == crate::retention::WorkerRemovalPlanState::Failed { match self.store.prepare_worker_removal_execution( @@ -409,6 +414,57 @@ impl WorkspaceWorkerRemoveExecutor { } else { prepared }; + if must_close_session { + let session = { + self.workdir_sessions + .lock() + .map_err(|_| "Workdir session registry was poisoned".to_string())? + .get(&target) + .cloned() + }; + if let Some(session) = session { + if session.close().await.is_err() { + let _ = self.store.fail_worker_removal( + &self.workspace_id, + &prepared.plan.operation_id, + &prepared.plan.input_fingerprint, + "workdir_session_close_failed", + ); + return Ok(worker_remove_error_response( + StatusCode::SERVICE_UNAVAILABLE, + "attachment_close_failed", + "Worker Workdir session could not be closed; removal can be retried", + )); + } + self.workdir_sessions + .lock() + .map_err(|_| "Workdir session registry was poisoned".to_string())? + .remove(&target); + } + } + if must_release_attachment + && self + .store + .detach_worker_workdir( + &self.workspace_id, + &target, + None, + &Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true), + ) + .is_err() + { + let _ = self.store.fail_worker_removal( + &self.workspace_id, + &prepared.plan.operation_id, + &prepared.plan.input_fingerprint, + "workdir_attachment_release_failed", + ); + return Ok(worker_remove_error_response( + StatusCode::SERVICE_UNAVAILABLE, + "attachment_release_failed", + "Worker Workdir attachment could not be released; removal can be retried", + )); + } return self .resume_worker_retention(&runtime, &target, prepared) .await;