Spaces:
Paused
Paused
| import { collections } from "$lib/server/database"; | |
| import { resolveGenerationStop } from "$lib/server/generation/stopGeneration"; | |
| import { logger } from "$lib/server/logger"; | |
| import { clampStoppedContent } from "$lib/server/stopTruncation"; | |
| import { downloadFile } from "$lib/server/files/downloadFile"; | |
| import type { Generation, GenerationStatus } from "$lib/types/Generation"; | |
| import { | |
| MessageUpdateStatus, | |
| MessageUpdateType, | |
| type MessageUpdate, | |
| } from "$lib/types/MessageUpdate"; | |
| import type { Workspace } from "$lib/types/Workspace"; | |
| import { | |
| AgentGenerationIncarnationMismatchError, | |
| AgentGenerationStore, | |
| AgentRelayLeaseLostError, | |
| AGENT_RELAY_LEASE_MS, | |
| AGENT_SANDBOX_CHANGED_MESSAGE, | |
| } from "./AgentGenerationStore"; | |
| import { AgentControlPump } from "./AgentControlPump"; | |
| import { AgentSessionManager } from "./AgentSessionManager"; | |
| import { prepareDetachedAgentPromptForSandbox } from "./agentPromptRecovery"; | |
| import { PI_RPC_REPLAY_GAP_MESSAGE, PiRpcReplayGapError } from "./PiRpcClient"; | |
| import type { SandboxClient } from "./SandboxClient"; | |
| import { agentUploadDirectory, isAgentUploadPath } from "./agentUploadPaths"; | |
| import { workspaceOperationAvailableFilter } from "./WorkspaceOperation"; | |
| interface ActiveAgentRelay { | |
| sandboxId: string; | |
| sandboxGeneration: string; | |
| promise: Promise<void>; | |
| } | |
| const activeRelays = new Map<string, ActiveAgentRelay>(); | |
| const ABORT_POLL_MS = 300; | |
| export interface EnsureAgentRelayOptions { | |
| generationId: string; | |
| workspace: Workspace; | |
| hfToken: string; | |
| bridgeUrl?: string; | |
| bridgeHeaders?: Record<string, string>; | |
| } | |
| /** | |
| * Opportunistically claim and resume a detached Agent run. The returned | |
| * promise only covers lease acquisition; the relay intentionally continues in | |
| * the background after the requesting SSE connection goes away. | |
| */ | |
| export async function ensureAgentGenerationRelay( | |
| options: EnsureAgentRelayOptions | |
| ): Promise<"started" | "already-running" | "gone"> { | |
| const sandboxId = options.workspace.sandbox.id; | |
| const sandboxGeneration = options.workspace.sandbox.generation; | |
| if (!sandboxId || !sandboxGeneration) { | |
| await AgentGenerationStore.interruptForSandboxChange({ | |
| generationId: options.generationId, | |
| workspaceId: options.workspace._id, | |
| }); | |
| return "gone"; | |
| } | |
| const workspaceOwner = options.workspace.userId | |
| ? { userId: options.workspace.userId } | |
| : options.workspace.sessionId | |
| ? { sessionId: options.workspace.sessionId } | |
| : undefined; | |
| if (!workspaceOwner) return "gone"; | |
| const admissionAt = new Date(); | |
| const dispatchableWorkspace = await collections.workspaces.findOneAndUpdate( | |
| { | |
| _id: options.workspace._id, | |
| conversationId: options.workspace.conversationId, | |
| ...workspaceOwner, | |
| retirement: { $exists: false }, | |
| ...workspaceOperationAvailableFilter(admissionAt), | |
| "sandbox.status": { $in: ["ready", "idle"] }, | |
| "sandbox.id": sandboxId, | |
| "sandbox.generation": sandboxGeneration, | |
| }, | |
| { $unset: { operation: "" } }, | |
| { projection: { _id: 1 }, returnDocument: "after", includeResultMetadata: false } | |
| ); | |
| if (!dispatchableWorkspace) return "already-running"; | |
| const active = activeRelays.get(options.generationId); | |
| if (active) { | |
| if (active.sandboxId === sandboxId && active.sandboxGeneration === sandboxGeneration) { | |
| return "already-running"; | |
| } | |
| await AgentGenerationStore.interruptForSandboxChange({ | |
| generationId: options.generationId, | |
| workspaceId: options.workspace._id, | |
| sandboxId, | |
| sandboxGeneration, | |
| }); | |
| return "gone"; | |
| } | |
| const promptPreparation = await prepareDetachedAgentPromptForSandbox({ | |
| generationId: options.generationId, | |
| workspaceId: options.workspace._id, | |
| sandbox: options.workspace.sandbox, | |
| }); | |
| if (promptPreparation === "busy") return "already-running"; | |
| if (promptPreparation === "gone") return "gone"; | |
| if (promptPreparation === "unrecoverable") { | |
| await AgentGenerationStore.interruptForSandboxChange({ | |
| generationId: options.generationId, | |
| workspaceId: options.workspace._id, | |
| sandboxId, | |
| sandboxGeneration, | |
| sandboxNonce: options.workspace.sandbox.nonce, | |
| }); | |
| return "gone"; | |
| } | |
| let lease: AgentGenerationStore | null; | |
| try { | |
| lease = await AgentGenerationStore.acquire({ | |
| generationId: options.generationId, | |
| workspaceId: options.workspace._id, | |
| sandboxId, | |
| sandboxGeneration, | |
| sandboxNonce: options.workspace.sandbox.nonce, | |
| }); | |
| } catch (cause) { | |
| if (!(cause instanceof AgentGenerationIncarnationMismatchError)) throw cause; | |
| await AgentGenerationStore.interruptForSandboxChange({ | |
| generationId: options.generationId, | |
| workspaceId: options.workspace._id, | |
| sandboxId, | |
| sandboxGeneration, | |
| }); | |
| return "gone"; | |
| } | |
| if (!lease) { | |
| const generation = await collections.generations.findOne( | |
| { generationId: options.generationId }, | |
| { projection: { status: 1 } } | |
| ); | |
| return generation?.status === "running" ? "already-running" : "gone"; | |
| } | |
| const relay = resumeWithLease(options, lease) | |
| .catch((cause) => | |
| logger.error( | |
| { cause, generationId: options.generationId }, | |
| "Detached Agent generation relay failed" | |
| ) | |
| ) | |
| .finally(() => { | |
| if (activeRelays.get(options.generationId)?.promise === relay) { | |
| activeRelays.delete(options.generationId); | |
| } | |
| }); | |
| activeRelays.set(options.generationId, { sandboxId, sandboxGeneration, promise: relay }); | |
| return "started"; | |
| } | |
| /** Test/diagnostic hook; production callers should use ensureAgentGenerationRelay. */ | |
| export function hasActiveAgentRelay(generationId: string): boolean { | |
| return activeRelays.has(generationId); | |
| } | |
| export function waitForActiveAgentRelay(generationId: string): Promise<void> | undefined { | |
| return activeRelays.get(generationId)?.promise; | |
| } | |
| async function readDurableLifecycleState(generationId: string): Promise<{ | |
| sawError: boolean; | |
| sawFinished: boolean; | |
| error?: string; | |
| }> { | |
| const events = await collections.generationEvents | |
| .find({ | |
| generationId, | |
| "event.type": MessageUpdateType.Status, | |
| "event.status": { | |
| $in: [MessageUpdateStatus.Error, MessageUpdateStatus.Finished], | |
| }, | |
| }) | |
| .sort({ seq: 1 }) | |
| .toArray(); | |
| let sawError = false; | |
| let sawFinished = false; | |
| let error: string | undefined; | |
| for (const record of events) { | |
| if (record.event.type !== MessageUpdateType.Status) continue; | |
| if (record.event.status === MessageUpdateStatus.Error) { | |
| sawError = true; | |
| error = record.event.message ?? error; | |
| } else if (record.event.status === MessageUpdateStatus.Finished) { | |
| sawFinished = true; | |
| } | |
| } | |
| return { sawError, sawFinished, ...(error ? { error } : {}) }; | |
| } | |
| async function resumeWithLease( | |
| options: EnsureAgentRelayOptions, | |
| store: AgentGenerationStore | |
| ): Promise<void> { | |
| let manager: AgentSessionManager | undefined; | |
| let controlPump: AgentControlPump | undefined; | |
| let controlPumpFailure: Error | undefined; | |
| let startControlPump: (() => AgentControlPump) | undefined; | |
| let leaseTimer: ReturnType<typeof setInterval> | undefined; | |
| let abortTimer: ReturnType<typeof setInterval> | undefined; | |
| let finished = false; | |
| let sawFinished = false; | |
| let sawError = false; | |
| let terminalError: string | undefined; | |
| const abortController = new AbortController(); | |
| const ownershipController = new AbortController(); | |
| try { | |
| // Start fencing renewal before WAL recovery, workspace reads, or attachment | |
| // restoration. Those operations can be slow enough to consume a short lease. | |
| await store.renew(); | |
| leaseTimer = setInterval( | |
| () => { | |
| void store.renew().catch(() => ownershipController.abort()); | |
| }, | |
| Math.max(1_000, Math.floor(AGENT_RELAY_LEASE_MS / 3)) | |
| ); | |
| leaseTimer.unref?.(); | |
| await store.resumePending(); | |
| const [generation, workspace, lifecycle] = await Promise.all([ | |
| collections.generations.findOne({ generationId: options.generationId }), | |
| collections.workspaces.findOne({ | |
| _id: options.workspace._id, | |
| conversationId: store.conversationId, | |
| retirement: { $exists: false }, | |
| "sandbox.id": store.sandboxId, | |
| "sandbox.generation": store.sandboxGeneration, | |
| "sandbox.status": { $in: ["ready", "idle"] }, | |
| }), | |
| readDurableLifecycleState(options.generationId), | |
| ]); | |
| if (!generation || generation.status !== "running") return; | |
| sawError = lifecycle.sawError; | |
| sawFinished = lifecycle.sawFinished; | |
| terminalError = lifecycle.error; | |
| if (!workspace) { | |
| // The reciprocal workspace lease can no longer be renewed, so normal | |
| // finish() cannot drain its WAL/mailbox preconditions. The exact relay | |
| // fence is still sufficient to terminate this now-unrecoverable run. | |
| finished = await store.finishUnrecoverable("interrupted", AGENT_SANDBOX_CHANGED_MESSAGE); | |
| return; | |
| } | |
| const lastSequence = Math.max( | |
| generation.agent?.sourceSequence ?? 0, | |
| workspace.piSession?.lastSequence ?? 0 | |
| ); | |
| manager = AgentSessionManager.connect({ | |
| workspace, | |
| hfToken: options.hfToken, | |
| bridgeUrl: options.bridgeUrl, | |
| bridgeHeaders: options.bridgeHeaders, | |
| lastSequence, | |
| }); | |
| const connectedManager = manager; | |
| startControlPump = () => { | |
| controlPump ??= new AgentControlPump({ | |
| store, | |
| manager: connectedManager, | |
| workspace, | |
| signal: ownershipController.signal, | |
| onFailure: (cause) => { | |
| controlPumpFailure = cause; | |
| ownershipController.abort(); | |
| }, | |
| }); | |
| return controlPump; | |
| }; | |
| const startControlPumpAfterDispatch = startControlPump; | |
| const existingAbort = await resolveGenerationStop({ | |
| conversationId: generation.conversationId, | |
| generationId: generation.generationId, | |
| }).catch(() => null); | |
| if (existingAbort) abortController.abort(); | |
| if (!existingAbort) { | |
| await restorePromptAttachments( | |
| manager.sandbox, | |
| generation.agent?.prompt?.attachments, | |
| generation.conversationId, | |
| generation.generationId, | |
| { | |
| signal: ownershipController.signal, | |
| assertOwnership: () => store.assertPromptDispatch(workspace), | |
| } | |
| ); | |
| } | |
| abortTimer = setInterval(() => { | |
| resolveGenerationStop({ | |
| conversationId: generation.conversationId, | |
| generationId: generation.generationId, | |
| }) | |
| .then((marker) => { | |
| if (marker && !abortController.signal.aborted) abortController.abort(); | |
| }) | |
| .catch(() => undefined); | |
| }, ABORT_POLL_MS); | |
| abortTimer.unref?.(); | |
| if (ownershipController.signal.aborted) throw new AgentRelayLeaseLostError(); | |
| const batches = abortController.signal.aborted | |
| ? manager.resumeBatches({ | |
| signal: abortController.signal, | |
| detachSignal: ownershipController.signal, | |
| expectedPromptId: generation.agent?.prompt?.id, | |
| beforeAbortDispatch: () => store.renew(), | |
| keepRpcOpen: true, | |
| }) | |
| : generation.agent?.prompt | |
| ? manager.runBatches(generation.agent.prompt.message, { | |
| signal: abortController.signal, | |
| detachSignal: ownershipController.signal, | |
| commandId: generation.agent.prompt.id, | |
| beforeAbortDispatch: () => store.renew(), | |
| beforePromptDispatch: () => store.assertPromptDispatch(workspace), | |
| afterPromptDispatch: () => { | |
| startControlPumpAfterDispatch(); | |
| }, | |
| keepRpcOpen: true, | |
| }) | |
| : manager.resumeBatches({ | |
| signal: abortController.signal, | |
| detachSignal: ownershipController.signal, | |
| beforeAbortDispatch: () => store.renew(), | |
| keepRpcOpen: true, | |
| }); | |
| if (abortController.signal.aborted || !generation.agent?.prompt) startControlPump(); | |
| for await (const batch of batches) { | |
| for (const event of batch.updates) { | |
| if ( | |
| event.type === MessageUpdateType.Status && | |
| event.status === MessageUpdateStatus.Finished | |
| ) { | |
| sawFinished = true; | |
| } | |
| if (event.type === MessageUpdateType.Status && event.status === MessageUpdateStatus.Error) { | |
| sawError = true; | |
| terminalError = event.message ?? terminalError; | |
| } | |
| } | |
| await store.commitBridgeBatch(batch.sequence, batch.updates); | |
| } | |
| if (controlPumpFailure) throw controlPumpFailure; | |
| await startControlPump().sealAndStop(); | |
| const stop = await resolveGenerationStop({ | |
| conversationId: generation.conversationId, | |
| generationId: generation.generationId, | |
| }).catch(() => null); | |
| const wasAborted = abortController.signal.aborted || Boolean(stop); | |
| const terminalEvents: MessageUpdate[] = []; | |
| if (wasAborted) { | |
| const conversation = await collections.conversations.findOne( | |
| { _id: generation.conversationId }, | |
| { projection: { messages: 1 } } | |
| ); | |
| const message = conversation?.messages.find( | |
| (item) => item.id === generation.messageId && item.generationId === generation.generationId | |
| ); | |
| const content = message?.content ?? ""; | |
| const clamped = clampStoppedContent({ | |
| content, | |
| initialLength: 0, | |
| generationId: generation.generationId, | |
| marker: stop, | |
| }); | |
| terminalEvents.push({ | |
| type: MessageUpdateType.FinalAnswer, | |
| text: clamped, | |
| interrupted: true, | |
| }); | |
| } | |
| if (!sawFinished) { | |
| terminalEvents.push({ | |
| type: MessageUpdateType.Status, | |
| status: MessageUpdateStatus.Finished, | |
| }); | |
| } | |
| terminalEvents.push({ | |
| type: MessageUpdateType.AgentSandbox, | |
| state: "idle", | |
| message: wasAborted ? "Agent stopped" : "Agent is ready for the next instruction", | |
| sandboxId: workspace.sandbox.id, | |
| flavor: workspace.flavor, | |
| }); | |
| await store.commitSystemBatch("relay-terminal", terminalEvents); | |
| const status: GenerationStatus = sawError ? "error" : wasAborted ? "interrupted" : "completed"; | |
| const now = new Date(); | |
| const idled = await collections.workspaces.updateOne( | |
| { | |
| _id: workspace._id, | |
| conversationId: store.conversationId, | |
| retirement: { $exists: false }, | |
| "sandbox.id": store.sandboxId, | |
| "sandbox.nonce": workspace.sandbox.nonce, | |
| "sandbox.generation": store.sandboxGeneration, | |
| "sandbox.status": { $in: ["ready", "idle"] }, | |
| "execution.generationId": store.generationId, | |
| "execution.phase": "relay", | |
| "execution.owner": store.owner, | |
| "execution.fence": store.fence, | |
| "execution.expiresAt": { $gt: now }, | |
| }, | |
| { | |
| $set: { | |
| "sandbox.status": "idle", | |
| "sandbox.lastSeenAt": now, | |
| updatedAt: now, | |
| }, | |
| } | |
| ); | |
| if (idled.matchedCount !== 1) throw new AgentGenerationIncarnationMismatchError(); | |
| finished = await store.finish(status, status === "error" ? terminalError : undefined); | |
| if (!finished) throw new AgentRelayLeaseLostError(); | |
| if (wasAborted) { | |
| await collections.abortedGenerations | |
| .deleteOne({ | |
| conversationId: generation.conversationId, | |
| generationId: generation.generationId, | |
| }) | |
| .catch(() => undefined); | |
| } | |
| } catch (cause) { | |
| if (controlPumpFailure) throw controlPumpFailure; | |
| if ( | |
| !(cause instanceof AgentGenerationIncarnationMismatchError) && | |
| !ownershipController.signal.aborted && | |
| startControlPump | |
| ) { | |
| await startControlPump().sealAndStop(); | |
| } | |
| if (cause instanceof AgentGenerationIncarnationMismatchError) { | |
| finished = await store.finishUnrecoverable("interrupted", AGENT_SANDBOX_CHANGED_MESSAGE); | |
| throw cause; | |
| } | |
| if (cause instanceof PiRpcReplayGapError) { | |
| try { | |
| await store.commitSystemBatch("relay-replay-gap", [ | |
| { | |
| type: MessageUpdateType.Status, | |
| status: MessageUpdateStatus.Error, | |
| message: PI_RPC_REPLAY_GAP_MESSAGE, | |
| }, | |
| ]); | |
| } finally { | |
| finished = await finishRelayFailure(store, "error", PI_RPC_REPLAY_GAP_MESSAGE); | |
| } | |
| throw cause; | |
| } | |
| if (abortController.signal.aborted) { | |
| // The abort timeout/path above normally reaches terminal persistence. | |
| // If the bridge itself disappeared, preserve a durable error instead of | |
| // allowing the generic reaper to silently strand the run. | |
| await store | |
| .commitSystemBatch("relay-abort-error", [ | |
| { | |
| type: MessageUpdateType.Status, | |
| status: MessageUpdateStatus.Error, | |
| message: "Agent stopped before the sandbox acknowledged the abort", | |
| }, | |
| ]) | |
| .catch(() => undefined); | |
| } | |
| finished = await finishRelayFailure( | |
| store, | |
| "error", | |
| cause instanceof Error ? cause.message : "Agent relay failed" | |
| ); | |
| throw cause; | |
| } finally { | |
| if (leaseTimer) clearInterval(leaseTimer); | |
| if (abortTimer) clearInterval(abortTimer); | |
| controlPump?.requestStop(); | |
| manager?.close(); | |
| await controlPump?.waitForStop().catch(() => undefined); | |
| if (!finished) await store.release().catch(() => undefined); | |
| } | |
| } | |
| async function finishRelayFailure( | |
| store: AgentGenerationStore, | |
| status: "interrupted" | "error", | |
| error: string | |
| ): Promise<boolean> { | |
| try { | |
| if (await store.finish(status, error)) return true; | |
| } catch { | |
| // A fatal bridge/incarnation failure can also make the reciprocal | |
| // workspace renewal inside normal finish() impossible. The exact | |
| // generation relay fence below is the final authority in that case. | |
| } | |
| return store.finishUnrecoverable(status, error); | |
| } | |
| async function restorePromptAttachments( | |
| sandbox: SandboxClient, | |
| attachments: NonNullable<NonNullable<Generation["agent"]>["prompt"]>["attachments"], | |
| conversationId: Generation["conversationId"], | |
| generationId: string, | |
| options: { signal: AbortSignal; assertOwnership: () => Promise<void> } | |
| ): Promise<void> { | |
| if (!attachments?.length) return; | |
| await options.assertOwnership(); | |
| const mkdir = await sandbox.exec( | |
| ["mkdir", "-p", `/workspace/${agentUploadDirectory(generationId)}`], | |
| { | |
| signal: options.signal, | |
| } | |
| ); | |
| if (mkdir.exitCode !== 0 || mkdir.timedOut) { | |
| throw new Error("Unable to restore Agent attachments"); | |
| } | |
| for (const attachment of attachments) { | |
| await options.assertOwnership(); | |
| if (!isAgentUploadPath(attachment.path, generationId)) { | |
| throw new Error("Stored Agent attachment path is invalid"); | |
| } | |
| const bytes = | |
| attachment.file.type === "base64" | |
| ? Buffer.from(attachment.file.value, "base64") | |
| : Buffer.from((await downloadFile(attachment.file.value, conversationId)).value, "base64"); | |
| await options.assertOwnership(); | |
| await sandbox.writeFile( | |
| `/workspace/${attachment.path}`, | |
| Uint8Array.from(bytes).buffer, | |
| undefined, | |
| { signal: options.signal } | |
| ); | |
| await options.assertOwnership(); | |
| } | |
| } | |