import { afterEach, beforeAll, describe, expect, it, vi } from "vitest";
import { ObjectId } from "mongodb";
import { randomUUID } from "node:crypto";
import { collections, ready } from "$lib/server/database";
import {
	MessageUpdateStatus,
	MessageUpdateType,
	type MessageUpdate,
} from "$lib/types/MessageUpdate";
import type { AgentPendingBatch, Generation } from "$lib/types/Generation";
import type { Workspace } from "$lib/types/Workspace";
import {
	claimWorkspaceOperation,
	releaseWorkspaceOperation,
	WORKSPACE_OPERATION_LEASE_MS,
} from "./WorkspaceOperation";
import { recordGenerationStop } from "$lib/server/generation/stopGeneration";
import {
	AgentGenerationIncarnationMismatchError,
	AgentGenerationStore,
	AgentRelayLeaseLostError,
} from "./AgentGenerationStore";
import { reapStaleGenerations } from "$lib/server/generation/reaper";

beforeAll(async () => {
	await ready;
});

afterEach(async () => {
	vi.restoreAllMocks();
	await Promise.all([
		collections.conversations.deleteMany({}),
		collections.workspaces.deleteMany({}),
		collections.generations.deleteMany({}),
		collections.generationEvents.deleteMany({}),
		collections.abortedGenerations.deleteMany({}),
	]);
});

