webavatar-core/src/lib/server/core/session.ts

import { randomUUID } from 'node:crypto';
import WebSocket from 'ws';
import type { ClientEvent, ServerEvent, SessionStartConfig } from '../../core/protocol.js';
import { formatToolError } from '../../core/tools.js';
import type { AnalyticsRecorder } from './analytics.js';
import type { ProviderSession, RealtimeProvider, ResolvedSessionConfig } from './provider.js';
import type { ToolRegistry } from './tools.js';

export type RealtimeSessionState = 'idle' | 'connecting' | 'active' | 'closing' | 'closed';

/** Server-owned policy for a realtime session. Never client-controlled. */
export interface SessionDefaults {
	model: string;
	instructions: string;
	/** Used only when the browser's `session.start` doesn't specify a voice. */
	defaultVoice?: string;
	toolRegistry: ToolRegistry;
}

/**
 * Merges the browser's (deliberately tiny) client-visible config into the
 * server-owned defaults. `instructions`/`model`/`tools` always come from
 * `defaults` — this is the trust boundary that keeps a shared relay safe to
 * expose to untrusted browsers.
 */
export function resolveSessionConfig(
	defaults: SessionDefaults,
	clientConfig?: SessionStartConfig
): ResolvedSessionConfig {
	return {
		model: defaults.model,
		instructions: defaults.instructions,
		voice: clientConfig?.voice ?? defaults.defaultVoice,
		tools: defaults.toolRegistry.definitions
	};
}

/**
 * Owns the lifecycle of one browser WS connection: turns its first
 * `session.start` into a `ProviderSession` via the injected `RealtimeProvider`,
 * pipes events both directions afterward, and guarantees that either side
 * closing/erroring tears down the other.
 */
export class RealtimeSession {
	private state: RealtimeSessionState = 'idle';
	private providerSession: ProviderSession | null = null;
	/** Tool calls currently awaiting resolution in the current turn (server executor pending or client round-trip pending). */
	private outstandingToolCalls = 0;
	/** Client-executed tool calls forwarded to the browser, awaiting `tool.result`. */
	private readonly pendingClientToolCalls = new Map<string, NodeJS.Timeout>();
	/** Set once `start()` reaches 'active'; null before then and stays that way if the session never gets that far. */
	private sessionId: string | null = null;
	/** Populated at the top of handleToolCall (before any branching) so even the unknown-tool path is covered; consumed in resolveToolCall. */
	private readonly toolCallMeta = new Map<
		string,
		{ name: string; arguments: string; startedAt: number }
	>();
	/** Assistant voice replies stream as many transcript.delta chunks then one text-less transcript.done — see docs/ANALYTICS.md. */
	private assistantTranscriptBuffer = '';

	constructor(
		private readonly browserSocket: WebSocket,
		private readonly provider: RealtimeProvider,
		private readonly defaults: SessionDefaults,
		private readonly recorder: AnalyticsRecorder,
		/** How long a client-executed tool call may go unanswered before we give up on it. */
		private readonly clientToolTimeoutMs: number
	) {
		this.browserSocket.on('message', (data, isBinary) => {
			void this.handleBrowserMessage(data, isBinary);
		});
		this.browserSocket.on('close', () => this.close('client_disconnected'));
		this.browserSocket.on('error', () => this.close('client_disconnected'));
	}

	getState(): RealtimeSessionState {
		return this.state;
	}

	private async handleBrowserMessage(data: WebSocket.RawData, isBinary: boolean): Promise<void> {
		if (isBinary) {
			this.sendToBrowser({
				type: 'error',
				code: 'invalid_client_event',
				message: 'Binary frames are not supported yet.'
			});
			return;
		}

		let event: ClientEvent;
		try {
			event = JSON.parse(data.toString());
		} catch {
			this.sendToBrowser({
				type: 'error',
				code: 'invalid_client_event',
				message: 'Malformed JSON.'
			});
			return;
		}

		if (event.type === 'session.start') {
			await this.start(event.config);
			return;
		}
		if (event.type === 'session.end') {
			this.close('client_requested');
			return;
		}
		if (this.state !== 'active' || !this.providerSession) {
			this.sendToBrowser({
				type: 'error',
				code: 'invalid_client_event',
				message: `Cannot handle '${event.type}' before session.start.`
			});
			return;
		}
		if (event.type === 'tool.result') {
			const timeout = this.pendingClientToolCalls.get(event.callId);
			if (!timeout) return; // stale/unknown callId (e.g. already timed out) — ignore
			clearTimeout(timeout);
			this.pendingClientToolCalls.delete(event.callId);
			this.resolveToolCall(event.callId, event.output);
			return;
		}
		if (event.type === 'input_text.send') {
			// Typed text is always already a complete string — unlike the voice
			// path, no buffering is needed.
			this.recorder.messageLogged({
				sessionId: this.sessionId!,
				role: 'user',
				text: event.text,
				loggedAt: Date.now()
			});
		}
		this.providerSession.send(event);
	}

