src/lib/server/agent/AgentGenerationStore.ts

import { randomUUID } from "node:crypto";
import { ObjectId, type Filter, type UpdateFilter } from "mongodb";
import { collections } from "$lib/server/database";
import type { AgentPendingBatch, Generation, GenerationStatus } from "$lib/types/Generation";
import type { Message } from "$lib/types/Message";
import { MessageUpdateType, type MessageUpdate } from "$lib/types/MessageUpdate";
import type { Workspace } from "$lib/types/Workspace";
import { reduceAgentMessageUpdates } from "./agentMessageReducer";
import { workspaceOperationAvailableFilter } from "./WorkspaceOperation";
import {
	generationTerminalCandidates,
	markGenerationMessageInterrupted,
} from "$lib/server/generation/terminalClaim";

export const AGENT_RELAY_LEASE_MS = 8_000;
export const AGENT_TERMINALIZATION_LEASE_MS = 15_000;
export const AGENT_SANDBOX_CHANGED_MESSAGE = "Agent workspace restarted before this run completed";

export class AgentRelayLeaseLostError extends Error {
	constructor() {
		super("Agent generation relay lease was lost");
		this.name = "AgentRelayLeaseLostError";
	}
}

export class AgentGenerationIncarnationMismatchError extends Error {
	constructor() {
		super(AGENT_SANDBOX_CHANGED_MESSAGE);
		this.name = "AgentGenerationIncarnationMismatchError";
	}
}

export interface AgentBatchCommit {
	inserted: boolean;
	endSeq: number;
	sourceSequence?: number;
}

interface AcquireOptions {
	generationId: string;
	workspaceId: Workspace["_id"];
	sandboxId: string;
	sandboxGeneration: string;
	sandboxNonce?: string;
	owner?: string;
	/** Exact pre-dispatch producer lease owned by the originating request. */
	producerOwner?: string;
	leaseMs?: number;
	now?: () => Date;
}

export type AgentGenerationOwner =
	| { userId: NonNullable<Generation["userId"]> }
	| { sessionId: NonNullable<Generation["sessionId"]> };

/**
 * Fenced, crash-recoverable writer for one Agent generation.
 *
 * A batch is first reserved in `generations.agent.pending` together with its
 * exact event-log sequence range. Only then are the append-only events and the
 * materialized message written. Clearing pending advances the bridge cursor.
 * A new process therefore either skips a completed batch or finishes the same
 * reservation; there is no prompt replay and no duplicate visible update.
 */
export class AgentGenerationStore {
	readonly generationId: string;
	readonly conversationId: Generation["conversationId"];
	readonly messageId: Generation["messageId"];
	readonly workspaceId: Workspace["_id"];
	readonly sandboxId: string;
	readonly sandboxGeneration: string;
	readonly sandboxNonce?: string;
	readonly owner: string;
	readonly fence: number;

	private constructor(
		options: Required<
			Pick<AcquireOptions, "generationId" | "workspaceId" | "sandboxId" | "sandboxGeneration">
		> & {
			conversationId: Generation["conversationId"];
			messageId: Generation["messageId"];
			sandboxNonce?: string;
			owner: string;
			leaseMs: number;
			now: () => Date;
			fence: number;
		}
	) {
		this.generationId = options.generationId;
		this.conversationId = options.conversationId;
		this.messageId = options.messageId;
		this.workspaceId = options.workspaceId;
		this.sandboxId = options.sandboxId;
		this.sandboxGeneration = options.sandboxGeneration;
		this.sandboxNonce = options.sandboxNonce;
		this.owner = options.owner;
		this.fence = options.fence;
		this.leaseMs = options.leaseMs;
		this.now = options.now;
	}

	private readonly leaseMs: number;
	private readonly now: () => Date;

