Spaces:
Paused
Paused
| import type { ObjectId } from "mongodb"; | |
| import { AbortRegistry } from "$lib/server/abortRegistry"; | |
| import { authCondition } from "$lib/server/auth"; | |
| import { collections } from "$lib/server/database"; | |
| import type { Conversation } from "$lib/types/Conversation"; | |
| import type { Generation } from "$lib/types/Generation"; | |
| export interface GenerationStopMarker { | |
| generationId: string; | |
| seenContentLength?: number; | |
| } | |
| /** Return the durable stop point carried by a registered generation. */ | |
| export function generationStopMarker( | |
| generation: Pick<Generation, "generationId" | "stopRequested"> | |
| ): GenerationStopMarker | null { | |
| if (!generation.stopRequested) return null; | |
| return { | |
| generationId: generation.generationId, | |
| ...(generation.stopRequested.seenContentLength === undefined | |
| ? {} | |
| : { seenContentLength: generation.stopRequested.seenContentLength }), | |
| }; | |
| } | |
| /** | |
| * Resolve a stop intent without depending on the short-lived wake-up marker. | |
| * | |
| * Once a Generation exists, its stopRequested field is authoritative. The | |
| * aborted-generations collection is only a pre-registration/cross-pod wake-up | |
| * handshake and may expire while a detached Agent run is still recoverable. | |
| */ | |
| export async function resolveGenerationStop(options: { | |
| conversationId: ObjectId; | |
| generationId: string; | |
| }): Promise<GenerationStopMarker | null> { | |
| const generation = await collections.generations.findOne( | |
| { | |
| conversationId: options.conversationId, | |
| generationId: options.generationId, | |
| }, | |
| { projection: { generationId: 1, stopRequested: 1 } } | |
| ); | |
| if (generation) return generationStopMarker(generation); | |
| const marker = await collections.abortedGenerations.findOne( | |
| { | |
| conversationId: options.conversationId, | |
| generationId: options.generationId, | |
| }, | |
| { projection: { seenContentLength: 1 } } | |
| ); | |
| if (!marker) return null; | |
| return { | |
| generationId: options.generationId, | |
| ...(marker.seenContentLength === undefined | |
| ? {} | |
| : { seenContentLength: marker.seenContentLength }), | |
| }; | |
| } | |
| interface ResolveStopTargetParams { | |
| conversation: Conversation; | |
| locals: App.Locals; | |
| requestedGenerationId?: string; | |
| } | |
| /** | |
| * Resolve a stop request to one exact, still-running generation owned by this | |
| * caller. Agent legacy requests have no generation id, so their only safe | |
| * target is the generation named by the workspace's current live execution | |
| * claim. Chat legacy requests retain their historical newest-running fallback. | |
| */ | |
| export async function resolveStopGeneration({ | |
| conversation, | |
| locals, | |
| requestedGenerationId, | |
| }: ResolveStopTargetParams): Promise<string | undefined> { | |
| const owner = authCondition(locals); | |
| const base = { | |
| conversationId: conversation._id, | |
| status: "running" as const, | |
| ...owner, | |
| }; | |
| if (requestedGenerationId) { | |
| const generation = await collections.generations.findOne( | |
| { generationId: requestedGenerationId }, | |
| { | |
| projection: { | |
| generationId: 1, | |
| conversationId: 1, | |
| userId: 1, | |
| sessionId: 1, | |
| status: 1, | |
| }, | |
| } | |
| ); | |
| // The client chooses the UUID before the producer registers its writer. | |
| // Preserve an exact pre-registration stop so the eventual owner sees the | |
| // tombstone before dispatch. An existing row, however, must be the caller's | |
| // still-running row; terminal or foreign IDs are safe no-ops. | |
| if (!generation) return requestedGenerationId; | |
| const sameOwner = locals.user | |
| ? generation.userId?.equals(locals.user._id) === true | |
| : generation.sessionId === locals.sessionId && generation.userId === undefined; | |
| if ( | |
| generation.status !== "running" || | |
| !generation.conversationId.equals(conversation._id) || | |
| !sameOwner | |
| ) { | |
| return undefined; | |
| } | |
| return generation.generationId; | |
| } | |
| if (conversation.mode === "agent") { | |
| if (!conversation.workspaceId) return undefined; | |
| const workspace = await collections.workspaces.findOne( | |
| { | |
| _id: conversation.workspaceId, | |
| conversationId: conversation._id, | |
| ...owner, | |
| retirement: { $exists: false }, | |
| "execution.expiresAt": { $gt: new Date() }, | |
| }, | |
| { projection: { execution: 1 } } | |
| ); | |
| if (!workspace?.execution) return undefined; | |
| const generation = await collections.generations.findOne( | |
| { | |
| ...base, | |
| generationId: workspace.execution.generationId, | |
| mode: "agent", | |
| workspaceId: conversation.workspaceId, | |
| }, | |
| { projection: { generationId: 1 } } | |
| ); | |
| return generation?.generationId; | |
| } | |
| const generation = await collections.generations.findOne(base, { | |
| projection: { generationId: 1 }, | |
| sort: { startedAt: -1, _id: -1 }, | |
| }); | |
| return generation?.generationId; | |
| } | |
| interface RecordStopParams { | |
| conversationId: ObjectId; | |
| generationId: string; | |
| seenContentLength?: number; | |
| } | |
| function stopIntentUpdate(now: Date, seenContentLength?: number) { | |
| return { | |
| $set: { | |
| "stopRequested.requestedAt": now, | |
| updatedAt: now, | |
| }, | |
| ...(seenContentLength === undefined | |
| ? {} | |
| : { $min: { "stopRequested.seenContentLength": seenContentLength } }), | |
| }; | |
| } | |
| async function upsertStopMarker({ | |
| conversationId, | |
| generationId, | |
| seenContentLength, | |
| now, | |
| }: RecordStopParams & { now: Date }): Promise<void> { | |
| await collections.abortedGenerations.updateOne( | |
| { conversationId, generationId }, | |
| { | |
| $set: { updatedAt: now }, | |
| ...(seenContentLength === undefined ? {} : { $min: { seenContentLength } }), | |
| $setOnInsert: { createdAt: now }, | |
| }, | |
| { upsert: true } | |
| ); | |
| } | |
| /** Complete the registration side of the pre-registration stop handshake. */ | |
| export async function reconcileRegisteredGenerationStop(options: { | |
| conversationId: ObjectId; | |
| generationId: string; | |
| }): Promise<boolean> { | |
| const marker = await collections.abortedGenerations.findOne( | |
| { conversationId: options.conversationId, generationId: options.generationId }, | |
| { projection: { seenContentLength: 1, updatedAt: 1 } } | |
| ); | |
| if (!marker) return false; | |
| const requestedAt = marker.updatedAt ?? new Date(); | |
| const result = await collections.generations.updateOne( | |
| { | |
| conversationId: options.conversationId, | |
| generationId: options.generationId, | |
| status: "running", | |
| }, | |
| stopIntentUpdate(requestedAt, marker.seenContentLength) | |
| ); | |
| return result.matchedCount === 1; | |
| } | |
| /** | |
| * Atomically compete with terminalization on the exact Generation row, then | |
| * write the cross-pod wake-up marker and signal this pod's exact owner. | |
| * | |
| * A client can stop before its POST has inserted the Generation. The marker and | |
| * the post-marker retry form a two-sided handshake with writer registration: | |
| * whichever side inserts second observes the other, so a pre-registration stop | |
| * cannot fall into the check/insert gap. | |
| */ | |
| export async function recordGenerationStop({ | |
| conversationId, | |
| generationId, | |
| seenContentLength, | |
| }: RecordStopParams): Promise<boolean> { | |
| const now = new Date(); | |
| const claimed = await collections.generations.updateOne( | |
| { conversationId, generationId, status: "running" }, | |
| stopIntentUpdate(now, seenContentLength) | |
| ); | |
| if (claimed.matchedCount === 0) { | |
| // An existing row means terminalization (or a foreign/conflicting reuse) | |
| // already won. Do not leave a marker that can mislabel a terminal answer. | |
| const existing = await collections.generations.findOne( | |
| { generationId }, | |
| { projection: { _id: 1 } } | |
| ); | |
| if (existing) return false; | |
| } | |
| await upsertStopMarker({ conversationId, generationId, seenContentLength, now }); | |
| if (claimed.matchedCount === 0) { | |
| // Close the no-row -> insert race. If registration appeared after our first | |
| // CAS, this retry installs the authoritative intent before returning. | |
| const raced = await collections.generations.updateOne( | |
| { conversationId, generationId, status: "running" }, | |
| stopIntentUpdate(now, seenContentLength) | |
| ); | |
| if (raced.matchedCount === 0) { | |
| const terminal = await collections.generations.findOne( | |
| { generationId }, | |
| { projection: { _id: 1 } } | |
| ); | |
| if (terminal) { | |
| await collections.abortedGenerations.deleteOne({ conversationId, generationId }); | |
| return false; | |
| } | |
| } | |
| } | |
| AbortRegistry.getInstance().abort(conversationId.toString(), generationId); | |
| return true; | |
| } | |