src/lib/server/agent/AgentSessionManager.ts

import type { MessageUpdate } from "$lib/types/MessageUpdate";
import { MessageUpdateStatus, MessageUpdateType } from "$lib/types/MessageUpdate";
import type { Workspace } from "$lib/types/Workspace";
import { PiEventTranslator } from "./piEvents";
import { PiRpcClient, type PiRpcEvent, type PiRpcResponse } from "./PiRpcClient";
import { SandboxClient, type SandboxConnection } from "./SandboxClient";

const BRIDGE_PORT = 4_777;
const EVENT_QUEUE_LIMIT = 5_000;
const ABORT_SETTLE_TIMEOUT_MS = 15_000;

export interface AgentSessionManagerOptions {
	workspace: Workspace;
	hfToken: string;
	fetch?: typeof fetch;
	onSequence?: (sequence: number, final: boolean) => void | Promise<void>;
	onTransportStatus?: (status: "connecting" | "connected" | "reconnecting" | "closed") => void;
	connectTimeoutMs?: number;
	commandTimeoutMs?: number;
	/** Dependency injection for the local bridge integration test. */
	bridgeUrl?: string;
	bridgeHeaders?: Record<string, string>;
	/** Override the workspace cursor when a generation-specific cursor is newer. */
	lastSequence?: number;
}

export type AgentApprovalResponse =
	{ confirmed: boolean } | { value: string } | { cancelled: true };

export interface AgentEventBatch {
	/** One sequence in the sandbox bridge's append-only replay log. */
	sequence: number;
	/** All chat-ui updates translated from that one raw pi event. */
	updates: MessageUpdate[];
	rawType?: string;
}

/**
 * Stateless server-side controller for one durable pi process. Every request
 * can reconstruct it from the persisted workspace and bridge replay cursor;
 * no process-local ownership is required for steer/stop or reconnect.
 */
export class AgentSessionManager {
	readonly sandbox: SandboxClient;
	readonly rpc: PiRpcClient;
	readonly translator: PiEventTranslator;

	private constructor(
		private readonly workspace: Workspace,
		private readonly options: AgentSessionManagerOptions,
		sandbox: SandboxClient,
		rpc: PiRpcClient
	) {
		this.sandbox = sandbox;
		this.rpc = rpc;
		this.translator = new PiEventTranslator({ model: workspace.model });
	}

	static connect(options: AgentSessionManagerOptions): AgentSessionManager {
		const connection = connectionFromWorkspace(options.workspace);
		const sandbox = SandboxClient.connect(connection, options.hfToken, { fetch: options.fetch });
		const rpc = new PiRpcClient(
			options.bridgeUrl ?? sandbox.proxyUrlFor(BRIDGE_PORT, "/ws", "wss"),
			{
				headers: options.bridgeHeaders ?? sandbox.proxyHeaders,
				lastSequence: options.lastSequence ?? options.workspace.piSession?.lastSequence ?? 0,
				connectTimeoutMs: options.connectTimeoutMs,
				commandTimeoutMs: options.commandTimeoutMs,
				onStatus: options.onTransportStatus,
			}
		);
		return new AgentSessionManager(options.workspace, options, sandbox, rpc);
	}

	get lastSequence(): number {
		return this.rpc.lastSequence;
	}

	async *run(
		message: string,
		options: {
			signal?: AbortSignal;
			detachSignal?: AbortSignal;
			expectedPromptId?: string;
			beforeAbortDispatch?: () => void | Promise<void>;
			images?: Array<Record<string, unknown>>;
		} = {}
	): AsyncGenerator<MessageUpdate> {
		for await (const batch of this.runBatches(message, options)) {
			for (const update of batch.updates) yield update;
		}
	}

	/** Start a new pi turn and retain the bridge sequence boundary for durable mirroring. */
	async *runBatches(
		message: string,
		options: {
			signal?: AbortSignal;
			/** Stops only this relay; unlike `signal`, never aborts the pi turn. */
			detachSignal?: AbortSignal;
			expectedPromptId?: string;
			beforeAbortDispatch?: () => void | Promise<void>;
			images?: Array<Record<string, unknown>>;
			/** Stable id lets the sandbox bridge suppress a retry after server death. */
			commandId?: string;
			/** Last-moment durable ownership/incarnation check after the socket connects. */
			beforePromptDispatch?: () => void | Promise<void>;
			/** Starts relay-owned side work only after the bridge accepted the prompt. */
			afterPromptDispatch?: () => void | Promise<void>;
			/** Keep the shared socket alive while the durable control mailbox is sealed. */
			keepRpcOpen?: boolean;
		} = {}
	): AsyncGenerator<AgentEventBatch> {
		if (!message.trim()) throw new Error("Agent prompt cannot be empty");
		yield* this.consumeBatches({
			...options,
			prompt: message,
			expectedPromptId: options.expectedPromptId ?? options.commandId,
		});
	}

