Implements address space cleanup when a workspace service is uninstalled

[James Chapman (JC-wk)]

Resolves #4727

PR

• unit testing

What is being addressed

• Deletes allocated address_space when a workspace service is uninstalled, preventing IP range exhaustion.
• Adds workspace-level atomic operation leasing to serialize concurrent workspace mutations and prevent state/address corruption.

How is this addressed

• Address Cleanup: Adds a pipeline step to remove the service's address space and triggers a workspace upgrade.
• Workspace Leasing: Implements atomic workspace-level leasing with automatic stale-lease recovery and rollback guards.
• Operation Serialization: Serializes multi-step pipelines and background workflows (such as Airlock review VM cleanup) under the workspace lease.
• Tests & Docs: Added unit tests, updated CHANGELOG.md, and incremented template versions.

[github-actions]

Unit Test Results

886 tests   886 ✅  11s ⏱️
  1 suites    0 💤
  1 files      0 ❌

Results for commit f17e48d.

♻️ This comment has been updated with latest results.

[James Chapman (JC-wk)]

I have been testing this for a few days, I am not sure if unit tests are needed and how best to write them if anyone wants to assist.

[Copilot]

Pull request overview

This PR addresses IP range exhaustion risk by ensuring workspace address spaces allocated by workspace services are freed on successful uninstall, and by triggering a workspace upgrade so downstream infra reflects the removal.

Changes:

• Add post-uninstall cleanup in the service bus deployment status handler to remove a workspace-service address_space from the parent workspace’s address_spaces.
• Update AzureML and Databricks workspace-service templates to run a workspace upgrade step after uninstall.
• Bump API + template versions and add a changelog entry.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
templates/workspace_services/databricks/template_schema.json	Adds a workspace upgrade step after uninstall (and JSON formatting changes).
templates/workspace_services/databricks/porter.yaml	Patch version bump.
templates/workspace_services/azureml/template_schema.json	Adds a workspace upgrade step after uninstall.
templates/workspace_services/azureml/porter.yaml	Patch version bump.
api_app/service_bus/deployment_status_updater.py	Implements address space cleanup after successful uninstall main step.
api_app/_version.py	API patch version bump.
CHANGELOG.md	Adds an Unreleased entry describing the change.

[Copilot]

🟡 Changes recommended

Critical lease and reconciliation defects can allow concurrent workspace operations.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

api_app/db/repositories/operations.py:159

• If persisting the stale operation's terminal state fails, this exception is swallowed and execution continues to replace the lease for the new operation. That leaves the old operation active while admitting a concurrent one. Fail closed (for example, return a 409) unless the terminal update succeeds.

                            except Exception:
                                pass

• Files reviewed: 23/24 changed files
• Comments generated: 3
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

Four moderate correctness and concurrency issues must be resolved before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (2)

api_app/db/repositories/operations.py:430

• The five-minute staleness check can invalidate a deployment that is still running. The resource processor emits one in-progress update and then waits synchronously for Porter while renewing the Service Bus session for up to an hour (resource_processor/vmss_porter/runner.py:67-68,179-189), so updatedWhen receives no heartbeat during that interval. If another request arrives after five minutes, this branch marks the live operation failed and releases its lease, allowing concurrent Terraform work in the same workspace. Please base expiry on a durable execution heartbeat or on the supported deployment timeout rather than this operation timestamp.

            if op_time is not None and (timestamp - op_time >= WORKSPACE_LEASE_EXPIRY_SECONDS):

api_app/db/repositories/operations.py:182

• When a stale active operation is marked terminal above, update_item also releases and deletes this workspace lease. The subsequent replace_item therefore uses the ETag of a deleted item and raises a not-found exception, but this handler only retries conflicts; callers can receive a 500 instead of acquiring the newly freed lease. Treat both Cosmos and Azure Core not-found errors as retryable here.

                except (CosmosAccessConditionFailedError, ResourceExistsError):

• Files reviewed: 24/25 changed files
• Comments generated: 2
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

