@baochunli/flakes 0.0.5 → 0.0.6

package/dist/node_modules/@flakes/store/dist/managed-sandbox-instances-store.js

@@ -4,7 +4,17 @@ import { HOST_PROVISIONING_QUOTAS, HostProvisioningQuotaError, jsonByteLength, }
 import { addMs, fromJson, toJson } from "./json.js";
 import { mapHost, mapHostCredential, mapHostPool, mapManagedSandboxInstance, mapProvisioningLease, mapSandboxCredential, mapSandboxIdentity, } from "./mappers.js";
 const nonTerminalManagedInstanceStatuses = ["pending_launch", "launching", "running", "draining", "stopping"];
-const cleanupManagedInstanceStatuses = ["stopped", "failed", "expired", "cancelled"];
+const cleanupManagedInstanceStatuses = ["stopped", "failed", "expired", "cancelled", "deleting"];
+const noProviderDeleteManagedInstanceStatuses = [
+    "stopped",
+    "failed",
+    "expired",
+    "cancelled",
+    "running",
+    "draining",
+    "stopping",
+];
+const providerDeleteManagedInstanceStatuses = [...noProviderDeleteManagedInstanceStatuses, "deleting"];
 const nonTerminalProvisioningLeaseStates = [
     "pending",
     "claimed",
@@ -1061,6 +1071,261 @@ export class ManagedSandboxInstancesStore extends ArtifactsEventsStore {
             clearedAssignments: cleared.cleared,
         };
     }
