@agentuity/core 1.0.32 → 1.0.34

package/src/services/queue/websocket.ts

@@ -0,0 +1,554 @@
+/**
+ * @module websocket
+ *
+ * WebSocket client for real-time queue message subscriptions.
+ *
+ * Provides both a callback-based API ({@link createQueueWebSocket}) and an
+ * async iterator API ({@link subscribeToQueue}) for receiving messages from
+ * a queue in real time via WebSocket.
+ *
+ * @example Callback-based API
+ * ```typescript
+ * import { createQueueWebSocket } from '@agentuity/server';
+ *
+ * const connection = createQueueWebSocket({
+ *     queueName: 'order-processing',
+ *     baseUrl: 'https://catalyst.agentuity.cloud',
+ *     onMessage: (message) => {
+ *         console.log('Received:', message.id, message.payload);
+ *     },
+ *     onOpen: () => console.log('Connected'),
+ *     onClose: (code, reason) => console.log('Closed:', code, reason),
+ *     onError: (error) => console.error('Error:', error),
+ * });
+ *
+ * // Later: close the connection
+ * connection.close();
+ * ```
+ *
+ * @example Resuming from a previous session
+ * ```typescript
+ * import { createQueueWebSocket } from '@agentuity/server';
+ *
+ * // Use a previously obtained clientId and lastOffset to resume
+ * const connection = createQueueWebSocket({
+ *     queueName: 'order-processing',
+ *     baseUrl: 'https://catalyst.agentuity.cloud',
+ *     clientId: previousClientId,
+ *     lastOffset: previousOffset,
+ *     onMessage: (message) => {
+ *         console.log('Received:', message.id, message.payload);
+ *     },
+ * });
+ * ```
+ *
+ * @example Async iterator API
+ * ```typescript
+ * import { subscribeToQueue } from '@agentuity/server';
+ *
+ * const controller = new AbortController();
+ * for await (const message of subscribeToQueue({
+ *     queueName: 'order-processing',
+ *     baseUrl: 'https://catalyst.agentuity.cloud',
+ *     signal: controller.signal,
+ * })) {
+ *     console.log('Received:', message.id, message.payload);
+ * }
+ * ```
+ */
+
+import { z } from 'zod';
+import type { Message } from './types.ts';
+import { WebSocketAuthResponseSchema, WebSocketMessageSchema } from './types.ts';
+import { getEnv } from '../env.ts';
+import { QueueError } from './util.ts';
+import { validateQueueName } from './validation.ts';
+
+// ============================================================================
+// Types
+// ============================================================================
+
+/** Connection state for a queue WebSocket connection. */
+export type QueueWebSocketState =
+	| 'connecting'
+	| 'authenticating'
+	| 'connected'
+	| 'reconnecting'
+	| 'closed';
+
+/** Options for creating a queue WebSocket subscription. */
+export const QueueWebSocketOptionsSchema = z.object({
+	queueName: z.string().describe('Queue name to subscribe to'),
+	apiKey: z
+		.string()
+		.optional()
+		.describe('Optional API key override; falls back to AGENTUITY_SDK_KEY'),
+	baseUrl: z.string().describe('Base Catalyst URL used to construct the WebSocket endpoint'),
+	onMessage: z
+		.custom<(message: Message) => void>()
+		.describe('Callback invoked for each received queue message'),
+	onOpen: z.custom<() => void>().optional().describe('Callback invoked after auth succeeds'),
+	onClose: z
+		.custom<(code: number, reason: string) => void>()
+		.optional()
+		.describe('Callback invoked when the socket closes'),
+	onError: z
+		.custom<(error: Error) => void>()
+		.optional()
+		.describe('Callback invoked when an error occurs'),
+	autoReconnect: z
+		.boolean()
+		.optional()
+		.describe('Whether to automatically reconnect after disconnection'),
+	maxReconnectAttempts: z
+		.number()
+		.optional()
+		.describe('Maximum reconnect attempts before giving up'),
+	reconnectDelayMs: z.number().optional().describe('Initial reconnect delay in milliseconds'),
+	maxReconnectDelayMs: z.number().optional().describe('Maximum reconnect delay in milliseconds'),
+	clientId: z.string().optional().describe('Optional prior client id used for resume semantics'),
+	lastOffset: z
+		.number()
+		.optional()
+		.describe('Optional last processed offset used to resume stream consumption'),
+	orgId: z
+		.string()
+		.optional()
+		.describe('Optional org id for API keys requiring explicit org scope'),
+});
+
+export type QueueWebSocketOptions = z.infer<typeof QueueWebSocketOptionsSchema>;
+
+/** Return type from {@link createQueueWebSocket}. */
+export interface QueueWebSocketConnection {
+	/** Close the WebSocket connection. Disables auto-reconnect. */
+	close(): void;
+	/** The current connection state. */
+	readonly state: QueueWebSocketState;
+	/** The client/subscription ID assigned by the server. Stable across reconnections. */
+	readonly clientId: string | undefined;
+	/** The offset of the last message processed. */
+	readonly lastOffset: number | undefined;
+}
+
+/** Options for the async iterator queue subscription. */
+export const SubscribeToQueueOptionsSchema = z.object({
+	queueName: z.string().describe('Queue name to subscribe to'),
+	apiKey: z
+		.string()
+		.optional()
+		.describe('Optional API key override; falls back to AGENTUITY_SDK_KEY'),
+	baseUrl: z.string().describe('Base Catalyst URL used to construct the WebSocket endpoint'),
+	signal: z.custom<AbortSignal>().optional().describe('AbortSignal used to stop the subscription'),
+	clientId: z.string().optional().describe('Optional prior client id used for resume semantics'),
+	lastOffset: z
+		.number()
+		.optional()
+		.describe('Optional last processed offset used to resume stream consumption'),
+	orgId: z
+		.string()
+		.optional()
+		.describe('Optional org id for API keys requiring explicit org scope'),
+});
+
+export type SubscribeToQueueOptions = z.infer<typeof SubscribeToQueueOptionsSchema>;
+
+// ============================================================================
+// Internal Helpers
+// ============================================================================
+
+/**
+ * Resolve the API key from the options or the AGENTUITY_SDK_KEY environment variable.
+ * Throws a {@link QueueError} if no API key is available.
+ */
+function resolveApiKey(apiKey?: string): string {
+	const key = apiKey ?? getEnv('AGENTUITY_SDK_KEY');
+	if (!key) {
+		throw new QueueError({
+			message:
+				'No API key provided. Pass apiKey in options or set the AGENTUITY_SDK_KEY environment variable.',
+		});
+	}
+	return key;
+}
+
+/**
+ * Convert an HTTP(S) base URL to a WebSocket URL and append the queue path.
+ *
+ * The WebSocket route is registered at `/queue/ws/{name}` (not versioned).
+ */
+function buildWebSocketUrl(baseUrl: string, queueName: string): string {
+	const wsUrl = baseUrl.replace(/^https:\/\//, 'wss://').replace(/^http:\/\//, 'ws://');
+	// Remove trailing slash if present
+	const base = wsUrl.replace(/\/$/, '');
+	return `${base}/queue/ws/${encodeURIComponent(queueName)}`;
+}
+
+// ============================================================================
+// Callback-based API
+// ============================================================================
+
+/**
+ * Create a WebSocket connection to receive real-time messages from a queue.
+ *
+ * The connection handles authentication, automatic reconnection with exponential
+ * backoff, and ping/pong keep-alive (handled automatically by the WebSocket
+ * implementation).
+ *
+ * @param options - Configuration for the WebSocket connection
+ * @returns A {@link QueueWebSocketConnection} handle for managing the connection
+ * @throws {QueueError} If no API key is available
+ * @throws {QueueValidationError} If the queue name is invalid
+ *
+ * @example
+ * ```typescript
+ * const connection = createQueueWebSocket({
+ *     queueName: 'order-processing',
+ *     baseUrl: 'https://catalyst.agentuity.cloud',
+ *     onMessage: (message) => {
+ *         console.log('Received:', message.id, message.payload);
+ *     },
+ * });
+ *
+ * // Later: close the connection
+ * connection.close();
+ * ```
+ */
+export function createQueueWebSocket(options: QueueWebSocketOptions): QueueWebSocketConnection {
+	// Validate inputs eagerly so callers get immediate feedback.
+	validateQueueName(options.queueName);
+	const apiKey = resolveApiKey(options.apiKey);
+
+	const {
+		queueName,
+		baseUrl,
+		onMessage,
+		onOpen,
+		onClose,
+		onError,
+		autoReconnect = true,
+		maxReconnectAttempts = Infinity,
+		reconnectDelayMs = 1000,
+		maxReconnectDelayMs = 30000,
+		orgId,
+	} = options;
+
+	let state: QueueWebSocketState = 'connecting';
+	let ws: WebSocket | null = null;
+	let intentionallyClosed = false;
+	let reconnectAttempts = 0;
+	let reconnectTimer: ReturnType<typeof setTimeout> | null = null;
+	let clientId: string | undefined = options.clientId;
+	let lastProcessedOffset: number | undefined = options.lastOffset;
+
+	function connect() {
+		if (intentionallyClosed) return;
+
+		const url = buildWebSocketUrl(baseUrl, queueName);
+		state = reconnectAttempts > 0 ? 'reconnecting' : 'connecting';
+
+		try {
+			ws = new WebSocket(url);
+		} catch (err) {
+			state = 'closed';
+			onError?.(
+				new QueueError({
+					message: `Failed to create WebSocket connection: ${err instanceof Error ? err.message : String(err)}`,
+					queueName,
+					cause: err instanceof Error ? err : undefined,
+				})
+			);
+			scheduleReconnect();
+			return;
+		}
+
+		ws.onopen = () => {
+			state = 'authenticating';
+			// Send auth message — raw API key, no "Bearer " prefix.
+			// Include client_id and last_offset on reconnect for resumption.
+			const authPayload: Record<string, unknown> = { authorization: apiKey, orgId };
+			if (clientId) {
+				authPayload.client_id = clientId;
+			}
+			if (lastProcessedOffset !== undefined) {
+				authPayload.last_offset = lastProcessedOffset;
+			}
+			ws!.send(JSON.stringify(authPayload));
+		};
+
+		/** Whether the auth handshake has completed successfully. */
+		let authenticated = false;
+
+		ws.onmessage = (event: MessageEvent) => {
+			const raw = typeof event.data === 'string' ? event.data : String(event.data);
+
+			if (!authenticated) {
+				// First message after open must be the auth response.
+				try {
+					const parsed = JSON.parse(raw);
+					const authResult = WebSocketAuthResponseSchema.safeParse(parsed);
+					if (!authResult.success) {
+						const err = new QueueError({
+							message: `Unexpected auth response from server: ${raw}`,
+							queueName,
+						});
+						onError?.(err);
+						ws?.close(4000, 'Invalid auth response');
+						return;
+					}
+
+					if (!authResult.data.success) {
+						const err = new QueueError({
+							message: `Authentication failed: ${authResult.data.error ?? 'Unknown error'}`,
+							queueName,
+						});
+						onError?.(err);
+						// Auth rejection is terminal — do not reconnect with the same bad credentials.
+						intentionallyClosed = true;
+						ws?.close(4001, 'Auth failed');
+						return;
+					}
+
+					authenticated = true;
+					reconnectAttempts = 0; // Reset on successful auth.
+					state = 'connected';
+					if (authResult.data.client_id) {
+						clientId = authResult.data.client_id;
+					}
+					onOpen?.();
+				} catch {
+					const err = new QueueError({
+						message: `Failed to parse auth response: ${raw}`,
+						queueName,
+					});
+					onError?.(err);
+					ws?.close(4000, 'Invalid auth response');
+				}
+				return;
+			}
+
+			// Normal message after authentication.
+			try {
+				const parsed = JSON.parse(raw);
+				const msgResult = WebSocketMessageSchema.safeParse(parsed);
+				if (msgResult.success && msgResult.data.messages.length > 0) {
+					for (const msg of msgResult.data.messages) {
+						onMessage(msg);
+						if (msg.offset !== undefined) {
+							lastProcessedOffset = msg.offset;
+						}
+					}
+				}
+			} catch {
+				// Non-JSON frames are silently ignored; the server may send
+				// ping text frames that are not JSON.
+			}
+		};
+
+		ws.onclose = (event: CloseEvent) => {
+			state = 'closed';
+			ws = null;
+
+			onClose?.(event.code, event.reason);
+
+			// Reconnect on any unintentional close — whether we were fully
+			// connected, mid-auth, or never authenticated (transient network issue).
+			if (!intentionallyClosed) {
+				scheduleReconnect();
+			}
+		};
+
+		ws.onerror = () => {
+			// The browser/Node WebSocket fires `error` then `close`.
+			// We report the error but let `onclose` handle reconnection.
+			onError?.(
+				new QueueError({
+					message: 'WebSocket connection error',
+					queueName,
+				})
+			);
+		};
+	}
+
+	function scheduleReconnect() {
+		if (intentionallyClosed || !autoReconnect) return;
+		if (reconnectAttempts >= maxReconnectAttempts) {
+			onError?.(
+				new QueueError({
+					message: `Exceeded maximum reconnection attempts (${maxReconnectAttempts})`,
+					queueName,
+				})
+			);
+			return;
+		}
+
+		// Exponential backoff with jitter, capped at maxReconnectDelayMs.
+		const baseDelay = reconnectDelayMs * Math.pow(2, reconnectAttempts);
+		const jitter = 0.5 + Math.random() * 0.5;
+		const delay = Math.min(Math.floor(baseDelay * jitter), maxReconnectDelayMs);
+
+		reconnectAttempts++;
+		state = 'reconnecting';
+		reconnectTimer = setTimeout(() => {
+			reconnectTimer = null;
+			connect();
+		}, delay);
+	}
+
+	// Kick off the initial connection.
+	connect();
+
+	return {
+		close() {
+			intentionallyClosed = true;
+			if (reconnectTimer !== null) {
+				clearTimeout(reconnectTimer);
+				reconnectTimer = null;
+			}
+			if (ws) {
+				ws.close(1000, 'Client closed');
+				ws = null;
+			}
+			state = 'closed';
+		},
+		get state() {
+			return state;
+		},
+		get clientId() {
+			return clientId;
+		},
+		get lastOffset() {
+			return lastProcessedOffset;
+		},
+	};
+}
+
+// ============================================================================
+// Async Iterator API
+// ============================================================================
+
+/**
+ * Subscribe to real-time messages from a queue via WebSocket.
+ *
+ * Returns an async iterator that yields messages as they arrive.
+ * The connection is automatically managed (auth, reconnection, cleanup).
+ *
+ * @param options - Configuration for the subscription
+ * @returns An async generator that yields {@link Message} objects
+ * @throws {QueueError} If no API key is available
+ * @throws {QueueValidationError} If the queue name is invalid
+ *
+ * @example
+ * ```typescript
+ * const controller = new AbortController();
+ * for await (const message of subscribeToQueue({
+ *     queueName: 'order-processing',
+ *     baseUrl: 'https://catalyst.agentuity.cloud',
+ *     signal: controller.signal,
+ * })) {
+ *     console.log('Received:', message.id, message.payload);
+ * }
+ * ```
+ */
+export async function* subscribeToQueue(
+	options: SubscribeToQueueOptions
+): AsyncGenerator<Message, void, unknown> {
+	const { signal } = options;
+
+	// Check if already aborted.
+	if (signal?.aborted) return;
+
+	// A queue for buffering messages between the WebSocket callbacks and the
+	// async iterator consumer.
+	const buffer: Message[] = [];
+	let resolve: (() => void) | null = null;
+	let done = false;
+	let lastError: Error | null = null;
+
+	function push(message: Message) {
+		buffer.push(message);
+		if (resolve) {
+			resolve();
+			resolve = null;
+		}
+	}
+
+	function finish(error?: Error) {
+		done = true;
+		if (error) lastError = error;
+		if (resolve) {
+			resolve();
+			resolve = null;
+		}
+	}
+
+	const connection = createQueueWebSocket({
+		queueName: options.queueName,
+		apiKey: options.apiKey,
+		baseUrl: options.baseUrl,
+		clientId: options.clientId,
+		lastOffset: options.lastOffset,
+		orgId: options.orgId,
+		onMessage: push,
+		onError: (err) => {
+			// Terminal errors should stop the iterator.
+			// Auth failures and max-reconnect-exceeded are terminal.
+			// Transient errors (connection drops, server shutdown) are handled
+			// by the reconnection logic in createQueueWebSocket.
+			const msg = err instanceof Error ? err.message : String(err);
+			if (
+				msg.includes('Authentication failed') ||
+				msg.includes('Exceeded maximum reconnection attempts') ||
+				msg.includes('No API key provided')
+			) {
+				finish(err);
+			} else {
+				// Transient error — don't store it; the reconnection logic in
+				// createQueueWebSocket will handle retry.  Storing it would cause
+				// it to be thrown on clean shutdown / abort.
+			}
+		},
+		onClose: () => {
+			// Only finish if the connection is intentionally closed (signal aborted).
+			// Otherwise, the callback-based API handles reconnection.
+		},
+		autoReconnect: true,
+	});
+
+	// Wire up the abort signal to close the connection.
+	const onAbort = () => {
+		connection.close();
+		finish();
+	};
+	signal?.addEventListener('abort', onAbort, { once: true });
+
+	try {
+		while (!done) {
+			// Drain buffered messages.
+			while (buffer.length > 0) {
+				yield buffer.shift()!;
+				// Re-check after yield in case signal was aborted.
+				if (done || signal?.aborted) return;
+			}
+
+			if (done || signal?.aborted) return;
+
+			// Wait for the next message or completion.
+			await new Promise<void>((r) => {
+				resolve = r;
+			});
+		}
+
+		// Drain any remaining messages.
+		while (buffer.length > 0) {
+			yield buffer.shift()!;
+		}
+
+		if (lastError) {
+			throw lastError;
+		}
+	} finally {
+		signal?.removeEventListener('abort', onAbort);
+		connection.close();
+	}
+}