	static async acquire(options: AcquireOptions): Promise<AgentGenerationStore | null> {
		const owner = options.owner ?? randomUUID();
		const now = (options.now ?? (() => new Date()))();
		const leaseMs = options.leaseMs ?? AGENT_RELAY_LEASE_MS;
		const generation = await collections.generations.findOneAndUpdate(
			{
				generationId: options.generationId,
				workspaceId: options.workspaceId,
				mode: "agent",
				status: "running",
				"agent.terminalization": { $exists: false },
				$and: [
					{
						$or: [
							{ "agent.relay": { $exists: false } },
							{ "agent.relay.expiresAt": { $lte: now } },
						],
					},
					{
						$or: [
							{
								"agent.sandboxId": options.sandboxId,
								"agent.sandboxGeneration": options.sandboxGeneration,
							},
							{
								$and: [
									{ "agent.sandboxId": { $exists: false } },
									{ "agent.sandboxGeneration": { $exists: false } },
								],
							},
						],
					},
					options.producerOwner
						? {
								"agent.producer.owner": options.producerOwner,
								"agent.producer.expiresAt": { $gt: now },
							}
						: {
								$or: [
									{ "agent.producer": { $exists: false } },
									{ "agent.producer.expiresAt": { $lte: now } },
								],
							},
				],
			},
			{
				$set: {
					"agent.sandboxId": options.sandboxId,
					"agent.sandboxGeneration": options.sandboxGeneration,
					"agent.relay.owner": owner,
					"agent.relay.expiresAt": new Date(now.getTime() + leaseMs),
					lastHeartbeatAt: now,
					updatedAt: now,
				},
				$inc: { "agent.relay.fence": 1 },
				$unset: { "agent.producer": "", "agent.prompt.recovery": "" },
			},
			{ returnDocument: "after", includeResultMetadata: false }
		);

		const fence = generation?.agent?.relay?.fence;
		if (!generation || fence === undefined) {
			const current = await collections.generations.findOne(
				{
					generationId: options.generationId,
					workspaceId: options.workspaceId,
					mode: "agent",
					status: "running",
				},
				{
					projection: {
						"agent.sandboxId": 1,
						"agent.sandboxGeneration": 1,
						"agent.producer": 1,
					},
				}
			);
			const producer = current?.agent?.producer;
			if (producer && producer.expiresAt > now && producer.owner !== options.producerOwner) {
				return null;
			}
			const boundSandboxId = current?.agent?.sandboxId;
			const boundSandboxGeneration = current?.agent?.sandboxGeneration;
			if (
				current &&
				(boundSandboxId !== undefined || boundSandboxGeneration !== undefined) &&
				(boundSandboxId !== options.sandboxId ||
					boundSandboxGeneration !== options.sandboxGeneration)
			) {
				throw new AgentGenerationIncarnationMismatchError();
			}
			return null;
		}
		const generationOwner = generation.userId
			? { userId: generation.userId }
			: generation.sessionId
				? { sessionId: generation.sessionId }
				: undefined;
		if (!generationOwner) {
			await clearAcquiredRelay(options.generationId, owner, fence);
			throw new AgentGenerationIncarnationMismatchError();
		}
		const execution = await collections.workspaces.updateOne(
			{
				_id: options.workspaceId,
				conversationId: generation.conversationId,
				...generationOwner,
				retirement: { $exists: false },
				"sandbox.status": { $in: ["ready", "idle"] },
				"sandbox.id": options.sandboxId,
				"sandbox.generation": options.sandboxGeneration,
				...(options.sandboxNonce ? { "sandbox.nonce": options.sandboxNonce } : {}),
				$and: [
					workspaceOperationAvailableFilter(now),
					options.producerOwner
						? {
								"execution.generationId": options.generationId,
								"execution.phase": "producer",
								"execution.owner": options.producerOwner,
								"execution.expiresAt": { $gt: now },
							}
						: {
								$or: [
									{ "execution.generationId": options.generationId },
									{ execution: { $exists: false } },
									{ "execution.expiresAt": { $lte: now } },
								],
							},
				],
			},
			{
				$set: {
					execution: {
						generationId: options.generationId,
						phase: "relay",
						owner,
						fence,
						expiresAt: new Date(now.getTime() + leaseMs),
						sandboxId: options.sandboxId,
						sandboxGeneration: options.sandboxGeneration,
						...(options.sandboxNonce ? { sandboxNonce: options.sandboxNonce } : {}),
					},
				},
				$unset: { operation: "" },
			}
		);
		if (execution.matchedCount !== 1) {
			await clearAcquiredRelay(options.generationId, owner, fence);
			// An exact sandbox can be temporarily fenced by a workspace operation
			// (configuration, deliverables, reset) or another execution claimant.
			// That is contention, not an incarnation change, and must never cause the
			// caller to interrupt the durable generation.
			const contended = await collections.workspaces.findOne(
				{
					_id: options.workspaceId,
					conversationId: generation.conversationId,
					...generationOwner,
					retirement: { $exists: false },
					"sandbox.status": { $in: ["ready", "idle"] },
					"sandbox.id": options.sandboxId,
					"sandbox.generation": options.sandboxGeneration,
					...(options.sandboxNonce ? { "sandbox.nonce": options.sandboxNonce } : {}),
				},
				{ projection: { _id: 1 } }
			);
			if (contended) return null;
			throw new AgentGenerationIncarnationMismatchError();
		}
		return new AgentGenerationStore({
			generationId: options.generationId,
			conversationId: generation.conversationId,
			messageId: generation.messageId,
			workspaceId: options.workspaceId,
			sandboxId: options.sandboxId,
			sandboxGeneration: options.sandboxGeneration,
			...(options.sandboxNonce ? { sandboxNonce: options.sandboxNonce } : {}),
			owner,
			leaseMs,
			now: options.now ?? (() => new Date()),
			fence,
		});
	}

