Spaces:
Paused
Paused
| import type { ObjectId } from "mongodb"; | |
| import { collections } from "$lib/server/database"; | |
| import { logger } from "$lib/server/logger"; | |
| import { markGenerationMessageInterrupted } from "./terminalClaim"; | |
| export interface ActiveRun { | |
| conversationId: ObjectId; | |
| messageId: string; | |
| mode?: "chat" | "agent"; | |
| } | |
| // This pod's in-flight runs, so a graceful shutdown can finalize them at once | |
| // instead of leaving each to the reaper's ~90s sweep. | |
| const activeRuns = new Map<string, ActiveRun>(); | |
| export function registerActiveRun(generationId: string, run: ActiveRun): void { | |
| activeRuns.set(generationId, run); | |
| } | |
| export function unregisterActiveRun(generationId: string): void { | |
| activeRuns.delete(generationId); | |
| } | |
| // The message flag — not just the generation status — is what every existing reader | |
| // treats as terminal, so both must move together. Flip the generation first, guarded on | |
| // `status: "running"`, so only one caller (a reaper on any pod, or the owner's shutdown) | |
| // wins the claim; mark the message only on that win. | |
| export async function markGenerationInterrupted( | |
| generationId: string, | |
| _run: ActiveRun | |
| ): Promise<void> { | |
| const now = new Date(); | |
| const generation = await collections.generations.findOneAndUpdate( | |
| { generationId, status: "running" }, | |
| { $set: { status: "interrupted", endedAt: now, updatedAt: now } }, | |
| { returnDocument: "before", includeResultMetadata: false } | |
| ); | |
| // Lost the claim: the run completed or errored elsewhere first. Marking the message now | |
| // would flag a finished answer as stopped, so leave it untouched. | |
| if (!generation) return; | |
| await collections.workspaces.updateOne( | |
| { "execution.generationId": generationId }, | |
| { $unset: { execution: "" } } | |
| ); | |
| await markGenerationMessageInterrupted(generation, now); | |
| } | |
| export async function finalizeActiveRunsOnExit(): Promise<void> { | |
| // A pi process runs in its HF Sandbox, not in this Node process. Leaving an | |
| // Agent generation running is intentional: another pod will acquire its | |
| // durable relay lease and continue it. Chat generations still need the old | |
| // interrupted finalization because their upstream request dies with this pod. | |
| const runs = [...activeRuns.entries()].filter(([, run]) => run.mode !== "agent"); | |
| if (runs.length === 0) return; | |
| logger.info({ count: runs.length }, "[generation] finalizing in-flight runs on shutdown"); | |
| await Promise.all( | |
| runs.map(([generationId, run]) => | |
| markGenerationInterrupted(generationId, run).catch((err) => | |
| logger.error({ err, generationId }, "[generation] failed to finalize run on shutdown") | |
| ) | |
| ) | |
| ); | |
| } | |