describe.sequential("AgentGenerationStore", () => {
	it("commits a replayed bridge sequence exactly once", async () => {
		const seeded = await seedAgentRun();
		const store = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		);
		const update: MessageUpdate = { type: MessageUpdateType.Stream, token: "once" };
		const bound = await collections.generations.findOne({
			generationId: seeded.generationId,
		});
		expect(bound?.agent).toMatchObject(sandboxIncarnation(seeded.workspace));

		const first = await store.commitBridgeBatch(41, [update]);
		const replay = await store.commitBridgeBatch(41, [update]);

		expect(first.inserted).toBe(true);
		expect(replay.inserted).toBe(false);
		expect(
			await collections.generationEvents.countDocuments({ generationId: seeded.generationId })
		).toBe(1);
		const conversation = await collections.conversations.findOne({ _id: seeded.conversationId });
		const message = conversation?.messages.find((item) => item.id === seeded.messageId);
		expect(message?.content).toBe("once");
		expect(message?.agentSequence).toBe(41);
		expect(message?.materializedSeq).toBe(1);
		expect(
			(await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.piSession?.lastSequence
		).toBe(41);
	});

	it("binds once and refuses to acquire the same prompt for a replacement sandbox", async () => {
		const seeded = await seedAgentRun();
		const original = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		);
		await original.release();

		await expect(
			AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				sandboxId: "sandbox-replacement",
				sandboxGeneration: "generation-replacement",
			})
		).rejects.toBeInstanceOf(AgentGenerationIncarnationMismatchError);

		const generation = await collections.generations.findOne({
			generationId: seeded.generationId,
		});
		expect(generation?.agent).toMatchObject(sandboxIncarnation(seeded.workspace));
		expect(generation?.agent?.relay).toBeUndefined();
		expect(generation?.status).toBe("running");
	});

	it("continues after a durable provisioning event even when its heartbeat cursor lagged", async () => {
		const seeded = await seedAgentRun();
		await collections.generationEvents.insertOne({
			_id: new ObjectId(),
			generationId: seeded.generationId,
			seq: 1,
			event: { type: MessageUpdateType.AgentSandbox, state: "ready" },
			createdAt: new Date(),
		});
		const store = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		);
		const committed = await store.commitBridgeBatch(1, [
			{ type: MessageUpdateType.Stream, token: "after provisioning" },
		]);

		expect(committed.endSeq).toBe(2);
		expect(
			(
				await collections.generationEvents
					.find({ generationId: seeded.generationId })
					.sort({ seq: 1 })
					.toArray()
			).map((event) => event.event.type)
		).toEqual([MessageUpdateType.AgentSandbox, MessageUpdateType.Stream]);
		const conversation = await collections.conversations.findOne({ _id: seeded.conversationId });
		const message = conversation?.messages.find((item) => item.id === seeded.messageId);
		expect(message?.materializedSeq).toBe(2);
		expect(message?.updates?.map((event) => event.type)).toEqual([
			MessageUpdateType.AgentSandbox,
			MessageUpdateType.Stream,
		]);
	});

	it("tracks native approval requests as live generation state and clears them at terminal", async () => {
		const seeded = await seedAgentRun();
		const store = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		);
		await store.commitBridgeBatch(1, [
			{
				type: MessageUpdateType.AgentApproval,
				requestId: "approval-live",
				method: "confirm",
				title: "Allow command?",
			},
		]);

		expect(
			(await collections.generations.findOne({ generationId: seeded.generationId }))?.agent
				?.pendingApprovalIds
		).toEqual(["approval-live"]);
		expect(await store.finish("completed")).toBe(true);
		expect(
			(await collections.generations.findOne({ generationId: seeded.generationId }))?.agent
				?.pendingApprovalIds
		).toBeUndefined();
	});

	it("keeps one exact relay authoritative through every login ownership handoff phase", async () => {
		const seeded = await seedAgentRun();
		const userId = new ObjectId();
		const store = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		);

		await collections.workspaces.updateOne({ _id: seeded.workspace._id }, { $set: { userId } });
		expect(await store.authoritativeOwner()).toEqual({ sessionId: seeded.workspace.sessionId });

		await collections.generations.updateOne(
			{ generationId: seeded.generationId },
			{ $set: { userId } }
		);
		expect(await store.authoritativeOwner()).toEqual({ sessionId: seeded.workspace.sessionId });

		await collections.conversations.updateOne(
			{ _id: seeded.conversationId },
			{ $set: { userId }, $unset: { sessionId: "" } }
		);
		expect(await store.authoritativeOwner()).toEqual({ userId });

		await collections.workspaces.updateOne(
			{ _id: seeded.workspace._id },
			{ $unset: { sessionId: "" } }
		);
		await collections.generations.updateOne(
			{ generationId: seeded.generationId },
			{ $unset: { sessionId: "" } }
		);
		await store.assertPromptDispatch(seeded.workspace);
		expect(await store.authoritativeOwner()).toEqual({ userId });
		await store.release();
	});

	it("cannot complete after an exact stop intent wins the relay terminal barrier", async () => {
		const seeded = await seedAgentRun();
		const store = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		);
		await store.commitBridgeBatch(1, [
			{ type: MessageUpdateType.Stream, token: "stopped-late-output" },
		]);

		expect(
			await recordGenerationStop({
				conversationId: seeded.conversationId,
				generationId: seeded.generationId,
				seenContentLength: 4,
			})
		).toBe(true);
		expect(await store.finish("completed")).toBe(true);

		expect(
			await collections.generations.findOne({ generationId: seeded.generationId })
		).toMatchObject({
			status: "interrupted",
			stopRequested: { seenContentLength: 4 },
		});
		const message = (await collections.conversations.findOne({ _id: seeded.conversationId }))
			?.messages[0];
		expect(message?.content).toBe("stop");
		expect(message?.interrupted).toBe(true);
	});

	it("finishes a crash-partial write-ahead batch without gaps or duplication", async () => {
		const seeded = await seedAgentRun();
		const store = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		);
		const events: MessageUpdate[] = [
			{ type: MessageUpdateType.Stream, token: "alpha " },
			{ type: MessageUpdateType.Stream, token: "beta" },
		];
		await collections.generations.updateOne(
			{ generationId: seeded.generationId },
			{
				$set: {
					seq: 2,
					"agent.pending": {
						kind: "bridge",
						sourceSequence: 7,
						startSeq: 1,
						events,
					},
				},
			}
		);
		await collections.generationEvents.insertOne({
			_id: new ObjectId(),
			generationId: seeded.generationId,
			seq: 1,
			event: events[0],
			sourceSequence: 7,
			sourceIndex: 0,
			createdAt: new Date(),
		});

		await store.resumePending();

		const persisted = await collections.generationEvents
			.find({ generationId: seeded.generationId })
			.sort({ seq: 1 })
			.toArray();
		expect(persisted.map((item) => item.seq)).toEqual([1, 2]);
		expect(persisted.map((item) => item.sourceIndex)).toEqual([0, 1]);
		const conversation = await collections.conversations.findOne({ _id: seeded.conversationId });
		expect(conversation?.messages.find((item) => item.id === seeded.messageId)?.content).toBe(
			"alpha beta"
		);
		const generation = await collections.generations.findOne({
			generationId: seeded.generationId,
		});
		expect(generation?.agent?.pending).toBeUndefined();
		expect(generation?.agent?.sourceSequence).toBe(7);
	});

	it("keeps the WAL pending when a relay crashes after message materialization", async () => {
		const seeded = await seedAgentRun();
		let timestamp = 0;
		const now = () => new Date(timestamp);
		const first = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
				owner: "crashing-relay",
				leaseMs: 10,
				now,
			})
		);
		timestamp = 5;
		const restoreMaterialization = failAfterNextPendingMaterialization(first);
		try {
			await expect(
				first.commitBridgeBatch(1, [
					{ type: MessageUpdateType.Stream, token: "materialized exactly once" },
				])
			).rejects.toThrow("simulated process crash after Agent WAL materialization");
		} finally {
			restoreMaterialization();
		}

		let generation = await collections.generations.findOne({
			generationId: seeded.generationId,
		});
		expect(generation?.agent?.pending).toMatchObject({
			kind: "bridge",
			sourceSequence: 1,
		});
		expect(generation?.agent?.sourceSequence).toBe(0);
		expect(generation?.agent?.relay?.expiresAt).toEqual(new Date(15));
		let workspace = await collections.workspaces.findOne({ _id: seeded.workspace._id });
		expect(workspace?.execution?.expiresAt).toEqual(new Date(15));
		let conversation = await collections.conversations.findOne({ _id: seeded.conversationId });
		expect(conversation?.messages[0]?.content).toBe("materialized exactly once");

		await first.release();
		const replacement = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
				owner: "replacement-relay",
				leaseMs: 10,
				now,
			})
		);
		await replacement.resumePending();

		generation = await collections.generations.findOne({ generationId: seeded.generationId });
		expect(generation?.agent?.pending).toBeUndefined();
		expect(generation?.agent?.sourceSequence).toBe(1);
		workspace = await collections.workspaces.findOne({ _id: seeded.workspace._id });
		expect(workspace?.piSession?.lastSequence).toBe(1);
		conversation = await collections.conversations.findOne({ _id: seeded.conversationId });
		expect(conversation?.messages[0]?.content).toBe("materialized exactly once");
		expect(
			await collections.generationEvents.countDocuments({ generationId: seeded.generationId })
		).toBe(1);
	});

	it("recovers a sandbox-change crash after pending message materialization", async () => {
		const seeded = await seedAgentRun();
		const store = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		);
		const pending: AgentPendingBatch = {
			kind: "bridge",
			sourceSequence: 1,
			startSeq: 1,
			events: [{ type: MessageUpdateType.Stream, token: "terminal batch exactly once" }],
		};
		await collections.generations.updateOne(
			{ generationId: seeded.generationId },
			{ $set: { seq: 1, "agent.pending": pending, "agent.controlsOpen": true } }
		);

		const originalUpdateOne = collections.conversations.updateOne.bind(collections.conversations);
		let crashed = false;
		const updateSpy = vi
			.spyOn(collections.conversations, "updateOne")
			.mockImplementation(async (filter, update, options) => {
				const result = await originalUpdateOne(filter, update, options);
				if (!crashed && "$set" in update && update.$set && "messages.$" in update.$set) {
					crashed = true;
					throw new Error("simulated sandbox terminalizer crash after materialization");
				}
				return result;
			});

		await expect(
			AgentGenerationStore.interruptForSandboxChange({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				sandboxId: "replacement-sandbox",
				sandboxGeneration: "replacement-generation",
			})
		).rejects.toThrow("simulated sandbox terminalizer crash after materialization");

		let generation = await collections.generations.findOne({
			generationId: seeded.generationId,
		});
		expect(generation).toMatchObject({
			status: "running",
			agent: {
				controlsOpen: false,
				pending,
				terminalization: { kind: "sandbox-change" },
			},
		});
		expect(generation?.agent?.relay).toBeUndefined();
		expect(
			(await collections.conversations.findOne({ _id: seeded.conversationId }))?.messages[0]
				?.content
		).toBe("terminal batch exactly once");
		expect(
			await collections.generationEvents.countDocuments({ generationId: seeded.generationId })
		).toBe(1);
		await expect(store.commitBridgeBatch(2, [])).rejects.toBeInstanceOf(AgentRelayLeaseLostError);
		expect(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		).toBeNull();

		updateSpy.mockRestore();
		await collections.generations.updateOne(
			{ generationId: seeded.generationId },
			{ $set: { "agent.terminalization.expiresAt": new Date(0) } }
		);
		expect(
			await AgentGenerationStore.interruptForSandboxChange({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				sandboxId: "replacement-sandbox",
				sandboxGeneration: "replacement-generation",
			})
		).toBe(true);

		generation = await collections.generations.findOne({ generationId: seeded.generationId });
		expect(generation).toMatchObject({
			status: "interrupted",
			agent: { sourceSequence: 1, controlsOpen: false },
		});
		expect(generation?.agent?.pending).toBeUndefined();
		expect(generation?.agent?.terminalization).toBeUndefined();
		expect(
			(await collections.conversations.findOne({ _id: seeded.conversationId }))?.messages[0]
		).toMatchObject({ content: "terminal batch exactly once", interrupted: true });
		expect(
			await collections.generationEvents.countDocuments({ generationId: seeded.generationId })
		).toBe(1);
	});

	it("reaper resumes an expired post-terminal marker before clearing side effects", async () => {
		const seeded = await seedAgentRun();
		await AgentGenerationStore.acquire({
			generationId: seeded.generationId,
			workspaceId: seeded.workspace._id,
			...sandboxIncarnation(seeded.workspace),
		});
		const originalFindOne = collections.conversations.findOne.bind(collections.conversations);
		let crashed = false;
		const findSpy = vi
			.spyOn(collections.conversations, "findOne")
			.mockImplementation(async (filter, options) => {
				const generation = await collections.generations.findOne({
					generationId: seeded.generationId,
				});
				if (!crashed && generation?.status === "interrupted" && generation.agent?.terminalization) {
					crashed = true;
					throw new Error("simulated crash after terminal CAS");
				}
				return originalFindOne(filter, options);
			});

		await expect(
			AgentGenerationStore.interruptForSandboxChange({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				sandboxId: "replacement-sandbox",
				sandboxGeneration: "replacement-generation",
			})
		).rejects.toThrow("simulated crash after terminal CAS");
		let generation = await collections.generations.findOne({
			generationId: seeded.generationId,
		});
		expect(generation).toMatchObject({
			status: "interrupted",
			agent: { terminalization: { kind: "sandbox-change" } },
		});
		expect(
			(await collections.conversations.findOne({ _id: seeded.conversationId }))?.messages[0]
				?.interrupted
		).not.toBe(true);
		expect(
			(await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.execution
		).toBeDefined();

		findSpy.mockRestore();
		await collections.generations.updateOne(
			{ generationId: seeded.generationId },
			{ $set: { "agent.terminalization.expiresAt": new Date(0) } }
		);
		await reapStaleGenerations();

		generation = await collections.generations.findOne({ generationId: seeded.generationId });
		expect(generation?.agent?.terminalization).toBeUndefined();
		expect(
			(await collections.conversations.findOne({ _id: seeded.conversationId }))?.messages[0]
				?.interrupted
		).toBe(true);
		expect(
			(await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.execution
		).toBeUndefined();
	});

	it("does not materialize a captured WAL batch after its relay fence moves", async () => {
		const seeded = await seedAgentRun();
		let timestamp = 0;
		const now = () => new Date(timestamp);
		const first = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
				owner: "stale-relay",
				leaseMs: 10,
				now,
			})
		);
		const pending: AgentPendingBatch = {
			kind: "bridge",
			sourceSequence: 1,
			startSeq: 1,
			events: [{ type: MessageUpdateType.Stream, token: "authoritative output" }],
		};
		await collections.generations.updateOne(
			{ generationId: seeded.generationId },
			{ $set: { seq: 1, "agent.pending": pending } }
		);
		const captured = await collections.generations.findOne({
			generationId: seeded.generationId,
		});
		if (!captured) throw new Error("Missing captured generation WAL");

		timestamp = 11;
		const replacement = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
				owner: "replacement-relay",
				leaseMs: 10,
				now,
			})
		);
		const capturedFlush = first as unknown as {
			flushPending(generation: Generation, batch: AgentPendingBatch): Promise<unknown>;
		};
		await expect(capturedFlush.flushPending(captured, pending)).rejects.toBeInstanceOf(
			AgentRelayLeaseLostError
		);
		let conversation = await collections.conversations.findOne({ _id: seeded.conversationId });
		expect(conversation?.messages[0]?.content).toBe("");

		await replacement.resumePending();
		conversation = await collections.conversations.findOne({ _id: seeded.conversationId });
		expect(conversation?.messages[0]?.content).toBe("authoritative output");
		expect(
			await collections.generationEvents.countDocuments({ generationId: seeded.generationId })
		).toBe(1);
	});

	it("fences an expired producer and lets only the replacement finish", async () => {
		const seeded = await seedAgentRun();
		let timestamp = 0;
		const now = () => new Date(timestamp);
		const first = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
				owner: "first",
				leaseMs: 10,
				now,
			})
		);
		timestamp = 11;
		const replacement = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
				owner: "replacement",
				leaseMs: 10,
				now,
			})
		);

		await expect(first.renew()).rejects.toBeInstanceOf(AgentRelayLeaseLostError);
		await replacement.commitBridgeBatch(1, [
			{ type: MessageUpdateType.Stream, token: "replacement" },
		]);
		expect(await replacement.finish("completed")).toBe(true);
		expect(
			(await collections.generations.findOne({ generationId: seeded.generationId }))?.status
		).toBe("completed");
	});

	it("never lets advisory pi stats lower authoritative proxy or sandbox spend", async () => {
		const seeded = await seedAgentRun();
		await collections.workspaces.updateOne(
			{ _id: seeded.workspace._id },
			{
				$set: {
					spend: {
						accruedUsd: 12,
						inferenceUsd: 7,
						sandboxUsd: 5,
						lastCalculatedAt: new Date(),
					},
				},
			}
		);
		const store = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		);
		await store.commitBridgeBatch(1, [
			{
				type: MessageUpdateType.AgentStats,
				scope: "session",
				tokens: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, total: 2 },
				costUsd: 0.01,
			},
		]);

		const spend = (await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.spend;
		expect(spend).toMatchObject({ accruedUsd: 12, inferenceUsd: 7, sandboxUsd: 5 });
	});

	it("does not advance the pi cursor on a replacement sandbox incarnation", async () => {
		const seeded = await seedAgentRun();
		const store = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		);
		await collections.workspaces.updateOne(
			{ _id: seeded.workspace._id },
			{
				$set: {
					"sandbox.id": "sandbox-replacement",
					"sandbox.generation": "generation-replacement",
					"piSession.lastSequence": 0,
				},
			}
		);

		await expect(
			store.commitBridgeBatch(41, [{ type: MessageUpdateType.Stream, token: "old incarnation" }])
		).rejects.toBeInstanceOf(AgentRelayLeaseLostError);

		const workspace = await collections.workspaces.findOne({ _id: seeded.workspace._id });
		expect(workspace?.sandbox.id).toBe("sandbox-replacement");
		expect(workspace?.piSession?.lastSequence).toBe(0);
	});

	it("treats an exact workspace operation as contention instead of a sandbox mismatch", async () => {
		const seeded = await seedAgentRun();
		const claimed = await claimWorkspaceOperation({
			conversationId: seeded.conversationId,
			workspace: seeded.workspace,
			owner: { sessionId: seeded.workspace.sessionId },
			kind: "configuration",
			action: "compact",
		});

		expect(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		).toBeNull();
		const generation = await collections.generations.findOne({
			generationId: seeded.generationId,
		});
		expect(generation?.status).toBe("running");
		expect(generation?.agent?.relay).toBeUndefined();
		await releaseWorkspaceOperation(
			seeded.workspace._id,
			{ sessionId: seeded.workspace.sessionId },
			claimed.lease.id
		);
	});

	it("atomically reclaims orphaned operation leases on relay acquire and renew", async () => {
		const seeded = await seedAgentRun();
		const orphanedOperation = (id: string) => ({
			id,
			kind: "configuration" as const,
			action: "set-model",
			claimedAt: new Date(Date.now() - WORKSPACE_OPERATION_LEASE_MS - 1),
			workspaceUpdatedAt: seeded.workspace.updatedAt,
		});
		await collections.workspaces.updateOne(
			{ _id: seeded.workspace._id },
			{ $set: { operation: orphanedOperation("orphaned-acquire") } }
		);

		const store = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		);
		expect(
			(await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.operation
		).toBeUndefined();

		await collections.workspaces.updateOne(
			{ _id: seeded.workspace._id },
			{ $set: { operation: orphanedOperation("orphaned-renew") } }
		);
		await expect(store.renew()).resolves.toBeUndefined();
		expect(
			(await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.operation
		).toBeUndefined();
		await store.finish("completed");
	});

	it("keeps a live operation lease blocking relay renewal", async () => {
		const seeded = await seedAgentRun();
		const store = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		);
		await collections.workspaces.updateOne(
			{ _id: seeded.workspace._id },
			{
				$set: {
					operation: {
						id: "live-renew-operation",
						kind: "configuration",
						action: "set-policy",
						claimedAt: new Date(),
						workspaceUpdatedAt: seeded.workspace.updatedAt,
					},
				},
			}
		);

		await expect(store.renew()).rejects.toBeInstanceOf(AgentRelayLeaseLostError);
		expect(
			(await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.operation?.id
		).toBe("live-renew-operation");
	});

	it("clears an unrecoverable pending WAL while preserving its terminal error", async () => {
		const seeded = await seedAgentRun();
		const store = requireStore(
			await AgentGenerationStore.acquire({
				generationId: seeded.generationId,
				workspaceId: seeded.workspace._id,
				...sandboxIncarnation(seeded.workspace),
			})
		);
		await collections.generations.updateOne(
			{
				generationId: seeded.generationId,
				"agent.relay.owner": store.owner,
				"agent.relay.fence": store.fence,
			},
			{
				$set: {
					"agent.controlsOpen": true,
					"agent.control": {
						id: "undeliverable-control",
						fingerprint: "fingerprint",
						expectedPromptId: "prompt",
						payload: { kind: "compact" },
						requestedAt: new Date(),
						state: "pending",
					},
				},
			}
		);
		await collections.workspaces.deleteOne({ _id: seeded.workspace._id });

		await expect(store.finish("interrupted")).rejects.toBeInstanceOf(AgentRelayLeaseLostError);
		await collections.generations.updateOne(
			{
				generationId: seeded.generationId,
				"agent.relay.owner": store.owner,
				"agent.relay.fence": store.fence,
			},
			{
				$set: {
					seq: 1,
					"agent.pending": {
						kind: "system",
						key: "unrecoverable-pending",
						startSeq: 1,
						events: [
							{
								type: MessageUpdateType.Status,
								status: MessageUpdateStatus.Error,
								message: "historical terminal failure",
							},
						],
					},
				},
			}
		);
		const restoreMaterialization = failAfterNextPendingMaterialization(store);
		try {
			await expect(
				store.finishUnrecoverable("interrupted", "Agent workspace is no longer available")
			).rejects.toThrow("simulated process crash after Agent WAL materialization");
		} finally {
			restoreMaterialization();
		}
		let generation = await collections.generations.findOne({
			generationId: seeded.generationId,
		});
		expect(generation?.status).toBe("running");
		expect(generation?.agent?.pending).toMatchObject({
			kind: "system",
			key: "unrecoverable-pending",
		});
		let conversation = await collections.conversations.findOne({ _id: seeded.conversationId });
		expect(conversation?.messages[0]?.updates).toContainEqual({
			type: MessageUpdateType.Status,
			status: MessageUpdateStatus.Error,
			message: "historical terminal failure",
		});

		expect(
			await store.finishUnrecoverable("interrupted", "Agent workspace is no longer available")
		).toBe(true);

		generation = await collections.generations.findOne({
			generationId: seeded.generationId,
		});
		expect(generation).toMatchObject({
			status: "interrupted",
			error: "Agent workspace is no longer available",
			agent: { controlsOpen: false },
		});
		expect(generation?.agent?.relay).toBeUndefined();
		expect(generation?.agent?.control).toBeUndefined();
		expect(generation?.agent?.pending).toBeUndefined();
		expect(generation?.agent?.systemKeys).toContain("unrecoverable-pending");
		expect(
			await collections.generationEvents.countDocuments({ generationId: seeded.generationId })
		).toBe(1);
		conversation = await collections.conversations.findOne({ _id: seeded.conversationId });
		const message = conversation?.messages.find((item) => item.id === seeded.messageId);
		expect(message?.interrupted).toBe(true);
		expect(message?.updates).toContainEqual({
			type: MessageUpdateType.Status,
			status: MessageUpdateStatus.Error,
			message: "historical terminal failure",
		});
	});
});