	/**
	 * Atomically terminate a run that is bound to a different sandbox than the
	 * workspace's current incarnation. Missing current identity is also unsafe
	 * and terminates the run rather than permitting an unbound resume.
	 */
	static async interruptForSandboxChange(options: {
		generationId: string;
		workspaceId: Workspace["_id"];
		sandboxId?: string;
		sandboxGeneration?: string;
		sandboxNonce?: string;
	}): Promise<boolean> {
		const claimAt = new Date();
		const terminalizationId = randomUUID();
		const exactCurrentIncarnation = options.sandboxId && options.sandboxGeneration;
		const mismatchFilter: Filter<Generation> = exactCurrentIncarnation
			? {
					$and: [
						{
							$or: [
								{ "agent.sandboxId": { $exists: true } },
								{ "agent.sandboxGeneration": { $exists: true } },
							],
						},
						{
							$or: [
								{ "agent.sandboxId": { $ne: options.sandboxId } },
								{ "agent.sandboxGeneration": { $ne: options.sandboxGeneration } },
							],
						},
					],
				}
			: {};
		const generation = await collections.generations.findOneAndUpdate(
			{
				generationId: options.generationId,
				workspaceId: options.workspaceId,
				mode: "agent",
				$or: [
					{
						status: "running",
						"agent.terminalization": { $exists: false },
						...mismatchFilter,
					},
					{
						status: { $in: ["running", "interrupted"] },
						"agent.terminalization.kind": "sandbox-change",
						"agent.terminalization.expiresAt": { $lte: claimAt },
					},
				],
			},
			{
				$set: {
					"agent.terminalization": {
						id: terminalizationId,
						kind: "sandbox-change",
						claimedAt: claimAt,
						expiresAt: new Date(claimAt.getTime() + AGENT_TERMINALIZATION_LEASE_MS),
					},
					"agent.controlsOpen": false,
					lastHeartbeatAt: claimAt,
					updatedAt: claimAt,
				},
				$unset: {
					"agent.relay": "",
					"agent.producer": "",
					"agent.prompt.recovery": "",
					"agent.control": "",
				},
			},
			{ returnDocument: "after", includeResultMetadata: false }
		);
		if (!generation) return false;

		let terminalGeneration = generation;
		if (generation.status === "running") {
			const pending = generation.agent?.pending;
			if (pending) {
				await persistAgentPendingEvents(generation.generationId, pending, () => new Date());
				const renewed = await renewAgentTerminalization(
					generation.generationId,
					options.workspaceId,
					terminalizationId
				);
				if (!renewed) return false;
				const endSeq = pending.startSeq + pending.events.length - 1;
				await applyAgentPendingToMessage(renewed, pending, endSeq, () => new Date());
			}

			const terminalAt = new Date();
			const renewed = await renewAgentTerminalization(
				generation.generationId,
				options.workspaceId,
				terminalizationId,
				terminalAt
			);
			if (!renewed) return false;
			const terminalUpdate: UpdateFilter<Generation> = {
				...(pending?.kind === "system" ? { $addToSet: { "agent.systemKeys": pending.key } } : {}),
				$set: {
					status: "interrupted",
					"agent.controlsOpen": false,
					...(pending?.kind === "bridge" ? { "agent.sourceSequence": pending.sourceSequence } : {}),
					error: AGENT_SANDBOX_CHANGED_MESSAGE,
					endedAt: terminalAt,
					lastHeartbeatAt: terminalAt,
					updatedAt: terminalAt,
				},
				$unset: {
					"agent.relay": "",
					"agent.producer": "",
					"agent.prompt.recovery": "",
					"agent.control": "",
					"agent.pending": "",
					"agent.pendingApprovalIds": "",
				},
			};
			const terminalized = await collections.generations.findOneAndUpdate(
				{
					generationId: generation.generationId,
					workspaceId: options.workspaceId,
					mode: "agent",
					status: "running",
					"agent.terminalization.id": terminalizationId,
					"agent.terminalization.expiresAt": { $gt: terminalAt },
					...(pending ? { "agent.pending": pending } : { "agent.pending": { $exists: false } }),
				},
				terminalUpdate,
				{ returnDocument: "before", includeResultMetadata: false }
			);
			if (!terminalized) return false;
			terminalGeneration = terminalized;
		}

		// Keep the dedicated marker through both post-terminal side effects. If the
		// process dies here, the reaper can discover and take over the terminal row.
		const sideEffectAt = new Date();
		if (
			!(await renewAgentTerminalization(
				generation.generationId,
				options.workspaceId,
				terminalizationId,
				sideEffectAt
			))
		) {
			return false;
		}
		await markGenerationMessageInterrupted(terminalGeneration, sideEffectAt);
		await collections.workspaces.updateOne(
			{
				_id: options.workspaceId,
				"execution.generationId": options.generationId,
			},
			{ $unset: { execution: "" } }
		);
		await collections.generations.updateOne(
			{
				generationId: generation.generationId,
				workspaceId: options.workspaceId,
				"agent.terminalization.id": terminalizationId,
			},
			{ $unset: { "agent.terminalization": "" } }
		);
		return true;
	}