Lease-release races can permit overlapping workspace operations and inconsistent address state.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

api_app/db/repositories/operations.py:453

• The age check can release serialization while the original request is still queued or executing. updatedWhen is not a worker heartbeat, and the one-hour Service Bus lock-renewal limit in resource_processor/vmss_porter/runner.py:68 does not cancel Porter when it expires. Marking the operation failed here releases its workspace lease, so a new Terraform operation can start and the old queued/running request can later execute against the same workspace. Please fail closed until there is authoritative evidence that the old request cannot run (or explicitly cancel/dead-letter it before reclaiming the lease).

            # If timestamp is present and exceeds expiry, it is stale
            if op_time is not None and (timestamp - op_time >= WORKSPACE_LEASE_EXPIRY_SECONDS):
                # Stale active operation from interrupted dispatch; reconcile to failure status
                try:
                    op = TypeAdapter(Operation).validate_python(op_dict)
                    op.status = get_failure_status_for_action(op.action)
                    op.message = "Operation timed out or was interrupted before completion"
                    op.updatedWhen = timestamp
                    update_call = self.update_item(op)
                    if hasattr(update_call, "__await__"):
                        await update_call

• Files reviewed: 24/25 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

Two critical lease-safety issues and one failure-recovery issue must be resolved before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

api_app/service_bus/resource_request_sender.py:80

• If dispatch fails and this persistence write also fails, the exception is only logged and the operation remains in its active initial status. For workspace-scoped resources, create_operation_item has already acquired the workspace lease, so subsequent mutations are rejected until the one-hour stale-operation timeout. Please add a fallback that removes/reconciles the orphan operation and releases its workspace lease before re-raising when the terminal failure state cannot be saved.

• Files reviewed: 24/25 changed files
• Comments generated: 2
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

The critical airlock serialization and moderate lease-release issues must be resolved before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 25/26 changed files
• Comments generated: 2
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

One critical and three moderate issues remain unresolved.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

api_app/api/routes/workspaces.py:320

• The rollback is gated only by whether the service object has an address_space, not whether this request actually appended that address to the workspace. Because address_space is an accepted (UI-hidden) input property, if address allocation fails before line 293 overwrites it, a caller-supplied value matching another service's range will be removed here. Track whether patch_workspace successfully added the newly allocated range and only run rollback in that case.

• Files reviewed: 26/27 changed files
• Comments generated: 3
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

Moderate rollback, lease-release, and Airlock cleanup reliability issues remain unresolved.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (3)

api_app/api/routes/workspaces.py:331

• If this compensation write fails, the exception is only logged and the workspace lease is then released. The allocated CIDR remains in address_spaces even though resource creation failed, recreating the IP-range leak this change is intended to prevent. Retry the rollback with fresh ETags or persist a recoverable compensation job before releasing the lease.

            except Exception:
                logger.exception("Failed to rollback allocated address space on workspace")
        if hasattr(operations_repo, "release_workspace_lease"):
            await operations_repo.release_workspace_lease(workspace.id, operation_id)

api_app/service_bus/resource_request_sender.py:111

• When operation_id was supplied, the route owns compensation and lease release, but this fallback releases the lease immediately after deleting the operation. For workspace-service creation that opens a race where another mutation can acquire the workspace while the caller is still rolling back the allocated address space. Keep the lease for caller-owned operations, matching the should_release_lease behavior above.

                if lease_released:

api_app/services/airlock.py:568

• This cancellation path has the same untracked fire-and-forget behavior: cleanup exceptions cannot reach the caller or trigger a retry, and shutdown can leave review VMs undeleted after the request is marked cancelled. Schedule the deletion through a managed/durable background mechanism instead of discarding the task.

    asyncio.create_task(delete_all_review_user_resources(airlock_request, user_resource_repo, workspace_service_repo, resource_template_repo, operations_repo, resource_history_repo, user))
    await asyncio.sleep(0)

• Files reviewed: 26/27 changed files
• Comments generated: 2
• Review effort level: Balanced