+    /** Deletes or queues cleanup for an owner-scoped stale managed sandbox row. */
+    async deleteManagedSandboxInstance(input) {
+        return this.withTransaction(async (trx) => {
+            const liveness = await this.lockManagedSandboxLiveness(trx, input.sandboxId);
+            const instance = await trx
+                .selectFrom("managed_sandbox_instances as instance")
+                .innerJoin("sandbox_identities as identity", "identity.sandbox_id", "instance.sandbox_id")
+                .select([
+                "instance.instance_id as instance_id",
+                "instance.owner_kind as owner_kind",
+                "instance.owner_id as owner_id",
+                "instance.owner_key as owner_key",
+                "instance.project_id as project_id",
+                "instance.host_id as host_id",
+                "instance.host_pool_id as host_pool_id",
+                "instance.sandbox_id as sandbox_id",
+                "instance.status as status",
+                "instance.provider_resource_id as provider_resource_id",
+                "identity.sandbox_identity_id as sandbox_identity_id",
+                "identity.disabled_at as identity_disabled_at",
+            ])
+                .where("instance.owner_key", "=", input.ownerKey)
+                .where("instance.sandbox_id", "=", input.sandboxId)
+                .where("identity.owner_key", "=", input.ownerKey)
+                .forUpdate()
+                .executeTakeFirst();
+            if (!instance) {
+                return { status: "not_found" };
+            }
+            const blockers = [];
+            const providerResourceId = instance.provider_resource_id;
+            const eligibleStatuses = providerResourceId === null
+                ? noProviderDeleteManagedInstanceStatuses
+                : providerDeleteManagedInstanceStatuses;
+            if (!eligibleStatuses.includes(instance.status)) {
+                blockers.push("unsafe_status");
+            }
+            if (liveness?.status === "active") {
+                blockers.push("live_sandbox");
+            }
+            if (blockers.length > 0) {
+                return {
+                    status: "blocked",
+                    sandboxId: input.sandboxId,
+                    blockers,
+                };
+            }
+            const latestLiveness = liveness ??
+                (await this.lockManagedSandboxLiveness(trx, input.sandboxId));
+            if (latestLiveness?.status === "active") {
+                return {
+                    status: "blocked",
+                    sandboxId: input.sandboxId,
+                    blockers: ["live_sandbox"],
+                };
+            }
+            const hostPoolId = instance.host_pool_id;
+            if (providerResourceId !== null) {
+                if (hostPoolId === null) {
+                    return {
+                        status: "blocked",
+                        sandboxId: input.sandboxId,
+                        blockers: ["provider_cleanup_required"],
+                    };
+                }
+                if (await this.managedSandboxHasAcceptedRunningAssignments(trx, input.sandboxId)) {
+                    return {
+                        status: "blocked",
+                        sandboxId: input.sandboxId,
+                        blockers: ["accepted_running_assignments"],
+                    };
+                }
+                const clearedAssignments = await this.clearReservedSandboxAssignmentsForDelete(trx, input.sandboxId);
+                const now = this.nowIso();
+                const reason = redactReason(input.reason) ?? "managed sandbox deleted";
+                await this.cancelSupersededProviderStopLeasesForDelete(trx, {
+                    ownerKey: instance.owner_key,
+                    hostId: instance.host_id,
+                    instanceId: instance.instance_id,
+                }, now, reason);
+                const cleanupLease = await this.queueManagedSandboxDeleteCleanup(trx, {
+                    instanceId: instance.instance_id,
+                    ownerKind: instance.owner_kind,
+                    ownerId: instance.owner_id,
+                    ownerKey: instance.owner_key,
+                    projectId: instance.project_id,
+                    hostId: instance.host_id,
+                    hostPoolId,
+                    sandboxId: instance.sandbox_id,
+                    providerResourceId,
+                }, now);
+                await trx
+                    .updateTable("managed_sandbox_instances")
+                    .set({
+                    status: "deleting",
+                    last_error: reason,
+                    updated_at: now,
+                })
+                    .where("instance_id", "=", instance.instance_id)
+                    .where("owner_key", "=", input.ownerKey)
+                    .where("sandbox_id", "=", input.sandboxId)
+                    .where("provider_resource_id", "=", providerResourceId)
+                    .where("status", "in", [...providerDeleteManagedInstanceStatuses])
+                    .executeTakeFirst();
+                await this.invalidateManagedSandboxRuntimeAuthForDelete(trx, {
+                    sandboxIdentityId: instance.sandbox_identity_id,
+                    identityDisabledAt: instance.identity_disabled_at,
+                    reason,
+                    now,
+                });
+                return {
+                    status: "cleanup_queued",
+                    sandboxId: input.sandboxId,
+                    instanceId: instance.instance_id,
+                    leaseId: cleanupLease.leaseId,
+                    providerResourceId,
+                    clearedAssignments,
+                };
+            }
+            if (await this.managedSandboxHasAcceptedRunningAssignments(trx, input.sandboxId)) {
+                return {
+                    status: "blocked",
+                    sandboxId: input.sandboxId,
+                    blockers: ["accepted_running_assignments"],
+                };
+            }
+            const clearedAssignments = await this.clearReservedSandboxAssignmentsForDelete(trx, input.sandboxId);
+            const now = this.nowIso();
+            const reason = redactReason(input.reason) ?? "managed sandbox deleted";
+            await this.invalidateManagedSandboxRuntimeAuthForDelete(trx, {
+                sandboxIdentityId: instance.sandbox_identity_id,
+                identityDisabledAt: instance.identity_disabled_at,
+                reason,
+                now,
+            });
+            await trx
+                .deleteFrom("managed_sandbox_instances")
+                .where("instance_id", "=", instance.instance_id)
+                .where("owner_key", "=", input.ownerKey)
+                .where("sandbox_id", "=", input.sandboxId)
+                .executeTakeFirst();
+            return {
+                status: "deleted",
+                sandboxId: input.sandboxId,
+                instanceId: instance.instance_id,
+                clearedAssignments,
+            };
+        });
+    }
+    async invalidateManagedSandboxRuntimeAuthForDelete(executor, input) {
+        await executor
+            .updateTable("sandbox_credentials")
+            .set({
+            revoked_at: input.now,
+            revoked_reason: input.reason,
+            updated_at: input.now,
+        })
+            .where("sandbox_identity_id", "=", input.sandboxIdentityId)
+            .where("revoked_at", "is", null)
+            .execute();
+        await executor
+            .updateTable("sandbox_identities")
+            .set({
+            status: "disabled",
+            disabled_at: input.identityDisabledAt ?? input.now,
+            updated_at: input.now,
+        })
+            .where("sandbox_identity_id", "=", input.sandboxIdentityId)
+            .execute();
+    }
+    async queueManagedSandboxDeleteCleanup(executor, instance, now) {
+        const existing = await executor
+            .selectFrom("provisioning_leases")
+            .selectAll()
+            .where("owner_key", "=", instance.ownerKey)
+            .where("host_id", "=", instance.hostId)
+            .where("instance_id", "=", instance.instanceId)
+            .where("action", "=", "provider_cleanup")
+            .where("state", "in", [...nonTerminalProvisioningLeaseStates])
+            .forUpdate()
+            .executeTakeFirst();
+        if (existing) {
+            if (existing.state === "pending" && existing.lease_deadline_at <= now) {
+                const refreshed = await executor
+                    .updateTable("provisioning_leases")
+                    .set({
+                    lease_deadline_at: addMs(this.clock.now(), providerCleanupLeaseTtlMs),
+                    last_error: null,
+                    updated_at: now,
+                })
+                    .where("lease_id", "=", existing.lease_id)
+                    .where("owner_key", "=", instance.ownerKey)
+                    .where("host_id", "=", instance.hostId)
+                    .where("action", "=", "provider_cleanup")
+                    .where("state", "=", "pending")
+                    .returningAll()
+                    .executeTakeFirst();
+                return mapProvisioningLease(refreshed ?? existing);
+            }
+            return mapProvisioningLease(existing);
+        }
+        await this.assertProvisioningLeaseQuota(executor, instance.ownerKey);
+        const row = await executor
+            .insertInto("provisioning_leases")
+            .values({
+            lease_id: this.idGenerator.createId("provisioning_lease"),
+            owner_kind: instance.ownerKind,
+            owner_id: instance.ownerId,
+            owner_key: instance.ownerKey,
+            project_id: instance.projectId,
+            host_project_binding_id: null,
+            host_id: instance.hostId,
+            host_pool_id: instance.hostPoolId,
+            instance_id: instance.instanceId,
+            sandbox_id: instance.sandboxId,
+            demand_task_id: null,
+            action: "provider_cleanup",
+            state: "pending",
+            claim_id: null,
+            attempt: 0,
+            max_attempts: providerCleanupLeaseMaxAttempts,
+            lease_deadline_at: addMs(this.clock.now(), providerCleanupLeaseTtlMs),
+            claim_deadline_at: null,
+            hello_deadline_at: null,
+            accepted_at: null,
+            running_at: null,
+            completed_at: null,
+            credential_ref: null,
+            provider_resource_id: instance.providerResourceId,
+            provider_result_json: null,
+            last_error: null,
+            release_audit_json: null,
+            created_at: now,
+            updated_at: now,
+        })
+            .returningAll()
+            .executeTakeFirstOrThrow();
+        return mapProvisioningLease(row);
+    }
+    async cancelSupersededProviderStopLeasesForDelete(executor, input, now, reason) {
+        await executor
+            .updateTable("provisioning_leases")
+            .set({
+            state: "cancelled",
+            completed_at: now,
+            last_error: reason,
+            updated_at: now,
+        })
+            .where("owner_key", "=", input.ownerKey)
+            .where("host_id", "=", input.hostId)
+            .where("instance_id", "=", input.instanceId)
+            .where("action", "=", "provider_stop")
+            .where("state", "in", [...nonTerminalProvisioningLeaseStates])
+            .execute();
+    }
     /** Returns one owner/host-scoped provisioning lease. */
     async getProvisioningLease(input) {
         const row = await this.db
@@ -1675,7 +1940,8 @@ export class ManagedSandboxInstancesStore extends ArtifactsEventsStore {
                 "available") {
                 return [];
             }
-            const now = this.nowIso();
+            const nowDate = this.clock.now();
+            const now = nowDate.toISOString();
             const candidates = await trx
                 .selectFrom("provisioning_leases as lease")
                 .innerJoin("host_pools as pool", "pool.host_pool_id", "lease.host_pool_id")
@@ -1775,7 +2041,8 @@ export class ManagedSandboxInstancesStore extends ArtifactsEventsStore {
                 action: "provider_stop",
                 limit,
             });
