dev/third_party/pi/packages/server/src/session-router.ts

import { randomUUID } from "node:crypto";
import type { JsonValue, ServiceCall, ServiceProviderUpdate } from "@earendil-works/chord";
import { BACKGROUND_CONTEXT, type Context, type SessionMetadata } from "@earendil-works/pi-agent-core";
import type { RpcTarget, SessionTarget } from "@earendil-works/pi-protocol";
import { ServerDrainingError, SessionNotAttachedError } from "./errors.ts";
import type { RoutedSessionAttachment, RoutedSessionHandle, ServerHost } from "./types.ts";

class SessionCleanupError extends AggregateError {}

interface ClientAttachment {
	readonly id: string;
	readonly client: object;
	readonly session: HostedSession;
	readonly operations: Set<Promise<unknown>>;
	acquiring?: Promise<RoutedSessionAttachment>;
	lease?: RoutedSessionAttachment;
	releasing?: Promise<void>;
}

interface HostedSession {
	readonly id: string;
	readonly handle: RoutedSessionHandle;
	readonly attachments: Set<ClientAttachment>;
}

interface SessionRouterOptions<TMetadata extends SessionMetadata> {
	host: ServerHost<TMetadata>;
	serverId: string;
	isClosing: () => boolean;
	publishAttachment(client: object, attachment: SessionTarget | undefined, context: Context): Promise<void>;
	reportError: (error: unknown) => void;
}

export class SessionRouter<TMetadata extends SessionMetadata = SessionMetadata> {
	private readonly options: SessionRouterOptions<TMetadata>;
	private readonly hostedSessions = new Map<string, HostedSession>();
	private readonly openingSessions = new Map<string, Promise<HostedSession>>();
	private readonly attachmentsByClient = new Map<object, ClientAttachment>();
	private readonly disconnectedClients = new Set<object>();
	private readonly clientOperations = new Map<object, Promise<void>>();
	private closePromise?: Promise<void>;

	constructor(options: SessionRouterOptions<TMetadata>) {
		this.options = options;
	}

	async executeServiceCall(
		call: ServiceCall,
		target: RpcTarget,
		client: object,
		publish: (subscriptionId: string, update: ServiceProviderUpdate, context: Context) => Promise<void>,
		context: Context,
	): Promise<JsonValue | undefined> {
		const admitted = await this.runForClient(client, () =>
			this.startServiceCall(client, target, call, publish, context),
		);
		return admitted.result;
	}

	attachClient(client: object, sessionId: string, context: Context): Promise<void> {
		if (this.options.isClosing()) return Promise.reject(new ServerDrainingError());
		return this.runForClient(client, () => this.attachClientNow(client, sessionId, context));
	}

	detachClient(client: object, context: Context): Promise<void> {
		return this.runForClient(client, async () => {
			const attachment = this.attachmentsByClient.get(client);
			if (attachment) await this.releaseAttachment(attachment, context);
		});
	}

	async removeSession(sessionId: string, context: Context): Promise<void> {
		if (this.options.isClosing()) throw new ServerDrainingError();
		const hosted = this.hostedSessions.get(sessionId);
		if (hosted === undefined) return;
		const errors: unknown[] = [];
		const releases = await Promise.allSettled(
			[...hosted.attachments].map((attachment) => this.releaseAttachment(attachment, context)),
		);
		for (const result of releases) if (result.status === "rejected") errors.push(result.reason);
		try {
			await hosted.handle.close(context);
		} catch (error) {
			errors.push(error);
		}
		if (this.hostedSessions.get(sessionId) === hosted) this.hostedSessions.delete(sessionId);
		if (errors.length === 1) throw errors[0];
		if (errors.length > 1) throw new AggregateError(errors, `Failed to close Session ${sessionId}`);
	}

	async disconnect(client: object, context: Context): Promise<void> {
		this.disconnectedClients.add(client);
		try {
			await this.runForClient(client, async () => {
				const attachment = this.attachmentsByClient.get(client);
				if (attachment) await this.releaseAttachment(attachment, context, false);
			});
		} finally {
			this.disconnectedClients.delete(client);
		}
	}

