src/lib/server/agent/AgentControlMailbox.ts

import { createHash } from "node:crypto";
import type { ObjectId } from "mongodb";
import { collections } from "$lib/server/database";
import type {
	AgentControlMailboxEntry,
	AgentControlPayload,
	AgentControlReceipt,
	Generation,
} from "$lib/types/Generation";
import type { Message } from "$lib/types/Message";
import { MessageUpdateType } from "$lib/types/MessageUpdate";
import type { User } from "$lib/types/User";
import type { Workspace } from "$lib/types/Workspace";

const MAX_CONTROL_RECEIPTS = 64;
const DEFAULT_WAIT_TIMEOUT_MS = 30_000;
const DEFAULT_WAIT_INTERVAL_MS = 50;

export interface AgentControlOwner {
	userId?: User["_id"];
	sessionId?: string;
}

export interface AgentControlTarget {
	generationId: string;
	conversationId: ObjectId;
	workspaceId: Workspace["_id"];
	messageId: Message["id"];
	owner: AgentControlOwner;
}

export type EnqueueAgentControlResult =
	| { status: "queued" | "duplicate"; control: AgentControlMailboxEntry }
	| { status: "completed"; receipt: AgentControlReceipt }
	| { status: "busy" | "gone" | "invalid-approval" | "conflict" };

export interface AgentRelayControlOwner {
	owner: string;
	fence: number;
}

export type SealAgentControlsResult = "sealed" | "pending" | "lost";

export type WaitForAgentControlResult =
	| { status: "completed" | "rejected"; receipt: AgentControlReceipt }
	| { status: "timeout" | "gone" | "conflict" };

/**
 * Enqueue one command into the generation's bounded mailbox.
 *
 * Enqueue and terminal sealing contend on the same generation document. If an
 * enqueue wins, sealing must drain it; if sealing wins, no late command can be
 * dispatched after the workspace is released to a later generation.
 */
export async function enqueueAgentControl(options: {
	target: AgentControlTarget;
	id: string;
	payload: AgentControlPayload;
	now?: Date;
}): Promise<EnqueueAgentControlResult> {
	const { target, id, payload } = options;
	if (!id || id.length > 256) return { status: "conflict" };
	const owner = exactOwner(target.owner);
	if (!owner) return { status: "gone" };
	const now = options.now ?? new Date();
	const fingerprint = fingerprintAgentControlPayload(payload);
	const key = agentControlIdempotencyKey(id, payload);

	const filter = targetFilter(target, owner);
	const generation = await collections.generations.findOne(filter, {
		projection: {
			status: 1,
			"agent.prompt.id": 1,
			"agent.relay": 1,
			"agent.controlsOpen": 1,
			"agent.control": 1,
			"agent.controlReceipts": 1,
			"agent.pendingApprovalIds": 1,
		},
	});
	const existing = classifyExisting(generation, id, key, fingerprint, now);
	if (existing) return existing;
	const expectedPromptId = generation?.agent?.prompt?.id;
	if (!generation || !expectedPromptId) return { status: "gone" };

	if (payload.kind === "approval") {
		// The append-only message update is display history, not authorization.
		// Only an unconsumed request emitted by Pi for this live generation may be
		// answered. The enqueue below consumes this id in the same document CAS.
		if (!generation.agent?.pendingApprovalIds?.includes(payload.requestId)) {
			return { status: "invalid-approval" };
		}
		const approvalExists = await collections.conversations.findOne(
			{
				_id: target.conversationId,
				...owner,
				mode: "agent",
				workspaceId: target.workspaceId,
				messages: {
					$elemMatch: {
						id: target.messageId,
						generationId: target.generationId,
						updates: {
							$elemMatch: {
								type: MessageUpdateType.AgentApproval,
								requestId: payload.requestId,
							},
						},
					},
				},
			},
			{ projection: { _id: 1 } }
		);
		if (!approvalExists) return { status: "invalid-approval" };
	}

	const control: AgentControlMailboxEntry = {
		id,
		fingerprint,
		expectedPromptId,
		payload,
		requestedAt: now,
		state: "pending",
	};
	const queued = await collections.generations.updateOne(
		{
			...filter,
			status: "running",
			"agent.prompt.id": expectedPromptId,
			"agent.controlsOpen": { $ne: false },
			"agent.control": { $exists: false },
			"agent.controlReceipts.id": { $ne: id },
			"agent.controlReceipts.key": { $ne: key },
			"agent.relay.expiresAt": { $gt: now },
			...(payload.kind === "approval" ? { "agent.pendingApprovalIds": payload.requestId } : {}),
		},
		{
			$set: {
				"agent.control": control,
				updatedAt: now,
			},
			...(payload.kind === "approval"
				? { $pull: { "agent.pendingApprovalIds": payload.requestId } }
				: {}),
		}
	);
	if (queued.matchedCount === 1) return { status: "queued", control };

	const raced = await collections.generations.findOne(filter, {
		projection: {
			status: 1,
			"agent.prompt.id": 1,
			"agent.relay": 1,
			"agent.controlsOpen": 1,
			"agent.control": 1,
			"agent.controlReceipts": 1,
			"agent.pendingApprovalIds": 1,
		},
	});
	const racedExisting = classifyExisting(raced, id, key, fingerprint, now);
	if (racedExisting) return racedExisting;
	if (
		payload.kind === "approval" &&
		!raced?.agent?.pendingApprovalIds?.includes(payload.requestId)
	) {
		return { status: "invalid-approval" };
	}
	return { status: "busy" };
}

