src/lib/server/agent/AgentControlMailbox.spec.ts

import { randomUUID } from "node:crypto";
import { afterEach, beforeAll, describe, expect, it } from "vitest";
import { ObjectId } from "mongodb";
import { collections, ready } from "$lib/server/database";
import { createTestLocals } from "$lib/server/api/__tests__/testHelpers";
import type {
	AgentControlMailboxEntry,
	AgentControlPayload,
	Generation,
} from "$lib/types/Generation";
import { AGENT_APPROVAL_NOT_ACTIVE_RESPONSE } from "$lib/types/AgentControl";
import { MessageUpdateType } from "$lib/types/MessageUpdate";
import {
	agentControlBridgeId,
	agentControlRpcEnvelope,
	claimAgentControl,
	completeAgentControl,
	dispatchAgentControl,
	enqueueAgentControl,
	fingerprintAgentControlPayload,
	sealAgentControls,
	waitForAgentControlCompletion,
	type AgentControlTarget,
} from "./AgentControlMailbox";
import { POST as agentControlPOST } from "../../../routes/conversation/[id]/agent/control/+server";

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

afterEach(async () => {
	await Promise.all([
		collections.conversations.deleteMany({ sessionId: /^agent-control-/ }),
		collections.workspaces.deleteMany({ sessionId: /^agent-control-/ }),
		collections.generations.deleteMany({ sessionId: /^agent-control-/ }),
	]);
});