[Copilot]

🔵 Needs a closer look

A critical lease-release flaw and two durability issues remain unresolved.

Review details

• Files reviewed: 26/27 changed files
• Comments generated: 3
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

Unresolved rollback, lease-retention, and synchronous Airlock cleanup issues can cause orphaned addresses, concurrency hazards, and request timeouts.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (2)

api_app/api/routes/resource_helpers.py:110

• The lease_retained guard is lost when the dispatch exception is wrapped in a new HTTPException. Callers such as create_workspace_service and create_user_resource inspect that attribute before rolling back state and releasing the workspace lease; without propagating it, they can release the lease (and roll back an allocated address) while the concurrently advanced operation is still running. Preserve the flag on the translated exception.

        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail=strings.SERVICE_BUS_GENERAL_ERROR_MESSAGE,
        ) from ex

api_app/services/airlock.py:561

• This now waits synchronously for both the disable and uninstall operations, each with a default timeout of 7,200 seconds, and processes multiple review resources sequentially. Because this runs inside the Airlock review/cancel HTTP request, it can exceed Azure App Service's roughly 230-second request timeout, so the client will time out while the request is still in progress. Queue this cleanup as background work (while preserving lease sequencing), or return after scheduling it rather than polling deployment completion in the request.

                if operation and hasattr(operation, "id"):
                    await wait_for_successful_operation(operations_repo, operation.id)

• Files reviewed: 26/27 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

The unresolved critical operation-reconciliation issue and moderate lease and Airlock workflow issues must be addressed before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (5)

Previously missed (2) — in code that hasn't changed since the last review.

api_app/service_bus/resource_request_sender.py:110

• If deleting the orphaned operation fails for any reason other than the three concurrency exceptions, lease_retained remains false. Callers that supplied operation_id then release the workspace lease and may delete/roll back the resource even though the operation document may still exist and the dispatch outcome is unknown, allowing a second workspace mutation to start. Treat every unconfirmed deletion as lease-retained; only report the lease releasable after the conditional delete succeeds.
  api_app/services/airlock.py:465
• All unexpected repository errors are silently retried until the timeout, so a Cosmos authentication/outage error is converted into a misleading 504 after a potentially long wait. Only EntityDoesNotExist is an expected transient here; let other exceptions propagate immediately so the caller receives the actual service failure.

api_app/db/repositories/operations.py:460

• The active-operation reconciliation also releases the lease after changing only the operation status. Late deployment status messages can subsequently make this operation/resource successful again while another workspace operation is running, because generic terminal operations are not ignored by the status updater. Use the same persisted stale/aborted marker and consumer-side guard for this reconciliation path.

                    op.status = get_failure_status_for_action(op.action)
                    op.message = "Operation timed out or was interrupted before completion"
                    op.updatedWhen = timestamp

api_app/services/airlock.py:497

• This makes the synchronous airlock review/cancel request fail after only 30 seconds whenever disabling a review VM takes longer, even though deployment operations and their leases are explicitly allowed to run for up to two hours. The request status has already been persisted as ApprovalInProgress/RejectionInProgress (or Cancelled) before this wait, so the caller receives a 504 while the request is left partially transitioned and cannot safely retry. Run this serialized cleanup as a durable background workflow, or otherwise avoid a fixed 30-second wait inside the request.

    disable_op = await disable_user_resource(user_resource, user, workspace_service, user_resource_repo, resource_template_repo, operations_repo, resource_history_repo)
    if disable_op and hasattr(disable_op, "id"):
        await wait_for_successful_operation(operations_repo, disable_op.id, timeout=30.0)

api_app/services/airlock.py:541

• This retry runs inline in the Airlock review and cancel HTTP handlers, which await delete_all_review_user_resources. Lease contention can therefore hold a request for up to two hours per review resource (the timer resets for each item), likely exceeding gateway/request timeouts and leaving the Airlock request in an in-progress state. Move serialized cleanup retries to a background workflow and return from the endpoint after scheduling them.

        start_time = time.time()
        max_lease_wait = WORKSPACE_LEASE_EXPIRY_SECONDS
        poll_interval = 2.0