/** Claim or reclaim the current mailbox command for one exact live relay fence. */
export async function claimAgentControl(options: {
	generationId: string;
	relay: AgentRelayControlOwner;
	now?: Date;
}): Promise<AgentControlMailboxEntry | null> {
	const now = options.now ?? new Date();
	const leaseFilter = {
		generationId: options.generationId,
		mode: "agent" as const,
		status: "running" as const,
		"agent.controlsOpen": { $ne: false },
		"agent.relay.owner": options.relay.owner,
		"agent.relay.fence": options.relay.fence,
		"agent.relay.expiresAt": { $gt: now },
	};
	const generation = await collections.generations.findOne(leaseFilter, {
		projection: { "agent.prompt.id": 1, "agent.control": 1 },
	});
	const control = generation?.agent?.control;
	if (!control || control.expectedPromptId !== generation.agent?.prompt?.id) return null;
	if (
		control.state === "dispatching" &&
		control.dispatch?.owner === options.relay.owner &&
		control.dispatch.fence === options.relay.fence
	) {
		return control;
	}

	const claimed = await collections.generations.findOneAndUpdate(
		{
			...leaseFilter,
			"agent.prompt.id": control.expectedPromptId,
			"agent.control.id": control.id,
			"agent.control.fingerprint": control.fingerprint,
			$or: [
				{ "agent.control.state": "pending" },
				{
					"agent.control.state": "dispatching",
					$or: [
						{ "agent.control.dispatch.owner": { $ne: options.relay.owner } },
						{ "agent.control.dispatch.fence": { $ne: options.relay.fence } },
					],
				},
			],
		},
		{
			$set: {
				"agent.control.state": "dispatching",
				"agent.control.dispatch": {
					owner: options.relay.owner,
					fence: options.relay.fence,
					claimedAt: now,
				},
				updatedAt: now,
			},
		},
		{ returnDocument: "after", includeResultMetadata: false }
	);
	return claimed?.agent?.control ?? null;
}

/**
 * Clear a dispatched command and retain its bounded idempotency receipt in one
 * atomic write. Transport failures should not call this: leaving the command
 * dispatching lets a replacement relay retry its stable bridge id.
 */
export async function completeAgentControl(options: {
	generationId: string;
	relay: AgentRelayControlOwner;
	control: AgentControlMailboxEntry;
	outcome?: "completed" | "rejected";
	error?: string;
	now?: Date;
}): Promise<boolean> {
	const now = options.now ?? new Date();
	const receipt: AgentControlReceipt = {
		id: options.control.id,
		key: agentControlIdempotencyKey(options.control.id, options.control.payload),
		fingerprint: options.control.fingerprint,
		outcome: options.outcome ?? "completed",
		completedAt: now,
		...(options.error ? { error: options.error.slice(0, 2_000) } : {}),
	};
	const completed = await collections.generations.updateOne(
		{
			generationId: options.generationId,
			mode: "agent",
			status: "running",
			"agent.relay.owner": options.relay.owner,
			"agent.relay.fence": options.relay.fence,
			"agent.relay.expiresAt": { $gt: now },
			"agent.control.id": options.control.id,
			"agent.control.fingerprint": options.control.fingerprint,
			"agent.control.state": "dispatching",
			"agent.control.dispatch.owner": options.relay.owner,
			"agent.control.dispatch.fence": options.relay.fence,
		},
		{
			$unset: { "agent.control": "" },
			$push: {
				"agent.controlReceipts": {
					$each: [receipt],
					$slice: -MAX_CONTROL_RECEIPTS,
				},
			},
			$set: { updatedAt: now },
		}
	);
	return completed.matchedCount === 1;
}