function requireStore(store: AgentGenerationStore | null): AgentGenerationStore {
	if (!store) throw new Error("Expected to acquire the Agent generation relay lease");
	return store;
}

function failAfterNextPendingMaterialization(store: AgentGenerationStore): () => void {
	const internals = store as unknown as {
		applyPendingToMessage(
			generation: Generation,
			pending: AgentPendingBatch,
			endSeq: number
		): Promise<void>;
	};
	const applyPendingToMessage = internals.applyPendingToMessage.bind(store);
	internals.applyPendingToMessage = async (generation, pending, endSeq) => {
		await applyPendingToMessage(generation, pending, endSeq);
		throw new Error("simulated process crash after Agent WAL materialization");
	};
	return () => {
		internals.applyPendingToMessage = applyPendingToMessage;
	};
}

function sandboxIncarnation(workspace: Workspace): {
	sandboxId: string;
	sandboxGeneration: string;
} {
	if (!workspace.sandbox.id || !workspace.sandbox.generation) {
		throw new Error("Expected a complete sandbox incarnation");
	}
	return {
		sandboxId: workspace.sandbox.id,
		sandboxGeneration: workspace.sandbox.generation,
	};
}

async function seedAgentRun() {
	const now = new Date();
	const conversationId = new ObjectId();
	const workspaceId = new ObjectId();
	const messageId = randomUUID();
	const generationId = randomUUID();
	const sandboxId = `sandbox-${randomUUID()}`;
	const workspace: Workspace = {
		_id: workspaceId,
		conversationId,
		sessionId: "agent-store-owner",
		repo: { kind: "blank" },
		flavor: "cpu-basic",
		model: "fake/model",
		thinkingLevel: "medium",
		policy: "auto",
		sandbox: {
			status: "ready",
			generation: `generation-${randomUUID()}`,
			id: sandboxId,
			namespace: "tester",
			baseUrl: `https://${sandboxId}--49983.hf.jobs`,
			nonce: "nonce",
			image: "node:22",
			flavor: "cpu-basic",
			createdAt: now,
		},
		piSession: { id: "pi-session", lastSequence: 0 },
		spend: { accruedUsd: 0 },
		createdAt: now,
		updatedAt: now,
	};
	await collections.conversations.insertOne({
		_id: conversationId,
		sessionId: workspace.sessionId,
		mode: "agent",
		workspaceId,
		messages: [
			{
				id: messageId,
				from: "assistant",
				content: "",
				updates: [],
				generationId,
				createdAt: now,
				updatedAt: now,
			},
		],
		createdAt: now,
		updatedAt: now,
	} as never);
	await collections.workspaces.insertOne(workspace);
	await collections.generations.insertOne({
		_id: new ObjectId(),
		generationId,
		conversationId,
		messageId,
		sessionId: workspace.sessionId,
		mode: "agent",
		workspaceId,
		agent: { sourceSequence: 0 },
		status: "running",
		seq: 0,
		startedAt: now,
		lastHeartbeatAt: now,
		createdAt: now,
		updatedAt: now,
	} satisfies Generation);
	return { conversationId, messageId, generationId, workspace };
}