	async renew(): Promise<void> {
		const now = this.now();
		const expiresAt = new Date(now.getTime() + this.leaseMs);
		const workspace = await collections.workspaces.updateOne(
			{
				_id: this.workspaceId,
				conversationId: this.conversationId,
				retirement: { $exists: false },
				...workspaceOperationAvailableFilter(now),
				"execution.generationId": this.generationId,
				"execution.phase": "relay",
				"execution.owner": this.owner,
				"execution.fence": this.fence,
				"execution.sandboxId": this.sandboxId,
				"execution.sandboxGeneration": this.sandboxGeneration,
				...(this.sandboxNonce ? { "execution.sandboxNonce": this.sandboxNonce } : {}),
				"execution.expiresAt": { $gt: now },
				"sandbox.id": this.sandboxId,
				"sandbox.generation": this.sandboxGeneration,
				...(this.sandboxNonce ? { "sandbox.nonce": this.sandboxNonce } : {}),
			},
			{
				$set: { "execution.expiresAt": expiresAt, updatedAt: now },
				$unset: { operation: "" },
			}
		);
		if (workspace.matchedCount === 0) throw new AgentRelayLeaseLostError();
		const result = await collections.generations.updateOne(this.leaseFilter(true, now), {
			$set: {
				"agent.relay.expiresAt": expiresAt,
				lastHeartbeatAt: now,
				updatedAt: now,
			},
		});
		if (result.matchedCount === 0) {
			await this.releaseWorkspaceExecution();
			throw new AgentRelayLeaseLostError();
		}
	}

	/**
	 * Revalidate both fenced generation ownership and the exact live workspace
	 * incarnation after the bridge socket connects, immediately before prompt
	 * dispatch. A retirement claim is a hard stop even while the Job still says
	 * ready.
	 */
	async assertPromptDispatch(
		workspace: Pick<Workspace, "conversationId" | "sandbox">
	): Promise<void> {
		await this.renew();
		if (!workspace.conversationId.equals(this.conversationId)) {
			throw new AgentGenerationIncarnationMismatchError();
		}
		await this.resolveAuthoritativeOwner(workspace.sandbox.nonce);
	}

	/**
	 * Re-read ownership from the exact fenced generation, conversation, and
	 * workspace. Login moves these documents one collection at a time, so the
	 * temporary dual-owned generation/workspace may legitimately agree with the
	 * conversation through either principal. Returning only the principal shared
	 * by all three keeps this transition narrow and never treats the old session
	 * as an alias for the new account.
	 */
	async authoritativeOwner(): Promise<AgentGenerationOwner> {
		await this.renew();
		return this.resolveAuthoritativeOwner(this.sandboxNonce);
	}

	private async resolveAuthoritativeOwner(
		expectedSandboxNonce?: string
	): Promise<AgentGenerationOwner> {
		for (let attempt = 0; attempt < 4; attempt++) {
			const now = this.now();
			const [generation, conversation, workspace] = await Promise.all([
				collections.generations.findOne(this.leaseFilter(true, now), {
					projection: { userId: 1, sessionId: 1 },
				}),
				collections.conversations.findOne(
					{
						_id: this.conversationId,
						mode: "agent",
						workspaceId: this.workspaceId,
						messages: {
							$elemMatch: { id: this.messageId, generationId: this.generationId },
						},
					},
					{ projection: { userId: 1, sessionId: 1 } }
				),
				collections.workspaces.findOne(
					{
						_id: this.workspaceId,
						conversationId: this.conversationId,
						retirement: { $exists: false },
						"execution.generationId": this.generationId,
						"execution.phase": "relay",
						"execution.owner": this.owner,
						"execution.fence": this.fence,
						"execution.sandboxId": this.sandboxId,
						"execution.sandboxGeneration": this.sandboxGeneration,
						...(this.sandboxNonce ? { "execution.sandboxNonce": this.sandboxNonce } : {}),
						"execution.expiresAt": { $gt: now },
						"sandbox.status": { $in: ["ready", "idle"] },
						"sandbox.id": this.sandboxId,
						"sandbox.generation": this.sandboxGeneration,
						...(expectedSandboxNonce ? { "sandbox.nonce": expectedSandboxNonce } : {}),
					},
					{ projection: { userId: 1, sessionId: 1 } }
				),
			]);
			if (!generation) throw new AgentRelayLeaseLostError();
			if (!conversation || !workspace) {
				throw new AgentGenerationIncarnationMismatchError();
			}

			if (
				generation.userId &&
				conversation.userId?.equals(generation.userId) &&
				workspace.userId?.equals(generation.userId)
			) {
				return { userId: generation.userId };
			}
			if (
				generation.sessionId &&
				conversation.sessionId === generation.sessionId &&
				workspace.sessionId === generation.sessionId
			) {
				return { sessionId: generation.sessionId };
			}

			// The three primary reads can straddle one of the short login handoff
			// writes. Yield and re-read the exact fenced chain before declaring it
			// inconsistent.
			await Promise.resolve();
		}
		throw new AgentGenerationIncarnationMismatchError();
	}

