Spaces:
Paused
Paused
| import { randomUUID } from "node:crypto"; | |
| import { authCondition } from "$lib/server/auth"; | |
| import { collections } from "$lib/server/database"; | |
| import { getAgentJobsToken } from "$lib/server/agent/credentials"; | |
| import { logger } from "$lib/server/logger"; | |
| import { refreshWorkspaceSandboxCost } from "$lib/server/agent/sandboxCost"; | |
| import { workspaceProvisioner } from "$lib/server/agent/WorkspaceProvisioner"; | |
| import { workspaceOperationAvailableFilter } from "$lib/server/agent/WorkspaceOperation"; | |
| import type { Workspace } from "$lib/types/Workspace"; | |
| const RETIREMENT_LEASE_MS = 15 * 60 * 1_000; | |
| export class WorkspaceRetirementConflict extends Error { | |
| constructor() { | |
| super("Agent workspace changed while it was being retired"); | |
| this.name = "WorkspaceRetirementConflict"; | |
| } | |
| } | |
| export interface RetiredAgentWorkspace { | |
| workspace: Workspace; | |
| leaseId: string; | |
| } | |
| /** | |
| * Cancels Jobs before their owning records are deleted. This is deliberately | |
| * sequential so a bulk conversation delete cannot burst Jobs cancellation calls. | |
| */ | |
| export async function retireAgentWorkspaces( | |
| workspaces: Workspace[], | |
| locals: App.Locals | |
| ): Promise<RetiredAgentWorkspace[]> { | |
| const owner = authCondition(locals); | |
| const claimed: RetiredAgentWorkspace[] = []; | |
| // Claim every target before stopping the first Job. A multi-workspace delete | |
| // must either own the whole set or leave the records intact for a retry. | |
| try { | |
| for (const candidate of workspaces) { | |
| const leaseId = randomUUID(); | |
| const claimNow = new Date(); | |
| const retirementStaleBefore = new Date(claimNow.getTime() - RETIREMENT_LEASE_MS); | |
| const canReclaimLease = | |
| candidate.retirement && | |
| (candidate.sandbox.status === "error" || | |
| candidate.retirement.claimedAt < retirementStaleBefore); | |
| // A provisioner owns the interval between its database claim and publishing | |
| // the newly-created Job identity. Retiring that fresh claim could mark the | |
| // workspace dead (and later delete it) before the provisioner has a Job id | |
| // to stop. Let the in-flight provision finish or fail; its exact writes are | |
| // then fenced by the retirement claim on the next attempt. | |
| const provisioningClaimedAt = candidate.sandbox.claimedAt ?? candidate.updatedAt; | |
| if ( | |
| candidate.sandbox.status === "provisioning" && | |
| provisioningClaimedAt >= retirementStaleBefore && | |
| !canReclaimLease | |
| ) { | |
| throw new WorkspaceRetirementConflict(); | |
| } | |
| const reclaimFailedLease = | |
| canReclaimLease && candidate.retirement | |
| ? { | |
| "retirement.id": candidate.retirement.id, | |
| "retirement.claimedAt": candidate.retirement.claimedAt, | |
| } | |
| : { retirement: { $exists: false } }; | |
| const workspace = await collections.workspaces.findOneAndUpdate( | |
| { | |
| _id: candidate._id, | |
| ...owner, | |
| updatedAt: candidate.updatedAt, | |
| "sandbox.status": candidate.sandbox.status, | |
| ...(candidate.sandbox.id ? { "sandbox.id": candidate.sandbox.id } : {}), | |
| ...(candidate.sandbox.nonce ? { "sandbox.nonce": candidate.sandbox.nonce } : {}), | |
| ...(candidate.sandbox.generation | |
| ? { "sandbox.generation": candidate.sandbox.generation } | |
| : {}), | |
| $and: [ | |
| { | |
| $or: [ | |
| { execution: { $exists: false } }, | |
| { "execution.expiresAt": { $lte: claimNow } }, | |
| ], | |
| }, | |
| workspaceOperationAvailableFilter(claimNow), | |
| ], | |
| ...reclaimFailedLease, | |
| }, | |
| { | |
| $set: { | |
| retirement: { id: leaseId, claimedAt: claimNow }, | |
| }, | |
| $unset: { execution: "", operation: "" }, | |
| }, | |
| { returnDocument: "after", includeResultMetadata: false } | |
| ); | |
| if (!workspace) throw new WorkspaceRetirementConflict(); | |
| claimed.push({ workspace, leaseId }); | |
| } | |
| } catch (cause) { | |
| await releaseWorkspaceRetirements(claimed, locals); | |
| throw cause; | |
| } | |
| for (let index = 0; index < claimed.length; index += 1) { | |
| const retirement = claimed[index]; | |
| let workspace = retirement.workspace; | |
| try { | |
| try { | |
| workspace = await refreshWorkspaceSandboxCost(workspace, getAgentJobsToken(locals)); | |
| } catch (cause) { | |
| logger.warn( | |
| { cause, workspaceId: workspace._id.toString() }, | |
| "Unable to finalize sandbox cost before deletion" | |
| ); | |
| } | |
| const stopped = await workspaceProvisioner.kill({ | |
| workspace, | |
| requestHfToken: locals.token, | |
| }); | |
| claimed[index] = { workspace: stopped, leaseId: retirement.leaseId }; | |
| } catch (cause) { | |
| // Keep the failed target leased: its Job could not be confirmed stopped. | |
| // Mark only our exact lease stale so the next explicit teardown can retry | |
| // immediately, including failures that happened before kill() persisted an | |
| // error state. A process crash still observes the normal lease timeout. | |
| await collections.workspaces.updateOne( | |
| { _id: workspace._id, ...owner, "retirement.id": retirement.leaseId }, | |
| { $set: { "retirement.claimedAt": new Date(0) } } | |
| ); | |
| await releaseWorkspaceRetirements( | |
| claimed.filter((_, claimedIndex) => claimedIndex !== index), | |
| locals | |
| ); | |
| throw cause; | |
| } | |
| } | |
| return claimed; | |
| } | |
| /** Release leases only when the caller still owns them. */ | |
| export async function releaseWorkspaceRetirements( | |
| retired: RetiredAgentWorkspace[], | |
| locals: App.Locals | |
| ): Promise<void> { | |
| const owner = authCondition(locals); | |
| for (const { workspace, leaseId } of retired) { | |
| await collections.workspaces.updateOne( | |
| { _id: workspace._id, ...owner, "retirement.id": leaseId }, | |
| { $unset: { retirement: "" } } | |
| ); | |
| } | |
| } | |
| /** Exact filter for the stopped incarnation protected by a retirement lease. */ | |
| export function retiredWorkspaceFilter( | |
| retired: RetiredAgentWorkspace, | |
| locals: App.Locals | |
| ): Record<string, unknown> { | |
| const { workspace, leaseId } = retired; | |
| return { | |
| _id: workspace._id, | |
| ...authCondition(locals), | |
| "retirement.id": leaseId, | |
| updatedAt: workspace.updatedAt, | |
| "sandbox.status": workspace.sandbox.status, | |
| ...(workspace.sandbox.id ? { "sandbox.id": workspace.sandbox.id } : {}), | |
| ...(workspace.sandbox.nonce ? { "sandbox.nonce": workspace.sandbox.nonce } : {}), | |
| ...(workspace.sandbox.generation ? { "sandbox.generation": workspace.sandbox.generation } : {}), | |
| }; | |
| } | |