src/lib/server/agent/agentPromptRecovery.ts

import { collections } from "$lib/server/database";
import type { Workspace, WorkspaceSandbox } from "$lib/types/Workspace";

export interface AgentPromptRecovery {
	sandboxId: string;
	sandboxGeneration: string;
	seededMessage: string;
}

export type DetachedAgentPromptPreparation = "ready" | "busy" | "unrecoverable" | "gone";

/**
 * A prompt assembled for an existing pi session must be replaced with its
 * transcript-seeded form when provisioning returns a different sandbox. The
 * CAS only succeeds before any relay has claimed or dispatched the command.
 */
export async function reseedAgentPromptForReplacement(options: {
	generationId: string;
	conversationId: Workspace["conversationId"];
	workspaceId: Workspace["_id"];
	prompt: { id: string; message: string };
	recovery?: AgentPromptRecovery;
	sandbox: Pick<WorkspaceSandbox, "id" | "generation">;
	producerOwner?: string;
}): Promise<string> {
	const { recovery } = options;
	if (
		!recovery ||
		(options.sandbox.id === recovery.sandboxId &&
			options.sandbox.generation === recovery.sandboxGeneration)
	) {
		return options.prompt.message;
	}
	if (!options.sandbox.id || !options.sandbox.generation) {
		throw new Error("Replacement Agent sandbox has no durable incarnation");
	}

	const reseeded = await collections.generations.updateOne(
		{
			generationId: options.generationId,
			conversationId: options.conversationId,
			workspaceId: options.workspaceId,
			mode: "agent",
			status: "running",
			"agent.prompt.id": options.prompt.id,
			"agent.prompt.message": options.prompt.message,
			"agent.relay": { $exists: false },
			"agent.sandboxId": recovery.sandboxId,
			"agent.sandboxGeneration": recovery.sandboxGeneration,
			...(options.producerOwner
				? {
						"agent.producer.owner": options.producerOwner,
						"agent.producer.expiresAt": { $gt: new Date() },
					}
				: {
						$or: [
							{ "agent.producer": { $exists: false } },
							{ "agent.producer.expiresAt": { $lte: new Date() } },
						],
					}),
		},
		{
			$set: {
				"agent.prompt.message": recovery.seededMessage,
				"agent.sandboxId": options.sandbox.id,
				"agent.sandboxGeneration": options.sandbox.generation,
			},
			$unset: { "agent.prompt.recovery": "" },
		}
	);
	if (reseeded.matchedCount !== 1) {
		const current = await collections.generations.findOne(
			{
				generationId: options.generationId,
				conversationId: options.conversationId,
				workspaceId: options.workspaceId,
				mode: "agent",
				status: "running",
			},
			{
				projection: {
					"agent.prompt.id": 1,
					"agent.prompt.message": 1,
					"agent.sandboxId": 1,
					"agent.sandboxGeneration": 1,
				},
			}
		);
		if (
			current?.agent?.prompt?.id === options.prompt.id &&
			current.agent.prompt.message === recovery.seededMessage &&
			current.agent.sandboxId === options.sandbox.id &&
			current.agent.sandboxGeneration === options.sandbox.generation
		) {
			return recovery.seededMessage;
		}
		throw new Error("Agent prompt changed before replacement context was restored");
	}
	return recovery.seededMessage;
}

/**
 * Recover a replacement-bound prompt after its originating producer died.
 * The durable one-shot fallback means a detached relay never dispatches a
 * resume prompt into the wrong pi-session incarnation.
 */
export async function prepareDetachedAgentPromptForSandbox(options: {
	generationId: string;
	workspaceId: Workspace["_id"];
	sandbox: Pick<WorkspaceSandbox, "id" | "generation">;
	now?: Date;
}): Promise<DetachedAgentPromptPreparation> {
	const { id: sandboxId, generation: sandboxGeneration } = options.sandbox;
	if (!sandboxId || !sandboxGeneration) return "unrecoverable";
	const now = options.now ?? new Date();
	const generation = await collections.generations.findOne(
		{
			generationId: options.generationId,
			workspaceId: options.workspaceId,
			mode: "agent",
		},
		{
			projection: {
				status: 1,
				"agent.sandboxId": 1,
				"agent.sandboxGeneration": 1,
				"agent.prompt": 1,
				"agent.producer": 1,
				"agent.relay": 1,
			},
		}
	);
	if (!generation || generation.status !== "running") return "gone";

	const agent = generation.agent;
	const boundId = agent?.sandboxId;
	const boundGeneration = agent?.sandboxGeneration;
	if (
		(boundId === undefined && boundGeneration === undefined) ||
		(boundId === sandboxId && boundGeneration === sandboxGeneration)
	) {
		return "ready";
	}
	const producer = agent?.producer;
	if (producer && producer.expiresAt > now) return "busy";
	const prompt = agent?.prompt;
	const recovery = prompt?.recovery;
	if (
		!prompt ||
		!recovery ||
		agent?.relay ||
		recovery.sandboxId !== boundId ||
		recovery.sandboxGeneration !== boundGeneration
	) {
		return "unrecoverable";
	}

	const reseeded = await collections.generations.updateOne(
		{
			generationId: options.generationId,
			workspaceId: options.workspaceId,
			mode: "agent",
			status: "running",
			"agent.prompt.id": prompt.id,
			"agent.prompt.message": prompt.message,
			"agent.prompt.recovery.sandboxId": recovery.sandboxId,
			"agent.prompt.recovery.sandboxGeneration": recovery.sandboxGeneration,
			"agent.prompt.recovery.seededMessage": recovery.seededMessage,
			"agent.sandboxId": boundId,
			"agent.sandboxGeneration": boundGeneration,
			"agent.relay": { $exists: false },
			$or: [
				{ "agent.producer": { $exists: false } },
				{ "agent.producer.expiresAt": { $lte: now } },
			],
		},
		{
			$set: {
				"agent.prompt.message": recovery.seededMessage,
				"agent.sandboxId": sandboxId,
				"agent.sandboxGeneration": sandboxGeneration,
			},
			// Completing takeover recovery fences the expired producer. Its delayed
			// heartbeat must not resurrect the pre-dispatch lease between this CAS
			// and detached relay acquisition.
			$unset: { "agent.prompt.recovery": "", "agent.producer": "" },
		}
	);
	if (reseeded.matchedCount === 1) return "ready";

	const current = await collections.generations.findOne(
		{
			generationId: options.generationId,
			workspaceId: options.workspaceId,
			mode: "agent",
		},
		{
			projection: {
				status: 1,
				"agent.sandboxId": 1,
				"agent.sandboxGeneration": 1,
				"agent.prompt.message": 1,
				"agent.producer": 1,
			},
		}
	);
	if (!current || current.status !== "running") return "gone";
	if (
		current.agent?.sandboxId === sandboxId &&
		current.agent?.sandboxGeneration === sandboxGeneration &&
		current.agent?.prompt?.message === recovery.seededMessage
	) {
		return "ready";
	}
	if (current.agent?.producer && current.agent.producer.expiresAt > now) return "busy";
	return "unrecoverable";
}