Blame for src/lib/server/agent/AgentRunRelay.ts

import { collections } from "$lib/server/database";
import { resolveGenerationStop } from "$lib/server/generation/stopGeneration";
import { logger } from "$lib/server/logger";
import { clampStoppedContent } from "$lib/server/stopTruncation";
import { downloadFile } from "$lib/server/files/downloadFile";
import type { Generation, GenerationStatus } from "$lib/types/Generation";
import {
	MessageUpdateStatus,
	MessageUpdateType,
	type MessageUpdate,
} from "$lib/types/MessageUpdate";
import type { Workspace } from "$lib/types/Workspace";
import {
	AgentGenerationIncarnationMismatchError,
	AgentGenerationStore,
	AgentRelayLeaseLostError,
	AGENT_RELAY_LEASE_MS,
	AGENT_SANDBOX_CHANGED_MESSAGE,
} from "./AgentGenerationStore";
import { AgentControlPump } from "./AgentControlPump";
import { AgentSessionManager } from "./AgentSessionManager";
import { prepareDetachedAgentPromptForSandbox } from "./agentPromptRecovery";
import { PI_RPC_REPLAY_GAP_MESSAGE, PiRpcReplayGapError } from "./PiRpcClient";
import type { SandboxClient } from "./SandboxClient";
import { agentUploadDirectory, isAgentUploadPath } from "./agentUploadPaths";
import { workspaceOperationAvailableFilter } from "./WorkspaceOperation";

interface ActiveAgentRelay {
	sandboxId: string;
	sandboxGeneration: string;
	promise: Promise<void>;
}

const activeRelays = new Map<string, ActiveAgentRelay>();
const ABORT_POLL_MS = 300;

export interface EnsureAgentRelayOptions {
	generationId: string;
	workspace: Workspace;
	hfToken: string;
	bridgeUrl?: string;
	bridgeHeaders?: Record<string, string>;
}

/**
 * Opportunistically claim and resume a detached Agent run. The returned
 * promise only covers lease acquisition; the relay intentionally continues in
 * the background after the requesting SSE connection goes away.
 */
export async function ensureAgentGenerationRelay(
	options: EnsureAgentRelayOptions
): Promise<"started" | "already-running" | "gone"> {
	const sandboxId = options.workspace.sandbox.id;
	const sandboxGeneration = options.workspace.sandbox.generation;
	if (!sandboxId || !sandboxGeneration) {
		await AgentGenerationStore.interruptForSandboxChange({
			generationId: options.generationId,
			workspaceId: options.workspace._id,
		});
		return "gone";
	}
	const workspaceOwner = options.workspace.userId
		? { userId: options.workspace.userId }
		: options.workspace.sessionId
			? { sessionId: options.workspace.sessionId }
			: undefined;
	if (!workspaceOwner) return "gone";
	const admissionAt = new Date();
	const dispatchableWorkspace = await collections.workspaces.findOneAndUpdate(
		{
			_id: options.workspace._id,
			conversationId: options.workspace.conversationId,
			...workspaceOwner,
			retirement: { $exists: false },
			...workspaceOperationAvailableFilter(admissionAt),
			"sandbox.status": { $in: ["ready", "idle"] },
			"sandbox.id": sandboxId,
			"sandbox.generation": sandboxGeneration,
		},
		{ $unset: { operation: "" } },
		{ projection: { _id: 1 }, returnDocument: "after", includeResultMetadata: false }
	);
	if (!dispatchableWorkspace) return "already-running";

	const active = activeRelays.get(options.generationId);
	if (active) {
		if (active.sandboxId === sandboxId && active.sandboxGeneration === sandboxGeneration) {
			return "already-running";
		}
		await AgentGenerationStore.interruptForSandboxChange({
			generationId: options.generationId,
			workspaceId: options.workspace._id,
			sandboxId,
			sandboxGeneration,
		});
		return "gone";
	}

	const promptPreparation = await prepareDetachedAgentPromptForSandbox({
		generationId: options.generationId,
		workspaceId: options.workspace._id,
		sandbox: options.workspace.sandbox,
	});
	if (promptPreparation === "busy") return "already-running";
	if (promptPreparation === "gone") return "gone";
	if (promptPreparation === "unrecoverable") {
		await AgentGenerationStore.interruptForSandboxChange({
			generationId: options.generationId,
			workspaceId: options.workspace._id,
			sandboxId,
			sandboxGeneration,
			sandboxNonce: options.workspace.sandbox.nonce,
		});
		return "gone";
	}

	let lease: AgentGenerationStore | null;
	try {
		lease = await AgentGenerationStore.acquire({
			generationId: options.generationId,
			workspaceId: options.workspace._id,
			sandboxId,
			sandboxGeneration,
			sandboxNonce: options.workspace.sandbox.nonce,
		});
	} catch (cause) {
		if (!(cause instanceof AgentGenerationIncarnationMismatchError)) throw cause;
		await AgentGenerationStore.interruptForSandboxChange({
			generationId: options.generationId,
			workspaceId: options.workspace._id,
			sandboxId,
			sandboxGeneration,
		});
		return "gone";
	}
	if (!lease) {
		const generation = await collections.generations.findOne(
			{ generationId: options.generationId },
			{ projection: { status: 1 } }
		);
		return generation?.status === "running" ? "already-running" : "gone";
	}

	const relay = resumeWithLease(options, lease)
		.catch((cause) =>
			logger.error(
				{ cause, generationId: options.generationId },
				"Detached Agent generation relay failed"
			)
		)
		.finally(() => {
			if (activeRelays.get(options.generationId)?.promise === relay) {
				activeRelays.delete(options.generationId);
			}
		});
	activeRelays.set(options.generationId, { sandboxId, sandboxGeneration, promise: relay });
	return "started";
}