	async resumePending(): Promise<AgentBatchCommit | undefined> {
		await this.renew();
		let generation = await collections.generations.findOne(this.leaseFilter());
		if (!generation) throw new AgentRelayLeaseLostError();
		const pending = generation.agent?.pending;
		if (pending) return this.flushPending(generation, pending);
		await this.synchronizeEventCursor(generation);
		generation = await collections.generations.findOne(this.leaseFilter());
		if (!generation) throw new AgentRelayLeaseLostError();
		return undefined;
	}

	private async synchronizeEventCursor(generation: Generation): Promise<void> {
		const latest = await collections.generationEvents.findOne(
			{ generationId: this.generationId },
			{ sort: { seq: -1 }, projection: { seq: 1 } }
		);
		if (!latest) return;
		await this.materializeEventBacklog(generation, latest.seq);
		const result = await collections.generations.updateOne(this.leaseFilter(), {
			$max: { seq: latest.seq },
		});
		if (result.matchedCount === 0) throw new AgentRelayLeaseLostError();
	}

	private async materializeEventBacklog(generation: Generation, latestSeq: number): Promise<void> {
		for (let attempt = 0; attempt < 8; attempt++) {
			const conversation = await collections.conversations.findOne(
				{ _id: generation.conversationId },
				{ projection: { messages: 1 } }
			);
			const current = conversation?.messages.find(
				(message) =>
					message.id === generation.messageId && message.generationId === generation.generationId
			);
			if (!current) throw new Error("Agent generation message no longer exists");
			const cursor = current.materializedSeq ?? 0;
			if (cursor >= latestSeq) return;
			const backlog = await collections.generationEvents
				.find({ generationId: this.generationId, seq: { $gt: cursor, $lte: latestSeq } })
				.sort({ seq: 1 })
				.toArray();
			if (!backlog.length) return;
			let contiguousEnd = cursor;
			const updates: MessageUpdate[] = [];
			for (const record of backlog) {
				if (record.seq !== contiguousEnd + 1) break;
				updates.push(record.event);
				contiguousEnd = record.seq;
			}
			if (!updates.length) return;
			const reduced = reduceAgentMessageUpdates(current, updates);
			const result = await collections.conversations.updateOne(
				{
					_id: generation.conversationId,
					messages: {
						$elemMatch: {
							id: generation.messageId,
							generationId: generation.generationId,
							...(cursor === 0
								? {
										$or: [{ materializedSeq: { $exists: false } }, { materializedSeq: 0 }],
									}
								: { materializedSeq: cursor }),
						},
					},
				},
				{
					$set: {
						"messages.$": { ...reduced.message, materializedSeq: contiguousEnd },
						...(reduced.title ? { title: reduced.title } : {}),
						updatedAt: this.now(),
					},
				}
			);
			if (result.matchedCount === 1 && contiguousEnd >= latestSeq) return;
		}
		throw new Error("Agent event backlog changed too frequently to materialize");
	}

	async commitBridgeBatch(
		sourceSequence: number,
		events: readonly MessageUpdate[]
	): Promise<AgentBatchCommit> {
		if (!Number.isSafeInteger(sourceSequence) || sourceSequence < 1) {
			throw new Error("Agent bridge sequence must be a positive safe integer");
		}
		await this.resumePending();

		for (;;) {
			const generation = await collections.generations.findOne(this.leaseFilter());
			if (!generation) throw new AgentRelayLeaseLostError();
			if ((generation.agent?.sourceSequence ?? 0) >= sourceSequence) {
				return {
					inserted: false,
					endSeq: generation.seq,
					sourceSequence,
				};
			}
			if (generation.agent?.pending) {
				await this.flushPending(generation, generation.agent.pending);
				continue;
			}

			const pending: AgentPendingBatch = {
				kind: "bridge",
				sourceSequence,
				startSeq: generation.seq + 1,
				events: [...events],
			};
			const now = this.now();
			const reservation = await collections.generations.updateOne(
				{
					...this.leaseFilter(),
					seq: generation.seq,
					"agent.pending": { $exists: false },
					"agent.sourceSequence": { $lt: sourceSequence },
				},
				{
					$set: {
						"agent.pending": pending,
						seq: generation.seq + events.length,
						"agent.relay.expiresAt": new Date(now.getTime() + this.leaseMs),
						lastHeartbeatAt: now,
						updatedAt: now,
					},
				}
			);
			if (reservation.matchedCount === 0) continue;
			return this.flushPending(
				{
					...generation,
					seq: generation.seq + events.length,
					agent: {
						sourceSequence: generation.agent?.sourceSequence ?? 0,
						...generation.agent,
						pending,
					},
				},
				pending
			);
		}
	}