• Files reviewed: 26/27 changed files
• Comments generated: 2
• Review effort level: Balanced

[Copilot]

🔵 Needs a closer look

Airlock cleanup can block synchronous HTTP requests beyond the platform timeout.

Review details

Suppressed comments (1)

api_app/services/airlock.py:495

• Waiting for the disable operation here can block the Airlock review/cancel HTTP request for up to WORKSPACE_LEASE_EXPIRY_SECONDS (two hours). These functions are called directly by synchronous API routes, while Azure App Service requests time out after roughly 230 seconds, so a normal VM disable can make the client receive a timeout after the Airlock request status has already been changed. Move the wait-and-uninstall sequence to a background workflow (or persist a chained cleanup operation) and let the route return after scheduling it.

    disable_op = await disable_user_resource(user_resource, user, workspace_service, user_resource_repo, resource_template_repo, operations_repo, resource_history_repo)
    if disable_op and hasattr(disable_op, "id"):
        await wait_for_successful_operation(operations_repo, disable_op.id)

• Files reviewed: 26/27 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Copilot]

🔵 Needs a closer look

Two moderate issues remain in stale-step reconciliation and Airlock replacement deployment sequencing.

Review details

Suppressed comments (3)

Previously missed (1) — in code that hasn't changed since the last review.

api_app/db/repositories/operations.py:175

• Stale reconciliation assigns every unfinished pipeline step the failure status of the root operation. For an uninstall pipeline's workspace-upgrade cleanup step, this records DeletingFailed instead of UpdatingFailed, so the persisted step state no longer matches its resourceAction. Derive each step's failure status from step.resourceAction while keeping the operation status based on existing_op.action.

This issue also appears on line 472 of the same file.

api_app/db/repositories/operations.py:472

• This second stale-operation reconciliation path also derives each step's failure state from the root action. Multi-step pipelines can contain different actions, so this persists an incorrect status (for example, DeletingFailed for the cleanup workspace upgrade). Use the step action here.

                                step.status = get_failure_status_for_action(op.action)

api_app/services/airlock.py:508

• The new workspace lease makes the existing unhealthy-review-VM replacement path fail: _handle_existing_review_resource awaits this function and immediately calls _deploy_vm, but this function returns as soon as the uninstall is submitted, while that uninstall still owns the workspace lease. The replacement deployment then receives a 409. Wait for the uninstall operation to succeed before returning in that replacement flow (or make this helper wait before returning).

    disable_op = await disable_user_resource(user_resource, user, workspace_service, user_resource_repo, resource_template_repo, operations_repo, resource_history_repo)
    if disable_op and hasattr(disable_op, "id"):
        await wait_for_successful_operation(operations_repo, disable_op.id)

• Files reviewed: 27/28 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

The lease-release race, stale cascade snapshot, and non-durable Airlock cleanup must be resolved before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (3)

Previously missed (1) — in code that hasn't changed since the last review.

api_app/service_bus/resource_request_sender.py:50

• For cascade requests, resources_list is collected before create_operation_item acquires the workspace lease. A concurrent child-resource operation can therefore complete in that gap, after validation/snapshotting but before this operation obtains the lease, leaving the cascade to proceed from a stale dependency list. Acquire the workspace lease before get_resource_dependency_list (using a pre-created operation ID), and release it if dependency discovery or operation creation fails.

api_app/api/routes/workspaces.py:393

• This exception path bypasses the lease_retained guard below. send_resource_request_message can re-raise a CosmosAccessConditionFailedError with lease_retained=True when its operation was concurrently advanced; unconditionally releasing here then permits another workspace mutation while that operation remains active. Apply the same guard before releasing.

    except CosmosAccessConditionFailedError:
        if hasattr(operations_repo, "release_workspace_lease"):
            await operations_repo.release_workspace_lease(workspace_service.workspaceId, operation_id)
        raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=strings.ETAG_CONFLICT)