/** Test/diagnostic hook; production callers should use ensureAgentGenerationRelay. */
export function hasActiveAgentRelay(generationId: string): boolean {
	return activeRelays.has(generationId);
}

export function waitForActiveAgentRelay(generationId: string): Promise<void> | undefined {
	return activeRelays.get(generationId)?.promise;
}

async function readDurableLifecycleState(generationId: string): Promise<{
	sawError: boolean;
	sawFinished: boolean;
	error?: string;
}> {
	const events = await collections.generationEvents
		.find({
			generationId,
			"event.type": MessageUpdateType.Status,
			"event.status": {
				$in: [MessageUpdateStatus.Error, MessageUpdateStatus.Finished],
			},
		})
		.sort({ seq: 1 })
		.toArray();
	let sawError = false;
	let sawFinished = false;
	let error: string | undefined;
	for (const record of events) {
		if (record.event.type !== MessageUpdateType.Status) continue;
		if (record.event.status === MessageUpdateStatus.Error) {
			sawError = true;
			error = record.event.message ?? error;
		} else if (record.event.status === MessageUpdateStatus.Finished) {
			sawFinished = true;
		}
	}
	return { sawError, sawFinished, ...(error ? { error } : {}) };
}

async function resumeWithLease(
	options: EnsureAgentRelayOptions,
	store: AgentGenerationStore
): Promise<void> {
	let manager: AgentSessionManager | undefined;
	let controlPump: AgentControlPump | undefined;
	let controlPumpFailure: Error | undefined;
	let startControlPump: (() => AgentControlPump) | undefined;
	let leaseTimer: ReturnType<typeof setInterval> | undefined;
	let abortTimer: ReturnType<typeof setInterval> | undefined;
	let finished = false;
	let sawFinished = false;
	let sawError = false;
	let terminalError: string | undefined;
	const abortController = new AbortController();
	const ownershipController = new AbortController();

	try {
		// Start fencing renewal before WAL recovery, workspace reads, or attachment
		// restoration. Those operations can be slow enough to consume a short lease.
		await store.renew();
		leaseTimer = setInterval(
			() => {
				void store.renew().catch(() => ownershipController.abort());
			},
			Math.max(1_000, Math.floor(AGENT_RELAY_LEASE_MS / 3))
		);
		leaseTimer.unref?.();

		await store.resumePending();
		const [generation, workspace, lifecycle] = await Promise.all([
			collections.generations.findOne({ generationId: options.generationId }),
			collections.workspaces.findOne({
				_id: options.workspace._id,
				conversationId: store.conversationId,
				retirement: { $exists: false },
				"sandbox.id": store.sandboxId,
				"sandbox.generation": store.sandboxGeneration,
				"sandbox.status": { $in: ["ready", "idle"] },
			}),
			readDurableLifecycleState(options.generationId),
		]);
		if (!generation || generation.status !== "running") return;
		sawError = lifecycle.sawError;
		sawFinished = lifecycle.sawFinished;
		terminalError = lifecycle.error;
		if (!workspace) {
			// The reciprocal workspace lease can no longer be renewed, so normal
			// finish() cannot drain its WAL/mailbox preconditions. The exact relay
			// fence is still sufficient to terminate this now-unrecoverable run.
			finished = await store.finishUnrecoverable("interrupted", AGENT_SANDBOX_CHANGED_MESSAGE);
			return;
		}

		const lastSequence = Math.max(
			generation.agent?.sourceSequence ?? 0,
			workspace.piSession?.lastSequence ?? 0
		);
		manager = AgentSessionManager.connect({
			workspace,
			hfToken: options.hfToken,
			bridgeUrl: options.bridgeUrl,
			bridgeHeaders: options.bridgeHeaders,
			lastSequence,
		});
		const connectedManager = manager;
		startControlPump = () => {
			controlPump ??= new AgentControlPump({
				store,
				manager: connectedManager,
				workspace,
				signal: ownershipController.signal,
				onFailure: (cause) => {
					controlPumpFailure = cause;
					ownershipController.abort();
				},
			});
			return controlPump;
		};
		const startControlPumpAfterDispatch = startControlPump;
		const existingAbort = await resolveGenerationStop({
			conversationId: generation.conversationId,
			generationId: generation.generationId,
		}).catch(() => null);
		if (existingAbort) abortController.abort();
		if (!existingAbort) {
			await restorePromptAttachments(
				manager.sandbox,
				generation.agent?.prompt?.attachments,
				generation.conversationId,
				generation.generationId,
				{
					signal: ownershipController.signal,
					assertOwnership: () => store.assertPromptDispatch(workspace),
				}
			);
		}

		abortTimer = setInterval(() => {
			resolveGenerationStop({
				conversationId: generation.conversationId,
				generationId: generation.generationId,
			})
				.then((marker) => {
					if (marker && !abortController.signal.aborted) abortController.abort();
				})
				.catch(() => undefined);
		}, ABORT_POLL_MS);
		abortTimer.unref?.();

		if (ownershipController.signal.aborted) throw new AgentRelayLeaseLostError();
		const batches = abortController.signal.aborted
			? manager.resumeBatches({
					signal: abortController.signal,
					detachSignal: ownershipController.signal,
					expectedPromptId: generation.agent?.prompt?.id,
					beforeAbortDispatch: () => store.renew(),
					keepRpcOpen: true,
				})
			: generation.agent?.prompt
				? manager.runBatches(generation.agent.prompt.message, {
						signal: abortController.signal,
						detachSignal: ownershipController.signal,
						commandId: generation.agent.prompt.id,
						beforeAbortDispatch: () => store.renew(),
						beforePromptDispatch: () => store.assertPromptDispatch(workspace),
						afterPromptDispatch: () => {
							startControlPumpAfterDispatch();
						},
						keepRpcOpen: true,
					})
				: manager.resumeBatches({
						signal: abortController.signal,
						detachSignal: ownershipController.signal,
						beforeAbortDispatch: () => store.renew(),
						keepRpcOpen: true,
					});
		if (abortController.signal.aborted || !generation.agent?.prompt) startControlPump();
		for await (const batch of batches) {
			for (const event of batch.updates) {
				if (
					event.type === MessageUpdateType.Status &&
					event.status === MessageUpdateStatus.Finished
				) {
					sawFinished = true;
				}
				if (event.type === MessageUpdateType.Status && event.status === MessageUpdateStatus.Error) {
					sawError = true;
					terminalError = event.message ?? terminalError;
				}
			}
			await store.commitBridgeBatch(batch.sequence, batch.updates);
		}
		if (controlPumpFailure) throw controlPumpFailure;
		await startControlPump().sealAndStop();

		const stop = await resolveGenerationStop({
			conversationId: generation.conversationId,
			generationId: generation.generationId,
		}).catch(() => null);
		const wasAborted = abortController.signal.aborted || Boolean(stop);
		const terminalEvents: MessageUpdate[] = [];
		if (wasAborted) {
			const conversation = await collections.conversations.findOne(
				{ _id: generation.conversationId },
				{ projection: { messages: 1 } }
			);
			const message = conversation?.messages.find(
				(item) => item.id === generation.messageId && item.generationId === generation.generationId
			);
			const content = message?.content ?? "";
			const clamped = clampStoppedContent({
				content,
				initialLength: 0,
				generationId: generation.generationId,
				marker: stop,
			});
			terminalEvents.push({
				type: MessageUpdateType.FinalAnswer,
				text: clamped,
				interrupted: true,
			});
		}
		if (!sawFinished) {
			terminalEvents.push({
				type: MessageUpdateType.Status,
				status: MessageUpdateStatus.Finished,
			});
		}
		terminalEvents.push({
			type: MessageUpdateType.AgentSandbox,
			state: "idle",
			message: wasAborted ? "Agent stopped" : "Agent is ready for the next instruction",
			sandboxId: workspace.sandbox.id,
			flavor: workspace.flavor,
		});
		await store.commitSystemBatch("relay-terminal", terminalEvents);

		const status: GenerationStatus = sawError ? "error" : wasAborted ? "interrupted" : "completed";
		const now = new Date();
		const idled = await collections.workspaces.updateOne(
			{
				_id: workspace._id,
				conversationId: store.conversationId,
				retirement: { $exists: false },
				"sandbox.id": store.sandboxId,
				"sandbox.nonce": workspace.sandbox.nonce,
				"sandbox.generation": store.sandboxGeneration,
				"sandbox.status": { $in: ["ready", "idle"] },
				"execution.generationId": store.generationId,
				"execution.phase": "relay",
				"execution.owner": store.owner,
				"execution.fence": store.fence,
				"execution.expiresAt": { $gt: now },
			},
			{
				$set: {
					"sandbox.status": "idle",
					"sandbox.lastSeenAt": now,
					updatedAt: now,
				},
			}
		);
		if (idled.matchedCount !== 1) throw new AgentGenerationIncarnationMismatchError();
		finished = await store.finish(status, status === "error" ? terminalError : undefined);
		if (!finished) throw new AgentRelayLeaseLostError();
		if (wasAborted) {
			await collections.abortedGenerations
				.deleteOne({
					conversationId: generation.conversationId,
					generationId: generation.generationId,
				})
				.catch(() => undefined);
		}
	} catch (cause) {
		if (controlPumpFailure) throw controlPumpFailure;
		if (
			!(cause instanceof AgentGenerationIncarnationMismatchError) &&
			!ownershipController.signal.aborted &&
			startControlPump
		) {
			await startControlPump().sealAndStop();
		}
		if (cause instanceof AgentGenerationIncarnationMismatchError) {
			finished = await store.finishUnrecoverable("interrupted", AGENT_SANDBOX_CHANGED_MESSAGE);
			throw cause;
		}
		if (cause instanceof PiRpcReplayGapError) {
			try {
				await store.commitSystemBatch("relay-replay-gap", [
					{
						type: MessageUpdateType.Status,
						status: MessageUpdateStatus.Error,
						message: PI_RPC_REPLAY_GAP_MESSAGE,
					},
				]);
			} finally {
				finished = await finishRelayFailure(store, "error", PI_RPC_REPLAY_GAP_MESSAGE);
			}
			throw cause;
		}
		if (abortController.signal.aborted) {
			// The abort timeout/path above normally reaches terminal persistence.
			// If the bridge itself disappeared, preserve a durable error instead of
			// allowing the generic reaper to silently strand the run.
			await store
				.commitSystemBatch("relay-abort-error", [
					{
						type: MessageUpdateType.Status,
						status: MessageUpdateStatus.Error,
						message: "Agent stopped before the sandbox acknowledged the abort",
					},
				])
				.catch(() => undefined);
		}
		finished = await finishRelayFailure(
			store,
			"error",
			cause instanceof Error ? cause.message : "Agent relay failed"
		);
		throw cause;
	} finally {
		if (leaseTimer) clearInterval(leaseTimer);
		if (abortTimer) clearInterval(abortTimer);
		controlPump?.requestStop();
		manager?.close();
		await controlPump?.waitForStop().catch(() => undefined);
		if (!finished) await store.release().catch(() => undefined);
	}
}