	async commitSystemBatch(
		key: string,
		events: readonly MessageUpdate[]
	): Promise<AgentBatchCommit> {
		if (!key || key.length > 128) throw new Error("Invalid Agent system batch key");
		await this.resumePending();

		for (;;) {
			const generation = await collections.generations.findOne(this.leaseFilter());
			if (!generation) throw new AgentRelayLeaseLostError();
			if (generation.agent?.systemKeys?.includes(key)) {
				return { inserted: false, endSeq: generation.seq };
			}
			if (generation.agent?.pending) {
				await this.flushPending(generation, generation.agent.pending);
				continue;
			}

			const pending: AgentPendingBatch = {
				kind: "system",
				key,
				startSeq: generation.seq + 1,
				events: [...events],
			};
			const now = this.now();
			const reservation = await collections.generations.updateOne(
				{
					...this.leaseFilter(),
					seq: generation.seq,
					"agent.pending": { $exists: false },
					"agent.systemKeys": { $ne: key },
				},
				{
					$set: {
						"agent.pending": pending,
						seq: generation.seq + events.length,
						"agent.relay.expiresAt": new Date(now.getTime() + this.leaseMs),
						lastHeartbeatAt: now,
						updatedAt: now,
					},
				}
			);
			if (reservation.matchedCount === 0) continue;
			return this.flushPending(
				{
					...generation,
					seq: generation.seq + events.length,
					agent: {
						sourceSequence: generation.agent?.sourceSequence ?? 0,
						...generation.agent,
						pending,
					},
				},
				pending
			);
		}
	}

	async finish(status: GenerationStatus, error?: string): Promise<boolean> {
		await this.resumePending();
		const now = this.now();
		for (const candidate of generationTerminalCandidates(status)) {
			const stoppedFallback = candidate.status === "interrupted" && status !== "interrupted";
			const generation = await collections.generations.findOneAndUpdate(
				{
					...this.leaseFilter(),
					...candidate.filter,
					// Enqueue and terminalization contend on this same document. A control
					// which wins must be drained before finish may be retried.
					"agent.control": { $exists: false },
				},
				{
					$set: {
						status: candidate.status,
						"agent.controlsOpen": false,
						endedAt: now,
						lastHeartbeatAt: now,
						updatedAt: now,
						...(!stoppedFallback && error ? { error } : {}),
					},
					$unset: { "agent.relay": "", "agent.pendingApprovalIds": "" },
				},
				{ returnDocument: "before", includeResultMetadata: false }
			);
			if (!generation) continue;
			if (candidate.status === "interrupted") {
				await markGenerationMessageInterrupted(generation, now);
			}
			await this.releaseWorkspaceExecution();
			return true;
		}
		return false;
	}

	/**
	 * Fail an exact relay after recovery has become impossible.
	 *
	 * Normal completion must drain and seal the control mailbox through
	 * `finish()`. Fatal recovery paths are different: the workspace or bridge
	 * needed to drain that mailbox may no longer exist, and `resumePending()` can
	 * no longer renew the reciprocal workspace lease. Keeping the generation in
	 * `running` in that case only strands it for every later reaper sweep.
	 *
	 * This transition is deliberately fenced by the still-live generation relay
	 * owner/fence/incarnation. It closes and discards an undeliverable control in
	 * the same Mongo write, then releases only this relay's exact workspace
	 * execution claim. A replacement relay which has already taken ownership wins
	 * the fence and this method becomes a no-op.
	 */
	async finishUnrecoverable(status: "interrupted" | "error", error: string): Promise<boolean> {
		for (const candidate of generationTerminalCandidates(status)) {
			for (let attempt = 0; attempt < 8; attempt++) {
				const observedAt = this.now();
				const filter = { ...this.leaseFilter(true, observedAt), ...candidate.filter };
				const observed = await collections.generations.findOne(filter);
				if (!observed) break;
				const pending = observed.agent?.pending;
				if (pending) {
					// Keep the WAL reservation durable until its idempotent conversation
					// materialization succeeds. A crash after the apply therefore leaves a
					// retryable pending record instead of exposing event rows which the
					// durable conversation never incorporated.
					await this.persistPendingEvents(pending);
					const renewalAt = this.now();
					const renewed = await collections.generations.updateOne(
						{
							...this.leaseFilter(true, renewalAt),
							...candidate.filter,
							"agent.pending": pending,
						},
						{
							$set: {
								"agent.relay.expiresAt": new Date(renewalAt.getTime() + this.leaseMs),
								lastHeartbeatAt: renewalAt,
								updatedAt: renewalAt,
							},
						}
					);
					if (renewed.matchedCount === 0) continue;
					const endSeq = pending.startSeq + pending.events.length - 1;
					await this.applyPendingToMessage(observed, pending, endSeq);
				}

				const terminalAt = this.now();
				const stoppedFallback = candidate.status === "interrupted" && status !== "interrupted";
				const generation = await collections.generations.findOneAndUpdate(
					{
						...this.leaseFilter(true, terminalAt),
						...candidate.filter,
						...(pending ? { "agent.pending": pending } : { "agent.pending": { $exists: false } }),
					},
					{
						...(pending?.kind === "system"
							? { $addToSet: { "agent.systemKeys": pending.key } }
							: {}),
						$set: {
							status: candidate.status,
							"agent.controlsOpen": false,
							...(pending?.kind === "bridge"
								? { "agent.sourceSequence": pending.sourceSequence }
								: {}),
							...(!stoppedFallback ? { error } : {}),
							endedAt: terminalAt,
							lastHeartbeatAt: terminalAt,
							updatedAt: terminalAt,
						},
						$unset: {
							"agent.relay": "",
							"agent.control": "",
							"agent.pending": "",
							"agent.pendingApprovalIds": "",
						},
					},
					{ returnDocument: "before", includeResultMetadata: false }
				);
				if (!generation) continue;
				if (candidate.status === "interrupted") {
					await markGenerationMessageInterrupted(generation, terminalAt);
				}
				await this.releaseWorkspaceExecution();
				return true;
			}
		}
		return false;
	}

