Add AI Gateway WebSocket client with reconnect and draining support

packages/aigateway/package.json

@@ -26,6 +26,7 @@
 	},
 	"dependencies": {
 		"@agentuity/adapter": "workspace:*",
+		"@agentuity/core": "workspace:*",
 		"@agentuity/config": "workspace:*",
 		"@agentuity/client": "workspace:*",
 		"zod": "^4.3.5"

packages/aigateway/src/index.ts

@@ -1,4 +1,6 @@
 export * from './service.ts';
+export * from './protocol.ts';
+export * from './websocket.ts';
 
 import {
 	AIGatewayService,
@@ -24,6 +26,13 @@ import {
 	resolveServiceUrl,
 	type Logger,
 } from '@agentuity/client';
+import { isCliApiKey } from './protocol.ts';
+import {
+	createAIGatewayWebSocketClient,
+	AIGatewayWebSocketError,
+	type AIGatewayWebSocketClient,
+	type AIGatewayWebSocketOptions,
+} from './websocket.ts';
 import { z } from 'zod';
 
 function normalizeOrgId(orgId: string | undefined): string | undefined {
@@ -51,25 +60,59 @@ export type AIGatewayClientOptions = z.infer<typeof AIGatewayClientOptionsSchema
 
 export class AIGatewayClient {
 	readonly #service: AIGatewayService;
+	readonly #apiKey: string;
+	readonly #orgId: string | undefined;
+	readonly #url: string;
 
 	constructor(options: AIGatewayClientOptions = {}) {
 		const validatedOptions = AIGatewayClientOptionsSchema.parse(options);
 		const apiKey =
-			validatedOptions.apiKey || getEnv('AGENTUITY_AIGATEWAY_KEY') || resolveApiKey();
+			validatedOptions.apiKey || getEnv('AGENTUITY_AIGATEWAY_KEY') || resolveApiKey() || '';
 		const serviceUrls = getServiceUrls(resolveRegion());
 		const url = resolveServiceUrl({
 			url: validatedOptions.url,
 			envKey: 'AGENTUITY_AIGATEWAY_URL',
 			fallback: serviceUrls.aigateway,
 		});
+		const orgId = resolveOrgId(validatedOptions.orgId);
 		const { adapter } = createServiceAdapter({
 			apiKey,
-			orgId: resolveOrgId(validatedOptions.orgId),
+			orgId,
 			logger: validatedOptions.logger,
 		});
+		this.#apiKey = apiKey;
+		this.#orgId = orgId;
+		this.#url = url;
 		this.#service = new AIGatewayService(url, adapter);
 	}
 
+	createWebSocket(
+		options: Omit<AIGatewayWebSocketOptions, 'apiKey' | 'orgId' | 'url'> & {
+			url?: string;
+		} = {}
+	): AIGatewayWebSocketClient {
+		if (!this.#apiKey) {
+			throw new AIGatewayWebSocketError({
+				message:
+					'API key is required for AI Gateway WebSocket connections. Provide apiKey when constructing AIGatewayClient or set AGENTUITY_AIGATEWAY_KEY / AGENTUITY_SDK_KEY.',
+				code: 'connection_error',
+			});
+		}
+		if (isCliApiKey(this.#apiKey) && !this.#orgId) {
+			throw new AIGatewayWebSocketError({
+				message:
+					'Organization ID is required for AI Gateway WebSocket connections when using a CLI API key (ck_*). Provide orgId when constructing AIGatewayClient or set AGENTUITY_ORGID / AGENTUITY_ORG_ID / AGENTUITY_CLOUD_ORG_ID.',
+				code: 'connection_error',
+			});
+		}
+		return createAIGatewayWebSocketClient({
+			apiKey: this.#apiKey,
+			orgId: this.#orgId,
+			...options,
+			url: options.url ?? this.#url,
+		});
+	}
+
 	async listModels(): Promise<AIGatewayModels> {
 		return this.#service.listModels();
 	}

packages/aigateway/src/protocol.ts

@@ -0,0 +1,96 @@
+import { z } from 'zod';
+
+export function isCliApiKey(apiKey: string): boolean {
+	return apiKey.startsWith('ck_');
+}
+
+export const AIGatewayWSFrameType = {
+	request: 'request',
+	response: 'response',
+	error: 'error',
+	cancel: 'cancel',
+	draining: 'draining',
+} as const;
+
+export const AIGatewayWSResponseStatus = {
+	complete: 'complete',
+	delta: 'delta',
+	thinkingDelta: 'thinking_delta',
+} as const;
+
+export const AIGatewayWSUsageSchema = z.object({
+	prompt: z.number(),
+	completion: z.number(),
+	total: z.number(),
+	cached: z.number().optional(),
+});
+
+export type AIGatewayWSUsage = z.infer<typeof AIGatewayWSUsageSchema>;
+
+export const AIGatewayWSServerResponseSchema = z.looseObject({
+	type: z.literal(AIGatewayWSFrameType.response),
+	id: z.string(),
+	compact: z.boolean().optional(),
+	status: z.string(),
+	status_code: z.number().optional(),
+	content: z.string().optional(),
+	delta: z.string().optional(),
+	thinking: z.string().optional(),
+	usage: AIGatewayWSUsageSchema.optional(),
+	cost: z.number().optional(),
+	unit: z.string().optional(),
+	input_qty: z.number().optional(),
+	output_qty: z.number().optional(),
+	event: z.string().optional(),
+	data: z.unknown().optional(),
+});
+
+export type AIGatewayWSServerResponse = z.infer<typeof AIGatewayWSServerResponseSchema>;
+
+export const AIGatewayWSServerErrorSchema = z.object({
+	type: z.literal(AIGatewayWSFrameType.error),
+	id: z.string().optional(),
+	status_code: z.number(),
+	message: z.string(),
+});
+
+export type AIGatewayWSServerError = z.infer<typeof AIGatewayWSServerErrorSchema>;
+
+export const AIGatewayWSDrainingSchema = z.object({
+	type: z.literal(AIGatewayWSFrameType.draining),
+	message: z.string().optional(),
+});
+
+export type AIGatewayWSDraining = z.infer<typeof AIGatewayWSDrainingSchema>;
+
+export type AIGatewayWSServerFrame =
+	| AIGatewayWSServerResponse
+	| AIGatewayWSServerError
+	| AIGatewayWSDraining;
+
+export function parseAIGatewayWSServerFrame(raw: unknown): AIGatewayWSServerFrame | null {
+	if (!raw || typeof raw !== 'object') {
+		return null;
+	}
+	const frame = raw as Record<string, unknown>;
+	const type = frame.type;
+	if (type === AIGatewayWSFrameType.draining) {
+		const parsed = AIGatewayWSDrainingSchema.safeParse(raw);
+		return parsed.success ? parsed.data : null;
+	}
+	if (type === AIGatewayWSFrameType.error) {
+		const parsed = AIGatewayWSServerErrorSchema.safeParse(raw);
+		return parsed.success ? parsed.data : null;
+	}
+	if (type === AIGatewayWSFrameType.response) {
+		const parsed = AIGatewayWSServerResponseSchema.safeParse(raw);
+		return parsed.success ? parsed.data : null;
+	}
+	return null;
+}
+
+export function buildAIGatewayWebSocketUrl(baseUrl: string): string {
+	const trimmed = baseUrl.trim().replace(/\/$/, '');
+	const wsBase = trimmed.replace(/^https:\/\//, 'wss://').replace(/^http:\/\//, 'ws://');
+	return `${wsBase}/v1/ws`;
+}