	/**
	 * Attach to the already-running pi turn. No prompt is sent: bridge replay is
	 * consumed from the persisted sequence and then the same socket tails live
	 * events until pi settles.
	 */
	async *resumeBatches(
		options: {
			signal?: AbortSignal;
			detachSignal?: AbortSignal;
			expectedPromptId?: string;
			beforeAbortDispatch?: () => void | Promise<void>;
			/** Keep the shared socket alive while the durable control mailbox is sealed. */
			keepRpcOpen?: boolean;
		} = {}
	): AsyncGenerator<AgentEventBatch> {
		yield* this.consumeBatches(options);
	}

	async *resume(
		options: {
			signal?: AbortSignal;
			detachSignal?: AbortSignal;
			expectedPromptId?: string;
			beforeAbortDispatch?: () => void | Promise<void>;
		} = {}
	): AsyncGenerator<MessageUpdate> {
		for await (const batch of this.resumeBatches(options)) {
			for (const update of batch.updates) yield update;
		}
	}

	private async *consumeBatches(options: {
		signal?: AbortSignal;
		detachSignal?: AbortSignal;
		expectedPromptId?: string;
		beforeAbortDispatch?: () => void | Promise<void>;
		prompt?: string;
		images?: Array<Record<string, unknown>>;
		commandId?: string;
		beforePromptDispatch?: () => void | Promise<void>;
		afterPromptDispatch?: () => void | Promise<void>;
		keepRpcOpen?: boolean;
	}): AsyncGenerator<AgentEventBatch> {
		const queue = new BoundedAsyncQueue<AgentEventBatch>();
		let settled = false;
		let abortTimer: ReturnType<typeof setTimeout> | undefined;
		let abortStarted = false;
		let statsRequested = false;
		let drainedSequence = this.options.lastSequence ?? this.workspace.piSession?.lastSequence ?? 0;

		const finishAfterStats = () => {
			if (statsRequested) return;
			statsRequested = true;
			void this.rpc
				.command({ type: "get_session_stats" }, { timeoutMs: 10_000 })
				.catch(() => undefined)
				.finally(() => queue.end());
		};

		const unsubscribe = this.rpc.subscribe((event, sequence) => {
			const updates: MessageUpdate[] = [];
			if (event.type === "bridge_stderr") {
				updates.push(...this.translator.translateBridgeStderr(event.data));
			}
			if (event.type === "bridge_error") {
				updates.push({
					type: MessageUpdateType.Status,
					status: MessageUpdateStatus.Error,
					message:
						typeof event.message === "string" ? event.message : "Agent bridge reported an error",
				});
			}
			updates.push(...this.translator.translate(event));
			queue.push({
				updates,
				sequence,
				...(typeof event.type === "string" ? { rawType: event.type } : {}),
			});

			if (event.type === "agent_settled") {
				settled = true;
				if (abortTimer) clearTimeout(abortTimer);
				finishAfterStats();
			} else if (event.type === "bridge_exit") {
				queue.fail(new Error("The agent process stopped unexpectedly"));
			}
		});
		const unsubscribeReplayGap = this.rpc.subscribeReplayGap((error) => queue.fail(error));

		const abort = () => {
			if (settled || abortStarted) return;
			abortStarted = true;
			void (async () => {
				try {
					await options.beforeAbortDispatch?.();
					if (options.detachSignal?.aborted) throw agentDetachedError();
					await this.rpc.command(
						{
							type: "abort",
							...(options.expectedPromptId
								? { _chatUiExpectedPromptId: options.expectedPromptId }
								: {}),
						},
						{ timeoutMs: 5_000 }
					);
					abortTimer = setTimeout(() => queue.end(), ABORT_SETTLE_TIMEOUT_MS);
				} catch (cause) {
					if (
						cause instanceof Error &&
						cause.message.includes("no longer belongs to the active prompt")
					) {
						queue.end();
						return;
					}
					queue.fail(cause instanceof Error ? cause : new Error("Unable to stop Agent run"));
				}
			})();
		};
		const detach = () => queue.fail(agentDetachedError());

		try {
			await this.rpc.connect();
			if (options.detachSignal?.aborted) throw agentDetachedError();
			options.detachSignal?.addEventListener("abort", detach, { once: true });
			if (options.signal?.aborted) {
				if (options.prompt !== undefined) throw agentAbortError();
				abort();
			} else {
				options.signal?.addEventListener("abort", abort, { once: true });
			}

			if (options.prompt !== undefined) {
				await options.beforePromptDispatch?.();
				if (options.detachSignal?.aborted) throw agentDetachedError();
				if (options.signal?.aborted) throw agentAbortError();
				await this.rpc.command({
					type: "prompt",
					...(options.commandId ? { id: options.commandId } : {}),
					message: options.prompt,
					...(options.images?.length ? { images: options.images } : {}),
				});
				await options.afterPromptDispatch?.();
			} else if (!settled) {
				// If the persisted cursor already covered agent_settled there is no
				// terminal replay record left to close this consumer. Query pi only to
				// decide whether to tail; this is observational and never starts work.
				const state = await this.rpc
					.command<Record<string, unknown>>({ type: "get_state" }, { timeoutMs: 10_000 })
					.catch(() => undefined);
				if (state && !stateIndicatesRunning(state.data)) finishAfterStats();
			}

			for await (const batch of queue) {
				yield batch;
				drainedSequence = batch.sequence;
				await this.options.onSequence?.(drainedSequence, false);
			}
		} finally {
			if (abortTimer) clearTimeout(abortTimer);
			options.signal?.removeEventListener("abort", abort);
			options.detachSignal?.removeEventListener("abort", detach);
			unsubscribe();
			unsubscribeReplayGap();
			await this.options.onSequence?.(drainedSequence, true);
			if (!options.keepRpcOpen) this.rpc.close();
		}
	}

