Spaces:
Paused
Paused
| import { randomUUID } from "node:crypto"; | |
| import type { Filter, ObjectId } from "mongodb"; | |
| import { collections } from "$lib/server/database"; | |
| import type { Workspace, WorkspaceOperationLease, WorkspaceSandbox } from "$lib/types/Workspace"; | |
| export const WORKSPACE_OPERATION_LEASE_MS = 20 * 60 * 1_000; | |
| export class WorkspaceOperationConflict extends Error { | |
| constructor(message = "Agent workspace changed while the operation was starting") { | |
| super(message); | |
| this.name = "WorkspaceOperationConflict"; | |
| } | |
| } | |
| export interface ClaimedWorkspaceOperation { | |
| workspace: Workspace; | |
| lease: WorkspaceOperationLease; | |
| } | |
| /** | |
| * Admission predicate shared by every workspace producer/relay/RPC path. | |
| * Callers which mutate the workspace must pair it with `$unset: { operation: | |
| * "" }`; the predicate then makes reclaiming an orphaned lease atomic while a | |
| * lease younger than the timeout remains a hard fence. | |
| */ | |
| export function workspaceOperationAvailableFilter(now: Date = new Date()): Filter<Workspace> { | |
| const staleBefore = new Date(now.getTime() - WORKSPACE_OPERATION_LEASE_MS); | |
| return { | |
| $or: [{ operation: { $exists: false } }, { "operation.claimedAt": { $lte: staleBefore } }], | |
| }; | |
| } | |
| /** | |
| * Claims the exact workspace revision only while its conversation still owns it | |
| * in Agent mode. Retirement uses the reciprocal `operation` guard, so either | |
| * this claim or a mode change wins; neither can cross the other's async work. | |
| */ | |
| export async function claimWorkspaceOperation(options: { | |
| conversationId: ObjectId; | |
| workspace: Workspace; | |
| owner: Record<string, unknown>; | |
| kind: WorkspaceOperationLease["kind"]; | |
| action: string; | |
| now?: Date; | |
| }): Promise<ClaimedWorkspaceOperation> { | |
| const now = options.now ?? new Date(); | |
| const lease: WorkspaceOperationLease = { | |
| id: randomUUID(), | |
| kind: options.kind, | |
| action: options.action.slice(0, 64), | |
| claimedAt: now, | |
| workspaceUpdatedAt: options.workspace.updatedAt, | |
| }; | |
| const workspace = await collections.workspaces.findOneAndUpdate( | |
| { | |
| _id: options.workspace._id, | |
| conversationId: options.conversationId, | |
| ...options.owner, | |
| updatedAt: options.workspace.updatedAt, | |
| ...sandboxIdentityFilter(options.workspace.sandbox), | |
| retirement: { $exists: false }, | |
| $and: [ | |
| workspaceOperationAvailableFilter(now), | |
| { | |
| $or: [{ execution: { $exists: false } }, { "execution.expiresAt": { $lte: now } }], | |
| }, | |
| ], | |
| }, | |
| { | |
| $set: { operation: lease }, | |
| // The predicate admits only an absent/expired execution. Removing an | |
| // expired claim fences a delayed producer/relay from renewing into this | |
| // workspace mutation after the lease has been granted. | |
| $unset: { execution: "" }, | |
| }, | |
| { returnDocument: "after", includeResultMetadata: false } | |
| ); | |
| if (!workspace) throw new WorkspaceOperationConflict(); | |
| const stillOwned = await collections.conversations.findOne({ | |
| _id: options.conversationId, | |
| ...options.owner, | |
| mode: "agent", | |
| workspaceId: workspace._id, | |
| }); | |
| if (!stillOwned) { | |
| await releaseWorkspaceOperation(workspace._id, options.owner, lease.id); | |
| throw new WorkspaceOperationConflict( | |
| "The conversation changed before the workspace operation could start" | |
| ); | |
| } | |
| return { workspace, lease }; | |
| } | |
| /** Revalidates both collections and the exact workspace revision after an await. */ | |
| export async function requireWorkspaceOperation(options: { | |
| conversationId: ObjectId; | |
| workspace: Workspace; | |
| owner: Record<string, unknown>; | |
| leaseId: string; | |
| }): Promise<Workspace> { | |
| const conversation = await collections.conversations.findOne({ | |
| _id: options.conversationId, | |
| ...options.owner, | |
| mode: "agent", | |
| workspaceId: options.workspace._id, | |
| }); | |
| if (!conversation) { | |
| throw new WorkspaceOperationConflict("The conversation changed during this operation"); | |
| } | |
| const workspace = await collections.workspaces.findOne({ | |
| _id: options.workspace._id, | |
| conversationId: options.conversationId, | |
| ...options.owner, | |
| "operation.id": options.leaseId, | |
| retirement: { $exists: false }, | |
| updatedAt: options.workspace.updatedAt, | |
| ...sandboxIdentityFilter(options.workspace.sandbox), | |
| }); | |
| if (!workspace) | |
| throw new WorkspaceOperationConflict("The agent sandbox changed during this operation"); | |
| return workspace; | |
| } | |
| /** | |
| * Reload an operation after a provisioner deliberately changed the sandbox | |
| * state before reporting an error. The reciprocal lease still proves no | |
| * supported mode/reset path crossed the operation. | |
| */ | |
| export async function reloadWorkspaceOperation(options: { | |
| conversationId: ObjectId; | |
| workspaceId: ObjectId; | |
| owner: Record<string, unknown>; | |
| leaseId: string; | |
| }): Promise<Workspace> { | |
| const conversation = await collections.conversations.findOne({ | |
| _id: options.conversationId, | |
| ...options.owner, | |
| mode: "agent", | |
| workspaceId: options.workspaceId, | |
| }); | |
| if (!conversation) { | |
| throw new WorkspaceOperationConflict("The conversation changed during this operation"); | |
| } | |
| const workspace = await collections.workspaces.findOne({ | |
| _id: options.workspaceId, | |
| conversationId: options.conversationId, | |
| ...options.owner, | |
| "operation.id": options.leaseId, | |
| retirement: { $exists: false }, | |
| }); | |
| if (!workspace) throw new WorkspaceOperationConflict("The agent workspace changed"); | |
| return workspace; | |
| } | |
| export async function releaseWorkspaceOperation( | |
| workspaceId: ObjectId, | |
| owner: Record<string, unknown>, | |
| leaseId: string | |
| ): Promise<void> { | |
| await collections.workspaces.updateOne( | |
| { _id: workspaceId, ...owner, "operation.id": leaseId }, | |
| { $unset: { operation: "" } } | |
| ); | |
| } | |
| export function sandboxIdentityFilter(sandbox: WorkspaceSandbox): Record<string, unknown> { | |
| return { | |
| "sandbox.status": sandbox.status, | |
| ...(sandbox.id ? { "sandbox.id": sandbox.id } : { "sandbox.id": { $exists: false } }), | |
| ...(sandbox.nonce | |
| ? { "sandbox.nonce": sandbox.nonce } | |
| : { "sandbox.nonce": { $exists: false } }), | |
| ...(sandbox.generation | |
| ? { "sandbox.generation": sandbox.generation } | |
| : { "sandbox.generation": { $exists: false } }), | |
| }; | |
| } | |