	private async start(clientConfig?: SessionStartConfig): Promise<void> {
		if (this.state !== 'idle') {
			this.sendToBrowser({
				type: 'error',
				code: 'invalid_client_event',
				message: 'Session already started.'
			});
			return;
		}
		this.state = 'connecting';
		const config = resolveSessionConfig(this.defaults, clientConfig);

		let providerSession: ProviderSession;
		try {
			providerSession = await this.provider.createSession(config);
		} catch (err) {
			this.state = 'closed';
			this.sendToBrowser({
				type: 'error',
				code: 'upstream_unavailable',
				message: err instanceof Error ? err.message : 'Failed to start session.'
			});
			this.browserSocket.close();
			return;
		}

		// The browser may have disconnected (session.end / close) while we were
		// awaiting the provider handshake — don't leak the upstream connection.
		const stateAfterConnect: RealtimeSessionState = this.state;
		if (stateAfterConnect !== 'connecting') {
			providerSession.close('client_requested');
			return;
		}

		this.providerSession = providerSession;
		const sessionId = randomUUID();
		this.sessionId = sessionId;
		this.state = 'active';
		this.recorder.sessionStarted({
			sessionId,
			model: config.model,
			voice: config.voice,
			startedAt: Date.now()
		});
		// Recorded above before wiring onEvent, so there's no window where a
		// provider event could reach handleProviderEvent before sessionId is set.
		this.providerSession.onEvent((event) => this.handleProviderEvent(event));
		this.sendToBrowser({ type: 'session.ready', sessionId });
	}

	private handleProviderEvent(event: ServerEvent): void {
		if (event.type === 'tool.call') {
			void this.handleToolCall(event);
			return;
		}
		if (event.type === 'transcript.delta') {
			if (event.role === 'assistant') {
				this.assistantTranscriptBuffer += event.text;
			} else {
				// User voice transcripts arrive as a single already-complete chunk
				// (conversation.item.input_audio_transcription.completed in
				// providers/openai.ts) with no corresponding
				// transcript.done{role:'user'} ever emitted today — log
				// immediately rather than buffering. Rests on unverified live
				// behavior; see docs/ANALYTICS.md.
				this.recorder.messageLogged({
					sessionId: this.sessionId!,
					role: 'user',
					text: event.text,
					loggedAt: Date.now()
				});
			}
			this.sendToBrowser(event);
			return;
		}
		if (event.type === 'transcript.done' && event.role === 'assistant') {
			const text = this.assistantTranscriptBuffer;
			this.assistantTranscriptBuffer = '';
			this.recorder.messageLogged({
				sessionId: this.sessionId!,
				role: 'assistant',
				text,
				loggedAt: Date.now()
			});
			this.sendToBrowser(event);
			return;
		}
		if (event.type === 'session.closed') {
			// Forward first, then close: close() moves the browser socket out of
			// OPEN, and sendToBrowser no-ops once that happens, so closing first
			// would silently swallow this notification.
			this.sendToBrowser(event);
			this.close('upstream_disconnected');
			return;
		}
		this.sendToBrowser(event);
	}