-            const now = this.nowIso();
+            const nowDate = this.clock.now();
+            const now = nowDate.toISOString();
             const candidates = await trx
                 .selectFrom("provisioning_leases as lease")
                 .innerJoin("host_pools as pool", "pool.host_pool_id", "lease.host_pool_id")
@@ -1841,17 +2108,26 @@ export class ManagedSandboxInstancesStore extends ArtifactsEventsStore {
                 action: "provider_cleanup",
                 limit,
             });
-            const now = this.nowIso();
+            const nowDate = this.clock.now();
+            const now = nowDate.toISOString();
             const candidates = await trx
                 .selectFrom("provisioning_leases as lease")
                 .innerJoin("host_pools as pool", "pool.host_pool_id", "lease.host_pool_id")
                 .innerJoin("managed_sandbox_instances as instance", "instance.instance_id", "lease.instance_id")
-                .select(["lease.lease_id", "pool.provider"])
+                .select([
+                "lease.lease_id",
+                "lease.lease_deadline_at",
+                "pool.provider",
+                "instance.status as instance_status",
+            ])
                 .where("lease.owner_key", "=", input.ownerKey)
                 .where("lease.host_id", "=", input.hostId)
                 .where("lease.action", "=", "provider_cleanup")
                 .where("lease.state", "=", "pending")