describe.sequential("AgentControlMailbox", () => {
	it("preserves synchronous steer then follow-up semantics with one bounded slot", async () => {
		const seeded = await seedAgentControlRun();
		const steer = await enqueueAgentControl({
			target: seeded.target,
			id: "steer-1",
			payload: { kind: "steer", message: "change course" },
		});
		expect(steer.status).toBe("queued");
		if (steer.status !== "queued") throw new Error("expected queued steer");

		const waiting = waitForAgentControlCompletion({
			generationId: seeded.generationId,
			controlId: steer.control.id,
			fingerprint: steer.control.fingerprint,
			pollIntervalMs: 1,
			timeoutMs: 2_000,
		});
		const claimedSteer = await claimAgentControl({
			generationId: seeded.generationId,
			relay: seeded.relay,
		});
		expect(claimedSteer?.payload).toEqual({ kind: "steer", message: "change course" });
		expect(
			await completeAgentControl({
				generationId: seeded.generationId,
				relay: seeded.relay,
				control: requireControl(claimedSteer),
			})
		).toBe(true);
		expect((await waiting).status).toBe("completed");

		// The first HTTP request waits for its receipt, so the immediately
		// sequential second request sees an empty mailbox rather than a 409.
		const followUp = await enqueueAgentControl({
			target: seeded.target,
			id: "follow-up-1",
			payload: { kind: "follow-up", message: "then run the tests" },
		});
		expect(followUp.status).toBe("queued");
		const claimedFollowUp = await claimAgentControl({
			generationId: seeded.generationId,
			relay: seeded.relay,
		});
		expect(claimedFollowUp?.payload).toEqual({
			kind: "follow-up",
			message: "then run the tests",
		});
	});

	it("lets a replacement fence reclaim a crash-partial dispatch with the same bridge id", async () => {
		const seeded = await seedAgentControlRun();
		const queued = await enqueueAgentControl({
			target: seeded.target,
			id: "reclaim-me",
			payload: { kind: "steer", message: "safe retry" },
		});
		if (queued.status !== "queued") throw new Error("expected queued control");
		const first = requireControl(
			await claimAgentControl({ generationId: seeded.generationId, relay: seeded.relay })
		);
		const firstEnvelope = agentControlRpcEnvelope(first);

		const replacement = { owner: "relay-replacement", fence: seeded.relay.fence + 1 };
		await collections.generations.updateOne(
			{ generationId: seeded.generationId },
			{
				$set: {
					"agent.relay.owner": replacement.owner,
					"agent.relay.fence": replacement.fence,
					"agent.relay.expiresAt": new Date(Date.now() + 60_000),
				},
			}
		);
		const reclaimed = requireControl(
			await claimAgentControl({ generationId: seeded.generationId, relay: replacement })
		);

		expect(reclaimed.dispatch).toMatchObject(replacement);
		expect(agentControlRpcEnvelope(reclaimed)).toEqual(firstEnvelope);
		expect(firstEnvelope).toMatchObject({
			id: agentControlBridgeId("reclaim-me"),
			type: "steer",
			_chatUiExpectedPromptId: seeded.promptId,
			message: "safe retry",
		});
		const dispatched: Array<{ command: Record<string, unknown>; timeoutMs?: number }> = [];
		await dispatchAgentControl(
			{
				command: async (command, options) => {
					dispatched.push({ command, timeoutMs: options?.timeoutMs });
					return { success: true };
				},
			},
			reclaimed,
			{ timeoutMs: 1_234 }
		);
		expect(dispatched).toEqual([{ command: firstEnvelope, timeoutMs: 1_234 }]);
		expect(
			await completeAgentControl({
				generationId: seeded.generationId,
				relay: seeded.relay,
				control: first,
			})
		).toBe(false);
		expect(
			await completeAgentControl({
				generationId: seeded.generationId,
				relay: replacement,
				control: reclaimed,
			})
		).toBe(true);
	});

	it("atomically orders enqueue against terminal sealing", async () => {
		const enqueueWinner = await seedAgentControlRun();
		const queued = await enqueueAgentControl({
			target: enqueueWinner.target,
			id: "before-seal",
			payload: { kind: "compact" },
		});
		expect(queued.status).toBe("queued");
		expect(
			await sealAgentControls({
				generationId: enqueueWinner.generationId,
				relay: enqueueWinner.relay,
			})
		).toBe("pending");

		const sealWinner = await seedAgentControlRun();
		expect(
			await sealAgentControls({
				generationId: sealWinner.generationId,
				relay: sealWinner.relay,
			})
		).toBe("sealed");
		expect(
			await enqueueAgentControl({
				target: sealWinner.target,
				id: "after-seal",
				payload: { kind: "steer", message: "too late" },
			})
		).toEqual({ status: "gone" });
	});

	it("binds approval to its persisted request and makes the first response idempotent", async () => {
		const seeded = await seedAgentControlRun({ approvalRequestId: "approval-1" });
		const approve: AgentControlPayload = {
			kind: "approval",
			requestId: "approval-1",
			response: { confirmed: true, always: true },
		};
		const first = await enqueueAgentControl({
			target: seeded.target,
			id: "approval-command-a",
			payload: approve,
		});
		expect(first.status).toBe("queued");
		const duplicate = await enqueueAgentControl({
			target: seeded.target,
			id: "approval-command-b",
			payload: approve,
		});
		expect(duplicate.status).toBe("duplicate");
		const conflicting = await enqueueAgentControl({
			target: seeded.target,
			id: "approval-command-c",
			payload: {
				kind: "approval",
				requestId: "approval-1",
				response: { cancelled: true },
			},
		});
		expect(conflicting).toEqual({ status: "conflict" });

		const claimed = requireControl(
			await claimAgentControl({ generationId: seeded.generationId, relay: seeded.relay })
		);
		expect(agentControlRpcEnvelope(claimed)).toEqual({
			id: agentControlBridgeId("approval-command-a"),
			type: "extension_ui_response",
			_chatUiExpectedPromptId: seeded.promptId,
			requestId: "approval-1",
			confirmed: true,
		});
		expect(
			await completeAgentControl({
				generationId: seeded.generationId,
				relay: seeded.relay,
				control: claimed,
			})
		).toBe(true);
		expect(
			await enqueueAgentControl({
				target: seeded.target,
				id: "missing-approval",
				payload: {
					kind: "approval",
					requestId: "approval-does-not-exist",
					response: { confirmed: true },
				},
			})
		).toEqual({ status: "invalid-approval" });
		expect(
			await sealAgentControls({
				generationId: seeded.generationId,
				relay: seeded.relay,
			})
		).toBe("sealed");

		const completedRetry = await enqueueAgentControl({
			target: seeded.target,
			id: "approval-command-d",
			payload: approve,
		});
		expect(completedRetry.status).toBe("completed");
	});

	it("waits on the canonical mailbox id for a duplicate approval from another tab", async () => {
		const seeded = await seedAgentControlRun({ approvalRequestId: "approval-two-tabs" });
		const payload: AgentControlPayload = {
			kind: "approval",
			requestId: "approval-two-tabs",
			response: { confirmed: true },
		};
		const originalId = randomUUID();
		const original = await enqueueAgentControl({
			target: seeded.target,
			id: originalId,
			payload,
		});
		expect(original.status).toBe("queued");

		const duplicateResponse = postAgentControlRoute(
			seeded,
			{
				action: "approval",
				controlId: randomUUID(),
				generationId: seeded.generationId,
				messageId: seeded.target.messageId,
				requestId: payload.requestId,
				response: payload.response,
			},
			AbortSignal.timeout(2_000)
		);
		// Leave the original in the mailbox long enough for the second request to
		// classify it as a duplicate before publishing its canonical receipt.
		await new Promise((resolve) => setTimeout(resolve, 100));
		const claimed = requireControl(
			await claimAgentControl({ generationId: seeded.generationId, relay: seeded.relay })
		);
		expect(claimed.id).toBe(originalId);
		expect(
			await completeAgentControl({
				generationId: seeded.generationId,
				relay: seeded.relay,
				control: claimed,
			})
		).toBe(true);

		expect((await duplicateResponse).status).toBe(200);
	});

	it("returns an exact completed queue retry before rejecting its stale snapshot", async () => {
		const seeded = await seedAgentControlRun();
		const expected = { steering: ["old steer"], followUp: ["old follow-up"] };
		const replacement = { steering: ["new steer"], followUp: [] };
		await setAgentQueue(seeded, expected);
		const controlId = randomUUID();
		const payload: AgentControlPayload = {
			kind: "queue-replace",
			expected,
			replacement,
		};
		const queued = await enqueueAgentControl({ target: seeded.target, id: controlId, payload });
		expect(queued.status).toBe("queued");
		const claimed = requireControl(
			await claimAgentControl({ generationId: seeded.generationId, relay: seeded.relay })
		);
		expect(
			await completeAgentControl({
				generationId: seeded.generationId,
				relay: seeded.relay,
				control: claimed,
			})
		).toBe(true);
		await setAgentQueue(seeded, replacement);

		const completedRetry = await postAgentControlRoute(seeded, {
			action: "queue-replace",
			controlId,
			generationId: seeded.generationId,
			messageId: seeded.target.messageId,
			expected,
			replacement,
		});
		expect(completedRetry.status).toBe(200);

		await expect(
			postAgentControlRoute(seeded, {
				action: "queue-replace",
				controlId,
				generationId: seeded.generationId,
				messageId: seeded.target.messageId,
				expected,
				replacement: { steering: [], followUp: ["different"] },
			})
		).rejects.toMatchObject({
			status: 409,
			body: { message: "Agent control id was reused" },
		});

		await expect(
			postAgentControlRoute(seeded, {
				action: "queue-replace",
				controlId: randomUUID(),
				generationId: seeded.generationId,
				messageId: seeded.target.messageId,
				expected,
				replacement,
			})
		).rejects.toMatchObject({
			status: 409,
			body: { message: "The Agent queue changed; refresh before editing it" },
		});
	});

	it("rejects a consumed historical approval even after its bounded receipt is evicted", async () => {
		const seeded = await seedAgentControlRun({ approvalRequestId: "approval-stale" });
		const payload: AgentControlPayload = {
			kind: "approval",
			requestId: "approval-stale",
			response: { confirmed: true, always: true },
		};
		const queued = await enqueueAgentControl({
			target: seeded.target,
			id: "approval-first",
			payload,
		});
		expect(queued.status).toBe("queued");
		const claimed = requireControl(
			await claimAgentControl({ generationId: seeded.generationId, relay: seeded.relay })
		);
		expect(
			await completeAgentControl({
				generationId: seeded.generationId,
				relay: seeded.relay,
				control: claimed,
			})
		).toBe(true);
		await collections.generations.updateOne(
			{ generationId: seeded.generationId },
			{ $unset: { "agent.controlReceipts": "" } }
		);

		expect(
			await enqueueAgentControl({
				target: seeded.target,
				id: "approval-replay-after-eviction",
				payload,
			})
		).toEqual({ status: "invalid-approval" });
	});

	it("returns a typed public stale response when an approval expires before submission", async () => {
		const seeded = await seedAgentControlRun({ approvalRequestId: "approval-expired" });
		await collections.generations.updateOne(
			{ generationId: seeded.generationId },
			{
				$pull: { "agent.pendingApprovalIds": "approval-expired" },
				$set: { updatedAt: new Date(Date.now() + 1_000) },
			}
		);

		const response = await postAgentControlRoute(seeded, {
			action: "approval",
			controlId: randomUUID(),
			generationId: seeded.generationId,
			messageId: seeded.target.messageId,
			requestId: "approval-expired",
			response: { confirmed: true },
		});

		expect(response.status).toBe(409);
		expect(response.headers.get("cache-control")).toBe("no-store");
		expect(await response.json()).toEqual(AGENT_APPROVAL_NOT_ACTIVE_RESPONSE);
		expect(
			await collections.generations.findOne(
				{ generationId: seeded.generationId },
				{ projection: { "agent.control": 1 } }
			)
		).not.toHaveProperty("agent.control");
	});

	it("fingerprints equivalent payloads deterministically and conflicting approvals differently", () => {
		expect(
			fingerprintAgentControlPayload({
				kind: "approval",
				requestId: "same",
				response: { always: true, confirmed: true },
			})
		).toBe(
			fingerprintAgentControlPayload({
				kind: "approval",
				requestId: "same",
				response: { confirmed: true, always: true },
			})
		);
		expect(
			fingerprintAgentControlPayload({
				kind: "approval",
				requestId: "same",
				response: { confirmed: true },
			})
		).not.toBe(
			fingerprintAgentControlPayload({
				kind: "approval",
				requestId: "same",
				response: { cancelled: true },
			})
		);
	});

	it("dispatches queue replacement as a direct, semantic compare-and-swap command", async () => {
		const seeded = await seedAgentControlRun();
		const queued = await enqueueAgentControl({
			target: seeded.target,
			id: "replace-queue-1",
			payload: {
				kind: "queue-replace",
				expected: { steering: ["old steer"], followUp: ["old follow-up"] },
				replacement: { steering: ["new steer"], followUp: [] },
			},
		});
		if (queued.status !== "queued") throw new Error("expected queued control");
		const claimed = requireControl(
			await claimAgentControl({ generationId: seeded.generationId, relay: seeded.relay })
		);
		const envelope = agentControlRpcEnvelope(claimed);

		expect(envelope).toEqual({
			id: agentControlBridgeId("replace-queue-1"),
			type: "replace_queue",
			_chatUiExpectedPromptId: seeded.promptId,
			expected: { steering: ["old steer"], followUp: ["old follow-up"] },
			replacement: { steering: ["new steer"], followUp: [] },
		});
		const dispatched: Record<string, unknown>[] = [];
		await dispatchAgentControl(
			{
				command: async (command) => {
					dispatched.push(command);
					return { success: true };
				},
			},
			claimed
		);
		expect(dispatched).toEqual([envelope]);
	});

	it("prefers durable user ownership when a logged-in request also has a session id", async () => {
		const seeded = await seedAgentControlRun({ includeUserAndSession: true });
		expect(
			await enqueueAgentControl({
				target: seeded.target,
				id: "logged-in-control",
				payload: { kind: "steer", message: "keep the user owner" },
			})
		).toMatchObject({ status: "queued" });
	});
});