/** Seal enqueue atomically with the absence of a pending/dispatching command. */
export async function sealAgentControls(options: {
	generationId: string;
	relay: AgentRelayControlOwner;
	now?: Date;
}): Promise<SealAgentControlsResult> {
	const now = options.now ?? new Date();
	const leaseFilter = {
		generationId: options.generationId,
		mode: "agent" as const,
		status: "running" as const,
		"agent.relay.owner": options.relay.owner,
		"agent.relay.fence": options.relay.fence,
		"agent.relay.expiresAt": { $gt: now },
	};
	const sealed = await collections.generations.updateOne(
		{
			...leaseFilter,
			"agent.controlsOpen": { $ne: false },
			"agent.control": { $exists: false },
		},
		{ $set: { "agent.controlsOpen": false, updatedAt: now } }
	);
	if (sealed.matchedCount === 1) return "sealed";

	const current = await collections.generations.findOne(leaseFilter, {
		projection: { "agent.controlsOpen": 1, "agent.control": 1 },
	});
	if (!current) return "lost";
	if (current.agent?.controlsOpen === false) return "sealed";
	return current.agent?.control ? "pending" : "lost";
}

/** Wait for the exact durable receipt so sequential controls preserve old HTTP semantics. */
export async function waitForAgentControlCompletion(options: {
	generationId: string;
	controlId: string;
	fingerprint: string;
	timeoutMs?: number;
	pollIntervalMs?: number;
	signal?: AbortSignal;
}): Promise<WaitForAgentControlResult> {
	const timeoutMs = options.timeoutMs ?? DEFAULT_WAIT_TIMEOUT_MS;
	const pollIntervalMs = options.pollIntervalMs ?? DEFAULT_WAIT_INTERVAL_MS;
	const deadline = Date.now() + timeoutMs;
	for (;;) {
		if (options.signal?.aborted) return { status: "gone" };
		const generation = await collections.generations.findOne(
			{ generationId: options.generationId },
			{ projection: { status: 1, "agent.control": 1, "agent.controlReceipts": 1 } }
		);
		if (!generation) return { status: "gone" };
		const receipt = generation.agent?.controlReceipts?.find(
			(item) => item.id === options.controlId
		);
		if (receipt) {
			if (receipt.fingerprint !== options.fingerprint) return { status: "conflict" };
			return { status: receipt.outcome, receipt };
		}
		const control = generation.agent?.control;
		if (control?.id === options.controlId && control.fingerprint !== options.fingerprint) {
			return { status: "conflict" };
		}
		if (generation.status !== "running" && !control) return { status: "gone" };
		if (Date.now() >= deadline) return { status: "timeout" };
		await wait(Math.min(pollIntervalMs, Math.max(1, deadline - Date.now())), options.signal);
	}
}

/** Stable bridge id for commands which receive normal correlated pi responses. */
export function agentControlBridgeId(controlId: string): string {
	return `agent-control:${controlId}`;
}

/**
 * Produce the exact bridge envelope. `expectedPromptId` is bridge-only metadata
 * and must be checked then stripped before the command is written to pi.
 */
export function agentControlRpcEnvelope(
	control: AgentControlMailboxEntry
): Record<string, unknown> {
	const outerId = agentControlBridgeId(control.id);
	switch (control.payload.kind) {
		case "steer":
			return {
				id: outerId,
				type: "steer",
				_chatUiExpectedPromptId: control.expectedPromptId,
				message: control.payload.message,
			};
		case "follow-up":
			return {
				id: outerId,
				type: "follow_up",
				_chatUiExpectedPromptId: control.expectedPromptId,
				message: control.payload.message,
			};
		case "compact":
			return {
				id: outerId,
				type: "compact",
				_chatUiExpectedPromptId: control.expectedPromptId,
			};
		case "queue-replace":
			// Queue replacement needs pi's semantic compare-and-swap response, not
			// the bridge wrapper's transport-level write acknowledgement. A direct
			// stable id is still deduplicated by the bridge across relay retries.
			return {
				id: outerId,
				type: "replace_queue",
				_chatUiExpectedPromptId: control.expectedPromptId,
				expected: control.payload.expected,
				replacement: control.payload.replacement,
			};
		case "approval": {
			const response = control.payload.response;
			return {
				id: outerId,
				type: "extension_ui_response",
				_chatUiExpectedPromptId: control.expectedPromptId,
				requestId: control.payload.requestId,
				...("confirmed" in response
					? { confirmed: response.confirmed }
					: "value" in response
						? { value: response.value }
						: { cancelled: true }),
			};
		}
	}
}