async function finishRelayFailure(
	store: AgentGenerationStore,
	status: "interrupted" | "error",
	error: string
): Promise<boolean> {
	try {
		if (await store.finish(status, error)) return true;
	} catch {
		// A fatal bridge/incarnation failure can also make the reciprocal
		// workspace renewal inside normal finish() impossible. The exact
		// generation relay fence below is the final authority in that case.
	}
	return store.finishUnrecoverable(status, error);
}

async function restorePromptAttachments(
	sandbox: SandboxClient,
	attachments: NonNullable<NonNullable<Generation["agent"]>["prompt"]>["attachments"],
	conversationId: Generation["conversationId"],
	generationId: string,
	options: { signal: AbortSignal; assertOwnership: () => Promise<void> }
): Promise<void> {
	if (!attachments?.length) return;
	await options.assertOwnership();
	const mkdir = await sandbox.exec(
		["mkdir", "-p", `/workspace/${agentUploadDirectory(generationId)}`],
		{
			signal: options.signal,
		}
	);
	if (mkdir.exitCode !== 0 || mkdir.timedOut) {
		throw new Error("Unable to restore Agent attachments");
	}
	for (const attachment of attachments) {
		await options.assertOwnership();
		if (!isAgentUploadPath(attachment.path, generationId)) {
			throw new Error("Stored Agent attachment path is invalid");
		}
		const bytes =
			attachment.file.type === "base64"
				? Buffer.from(attachment.file.value, "base64")
				: Buffer.from((await downloadFile(attachment.file.value, conversationId)).value, "base64");
		await options.assertOwnership();
		await sandbox.writeFile(
			`/workspace/${attachment.path}`,
			Uint8Array.from(bytes).buffer,
			undefined,
			{ signal: options.signal }
		);
		await options.assertOwnership();
	}
}