api_app/services/airlock.py:608

• Cancellation cleanup is now process-local and starts only after the response is sent. A worker restart can discard this task after the request is marked cancelled, leaving review resources behind without a durable retry path. Use a durable queued job, or await cleanup until such a job has been persisted.

    if background_tasks is not None:
        background_tasks.add_task(
            delete_all_review_user_resources,

• Files reviewed: 28/29 changed files
• Comments generated: 2
• Review effort level: Balanced

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

🔵 Needs a closer look

Five unresolved moderate issues affect lease safety, stale resource reconciliation, and the blocking Airlock workflow.

Review details

Suppressed comments (5)

Previously missed (1) — in code that hasn't changed since the last review.

api_app/services/airlock.py:275

• Waiting for the uninstall operation inline can hold the create-review HTTP request for up to two hours here, after delete_review_user_resource may already have waited another two hours for disablement. Azure App Service requests time out at roughly 230 seconds, so callers can time out long before this workflow returns, while each wait also polls Cosmos every 0.5 seconds. Move this delete-then-redeploy sequence to a durable/background workflow (or resume deployment from the terminal operation event) instead of blocking the request.

api_app/api/routes/resource_helpers.py:102

• The HTTPException path deletes the newly saved resource unconditionally, unlike the generic path below. send_resource_request_message can re-raise an HTTPException carrying lease_retained=True when the operation was concurrently advanced; deleting the resource then leaves that live operation without its resource document. Honor the marker in this branch as well.

    except HTTPException:
        await resource_repo.delete_item(resource.id)
        raise

api_app/api/routes/workspaces.py:393

• This handler releases the workspace lease even when send_resource_request_message marks the exception with lease_retained=True, which specifically means the operation was concurrently advanced and its lease must remain held. The analogous workspace handler above preserves the lease in this case; apply the same guard here to avoid allowing a second workspace mutation to start.

    except CosmosAccessConditionFailedError:
        if hasattr(operations_repo, "release_workspace_lease"):
            await operations_repo.release_workspace_lease(workspace_service.workspaceId, operation_id)
        raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=strings.ETAG_CONFLICT)

api_app/db/repositories/operations.py:177

• Stale recovery only marks the operation and its steps failed, then releases the lease; it never changes the primary resource's deploymentStatus. Because the status updater now rejects messages for reconciled operations, a resource already marked Deploying/Updating/Deleting remains stranded in that active state even after a new workspace mutation is allowed. Reconcile the affected resource statuses atomically before releasing the lease.

                                existing_op.status = get_failure_status_for_action(existing_op.action)
                                existing_op.message = "Operation timed out or was interrupted before completion"
                                existing_op.updatedWhen = timestamp
                                existing_op.reconciled = True
                                if getattr(existing_op, "steps", None):
                                    for step in existing_op.steps:
                                        if not step.is_failure() and not step.is_success():
                                            step.status = get_failure_status_for_action(step.resourceAction or existing_op.action)
                                            step.message = "Operation timed out or was interrupted before completion"
                                            step.updatedWhen = timestamp

api_app/db/repositories/operations.py:474

• This second stale-operation path also persists only the terminal operation state. The corresponding resource can therefore keep an active deploymentStatus, and delayed status messages cannot repair it because reconciled operations are discarded by the updater. Update the resource state as part of reconciliation before reporting that no active operation remains.

                    op.status = get_failure_status_for_action(op.action)
                    op.message = "Operation timed out or was interrupted before completion"
                    op.updatedWhen = timestamp
                    op.reconciled = True
                    if getattr(op, "steps", None):
                        for step in op.steps:
                            if not step.is_failure() and not step.is_success():
                                step.status = get_failure_status_for_action(step.resourceAction or op.action)
                                step.message = "Operation timed out or was interrupted before completion"
                                step.updatedWhen = timestamp

• Files reviewed: 27/28 changed files
• Comments generated: 0 new
• Review effort level: Balanced