export interface AgentControlRpcTransport {
	command(
		command: Record<string, unknown> & { type: string },
		options?: { timeoutMs?: number }
	): Promise<unknown>;
}

/** Dispatch through the stable, turn-fenced bridge command for this control kind. */
export async function dispatchAgentControl(
	transport: AgentControlRpcTransport,
	control: AgentControlMailboxEntry,
	options: { timeoutMs?: number } = {}
): Promise<void> {
	const envelope = agentControlRpcEnvelope(control) as Record<string, unknown> & { type: string };
	await transport.command(envelope, { timeoutMs: options.timeoutMs ?? DEFAULT_WAIT_TIMEOUT_MS });
}

export function fingerprintAgentControlPayload(payload: AgentControlPayload): string {
	const canonical =
		payload.kind === "steer" || payload.kind === "follow-up"
			? JSON.stringify({ kind: payload.kind, message: payload.message })
			: payload.kind === "compact"
				? JSON.stringify({ kind: payload.kind })
				: payload.kind === "queue-replace"
					? JSON.stringify({
							kind: payload.kind,
							expected: payload.expected,
							replacement: payload.replacement,
						})
					: JSON.stringify({
							kind: payload.kind,
							requestId: payload.requestId,
							response:
								"confirmed" in payload.response
									? {
											confirmed: payload.response.confirmed,
											always: payload.response.always === true,
										}
									: "value" in payload.response
										? { value: payload.response.value }
										: { cancelled: true },
						});
	return createHash("sha256").update(canonical).digest("hex");
}

function classifyExisting(
	generation: Generation | null,
	id: string,
	key: string,
	fingerprint: string,
	now: Date
): EnqueueAgentControlResult | undefined {
	if (!generation) return { status: "gone" };
	const receipt = generation.agent?.controlReceipts?.find(
		(item) => item.id === id || item.key === key
	);
	if (receipt) {
		return receipt.fingerprint === fingerprint
			? { status: "completed", receipt }
			: { status: "conflict" };
	}
	if (generation.status !== "running" || generation.agent?.controlsOpen === false) {
		return { status: "gone" };
	}
	const control = generation.agent?.control;
	if (control) {
		const sameKey = agentControlIdempotencyKey(control.id, control.payload) === key;
		if (control.id === id || sameKey) {
			return control.fingerprint === fingerprint
				? { status: "duplicate", control }
				: { status: "conflict" };
		}
		return { status: "busy" };
	}
	if (!generation.agent?.prompt?.id) return { status: "gone" };
	const relay = generation.agent.relay;
	if (!relay || relay.expiresAt <= now) return { status: "busy" };
	return undefined;
}

function agentControlIdempotencyKey(id: string, payload: AgentControlPayload): string {
	return payload.kind === "approval" ? `approval:${payload.requestId}` : `command:${id}`;
}

function exactOwner(
	owner: AgentControlOwner
): { userId: User["_id"] } | { sessionId: string } | null {
	// Logged-in requests may retain an anonymous session id. Ownership throughout
	// chat-ui deliberately prefers the durable user id in that case.
	if (owner.userId) return { userId: owner.userId };
	if (owner.sessionId) return { sessionId: owner.sessionId };
	return null;
}

function targetFilter(
	target: AgentControlTarget,
	owner: { userId: User["_id"] } | { sessionId: string }
): Record<string, unknown> {
	return {
		generationId: target.generationId,
		conversationId: target.conversationId,
		workspaceId: target.workspaceId,
		messageId: target.messageId,
		mode: "agent",
		...owner,
	};
}

async function wait(ms: number, signal?: AbortSignal): Promise<void> {
	if (signal?.aborted) return;
	await new Promise<void>((resolve) => {
		const onAbort = () => {
			clearTimeout(timer);
			resolve();
		};
		const timer = setTimeout(() => {
			signal?.removeEventListener("abort", onAbort);
			resolve();
		}, ms);
		signal?.addEventListener("abort", onAbort, { once: true });
	});
}