-                .where("lease.lease_deadline_at", ">", now)
+                .where((eb) => eb.or([
+                eb("lease.lease_deadline_at", ">", now),
+                eb("instance.status", "=", "deleting"),
+            ]))
                 .where("lease.provider_resource_id", "is not", null)
                 .where("pool.owner_key", "=", input.ownerKey)
                 .where("pool.host_id", "=", input.hostId)
@@ -1868,6 +2144,25 @@ export class ManagedSandboxInstancesStore extends ArtifactsEventsStore {
                 .execute();
             const claimed = [];
             for (const candidate of candidates) {
+                if (candidate.instance_status === "deleting" &&
+                    candidate.lease_deadline_at <= now) {
+                    const refreshed = await trx
+                        .updateTable("provisioning_leases")
+                        .set({
+                        lease_deadline_at: addMs(nowDate, providerCleanupLeaseTtlMs),
+                        last_error: null,
+                        updated_at: now,
+                    })
+                        .where("lease_id", "=", candidate.lease_id)
+                        .where("owner_key", "=", input.ownerKey)
+                        .where("host_id", "=", input.hostId)
+                        .where("action", "=", "provider_cleanup")
+                        .where("state", "=", "pending")
+                        .executeTakeFirst();
+                    if (refreshed.numUpdatedRows === 0n) {
+                        continue;
+                    }
+                }
                 const result = await this.claimLeaseInTransaction(trx, {
                     ownerKey: input.ownerKey,
                     hostId: input.hostId,
@@ -2360,6 +2655,20 @@ export class ManagedSandboxInstancesStore extends ArtifactsEventsStore {
             if (!instanceRow) {
                 return { status: "rejected" };
             }
+            const deletingIdentity = input.status === "succeeded" && instanceRow.status === "deleting"
+                ? await trx
+                    .selectFrom("sandbox_identities")
+                    .select(["sandbox_identity_id", "disabled_at"])
+                    .where("owner_key", "=", current.owner_key)
+                    .where("sandbox_id", "=", current.sandbox_id ?? "")
+                    .forUpdate()
+                    .executeTakeFirst()
+                : undefined;
+            if (input.status === "succeeded" && instanceRow.status === "deleting") {
+                if (!deletingIdentity) {
+                    return { status: "rejected" };
+                }
+            }
             const row = await trx
                 .updateTable("provisioning_leases")
                 .set({
@@ -2385,6 +2694,22 @@ export class ManagedSandboxInstancesStore extends ArtifactsEventsStore {
             if (!row) {
                 throw new Error("failed to record provider cleanup lease result");
             }
+            if (deletingIdentity) {
+                await this.invalidateManagedSandboxRuntimeAuthForDelete(trx, {
+                    sandboxIdentityId: deletingIdentity.sandbox_identity_id,
+                    identityDisabledAt: deletingIdentity.disabled_at,
+                    reason: redactReason(input.reason) ?? "managed sandbox deleted",
+                    now,
+                });
+                await trx
+                    .deleteFrom("managed_sandbox_instances")
+                    .where("instance_id", "=", current.instance_id)
+                    .where("owner_key", "=", current.owner_key)
+                    .where("host_id", "=", current.host_id)
+                    .where("status", "=", "deleting")
+                    .where("provider_resource_id", "is", null)
+                    .executeTakeFirst();
+            }
             return {
                 status: "recorded",
                 lease: mapProvisioningLease(row),
@@ -2836,8 +3161,84 @@ export class ManagedSandboxInstancesStore extends ArtifactsEventsStore {
         if (nonterminalManagedInstances.count > 0) {
             blockers.push("nonterminal_managed_instances");
         }
+        const retryableDeletingProviderInstances = await executor
+            .selectFrom("managed_sandbox_instances")
+            .select(sql `count(*)::int`.as("count"))
+            .where("owner_key", "=", ownerKey)
+            .where("host_id", "=", hostId)
+            .where("status", "=", "deleting")
+            .where("provider_resource_id", "is not", null)
+            .executeTakeFirstOrThrow();
+        if (retryableDeletingProviderInstances.count > 0 &&
+            !blockers.includes("nonterminal_managed_instances")) {
+            blockers.push("nonterminal_managed_instances");
+        }
         return blockers;
     }
+    async lockManagedSandboxLiveness(executor, sandboxId) {
+        return executor
+            .selectFrom("sandboxes")
+            .select(["sandbox_id", "status"])
+            .where("sandbox_id", "=", sandboxId)
+            .forUpdate()
+            .executeTakeFirst();
+    }
+    async managedSandboxHasAcceptedRunningAssignments(executor, sandboxId) {
+        const row = await executor
+            .selectFrom("tasks")
+            .select("task_id")
+            .where("status", "=", "RUNNING")
+            .where("assigned_sandbox_id", "=", sandboxId)
+            .where("assignment_accepted_at", "is not", null)
+            .where("current_assignment_id", "is not", null)
+            .forUpdate()
+            .executeTakeFirst();
+        return row !== undefined;
+    }
+    async clearReservedSandboxAssignmentsForDelete(executor, sandboxId) {
+        const tasks = await executor
+            .selectFrom("tasks")
+            .select(["task_id", "run_id"])
+            .where("status", "=", "RUNNING")
+            .where("assigned_sandbox_id", "=", sandboxId)
+            .where("assignment_accepted_at", "is", null)
+            .where("current_assignment_id", "is not", null)
+            .forUpdate()
+            .execute();
+        if (tasks.length === 0) {
+            return 0;
+        }
+        const now = this.nowIso();
+        const taskIds = tasks.map((task) => task.task_id);
+        const cleared = await executor
+            .updateTable("tasks")
+            .set({
+            status: "QUEUED",
+            current_attempt_number: null,
+            current_assignment_id: null,
+            assigned_sandbox_id: null,
+            assignment_created_at: null,
+            assignment_sent_at: null,
+            assignment_delivery_deadline_at: null,
+            assignment_accepted_at: null,
+            assignment_expires_at: null,
+            assignment_heartbeat_at: null,
+            assignment_delivery_attempt_count: 0,
+            assignment_last_delivery_error: "managed_sandbox_deleted",
+            not_before_at: null,
+            updated_at: now,
+            queued_at: now,
+        })
+            .where("task_id", "in", taskIds)
+            .where("status", "=", "RUNNING")
+            .where("assigned_sandbox_id", "=", sandboxId)
+            .where("assignment_accepted_at", "is", null)
+            .executeTakeFirst();
+        for (const runId of new Set(tasks.map((task) => task.run_id))) {
+            await this.updateRunState(executor, runId);
+        }
+        return Number(cleared.numUpdatedRows);
+    }
     async requireHostPoolOwner(executor, hostPoolId, hostId, ownerKey) {
         const pool = await executor
             .selectFrom("host_pools")