	/**
	 * Routes a tool call by its registered `location`: a server tool is
	 * executed here directly (the browser never sees it); a client tool is
	 * forwarded to the browser and tracked with a timeout so an unresponsive
	 * tab can't hang the model's turn forever.
	 */
	private async handleToolCall(event: Extract<ServerEvent, { type: 'tool.call' }>): Promise<void> {
		this.outstandingToolCalls++;
		this.toolCallMeta.set(event.callId, {
			name: event.name,
			arguments: event.arguments,
			startedAt: Date.now()
		});
		const definition = this.defaults.toolRegistry.definitions.find((d) => d.name === event.name);
		if (!definition) {
			this.resolveToolCall(event.callId, formatToolError(`Unknown tool "${event.name}".`));
			return;
		}

		if (definition.location === 'client') {
			this.sendToBrowser(event);
			const timeout = setTimeout(() => {
				this.pendingClientToolCalls.delete(event.callId);
				this.resolveToolCall(event.callId, formatToolError('Client tool call timed out.'));
			}, this.clientToolTimeoutMs);
			this.pendingClientToolCalls.set(event.callId, timeout);
			return;
		}

		let args: Record<string, unknown>;
		try {
			args = JSON.parse(event.arguments);
		} catch {
			this.resolveToolCall(event.callId, formatToolError('Malformed tool arguments.'));
			return;
		}
		try {
			const executor = this.defaults.toolRegistry.executors.get(event.name);
			const result = await executor!(args);
			this.resolveToolCall(event.callId, JSON.stringify(result));
		} catch (err) {
			this.resolveToolCall(
				event.callId,
				formatToolError(err instanceof Error ? err.message : 'Tool execution failed.')
			);
		}
	}

	/**
	 * Sends a tool's result to the provider and, once every tool call
	 * outstanding in the current turn has resolved, explicitly triggers the
	 * model to continue — see the `response.trigger` note in protocol.ts for
	 * why this can't just be bundled with each individual result.
	 */
	private resolveToolCall(callId: string, output: string): void {
		this.sendToProvider({ type: 'tool.result', callId, output });
		const meta = this.toolCallMeta.get(callId);
		if (meta) {
			// Should always be present given handleToolCall sets it for every
			// path (including the unknown-tool branch) before this can run —
			// guarded defensively rather than asserted, so a future refactor
			// that breaks the invariant fails soft here instead of crashing.
			this.toolCallMeta.delete(callId);
			this.recorder.toolCallLogged({
				sessionId: this.sessionId!,
				callId,
				name: meta.name,
				arguments: meta.arguments,
				output,
				durationMs: Date.now() - meta.startedAt,
				loggedAt: Date.now()
			});
		}
		this.outstandingToolCalls--;
		if (this.outstandingToolCalls === 0) {
			this.sendToProvider({ type: 'response.trigger' });
		}
	}

	private close(reason: string): void {
		if (this.state === 'closing' || this.state === 'closed') return;
		this.state = 'closing';
		for (const timeout of this.pendingClientToolCalls.values()) clearTimeout(timeout);
		this.pendingClientToolCalls.clear();
		this.providerSession?.close(reason);
		this.providerSession = null;
		if (this.browserSocket.readyState === WebSocket.OPEN) {
			this.browserSocket.close();
		}
		this.state = 'closed';
		if (this.sessionId) {
			// Guarded: sessionId is still null if the browser disconnected
			// mid-handshake, before start() ever reached 'active' — no sessions
			// row exists yet to update in that case.
			this.recorder.sessionEnded({ sessionId: this.sessionId, endedAt: Date.now(), reason });
		}
	}

	private sendToBrowser(event: ServerEvent): void {
		if (event.type === 'error' || event.type === 'session.closed') {
			// Logged before the delivery guard below: the event "happened" from
			// the session's point of view even if the browser socket is already
			// dead by the time we try to deliver it.
			this.recorder.eventLogged({
				sessionId: this.sessionId,
				type: event.type,
				detail: JSON.stringify(event),
				loggedAt: Date.now()
			});
		}
		if (this.browserSocket.readyState !== WebSocket.OPEN) return;
		this.browserSocket.send(JSON.stringify(event));
	}

	// Guarded the same way sendToBrowser is: tool resolution can now reach
	// the provider from an async callback (executor await, timeout) rather
	// than only from a browser message already known to be in the 'active'
	// state, so providerSession may have gone null by the time this runs.
	private sendToProvider(event: ClientEvent): void {
		this.providerSession?.send(event);
	}
}

export function createRealtimeSession(
	browserSocket: WebSocket,
	provider: RealtimeProvider,
	defaults: SessionDefaults,
	recorder: AnalyticsRecorder,
	clientToolTimeoutMs: number
): RealtimeSession {
	return new RealtimeSession(browserSocket, provider, defaults, recorder, clientToolTimeoutMs);
}