	prompt(message: string, streamingBehavior?: "steer" | "followUp"): Promise<PiRpcResponse> {
		return this.rpc.command({
			type: "prompt",
			message,
			...(streamingBehavior ? { streamingBehavior } : {}),
		});
	}

	steer(message: string): Promise<PiRpcResponse> {
		return this.rpc.command({ type: "steer", message });
	}

	followUp(message: string): Promise<PiRpcResponse> {
		return this.rpc.command({ type: "follow_up", message });
	}

	abort(): Promise<PiRpcResponse> {
		return this.rpc.command({ type: "abort" });
	}

	setModel(provider: string, modelId: string): Promise<PiRpcResponse> {
		return this.rpc.command({ type: "set_model", provider, modelId });
	}

	setThinkingLevel(level: Workspace["thinkingLevel"]): Promise<PiRpcResponse> {
		return this.rpc.command({ type: "set_thinking_level", level });
	}

	compact(): Promise<PiRpcResponse> {
		return this.rpc.command({ type: "compact" }, { timeoutMs: 120_000 });
	}

	getSessionStats(): Promise<PiRpcResponse> {
		return this.rpc.command({ type: "get_session_stats" });
	}

	respondToApproval(requestId: string, response: AgentApprovalResponse): Promise<void> {
		return this.rpc.send({ type: "extension_ui_response", id: requestId, ...response });
	}

	close(): void {
		this.rpc.close();
	}
}

function agentAbortError(): Error {
	const error = new Error("Request was aborted.");
	error.name = "AbortError";
	return error;
}

function agentDetachedError(): Error {
	const error = new Error("Agent relay ownership moved to another server");
	error.name = "AgentSessionDetachedError";
	return error;
}

function connectionFromWorkspace(workspace: Workspace): SandboxConnection {
	const sandbox = workspace.sandbox;
	if (!sandbox.id || !sandbox.namespace || !sandbox.baseUrl || !sandbox.nonce || !sandbox.image) {
		throw new Error("Workspace sandbox connection is incomplete");
	}
	if (sandbox.status !== "ready" && sandbox.status !== "idle") {
		throw new Error(`Workspace sandbox is not available (${sandbox.status})`);
	}
	return {
		id: sandbox.id,
		namespace: sandbox.namespace,
		baseUrl: sandbox.baseUrl,
		nonce: sandbox.nonce,
		image: sandbox.image,
		flavor: sandbox.flavor,
		createdAt: sandbox.createdAt ?? workspace.createdAt,
	};
}

class BoundedAsyncQueue<T> implements AsyncIterable<T> {
	private readonly values: T[] = [];
	private readonly waiters: Array<{
		resolve: (value: IteratorResult<T>) => void;
		reject: (reason: Error) => void;
	}> = [];
	private ended = false;
	private failure?: Error;

	push(value: T): void {
		if (this.ended || this.failure) return;
		const waiter = this.waiters.shift();
		if (waiter) {
			waiter.resolve({ value, done: false });
			return;
		}
		if (this.values.length >= EVENT_QUEUE_LIMIT) {
			this.fail(new Error("Agent event queue overflowed"));
			return;
		}
		this.values.push(value);
	}

	end(): void {
		if (this.ended || this.failure) return;
		this.ended = true;
		for (const waiter of this.waiters.splice(0)) waiter.resolve({ value: undefined, done: true });
	}

	fail(error: Error): void {
		if (this.ended || this.failure) return;
		this.failure = error;
		for (const waiter of this.waiters.splice(0)) waiter.reject(error);
	}

	[Symbol.asyncIterator](): AsyncIterator<T> {
		return { next: () => this.next() };
	}

	private next(): Promise<IteratorResult<T>> {
		const value = this.values.shift();
		if (value !== undefined) return Promise.resolve({ value, done: false });
		if (this.failure) return Promise.reject(this.failure);
		if (this.ended) return Promise.resolve({ value: undefined, done: true });
		return new Promise((resolve, reject) => this.waiters.push({ resolve, reject }));
	}
}

export function isAgentSettledEvent(event: PiRpcEvent): boolean {
	return event.type === "agent_settled";
}

function stateIndicatesRunning(data: unknown): boolean {
	if (!data || typeof data !== "object" || Array.isArray(data)) return true;
	const state = data as Record<string, unknown>;
	for (const key of ["isStreaming", "isRunning", "agentRunning", "isAgentRunning"]) {
		if (state[key] === true) return true;
		if (state[key] === false) return false;
	}
	// Unknown pi versions fail open: continuing to tail is safer than declaring a
	// live turn complete just because a diagnostic field was renamed.
	return true;
}