	async release(): Promise<void> {
		await collections.generations.updateOne(this.leaseFilter(false), {
			$unset: { "agent.relay": "" },
			$set: { updatedAt: this.now() },
		});
		await this.releaseWorkspaceExecution();
	}

	private async releaseWorkspaceExecution(): Promise<void> {
		await collections.workspaces.updateOne(
			{
				_id: this.workspaceId,
				"execution.generationId": this.generationId,
				"execution.phase": "relay",
				"execution.owner": this.owner,
				"execution.fence": this.fence,
			},
			{ $unset: { execution: "" } }
		);
	}

	private async flushPending(
		generation: Generation,
		pending: AgentPendingBatch
	): Promise<AgentBatchCommit> {
		const endSeq = pending.startSeq + pending.events.length - 1;
		await this.persistPendingEvents(pending);

		// Do not materialize the captured batch until both reciprocal workspace
		// ownership and the exact generation relay fence have accepted it. A stale
		// relay may remain inside this method after its lease expires; durable event
		// upserts are idempotent, but mutating the conversation before the fence
		// checks would let that zombie publish output after replacement/terminality.
		const now = this.now();
		const expiresAt = new Date(now.getTime() + this.leaseMs);
		const approvalRequestIds = pendingApprovalRequestIds(pending);
		const workspace = await collections.workspaces.updateOne(
			{
				_id: this.workspaceId,
				conversationId: this.conversationId,
				retirement: { $exists: false },
				...workspaceOperationAvailableFilter(now),
				"sandbox.id": this.sandboxId,
				"sandbox.generation": this.sandboxGeneration,
				...(this.sandboxNonce ? { "sandbox.nonce": this.sandboxNonce } : {}),
				"execution.generationId": this.generationId,
				"execution.phase": "relay",
				"execution.owner": this.owner,
				"execution.fence": this.fence,
				"execution.sandboxId": this.sandboxId,
				"execution.sandboxGeneration": this.sandboxGeneration,
				...(this.sandboxNonce ? { "execution.sandboxNonce": this.sandboxNonce } : {}),
				"execution.expiresAt": { $gt: now },
			},
			{
				...(pending.kind === "bridge"
					? { $max: { "piSession.lastSequence": pending.sourceSequence } }
					: {}),
				$set: {
					"execution.expiresAt": expiresAt,
					updatedAt: now,
				},
				$unset: { operation: "" },
			}
		);
		if (workspace.matchedCount === 0) throw new AgentRelayLeaseLostError();
		const renewed = await collections.generations.updateOne(
			{
				...this.leaseFilter(true, now),
				"agent.pending": pending,
			},
			{
				$set: {
					"agent.relay.expiresAt": expiresAt,
					lastHeartbeatAt: now,
					updatedAt: now,
				},
			}
		);
		if (renewed.matchedCount === 0) throw new AgentRelayLeaseLostError();

		await this.applyPendingToMessage(generation, pending, endSeq);

		const completionAt = this.now();
		const completion: UpdateFilter<Generation> =
			pending.kind === "bridge"
				? {
						...(approvalRequestIds.length
							? {
									$addToSet: {
										"agent.pendingApprovalIds": { $each: approvalRequestIds },
									},
								}
							: {}),
						$set: {
							"agent.sourceSequence": pending.sourceSequence,
							"agent.relay.expiresAt": new Date(completionAt.getTime() + this.leaseMs),
							lastHeartbeatAt: completionAt,
							updatedAt: completionAt,
						},
						$unset: { "agent.pending": "" as const },
					}
				: {
						$addToSet: { "agent.systemKeys": pending.key },
						$set: {
							"agent.relay.expiresAt": new Date(completionAt.getTime() + this.leaseMs),
							lastHeartbeatAt: completionAt,
							updatedAt: completionAt,
						},
						$unset: { "agent.pending": "" as const },
					};
		const result = await collections.generations.updateOne(
			{
				...this.leaseFilter(true, completionAt),
				"agent.pending": pending,
			},
			completion
		);
		if (result.matchedCount === 0) throw new AgentRelayLeaseLostError();
		return {
			inserted: true,
			endSeq: Math.max(generation.seq, endSeq),
			...(pending.kind === "bridge" ? { sourceSequence: pending.sourceSequence } : {}),
		};
	}