	close(context: Context): Promise<void> {
		this.closePromise ??= this.closeInternal(context);
		return this.closePromise;
	}

	private async closeInternal(context: Context): Promise<void> {
		const operationPromises = [...this.clientOperations.values()];
		const openingPromises = [...this.openingSessions.values()];
		const [operationResults, openingResults] = await Promise.all([
			Promise.allSettled(operationPromises),
			Promise.allSettled(openingPromises),
		]);
		const closeErrors: unknown[] = [];
		for (const result of [...operationResults, ...openingResults]) {
			if (result.status !== "rejected") continue;
			this.options.reportError(result.reason);
			if (result.reason instanceof SessionCleanupError) closeErrors.push(result.reason);
		}
		const attachmentResults = await Promise.allSettled(
			[...this.hostedSessions.values()].flatMap((session) =>
				[...session.attachments].map((attachment) => this.releaseAttachment(attachment, context)),
			),
		);
		for (const result of attachmentResults) {
			if (result.status === "rejected") closeErrors.push(result.reason);
		}
		const hosted = [...this.hostedSessions.values()];
		const closeResults = await Promise.allSettled(hosted.map(({ handle }) => handle.close(context)));
		for (let index = 0; index < closeResults.length; index++) {
			const result = closeResults[index]!;
			const session = hosted[index]!;
			if (result.status === "fulfilled") {
				if (this.hostedSessions.get(session.id) === session) this.hostedSessions.delete(session.id);
				continue;
			}
			this.options.reportError(result.reason);
			closeErrors.push(result.reason);
		}
		this.attachmentsByClient.clear();
		this.clientOperations.clear();
		if (closeErrors.length > 0) throw new AggregateError(closeErrors, "Failed to close routed Sessions");
	}

	private runForClient<T>(client: object, operation: () => Promise<T>): Promise<T> {
		const previous = this.clientOperations.get(client) ?? Promise.resolve();
		const result = previous.catch(() => {}).then(operation);
		const tail = result.then(
			() => undefined,
			() => undefined,
		);
		this.clientOperations.set(client, tail);
		void tail.finally(() => {
			if (this.clientOperations.get(client) === tail) this.clientOperations.delete(client);
		});
		return result;
	}

	private async attachClientNow(client: object, sessionId: string, context: Context): Promise<void> {
		if (this.options.isClosing() || this.disconnectedClients.has(client)) throw new ServerDrainingError();
		const current = this.attachmentsByClient.get(client);
		if (current?.session.id === sessionId) return;
		const hosted = await this.acquire(sessionId, context);
		if (this.options.isClosing() || this.disconnectedClients.has(client)) throw new ServerDrainingError();
		if (current) await this.releaseAttachment(current, context, false);
		const attachment: ClientAttachment = {
			id: randomUUID(),
			client,
			session: hosted,
			operations: new Set(),
		};
		hosted.attachments.add(attachment);
		try {
			const acquiring = Promise.resolve(hosted.handle.attachClient(context));
			attachment.acquiring = acquiring;
			attachment.lease = await acquiring;
		} catch (error) {
			hosted.attachments.delete(attachment);
			throw error;
		}
		if (
			this.hostedSessions.get(hosted.id) !== hosted ||
			!hosted.attachments.has(attachment) ||
			this.disconnectedClients.has(client) ||
			this.options.isClosing()
		) {
			await this.releaseAttachment(attachment, context);
			throw new ServerDrainingError();
		}
		this.attachmentsByClient.set(client, attachment);
		await this.options.publishAttachment(
			client,
			{ serverId: this.options.serverId, sessionId, attachmentId: attachment.id },
			context,
		);
	}