function requireControl(control: AgentControlMailboxEntry | null): AgentControlMailboxEntry {
	if (!control) throw new Error("expected a claimed Agent control");
	return control;
}

async function seedAgentControlRun(
	options: { approvalRequestId?: string; includeUserAndSession?: boolean } = {}
) {
	const now = new Date();
	const conversationId = new ObjectId();
	const workspaceId = new ObjectId();
	const generationId = randomUUID();
	const messageId = randomUUID();
	const promptId = randomUUID();
	const sessionId = `agent-control-${randomUUID()}`;
	const userId = options.includeUserAndSession ? new ObjectId() : undefined;
	const relay = { owner: randomUUID(), fence: 7 };
	const target: AgentControlTarget = {
		generationId,
		conversationId,
		workspaceId,
		messageId,
		owner: { ...(userId ? { userId } : {}), sessionId },
	};
	await collections.conversations.insertOne({
		_id: conversationId,
		...(userId ? { userId } : {}),
		sessionId,
		mode: "agent",
		workspaceId,
		messages: [
			{
				id: messageId,
				from: "assistant",
				content: "",
				generationId,
				updates: options.approvalRequestId
					? [
							{
								type: MessageUpdateType.AgentApproval,
								requestId: options.approvalRequestId,
								method: "confirm",
								title: "Allow command?",
							},
						]
					: [],
				createdAt: now,
				updatedAt: now,
			},
		],
		createdAt: now,
		updatedAt: now,
	} as never);
	await collections.workspaces.insertOne({
		_id: workspaceId,
		conversationId,
		...(userId ? { userId } : {}),
		sessionId,
		repo: { kind: "blank" },
		flavor: "cpu-basic",
		model: "fake/model",
		thinkingLevel: "medium",
		policy: "ask",
		sandbox: {
			status: "ready",
			id: `sandbox-${workspaceId.toString()}`,
			generation: `sandbox-generation-${workspaceId.toString()}`,
			namespace: "tester",
			baseUrl: `https://sandbox-${workspaceId.toString()}.example.test`,
			nonce: "nonce",
			image: "node:22",
			flavor: "cpu-basic",
			createdAt: now,
		},
		spend: { accruedUsd: 0 },
		createdAt: now,
		updatedAt: now,
	});
	await collections.generations.insertOne({
		_id: new ObjectId(),
		generationId,
		conversationId,
		workspaceId,
		messageId,
		...(userId ? { userId } : {}),
		sessionId,
		mode: "agent",
		status: "running",
		seq: 0,
		lastHeartbeatAt: now,
		startedAt: now,
		createdAt: now,
		updatedAt: now,
		agent: {
			sourceSequence: 0,
			prompt: { id: promptId, message: "do the work" },
			relay: { ...relay, expiresAt: new Date(Date.now() + 60_000) },
			...(options.approvalRequestId ? { pendingApprovalIds: [options.approvalRequestId] } : {}),
		},
	} satisfies Generation);
	return { generationId, promptId, relay, target, sessionId };
}

function postAgentControlRoute(
	seeded: Awaited<ReturnType<typeof seedAgentControlRun>>,
	body: Record<string, unknown>,
	signal?: AbortSignal
): Promise<Response> {
	return agentControlPOST({
		params: { id: seeded.target.conversationId.toString() },
		locals: createTestLocals({ sessionId: seeded.sessionId }),
		request: new Request("http://localhost/agent/control", {
			method: "POST",
			headers: { "Content-Type": "application/json" },
			body: JSON.stringify(body),
			signal,
		}),
	} as never) as Promise<Response>;
}

async function setAgentQueue(
	seeded: Awaited<ReturnType<typeof seedAgentControlRun>>,
	queue: { steering: string[]; followUp: string[] }
): Promise<void> {
	const updated = await collections.conversations.updateOne(
		{ _id: seeded.target.conversationId, "messages.id": seeded.target.messageId },
		{
			$set: {
				"messages.$.updates": [{ type: MessageUpdateType.AgentQueue, ...queue }],
			},
		}
	);
	if (updated.matchedCount !== 1) throw new Error("Expected to update the Agent queue fixture");
}