	private async persistPendingEvents(pending: AgentPendingBatch): Promise<void> {
		await persistAgentPendingEvents(this.generationId, pending, this.now);
	}

	private async applyPendingToMessage(
		generation: Generation,
		pending: AgentPendingBatch,
		endSeq: number
	): Promise<void> {
		await applyAgentPendingToMessage(generation, pending, endSeq, this.now);
	}

	private leaseFilter(requireRunning = true, at = this.now()): Filter<Generation> {
		return {
			generationId: this.generationId,
			workspaceId: this.workspaceId,
			...(requireRunning ? { status: "running" as const } : {}),
			"agent.sandboxId": this.sandboxId,
			"agent.sandboxGeneration": this.sandboxGeneration,
			"agent.relay.owner": this.owner,
			"agent.relay.fence": this.fence,
			...(requireRunning ? { "agent.relay.expiresAt": { $gt: at } } : {}),
		};
	}
}

async function renewAgentTerminalization(
	generationId: string,
	workspaceId: Workspace["_id"],
	id: string,
	now = new Date()
): Promise<Generation | null> {
	return collections.generations.findOneAndUpdate(
		{
			generationId,
			workspaceId,
			"agent.terminalization.id": id,
		},
		{
			$set: {
				"agent.terminalization.expiresAt": new Date(now.getTime() + AGENT_TERMINALIZATION_LEASE_MS),
				lastHeartbeatAt: now,
				updatedAt: now,
			},
		},
		{ returnDocument: "after", includeResultMetadata: false }
	);
}

async function persistAgentPendingEvents(
	generationId: string,
	pending: AgentPendingBatch,
	now: () => Date
): Promise<void> {
	if (!pending.events.length) return;
	await collections.generationEvents.bulkWrite(
		pending.events.map((event, index) => ({
			updateOne: {
				filter: { generationId, seq: pending.startSeq + index },
				update: {
					$setOnInsert: {
						_id: new ObjectId(),
						generationId,
						seq: pending.startSeq + index,
						event,
						...(pending.kind === "bridge"
							? { sourceSequence: pending.sourceSequence, sourceIndex: index }
							: { systemKey: pending.key, systemIndex: index }),
						createdAt: now(),
					},
				},
				upsert: true,
			},
		})),
		{ ordered: true }
	);
}

async function applyAgentPendingToMessage(
	generation: Generation,
	pending: AgentPendingBatch,
	endSeq: number,
	now: () => Date
): Promise<void> {
	for (let attempt = 0; attempt < 8; attempt++) {
		const conversation = await collections.conversations.findOne(
			{ _id: generation.conversationId },
			{ projection: { messages: 1, title: 1 } }
		);
		const current = conversation?.messages.find(
			(message) =>
				message.id === generation.messageId && message.generationId === generation.generationId
		);
		if (!conversation || !current) throw new Error("Agent generation message no longer exists");

		if (pending.kind === "bridge" && (current.agentSequence ?? 0) >= pending.sourceSequence) {
			return;
		}
		if (pending.kind === "system" && current.agentSystemKeys?.includes(pending.key)) {
			return;
		}

		const reduced = reduceAgentMessageUpdates(current, pending.events);
		const next: Message = {
			...reduced.message,
			materializedSeq: Math.max(current.materializedSeq ?? 0, endSeq),
			...(pending.kind === "bridge"
				? { agentSequence: pending.sourceSequence }
				: { agentSystemKeys: [...(current.agentSystemKeys ?? []), pending.key] }),
		};
		const markerFilter: Record<string, unknown> =
			pending.kind === "bridge"
				? {
						$or: [
							{ agentSequence: { $exists: false } },
							{ agentSequence: { $lt: pending.sourceSequence } },
						],
					}
				: { agentSystemKeys: { $ne: pending.key } };
		const result = await collections.conversations.updateOne(
			{
				_id: generation.conversationId,
				messages: {
					$elemMatch: {
						id: generation.messageId,
						generationId: generation.generationId,
						...markerFilter,
					},
				},
			},
			{
				$set: {
					"messages.$": next,
					...(reduced.title ? { title: reduced.title } : {}),
					updatedAt: now(),
				},
			}
		);
		if (result.matchedCount === 1) return;
	}
	throw new Error("Agent message changed too frequently to materialize relay output");
}

function pendingApprovalRequestIds(pending: AgentPendingBatch): string[] {
	return [
		...new Set(
			pending.events.flatMap((event) =>
				event.type === MessageUpdateType.AgentApproval ? [event.requestId] : []
			)
		),
	];
}

async function clearAcquiredRelay(
	generationId: string,
	owner: string,
	fence: number
): Promise<void> {
	await collections.generations.updateOne(
		{
			generationId,
			status: "running",
			"agent.relay.owner": owner,
			"agent.relay.fence": fence,
		},
		{ $unset: { "agent.relay": "" } }
	);
}