	private async startServiceCall(
		client: object,
		target: RpcTarget,
		call: ServiceCall,
		publish: (subscriptionId: string, update: ServiceProviderUpdate, context: Context) => Promise<void>,
		context: Context,
	): Promise<{ result: Promise<JsonValue | undefined> }> {
		const attachment = this.requireAttachment(client, target);
		const result = attachment.lease!.invokeService(
			call,
			(subscriptionId, update, updateContext) => publish(subscriptionId, update, updateContext),
			context,
		);
		this.trackOperation(attachment, result);
		return { result };
	}

	private trackOperation(attachment: ClientAttachment, result: Promise<unknown>): void {
		attachment.operations.add(result);
		const remove = (): void => {
			attachment.operations.delete(result);
		};
		void result.then(remove, remove);
	}

	private requireAttachment(client: object, target: RpcTarget): ClientAttachment {
		if (this.options.isClosing() || this.disconnectedClients.has(client)) throw new ServerDrainingError();
		if (!("sessionId" in target)) throw new SessionNotAttachedError();
		const attachment = this.attachmentsByClient.get(client);
		if (!attachment || attachment.session.id !== target.sessionId || attachment.id !== target.attachmentId) {
			throw new SessionNotAttachedError();
		}
		return attachment;
	}

	private releaseAttachment(attachment: ClientAttachment, context: Context, publish = true): Promise<void> {
		attachment.releasing ??= (async () => {
			const errors: unknown[] = [];
			try {
				await Promise.allSettled(attachment.operations);
				try {
					const lease = attachment.lease ?? (await attachment.acquiring);
					if (lease) await lease.release(context);
				} catch (error) {
					errors.push(error);
				}
				if (errors.length === 1) throw errors[0];
				if (errors.length > 1) throw new AggregateError(errors, "Failed to release Session attachment");
			} finally {
				await this.clearAttachment(attachment, context, publish);
			}
		})();
		return attachment.releasing;
	}

	private async clearAttachment(attachment: ClientAttachment, context: Context, publish: boolean): Promise<void> {
		attachment.session.attachments.delete(attachment);
		if (this.attachmentsByClient.get(attachment.client) === attachment) {
			this.attachmentsByClient.delete(attachment.client);
			if (publish) await this.options.publishAttachment(attachment.client, undefined, context);
		}
	}

	private async acquire(sessionId: string, context: Context): Promise<HostedSession> {
		const existing = this.hostedSessions.get(sessionId);
		if (existing) return existing;
		const opening = this.openingSessions.get(sessionId);
		if (opening) return opening;
		const pending = this.open(sessionId, context);
		this.openingSessions.set(sessionId, pending);
		try {
			return await pending;
		} finally {
			if (this.openingSessions.get(sessionId) === pending) this.openingSessions.delete(sessionId);
		}
	}

	private async open(sessionId: string, context: Context): Promise<HostedSession> {
		const metadata = await this.options.host.resolveSession(sessionId, context);
		const handle = await this.options.host.openSession(metadata, context);
		if (this.options.isClosing()) {
			try {
				await handle.close(context);
			} catch (error) {
				this.options.reportError(error);
				throw new SessionCleanupError(
					[new ServerDrainingError(), error],
					"Failed to close routed Session acquired while draining",
				);
			}
			throw new ServerDrainingError();
		}
		const hosted: HostedSession = { id: metadata.id, handle, attachments: new Set() };
		this.hostedSessions.set(hosted.id, hosted);
		if (handle.terminated) {
			void handle.terminated.then(
				(error) => this.invalidate(hosted, error),
				(error: unknown) => this.invalidate(hosted, error instanceof Error ? error : new Error(String(error))),
			);
		}
		return hosted;
	}

	private invalidate(hosted: HostedSession, error: Error | undefined): void {
		if (this.hostedSessions.get(hosted.id) !== hosted) return;
		this.hostedSessions.delete(hosted.id);
		for (const attachment of hosted.attachments) {
			void this.releaseAttachment(attachment, BACKGROUND_CONTEXT).catch((releaseError: unknown) =>
				this.options.reportError(releaseError),
			);
		}
		if (error) this.options.reportError(error);
	}
}