src/lib/server/agent/AgentControlPump.ts

import { collections } from "$lib/server/database";
import type { AgentControlMailboxEntry } from "$lib/types/Generation";
import type { Workspace } from "$lib/types/Workspace";
import {
	claimAgentControl,
	completeAgentControl,
	dispatchAgentControl,
	sealAgentControls,
} from "./AgentControlMailbox";
import { AgentGenerationStore, AgentRelayLeaseLostError } from "./AgentGenerationStore";
import type { AgentSessionManager } from "./AgentSessionManager";
import { PiRpcCommandRejectedError } from "./PiRpcClient";
import { AGENT_POLICY_PATH } from "./WorkspaceProvisioner";

const CONTROL_POLL_MS = 50;

export interface AgentControlPumpOptions {
	store: AgentGenerationStore;
	manager: AgentSessionManager;
	workspace: Workspace;
	signal?: AbortSignal;
	onFailure?: (cause: Error) => void;
	pollIntervalMs?: number;
}

/**
 * Consumes the durable one-slot control mailbox for one exact relay fence.
 *
 * The bridge ACK and Mongo receipt intentionally form a retryable pair. If the
 * process disappears between them, the replacement relay resends the same
 * bridge id and receives the bridge's cached ACK before recording the receipt.
 */
export class AgentControlPump {
	private readonly done: Promise<void>;
	private readonly pollIntervalMs: number;
	private sealRequested = false;
	private stopRequested = false;
	private sealed = false;
	private failure: Error | undefined;
	private wakeWaiter: (() => void) | undefined;

	constructor(private readonly options: AgentControlPumpOptions) {
		this.pollIntervalMs = options.pollIntervalMs ?? CONTROL_POLL_MS;
		this.done = this.run().catch((cause) => {
			this.failure = asError(cause);
			options.onFailure?.(this.failure);
		});
		if (options.signal?.aborted) this.requestStop();
		else options.signal?.addEventListener("abort", this.handleAbort, { once: true });
	}

	/** Drain any enqueue which won the race, then atomically close the mailbox. */
	async sealAndStop(): Promise<void> {
		if (this.stopRequested && !this.sealed) {
			throw this.failure ?? new AgentRelayLeaseLostError();
		}
		this.sealRequested = true;
		this.wake();
		await this.done;
		if (this.failure) throw this.failure;
		if (!this.sealed) throw new AgentRelayLeaseLostError();
	}

	/** Stop without sealing when ownership moved; a replacement relay must drain. */
	requestStop(): void {
		this.stopRequested = true;
		this.wake();
	}

	async waitForStop(): Promise<void> {
		await this.done;
		if (this.failure) throw this.failure;
	}

	private readonly handleAbort = () => this.requestStop();

	private async run(): Promise<void> {
		try {
			while (!this.stopRequested) {
				const relay = { owner: this.options.store.owner, fence: this.options.store.fence };
				const control = await claimAgentControl({
					generationId: this.options.store.generationId,
					relay,
				});
				if (control) {
					await this.dispatch(control);
					continue;
				}

				if (this.sealRequested) {
					const result = await sealAgentControls({
						generationId: this.options.store.generationId,
						relay,
					});
					if (result === "sealed") {
						this.sealed = true;
						return;
					}
					if (result === "lost") throw new AgentRelayLeaseLostError();
					// An enqueue won the no-control/seal compare-and-swap race.
					continue;
				}

				await this.waitForWake();
			}
		} finally {
			this.options.signal?.removeEventListener("abort", this.handleAbort);
		}
	}

	private async dispatch(control: AgentControlMailboxEntry): Promise<void> {
		const relay = { owner: this.options.store.owner, fence: this.options.store.fence };
		try {
			// This check is deliberately adjacent to the irreversible bridge send.
			await this.options.store.renew();
			await dispatchAgentControl(this.options.manager.rpc, control);
			if (shouldPersistAlwaysPolicy(control)) await this.persistAlwaysPolicy();
		} catch (cause) {
			const error = asError(cause);
			if (!(error instanceof PiRpcCommandRejectedError) && !isStaleTurnRejection(error)) {
				throw error;
			}
			const completed = await completeAgentControl({
				generationId: this.options.store.generationId,
				relay,
				control,
				outcome: "rejected",
				error: error.message,
			});
			if (!completed) throw new AgentRelayLeaseLostError();
			return;
		}

		const completed = await completeAgentControl({
			generationId: this.options.store.generationId,
			relay,
			control,
		});
		if (!completed) throw new AgentRelayLeaseLostError();
	}

	private async persistAlwaysPolicy(): Promise<void> {
		const store = this.options.store;
		let persisted = false;
		for (let attempt = 0; attempt < 3 && !persisted; attempt++) {
			const owner = await store.authoritativeOwner();
			const now = new Date();
			const updated = await collections.workspaces.updateOne(
				{
					_id: store.workspaceId,
					conversationId: store.conversationId,
					...owner,
					retirement: { $exists: false },
					"sandbox.id": store.sandboxId,
					"sandbox.generation": store.sandboxGeneration,
					...(store.sandboxNonce ? { "sandbox.nonce": store.sandboxNonce } : {}),
					policy: { $ne: "read-only" },
					"execution.generationId": store.generationId,
					"execution.phase": "relay",
					"execution.owner": store.owner,
					"execution.fence": store.fence,
					"execution.expiresAt": { $gt: now },
				},
				{ $set: { policy: "auto", updatedAt: now } }
			);
			persisted = updated.matchedCount === 1;
		}
		if (!persisted) throw new AgentRelayLeaseLostError();

		// Durable auto is the safe side of this two-system transition. If the
		// renewed execution lease or sandbox write fails, runtime remains ask and
		// the stable approval control can be retried. Runtime auto with durable ask
		// would fail open after an ownership race, so the file is never written first.
		await store.renew();
		await this.options.manager.sandbox.writeFile(AGENT_POLICY_PATH, "auto\n", "0600", {
			signal: this.options.signal,
		});
	}

	private waitForWake(): Promise<void> {
		if (this.stopRequested || this.sealRequested) return Promise.resolve();
		return new Promise((resolve) => {
			const timer = setTimeout(() => {
				if (this.wakeWaiter === wake) this.wakeWaiter = undefined;
				resolve();
			}, this.pollIntervalMs);
			timer.unref?.();
			const wake = () => {
				clearTimeout(timer);
				if (this.wakeWaiter === wake) this.wakeWaiter = undefined;
				resolve();
			};
			this.wakeWaiter = wake;
		});
	}

	private wake(): void {
		this.wakeWaiter?.();
	}
}

function shouldPersistAlwaysPolicy(control: AgentControlMailboxEntry): boolean {
	return (
		control.payload.kind === "approval" &&
		"confirmed" in control.payload.response &&
		control.payload.response.confirmed === true &&
		control.payload.response.always === true
	);
}

function isStaleTurnRejection(error: Error): boolean {
	return error.message.includes("no longer belongs to the active prompt");
}

function asError(cause: unknown): Error {
	return cause instanceof Error ? cause : new Error("Agent control dispatch failed");
}