From bf795c56882abc8557a4ae7e5801a7544197d97c Mon Sep 17 00:00:00 2001 From: Jeff Haynie Date: Sat, 4 Jul 2026 09:45:17 -0500 Subject: [PATCH 1/7] Add reconnecting AI Gateway WebSocket client with v1 protocol support. Implements multiplexed request/response handling, draining-aware reconnect, and AIGatewayClient.createWebSocket() for shared auth configuration. Co-authored-by: Cursor --- packages/aigateway/src/index.ts | 41 +- packages/aigateway/src/protocol.ts | 94 +++ packages/aigateway/src/websocket.ts | 724 ++++++++++++++++++++++ packages/aigateway/test/websocket.test.ts | 382 ++++++++++++ 4 files changed, 1239 insertions(+), 2 deletions(-) create mode 100644 packages/aigateway/src/protocol.ts create mode 100644 packages/aigateway/src/websocket.ts create mode 100644 packages/aigateway/test/websocket.test.ts diff --git a/packages/aigateway/src/index.ts b/packages/aigateway/src/index.ts index dbb5091bb..ec5562fa7 100644 --- a/packages/aigateway/src/index.ts +++ b/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,11 @@ import { resolveServiceUrl, type Logger, } from '@agentuity/client'; +import { + createAIGatewayWebSocketClient, + type AIGatewayWebSocketClient, + type AIGatewayWebSocketOptions, +} from './websocket.ts'; import { z } from 'zod'; function normalizeOrgId(orgId: string | undefined): string | undefined { @@ -51,25 +58,55 @@ export type AIGatewayClientOptions = z.infer & { + url?: string; + } = {} + ): AIGatewayWebSocketClient { + if (!this.#apiKey) { + throw new Error( + 'API key is required for AI Gateway WebSocket connections. Provide apiKey when constructing AIGatewayClient or set AGENTUITY_AIGATEWAY_KEY / AGENTUITY_SDK_KEY.' + ); + } + if (!this.#orgId) { + throw new Error( + 'Organization ID is required for AI Gateway WebSocket connections. Provide orgId when constructing AIGatewayClient or set AGENTUITY_ORGID / AGENTUITY_ORG_ID / AGENTUITY_CLOUD_ORG_ID.' + ); + } + return createAIGatewayWebSocketClient({ + apiKey: this.#apiKey, + orgId: this.#orgId, + ...options, + url: options.url ?? this.#url, + }); + } + async listModels(): Promise { return this.#service.listModels(); } diff --git a/packages/aigateway/src/protocol.ts b/packages/aigateway/src/protocol.ts new file mode 100644 index 000000000..2851ae424 --- /dev/null +++ b/packages/aigateway/src/protocol.ts @@ -0,0 +1,94 @@ +import { z } from 'zod'; + +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; + +export const AIGatewayWSServerResponseSchema = z + .object({ + 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(), + }) + .passthrough(); + +export type AIGatewayWSServerResponse = z.infer; + +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; + +export const AIGatewayWSDrainingSchema = z.object({ + type: z.literal(AIGatewayWSFrameType.draining), + message: z.string().optional(), +}); + +export type AIGatewayWSDraining = z.infer; + +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; + 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`; +} diff --git a/packages/aigateway/src/websocket.ts b/packages/aigateway/src/websocket.ts new file mode 100644 index 000000000..a8d0ab780 --- /dev/null +++ b/packages/aigateway/src/websocket.ts @@ -0,0 +1,724 @@ +import { StructuredError } from '@agentuity/adapter'; +import { resolveRegion, resolveServiceUrl } from '@agentuity/client'; +import { getEnv, getServiceUrls } from '@agentuity/config'; +import { z } from 'zod'; +import { + AIGatewayWSFrameType, + AIGatewayWSResponseStatus, + buildAIGatewayWebSocketUrl, + parseAIGatewayWSServerFrame, + type AIGatewayWSUsage, + type AIGatewayWSServerError, + type AIGatewayWSServerResponse, +} from './protocol'; + +export const AIGatewayWebSocketErrorCode = { + connection_error: 'connection_error', + connection_draining: 'connection_draining', + connection_closed: 'connection_closed', + invalid_response: 'invalid_response', + max_reconnects_exceeded: 'max_reconnects_exceeded', + request_timeout: 'request_timeout', +} as const; + +export type AIGatewayWebSocketErrorCode = + (typeof AIGatewayWebSocketErrorCode)[keyof typeof AIGatewayWebSocketErrorCode]; + +export const AIGatewayWebSocketError = StructuredError('AIGatewayWebSocketError')<{ + code: AIGatewayWebSocketErrorCode; + statusCode?: number; + requestId?: string; + closeCode?: number; + closeReason?: string; +}>(); + +export type AIGatewayWebSocketErrorInstance = InstanceType; + +export type AIGatewayWebSocketState = + | 'connecting' + | 'connected' + | 'reconnecting' + | 'draining' + | 'closed'; + +export const AIGatewayWSRequestOptionsSchema = z.object({ + id: z.string().optional(), + compact: z.boolean().optional().default(true), + model: z.string().optional(), + prompt: z.string().optional(), + system: z.string().optional(), + thinking: z.union([z.boolean(), z.string()]).optional(), + temperature: z.number().optional(), + max_tokens: z.number().optional(), + stream: z.boolean().optional(), + data: z.record(z.string(), z.unknown()).optional(), + timeoutMs: z.number().optional(), +}); + +export type AIGatewayWSRequestOptions = z.infer; + +export interface AIGatewayWSResult { + id: string; + content?: string; + usage?: AIGatewayWSUsage; + cost?: number; + unit?: string; + inputQty?: number; + outputQty?: number; + statusCode?: number; + data?: unknown; +} + +export type AIGatewayWSStreamEvent = + | { type: 'delta'; delta: string } + | { type: 'thinking_delta'; thinking: string } + | { type: 'event'; event: string; data: unknown } + | { type: 'complete'; result: AIGatewayWSResult }; + +export interface AIGatewayWebSocketOptions { + apiKey: string; + orgId?: string; + url?: string; + autoReconnect?: boolean; + reconnectDelayMs?: number; + maxReconnectDelayMs?: number; + maxReconnectAttempts?: number; + defaultTimeoutMs?: number; + onOpen?: () => void; + onClose?: (code: number, reason: string) => void; + onError?: (error: AIGatewayWebSocketErrorInstance) => void; + onDraining?: (message: string) => void; + onReconnect?: (attempt: number) => void; + onStateChange?: (state: AIGatewayWebSocketState) => void; +} + +interface PendingRequest { + id: string; + stream: boolean; + resolve: (result: AIGatewayWSResult) => void; + reject: (error: Error) => void; + timeout: ReturnType; + push?: (event: AIGatewayWSStreamEvent) => void; + accumulated: string; + thinkingAccumulated: string; +} + +const DEFAULT_RECONNECT_DELAY_MS = 500; +const DEFAULT_MAX_RECONNECT_DELAY_MS = 30_000; +const DEFAULT_MAX_RECONNECT_ATTEMPTS = 10; +const DEFAULT_TIMEOUT_MS = 120_000; +const GOING_AWAY_CLOSE_CODE = 1001; + +function normalizeOrgId(orgId: string | undefined): string | undefined { + const trimmed = orgId?.trim(); + return trimmed ? trimmed : undefined; +} + +function resolveOrgId(options: { orgId?: string }): string { + const resolved = + normalizeOrgId(options.orgId) ?? + normalizeOrgId(getEnv('AGENTUITY_ORGID')) ?? + normalizeOrgId(getEnv('AGENTUITY_ORG_ID')) ?? + normalizeOrgId(getEnv('AGENTUITY_CLOUD_ORG_ID')); + if (!resolved) { + throw new AIGatewayWebSocketError({ + message: + 'Organization ID is required. Provide orgId in options or set AGENTUITY_ORGID / AGENTUITY_ORG_ID / AGENTUITY_CLOUD_ORG_ID.', + code: 'connection_error', + }); + } + return resolved; +} + +function resolveWebSocketUrl(options: { url?: string }): string { + const serviceUrls = getServiceUrls(resolveRegion()); + const baseUrl = resolveServiceUrl({ + url: options.url, + envKey: 'AGENTUITY_AIGATEWAY_URL', + fallback: serviceUrls.aigateway, + }); + return buildAIGatewayWebSocketUrl(baseUrl); +} + +function isTerminalCloseCode(code: number): boolean { + if (code === GOING_AWAY_CLOSE_CODE) { + return false; + } + return code >= 4000 && code < 5000; +} + +function createRequestId(): string { + if (typeof crypto !== 'undefined' && typeof crypto.randomUUID === 'function') { + return crypto.randomUUID(); + } + return `req_${Date.now()}_${Math.random().toString(36).slice(2, 10)}`; +} + +function buildRequestFrame( + options: AIGatewayWSRequestOptions, + id: string +): Record { + const parsed = AIGatewayWSRequestOptionsSchema.parse(options); + const frame: Record = { + type: AIGatewayWSFrameType.request, + id, + compact: parsed.compact ?? true, + }; + if (parsed.compact) { + if (parsed.model !== undefined) frame.model = parsed.model; + if (parsed.prompt !== undefined) frame.prompt = parsed.prompt; + if (parsed.system !== undefined) frame.system = parsed.system; + if (parsed.thinking !== undefined) frame.thinking = parsed.thinking; + if (parsed.temperature !== undefined) frame.temperature = parsed.temperature; + if (parsed.max_tokens !== undefined) frame.max_tokens = parsed.max_tokens; + if (parsed.stream !== undefined) frame.stream = parsed.stream; + } else if (parsed.data !== undefined) { + frame.data = parsed.data; + } + return frame; +} + +function mapCompleteResponse(response: AIGatewayWSServerResponse): AIGatewayWSResult { + return { + id: response.id, + content: response.content, + usage: response.usage, + cost: response.cost, + unit: response.unit, + inputQty: response.input_qty, + outputQty: response.output_qty, + statusCode: response.status_code, + data: response.data, + }; +} + +function mapServerError(error: AIGatewayWSServerError): AIGatewayWebSocketErrorInstance { + return new AIGatewayWebSocketError({ + message: error.message, + code: 'connection_error', + statusCode: error.status_code, + requestId: error.id, + }); +} + +export class AIGatewayWebSocketClient { + readonly #options: Required< + Pick< + AIGatewayWebSocketOptions, + | 'autoReconnect' + | 'reconnectDelayMs' + | 'maxReconnectDelayMs' + | 'maxReconnectAttempts' + | 'defaultTimeoutMs' + > + > & + AIGatewayWebSocketOptions; + + readonly #apiKey: string; + readonly #orgId: string; + readonly #url: string; + + #ws: WebSocket | null = null; + #state: AIGatewayWebSocketState = 'closed'; + #draining = false; + #intentionallyClosed = false; + #reconnectAttempts = 0; + #reconnectTimer: ReturnType | null = null; + #connectPromise: Promise | null = null; + #connectResolve: (() => void) | null = null; + #connectReject: ((error: Error) => void) | null = null; + #pending = new Map(); + #drainReconnectScheduled = false; + + constructor(options: AIGatewayWebSocketOptions) { + if (!options.apiKey) { + throw new AIGatewayWebSocketError({ + message: 'API key is required', + code: 'connection_error', + }); + } + this.#options = { + autoReconnect: options.autoReconnect ?? true, + reconnectDelayMs: options.reconnectDelayMs ?? DEFAULT_RECONNECT_DELAY_MS, + maxReconnectDelayMs: options.maxReconnectDelayMs ?? DEFAULT_MAX_RECONNECT_DELAY_MS, + maxReconnectAttempts: options.maxReconnectAttempts ?? DEFAULT_MAX_RECONNECT_ATTEMPTS, + defaultTimeoutMs: options.defaultTimeoutMs ?? DEFAULT_TIMEOUT_MS, + ...options, + }; + this.#apiKey = options.apiKey; + this.#orgId = resolveOrgId(options); + this.#url = resolveWebSocketUrl(options); + } + + get state(): AIGatewayWebSocketState { + return this.#state; + } + + get isDraining(): boolean { + return this.#draining; + } + + get url(): string { + return this.#url; + } + + get orgId(): string { + return this.#orgId; + } + + connect(): Promise { + if (this.#state === 'connected' && this.#ws?.readyState === WebSocket.OPEN) { + return Promise.resolve(); + } + if (this.#connectPromise) { + return this.#connectPromise; + } + this.#intentionallyClosed = false; + this.#connectPromise = new Promise((resolve, reject) => { + this.#connectResolve = resolve; + this.#connectReject = reject; + this.#connectInternal(); + }); + return this.#connectPromise; + } + + close(code = 1000, reason = 'client closed'): void { + this.#intentionallyClosed = true; + this.#clearReconnectTimer(); + this.#rejectAllPending( + new AIGatewayWebSocketError({ + message: reason || 'WebSocket closed by client', + code: 'connection_closed', + closeCode: code, + closeReason: reason, + }) + ); + if (this.#ws) { + const ws = this.#ws; + this.#ws = null; + ws.close(code, reason); + } + this.#setState('closed'); + this.#finishConnect( + new AIGatewayWebSocketError({ + message: reason || 'WebSocket closed by client', + code: 'connection_closed', + closeCode: code, + closeReason: reason, + }) + ); + } + + cancel(requestId: string): void { + this.#sendJson({ + type: AIGatewayWSFrameType.cancel, + id: requestId, + }); + } + + async complete(options: AIGatewayWSRequestOptions): Promise { + const streamOptions = { ...options, stream: false }; + const events: AIGatewayWSStreamEvent[] = []; + for await (const event of this.stream(streamOptions)) { + events.push(event); + } + const completeEvent = events.find((event) => event.type === 'complete'); + if (!completeEvent) { + throw new AIGatewayWebSocketError({ + message: 'Request completed without a final response frame', + code: 'invalid_response', + requestId: options.id, + }); + } + return completeEvent.result; + } + + async *stream(options: AIGatewayWSRequestOptions): AsyncGenerator { + this.#assertCanSendRequest(); + await this.connect(); + + const id = options.id ?? createRequestId(); + const timeoutMs = options.timeoutMs ?? this.#options.defaultTimeoutMs; + const queue: AIGatewayWSStreamEvent[] = []; + let notify: (() => void) | null = null; + let done = false; + let streamError: Error | null = null; + + const pending: PendingRequest = { + id, + stream: true, + accumulated: '', + thinkingAccumulated: '', + resolve: () => {}, + reject: (error) => { + streamError = error; + done = true; + notify?.(); + }, + timeout: setTimeout(() => { + this.cancel(id); + this.#rejectPending( + id, + new AIGatewayWebSocketError({ + message: `Request timed out after ${timeoutMs}ms`, + code: 'request_timeout', + requestId: id, + }) + ); + }, timeoutMs), + push: (event) => { + queue.push(event); + notify?.(); + }, + }; + + this.#pending.set(id, pending); + + try { + this.#sendJson(buildRequestFrame({ ...options, id }, id)); + + while (!done || queue.length > 0) { + if (queue.length === 0) { + await new Promise((resolve) => { + notify = resolve; + }); + notify = null; + if (streamError) { + throw streamError; + } + continue; + } + const event = queue.shift(); + if (!event) { + continue; + } + yield event; + if (event.type === 'complete') { + done = true; + } + } + } finally { + clearTimeout(pending.timeout); + this.#pending.delete(id); + this.#maybeScheduleDrainReconnect(); + } + } + + #assertCanSendRequest(): void { + if (this.#draining) { + throw new AIGatewayWebSocketError({ + message: + 'Connection is draining; new requests are not accepted until reconnect completes', + code: 'connection_draining', + }); + } + if (this.#intentionallyClosed) { + throw new AIGatewayWebSocketError({ + message: 'WebSocket client is closed', + code: 'connection_closed', + }); + } + } + + #connectInternal(): void { + this.#clearReconnectTimer(); + this.#setState(this.#reconnectAttempts > 0 ? 'reconnecting' : 'connecting'); + + const ws = new WebSocket(this.#url, { + headers: { + Authorization: `Bearer ${this.#apiKey}`, + 'x-agentuity-orgid': this.#orgId, + }, + }); + this.#ws = ws; + + ws.onopen = () => { + if (ws !== this.#ws) { + return; + } + this.#reconnectAttempts = 0; + this.#draining = false; + this.#drainReconnectScheduled = false; + this.#setState('connected'); + this.#options.onOpen?.(); + this.#finishConnect(); + }; + + ws.onmessage = (event) => { + if (ws !== this.#ws) { + return; + } + this.#handleMessage(event.data); + }; + + ws.onerror = () => { + if (ws !== this.#ws) { + return; + } + this.#options.onError?.( + new AIGatewayWebSocketError({ + message: 'WebSocket connection error', + code: 'connection_error', + }) + ); + }; + + ws.onclose = (event) => { + if (ws !== this.#ws) { + return; + } + this.#ws = null; + const wasDraining = this.#draining; + const terminalClose = isTerminalCloseCode(event.code); + if (terminalClose) { + this.#intentionallyClosed = true; + } + + if (this.#state !== 'closed') { + this.#setState(wasDraining ? 'draining' : 'closed'); + } + + this.#options.onClose?.(event.code, event.reason); + + if (!wasDraining) { + this.#rejectAllPending( + new AIGatewayWebSocketError({ + message: `WebSocket closed (code ${event.code})${event.reason ? `: ${event.reason}` : ''}`, + code: 'connection_closed', + closeCode: event.code, + closeReason: event.reason || undefined, + }) + ); + } + + this.#finishConnect( + new AIGatewayWebSocketError({ + message: `WebSocket closed before connection was ready (code ${event.code})`, + code: 'connection_error', + closeCode: event.code, + closeReason: event.reason || undefined, + }) + ); + + if (wasDraining) { + if (this.#pending.size > 0) { + this.#rejectAllPending( + new AIGatewayWebSocketError({ + message: + 'WebSocket closed during server drain before all in-flight requests completed', + code: 'connection_closed', + closeCode: event.code, + closeReason: event.reason || undefined, + }) + ); + } + this.#maybeScheduleDrainReconnect(); + return; + } + + if (!this.#intentionallyClosed && this.#options.autoReconnect) { + this.#scheduleReconnect(); + } + }; + } + + #handleMessage(raw: unknown): void { + let parsed: unknown; + try { + parsed = typeof raw === 'string' ? JSON.parse(raw) : raw; + } catch { + this.#options.onError?.( + new AIGatewayWebSocketError({ + message: 'Received invalid JSON from AI Gateway WebSocket', + code: 'invalid_response', + }) + ); + return; + } + + const frame = parseAIGatewayWSServerFrame(parsed); + if (!frame) { + this.#options.onError?.( + new AIGatewayWebSocketError({ + message: 'Received unrecognized WebSocket frame', + code: 'invalid_response', + }) + ); + return; + } + + if (frame.type === AIGatewayWSFrameType.draining) { + this.#handleDraining(frame.message ?? 'server is draining'); + return; + } + + if (frame.type === AIGatewayWSFrameType.error) { + this.#handleErrorFrame(frame); + return; + } + + this.#handleResponseFrame(frame); + } + + #handleDraining(message: string): void { + if (this.#draining) { + return; + } + this.#draining = true; + this.#setState('draining'); + this.#options.onDraining?.(message); + } + + #handleErrorFrame(error: AIGatewayWSServerError): void { + const mapped = mapServerError(error); + if (error.id) { + this.#rejectPending(error.id, mapped); + return; + } + this.#options.onError?.(mapped); + } + + #handleResponseFrame(response: AIGatewayWSServerResponse): void { + const pending = this.#pending.get(response.id); + if (!pending) { + return; + } + + if (response.status === AIGatewayWSResponseStatus.delta) { + const delta = response.delta ?? ''; + pending.accumulated += delta; + pending.push?.({ type: 'delta', delta }); + return; + } + + if (response.status === AIGatewayWSResponseStatus.thinkingDelta) { + const thinking = response.thinking ?? ''; + pending.thinkingAccumulated += thinking; + pending.push?.({ type: 'thinking_delta', thinking }); + return; + } + + if (response.event) { + pending.push?.({ + type: 'event', + event: response.event, + data: response.data, + }); + } + + if (response.status === AIGatewayWSResponseStatus.complete) { + const result = mapCompleteResponse({ + ...response, + content: response.content ?? pending.accumulated, + }); + clearTimeout(pending.timeout); + this.#pending.delete(response.id); + pending.push?.({ type: 'complete', result }); + pending.resolve(result); + this.#maybeScheduleDrainReconnect(); + } + } + + #rejectPending(id: string, error: Error): void { + const pending = this.#pending.get(id); + if (!pending) { + return; + } + clearTimeout(pending.timeout); + this.#pending.delete(id); + pending.reject(error); + this.#maybeScheduleDrainReconnect(); + } + + #rejectAllPending(error: AIGatewayWebSocketErrorInstance): void { + for (const [id, pending] of this.#pending.entries()) { + clearTimeout(pending.timeout); + pending.reject(error); + this.#pending.delete(id); + } + } + + #maybeScheduleDrainReconnect(): void { + if (!this.#draining || this.#drainReconnectScheduled || this.#intentionallyClosed) { + return; + } + if (this.#pending.size > 0 || this.#ws) { + return; + } + this.#drainReconnectScheduled = true; + this.#draining = false; + if (this.#options.autoReconnect) { + this.#scheduleReconnect(); + } + } + + #scheduleReconnect(): void { + if (this.#intentionallyClosed || !this.#options.autoReconnect) { + return; + } + if (this.#reconnectAttempts >= this.#options.maxReconnectAttempts) { + this.#options.onError?.( + new AIGatewayWebSocketError({ + message: `Exceeded maximum reconnection attempts (${this.#options.maxReconnectAttempts})`, + code: 'max_reconnects_exceeded', + }) + ); + return; + } + + const baseDelay = this.#options.reconnectDelayMs * 2 ** this.#reconnectAttempts; + const jitter = 0.5 + Math.random() * 0.5; + const delay = Math.min(Math.floor(baseDelay * jitter), this.#options.maxReconnectDelayMs); + this.#reconnectAttempts += 1; + this.#options.onReconnect?.(this.#reconnectAttempts); + this.#reconnectTimer = setTimeout(() => { + this.#reconnectTimer = null; + this.#connectInternal(); + }, delay); + } + + #clearReconnectTimer(): void { + if (this.#reconnectTimer) { + clearTimeout(this.#reconnectTimer); + this.#reconnectTimer = null; + } + } + + #sendJson(payload: Record): void { + if (!this.#ws || this.#ws.readyState !== WebSocket.OPEN) { + throw new AIGatewayWebSocketError({ + message: 'WebSocket is not connected', + code: 'connection_error', + }); + } + this.#ws.send(JSON.stringify(payload)); + } + + #setState(state: AIGatewayWebSocketState): void { + if (this.#state === state) { + return; + } + this.#state = state; + this.#options.onStateChange?.(state); + } + + #finishConnect(error?: AIGatewayWebSocketErrorInstance): void { + if (!this.#connectPromise) { + return; + } + if (error && this.#state !== 'connected') { + this.#connectReject?.(error); + } else { + this.#connectResolve?.(); + } + this.#connectPromise = null; + this.#connectResolve = null; + this.#connectReject = null; + } +} + +export function createAIGatewayWebSocketClient( + options: AIGatewayWebSocketOptions +): AIGatewayWebSocketClient { + return new AIGatewayWebSocketClient(options); +} diff --git a/packages/aigateway/test/websocket.test.ts b/packages/aigateway/test/websocket.test.ts new file mode 100644 index 000000000..8de29ed33 --- /dev/null +++ b/packages/aigateway/test/websocket.test.ts @@ -0,0 +1,382 @@ +import { afterEach, beforeEach, describe, expect, it } from 'bun:test'; +import { + AIGatewayWebSocketClient, + AIGatewayWebSocketError, + buildAIGatewayWebSocketUrl, + parseAIGatewayWSServerFrame, +} from '../src/index.ts'; + +const OriginalWebSocket = globalThis.WebSocket; + +class MockWebSocket { + static CONNECTING = 0; + static OPEN = 1; + static CLOSED = 3; + static instances: MockWebSocket[] = []; + + readyState = MockWebSocket.CONNECTING; + onopen: ((event: Event) => void) | null = null; + onmessage: ((event: MessageEvent) => void) | null = null; + onerror: ((event: Event) => void) | null = null; + onclose: ((event: CloseEvent) => void) | null = null; + readonly sent: string[] = []; + readonly url: string; + readonly headers?: Record; + + constructor(url: string, options?: { headers?: Record }) { + this.url = url; + this.headers = options?.headers; + MockWebSocket.instances.push(this); + } + + send(data: string) { + this.sent.push(data); + } + + close(code = 1000, reason = '') { + this.readyState = MockWebSocket.CLOSED; + this.onclose?.({ + code, + reason, + wasClean: true, + } as CloseEvent); + } + + open() { + this.readyState = MockWebSocket.OPEN; + this.onopen?.(new Event('open')); + } + + receive(payload: unknown) { + this.onmessage?.({ + data: JSON.stringify(payload), + } as MessageEvent); + } +} + +async function flushAsyncWork() { + await Promise.resolve(); + await Promise.resolve(); +} + +describe('AI Gateway WebSocket protocol', () => { + it('builds wss URLs from https base URLs', () => { + expect(buildAIGatewayWebSocketUrl('https://aigateway.example')).toBe( + 'wss://aigateway.example/v1/ws' + ); + }); + + it('parses draining, error, and response frames', () => { + expect( + parseAIGatewayWSServerFrame({ + type: 'draining', + message: 'shutting down', + }) + ).toEqual({ + type: 'draining', + message: 'shutting down', + }); + + expect( + parseAIGatewayWSServerFrame({ + type: 'error', + id: 'req_1', + status_code: 503, + message: 'draining', + }) + ).toEqual({ + type: 'error', + id: 'req_1', + status_code: 503, + message: 'draining', + }); + + expect( + parseAIGatewayWSServerFrame({ + type: 'response', + id: 'req_1', + status: 'complete', + content: 'hello', + usage: { prompt: 1, completion: 2, total: 3, cached: 1 }, + }) + ).toMatchObject({ + type: 'response', + id: 'req_1', + status: 'complete', + content: 'hello', + }); + }); +}); + +describe('AIGatewayWebSocketClient', () => { + beforeEach(() => { + MockWebSocket.instances = []; + globalThis.WebSocket = MockWebSocket as unknown as typeof WebSocket; + }); + + afterEach(() => { + globalThis.WebSocket = OriginalWebSocket; + MockWebSocket.instances = []; + }); + + it('connects with auth headers and completes compact requests', async () => { + const client = new AIGatewayWebSocketClient({ + apiKey: 'ag_test', + orgId: 'org_test', + url: 'https://aigateway.example', + }); + + const connectPromise = client.connect(); + const ws = MockWebSocket.instances[0]; + expect(ws?.url).toBe('wss://aigateway.example/v1/ws'); + expect(ws?.headers).toEqual({ + Authorization: 'Bearer ag_test', + 'x-agentuity-orgid': 'org_test', + }); + + ws?.open(); + await connectPromise; + + const resultPromise = client.complete({ + model: 'anthropic/claude-sonnet-4-20250514', + prompt: 'Hello', + }); + await flushAsyncWork(); + + const request = JSON.parse(ws?.sent[0] ?? '{}'); + expect(request).toMatchObject({ + type: 'request', + compact: true, + model: 'anthropic/claude-sonnet-4-20250514', + prompt: 'Hello', + stream: false, + }); + + ws?.receive({ + type: 'response', + id: request.id, + status: 'complete', + content: 'Hi there', + usage: { prompt: 10, completion: 5, total: 15, cached: 2 }, + }); + + await expect(resultPromise).resolves.toEqual({ + id: request.id, + content: 'Hi there', + usage: { prompt: 10, completion: 5, total: 15, cached: 2 }, + }); + }); + + it('streams delta and thinking frames before complete', async () => { + const client = new AIGatewayWebSocketClient({ + apiKey: 'ag_test', + orgId: 'org_test', + url: 'https://aigateway.example', + }); + + const connectPromise = client.connect(); + const ws = MockWebSocket.instances[0]; + ws?.open(); + await connectPromise; + + const events: string[] = []; + const streamPromise = (async () => { + for await (const event of client.stream({ + model: 'anthropic/claude-sonnet-4-20250514', + prompt: 'Hello', + stream: true, + })) { + events.push( + `${event.type}:${'delta' in event ? event.delta : 'thinking' in event ? event.thinking : event.result.content}` + ); + } + })(); + + await flushAsyncWork(); + const request = JSON.parse(MockWebSocket.instances[0]?.sent[0] ?? '{}'); + + MockWebSocket.instances[0]?.receive({ + type: 'response', + id: request.id, + status: 'thinking_delta', + thinking: 'hmm', + }); + MockWebSocket.instances[0]?.receive({ + type: 'response', + id: request.id, + status: 'delta', + delta: 'Hello', + }); + MockWebSocket.instances[0]?.receive({ + type: 'response', + id: request.id, + status: 'complete', + content: 'Hello world', + }); + + await streamPromise; + expect(events).toEqual(['thinking_delta:hmm', 'delta:Hello', 'complete:Hello world']); + }); + + it('blocks new requests while draining and reconnects after in-flight work completes', async () => { + let drainingMessage = ''; + let reconnectAttempt = 0; + const client = new AIGatewayWebSocketClient({ + apiKey: 'ag_test', + orgId: 'org_test', + url: 'https://aigateway.example', + reconnectDelayMs: 1, + onDraining: (message) => { + drainingMessage = message; + }, + onReconnect: (attempt) => { + reconnectAttempt = attempt; + }, + }); + + const connectPromise = client.connect(); + const firstWs = MockWebSocket.instances[0]; + firstWs?.open(); + await connectPromise; + + const inFlight = client.complete({ + model: 'anthropic/claude-sonnet-4-20250514', + prompt: 'Finish me', + }); + await flushAsyncWork(); + const request = JSON.parse(firstWs?.sent[0] ?? '{}'); + + firstWs?.receive({ + type: 'draining', + message: 'server rolling restart', + }); + expect(client.isDraining).toBe(true); + expect(drainingMessage).toBe('server rolling restart'); + expect(client.state).toBe('draining'); + + await expect( + client.complete({ + model: 'anthropic/claude-sonnet-4-20250514', + prompt: 'Should fail', + }) + ).rejects.toMatchObject({ + code: 'connection_draining', + }); + + firstWs?.receive({ + type: 'response', + id: request.id, + status: 'complete', + content: 'done', + }); + await expect(inFlight).resolves.toMatchObject({ content: 'done' }); + + firstWs?.close(1001, 'going away'); + await flushAsyncWork(); + await new Promise((resolve) => setTimeout(resolve, 5)); + + expect(reconnectAttempt).toBe(1); + expect(MockWebSocket.instances.length).toBe(2); + expect(client.isDraining).toBe(false); + + const secondWs = MockWebSocket.instances[1]; + secondWs?.open(); + await flushAsyncWork(); + expect(client.state).toBe('connected'); + }); + + it('multiplexes concurrent requests over one connection by request id', async () => { + const client = new AIGatewayWebSocketClient({ + apiKey: 'ag_test', + orgId: 'org_test', + url: 'https://aigateway.example', + }); + + const connectPromise = client.connect(); + const ws = MockWebSocket.instances[0]; + ws?.open(); + await connectPromise; + + const firstPromise = client.complete({ + id: 'req_alpha', + model: 'anthropic/claude-sonnet-4-20250514', + prompt: 'First prompt', + }); + const secondPromise = client.complete({ + id: 'req_beta', + model: 'anthropic/claude-sonnet-4-20250514', + prompt: 'Second prompt', + }); + + await flushAsyncWork(); + + expect(ws?.sent).toHaveLength(2); + expect(JSON.parse(ws?.sent[0] ?? '{}')).toMatchObject({ + type: 'request', + id: 'req_alpha', + prompt: 'First prompt', + }); + expect(JSON.parse(ws?.sent[1] ?? '{}')).toMatchObject({ + type: 'request', + id: 'req_beta', + prompt: 'Second prompt', + }); + + // Complete responses out of order to prove routing is per-id, not FIFO. + ws?.receive({ + type: 'response', + id: 'req_beta', + status: 'complete', + content: 'Beta reply', + }); + ws?.receive({ + type: 'response', + id: 'req_alpha', + status: 'complete', + content: 'Alpha reply', + }); + + await expect(firstPromise).resolves.toMatchObject({ + id: 'req_alpha', + content: 'Alpha reply', + }); + await expect(secondPromise).resolves.toMatchObject({ + id: 'req_beta', + content: 'Beta reply', + }); + }); + + it('surfaces server error frames for a request id', async () => { + const client = new AIGatewayWebSocketClient({ + apiKey: 'ag_test', + orgId: 'org_test', + url: 'https://aigateway.example', + }); + + const connectPromise = client.connect(); + const ws = MockWebSocket.instances[0]; + ws?.open(); + await connectPromise; + + const resultPromise = client.complete({ + id: 'req_error', + model: 'anthropic/claude-sonnet-4-20250514', + prompt: 'Hello', + }); + await flushAsyncWork(); + + ws?.receive({ + type: 'error', + id: 'req_error', + status_code: 503, + message: 'service draining', + }); + + await expect(resultPromise).rejects.toBeInstanceOf(AIGatewayWebSocketError); + await expect(resultPromise).rejects.toMatchObject({ + code: 'connection_error', + statusCode: 503, + requestId: 'req_error', + }); + }); +}); From d2a023c7387b308a9d29acf4e2f0f6a284c55539 Mon Sep 17 00:00:00 2001 From: Jeff Haynie Date: Sat, 4 Jul 2026 09:48:39 -0500 Subject: [PATCH 2/7] Hand off to a new WebSocket while retiring socket finishes in-flight work. On server drain, keep the old connection alive for pending requests and route new requests to a replacement active socket. Co-authored-by: Cursor --- packages/aigateway/src/websocket.ts | 302 ++++++++++++++-------- packages/aigateway/test/websocket.test.ts | 52 ++-- 2 files changed, 220 insertions(+), 134 deletions(-) diff --git a/packages/aigateway/src/websocket.ts b/packages/aigateway/src/websocket.ts index a8d0ab780..5e86ffcd8 100644 --- a/packages/aigateway/src/websocket.ts +++ b/packages/aigateway/src/websocket.ts @@ -103,6 +103,12 @@ interface PendingRequest { thinkingAccumulated: string; } +interface WebSocketLane { + ws: WebSocket; + role: 'active' | 'retiring'; + requestIds: Set; +} + const DEFAULT_RECONNECT_DELAY_MS = 500; const DEFAULT_MAX_RECONNECT_DELAY_MS = 30_000; const DEFAULT_MAX_RECONNECT_ATTEMPTS = 10; @@ -218,9 +224,9 @@ export class AIGatewayWebSocketClient { readonly #orgId: string; readonly #url: string; - #ws: WebSocket | null = null; + #activeLane: WebSocketLane | null = null; + #retiringLane: WebSocketLane | null = null; #state: AIGatewayWebSocketState = 'closed'; - #draining = false; #intentionallyClosed = false; #reconnectAttempts = 0; #reconnectTimer: ReturnType | null = null; @@ -228,7 +234,6 @@ export class AIGatewayWebSocketClient { #connectResolve: (() => void) | null = null; #connectReject: ((error: Error) => void) | null = null; #pending = new Map(); - #drainReconnectScheduled = false; constructor(options: AIGatewayWebSocketOptions) { if (!options.apiKey) { @@ -255,7 +260,7 @@ export class AIGatewayWebSocketClient { } get isDraining(): boolean { - return this.#draining; + return this.#retiringLane !== null; } get url(): string { @@ -267,7 +272,7 @@ export class AIGatewayWebSocketClient { } connect(): Promise { - if (this.#state === 'connected' && this.#ws?.readyState === WebSocket.OPEN) { + if (this.#activeLane?.ws.readyState === WebSocket.OPEN) { return Promise.resolve(); } if (this.#connectPromise) { @@ -277,7 +282,9 @@ export class AIGatewayWebSocketClient { this.#connectPromise = new Promise((resolve, reject) => { this.#connectResolve = resolve; this.#connectReject = reject; - this.#connectInternal(); + if (!this.#activeLane) { + this.#beginActiveConnection(); + } }); return this.#connectPromise; } @@ -293,11 +300,14 @@ export class AIGatewayWebSocketClient { closeReason: reason, }) ); - if (this.#ws) { - const ws = this.#ws; - this.#ws = null; - ws.close(code, reason); + for (const lane of [this.#activeLane, this.#retiringLane]) { + if (!lane) { + continue; + } + lane.ws.close(code, reason); } + this.#activeLane = null; + this.#retiringLane = null; this.#setState('closed'); this.#finishConnect( new AIGatewayWebSocketError({ @@ -310,7 +320,11 @@ export class AIGatewayWebSocketClient { } cancel(requestId: string): void { - this.#sendJson({ + const lane = this.#findLaneForRequest(requestId); + if (!lane) { + return; + } + this.#sendJsonOnLane(lane, { type: AIGatewayWSFrameType.cancel, id: requestId, }); @@ -375,7 +389,7 @@ export class AIGatewayWebSocketClient { this.#pending.set(id, pending); try { - this.#sendJson(buildRequestFrame({ ...options, id }, id)); + this.#sendJson(buildRequestFrame({ ...options, id }, id), id); while (!done || queue.length > 0) { if (queue.length === 0) { @@ -400,18 +414,11 @@ export class AIGatewayWebSocketClient { } finally { clearTimeout(pending.timeout); this.#pending.delete(id); - this.#maybeScheduleDrainReconnect(); + this.#clearRetiringLaneIfIdle(); } } #assertCanSendRequest(): void { - if (this.#draining) { - throw new AIGatewayWebSocketError({ - message: - 'Connection is draining; new requests are not accepted until reconnect completes', - code: 'connection_draining', - }); - } if (this.#intentionallyClosed) { throw new AIGatewayWebSocketError({ message: 'WebSocket client is closed', @@ -420,9 +427,15 @@ export class AIGatewayWebSocketClient { } } - #connectInternal(): void { + #beginActiveConnection(): void { + if (this.#activeLane) { + return; + } + this.#clearReconnectTimer(); - this.#setState(this.#reconnectAttempts > 0 ? 'reconnecting' : 'connecting'); + this.#setState( + this.#reconnectAttempts > 0 || this.#retiringLane ? 'reconnecting' : 'connecting' + ); const ws = new WebSocket(this.#url, { headers: { @@ -430,99 +443,134 @@ export class AIGatewayWebSocketClient { 'x-agentuity-orgid': this.#orgId, }, }); - this.#ws = ws; + const lane: WebSocketLane = { + ws, + role: 'active', + requestIds: new Set(), + }; + this.#activeLane = lane; + this.#attachLaneHandlers(lane); + } + + #attachLaneHandlers(lane: WebSocketLane): void { + const { ws } = lane; ws.onopen = () => { - if (ws !== this.#ws) { + if (!this.#isTrackedLane(lane)) { return; } - this.#reconnectAttempts = 0; - this.#draining = false; - this.#drainReconnectScheduled = false; - this.#setState('connected'); - this.#options.onOpen?.(); - this.#finishConnect(); + if (lane.role === 'active') { + this.#reconnectAttempts = 0; + this.#setState('connected'); + this.#options.onOpen?.(); + this.#finishConnect(); + } }; ws.onmessage = (event) => { - if (ws !== this.#ws) { + if (!this.#isTrackedLane(lane)) { return; } - this.#handleMessage(event.data); + this.#handleMessage(event.data, lane); }; ws.onerror = () => { - if (ws !== this.#ws) { + if (!this.#isTrackedLane(lane)) { return; } - this.#options.onError?.( - new AIGatewayWebSocketError({ - message: 'WebSocket connection error', - code: 'connection_error', - }) - ); + if (lane.role === 'active') { + this.#options.onError?.( + new AIGatewayWebSocketError({ + message: 'WebSocket connection error', + code: 'connection_error', + }) + ); + } }; - ws.onclose = (event) => { - if (ws !== this.#ws) { + ws.onclose = (event: CloseEvent) => { + if (!this.#isTrackedLane(lane)) { return; } - this.#ws = null; - const wasDraining = this.#draining; - const terminalClose = isTerminalCloseCode(event.code); - if (terminalClose) { - this.#intentionallyClosed = true; - } - - if (this.#state !== 'closed') { - this.#setState(wasDraining ? 'draining' : 'closed'); - } + this.#handleLaneClose(lane, event); + }; + } - this.#options.onClose?.(event.code, event.reason); + #handleLaneClose(lane: WebSocketLane, event: CloseEvent): void { + const terminalClose = isTerminalCloseCode(event.code); + if (terminalClose) { + this.#intentionallyClosed = true; + } - if (!wasDraining) { - this.#rejectAllPending( - new AIGatewayWebSocketError({ - message: `WebSocket closed (code ${event.code})${event.reason ? `: ${event.reason}` : ''}`, - code: 'connection_closed', - closeCode: event.code, - closeReason: event.reason || undefined, - }) - ); - } + this.#options.onClose?.(event.code, event.reason); - this.#finishConnect( + if (lane.role === 'retiring') { + this.#retiringLane = null; + this.#rejectLanePending( + lane, new AIGatewayWebSocketError({ - message: `WebSocket closed before connection was ready (code ${event.code})`, - code: 'connection_error', + message: + event.reason || + 'WebSocket closed during server drain before all in-flight requests completed', + code: 'connection_closed', closeCode: event.code, closeReason: event.reason || undefined, }) ); - - if (wasDraining) { - if (this.#pending.size > 0) { - this.#rejectAllPending( - new AIGatewayWebSocketError({ - message: - 'WebSocket closed during server drain before all in-flight requests completed', - code: 'connection_closed', - closeCode: event.code, - closeReason: event.reason || undefined, - }) - ); - } - this.#maybeScheduleDrainReconnect(); - return; + if (this.#activeLane?.ws.readyState === WebSocket.OPEN) { + this.#setState('connected'); } + return; + } - if (!this.#intentionallyClosed && this.#options.autoReconnect) { - this.#scheduleReconnect(); - } - }; + this.#activeLane = null; + this.#rejectLanePending( + lane, + new AIGatewayWebSocketError({ + message: `WebSocket closed (code ${event.code})${event.reason ? `: ${event.reason}` : ''}`, + code: 'connection_closed', + closeCode: event.code, + closeReason: event.reason || undefined, + }) + ); + + this.#finishConnect( + new AIGatewayWebSocketError({ + message: `WebSocket closed before connection was ready (code ${event.code})`, + code: 'connection_error', + closeCode: event.code, + closeReason: event.reason || undefined, + }) + ); + + if (!this.#intentionallyClosed && this.#options.autoReconnect) { + this.#scheduleReconnect(); + } else if (!this.#activeLane && !this.#retiringLane) { + this.#setState('closed'); + } } - #handleMessage(raw: unknown): void { + #handleDraining(message: string, lane: WebSocketLane): void { + if (lane.role !== 'active') { + return; + } + + this.#options.onDraining?.(message); + + if (this.#activeLane !== lane) { + return; + } + + lane.role = 'retiring'; + this.#retiringLane = lane; + this.#activeLane = null; + this.#setState('reconnecting'); + + this.#options.onReconnect?.(1); + this.#beginActiveConnection(); + } + + #handleMessage(raw: unknown, lane: WebSocketLane): void { let parsed: unknown; try { parsed = typeof raw === 'string' ? JSON.parse(raw) : raw; @@ -548,7 +596,7 @@ export class AIGatewayWebSocketClient { } if (frame.type === AIGatewayWSFrameType.draining) { - this.#handleDraining(frame.message ?? 'server is draining'); + this.#handleDraining(frame.message ?? 'server is draining', lane); return; } @@ -560,15 +608,6 @@ export class AIGatewayWebSocketClient { this.#handleResponseFrame(frame); } - #handleDraining(message: string): void { - if (this.#draining) { - return; - } - this.#draining = true; - this.#setState('draining'); - this.#options.onDraining?.(message); - } - #handleErrorFrame(error: AIGatewayWSServerError): void { const mapped = mapServerError(error); if (error.id) { @@ -613,9 +652,10 @@ export class AIGatewayWebSocketClient { }); clearTimeout(pending.timeout); this.#pending.delete(response.id); + this.#removeRequestFromLanes(response.id); pending.push?.({ type: 'complete', result }); pending.resolve(result); - this.#maybeScheduleDrainReconnect(); + this.#clearRetiringLaneIfIdle(); } } @@ -626,34 +666,56 @@ export class AIGatewayWebSocketClient { } clearTimeout(pending.timeout); this.#pending.delete(id); + this.#removeRequestFromLanes(id); pending.reject(error); - this.#maybeScheduleDrainReconnect(); + this.#clearRetiringLaneIfIdle(); + } + + #rejectLanePending(lane: WebSocketLane, error: AIGatewayWebSocketErrorInstance): void { + for (const id of [...lane.requestIds]) { + if (this.#pending.has(id)) { + this.#rejectPending(id, error); + } + } + lane.requestIds.clear(); } #rejectAllPending(error: AIGatewayWebSocketErrorInstance): void { - for (const [id, pending] of this.#pending.entries()) { - clearTimeout(pending.timeout); - pending.reject(error); - this.#pending.delete(id); + for (const id of [...this.#pending.keys()]) { + this.#rejectPending(id, error); } } - #maybeScheduleDrainReconnect(): void { - if (!this.#draining || this.#drainReconnectScheduled || this.#intentionallyClosed) { - return; + #removeRequestFromLanes(requestId: string): void { + this.#activeLane?.requestIds.delete(requestId); + this.#retiringLane?.requestIds.delete(requestId); + } + + #findLaneForRequest(requestId: string): WebSocketLane | null { + if (this.#activeLane?.requestIds.has(requestId)) { + return this.#activeLane; + } + if (this.#retiringLane?.requestIds.has(requestId)) { + return this.#retiringLane; } - if (this.#pending.size > 0 || this.#ws) { + return null; + } + + #clearRetiringLaneIfIdle(): void { + if (!this.#retiringLane || this.#retiringLane.requestIds.size > 0) { return; } - this.#drainReconnectScheduled = true; - this.#draining = false; - if (this.#options.autoReconnect) { - this.#scheduleReconnect(); + if (this.#retiringLane.ws.readyState === WebSocket.CLOSED) { + this.#retiringLane = null; } } + #isTrackedLane(lane: WebSocketLane): boolean { + return this.#activeLane === lane || this.#retiringLane === lane; + } + #scheduleReconnect(): void { - if (this.#intentionallyClosed || !this.#options.autoReconnect) { + if (this.#intentionallyClosed || !this.#options.autoReconnect || this.#activeLane) { return; } if (this.#reconnectAttempts >= this.#options.maxReconnectAttempts) { @@ -673,7 +735,7 @@ export class AIGatewayWebSocketClient { this.#options.onReconnect?.(this.#reconnectAttempts); this.#reconnectTimer = setTimeout(() => { this.#reconnectTimer = null; - this.#connectInternal(); + this.#beginActiveConnection(); }, delay); } @@ -684,14 +746,26 @@ export class AIGatewayWebSocketClient { } } - #sendJson(payload: Record): void { - if (!this.#ws || this.#ws.readyState !== WebSocket.OPEN) { + #sendJson(payload: Record, requestId: string): void { + const lane = this.#activeLane; + if (!lane || lane.ws.readyState !== WebSocket.OPEN) { + throw new AIGatewayWebSocketError({ + message: 'WebSocket is not connected', + code: 'connection_error', + }); + } + lane.requestIds.add(requestId); + this.#sendJsonOnLane(lane, payload); + } + + #sendJsonOnLane(lane: WebSocketLane, payload: Record): void { + if (lane.ws.readyState !== WebSocket.OPEN) { throw new AIGatewayWebSocketError({ message: 'WebSocket is not connected', code: 'connection_error', }); } - this.#ws.send(JSON.stringify(payload)); + lane.ws.send(JSON.stringify(payload)); } #setState(state: AIGatewayWebSocketState): void { diff --git a/packages/aigateway/test/websocket.test.ts b/packages/aigateway/test/websocket.test.ts index 8de29ed33..85bfb5d8e 100644 --- a/packages/aigateway/test/websocket.test.ts +++ b/packages/aigateway/test/websocket.test.ts @@ -218,14 +218,13 @@ describe('AIGatewayWebSocketClient', () => { expect(events).toEqual(['thinking_delta:hmm', 'delta:Hello', 'complete:Hello world']); }); - it('blocks new requests while draining and reconnects after in-flight work completes', async () => { + it('handoffs to a new websocket while the retiring socket completes in-flight requests', async () => { let drainingMessage = ''; let reconnectAttempt = 0; const client = new AIGatewayWebSocketClient({ apiKey: 'ag_test', orgId: 'org_test', url: 'https://aigateway.example', - reconnectDelayMs: 1, onDraining: (message) => { drainingMessage = message; }, @@ -240,49 +239,62 @@ describe('AIGatewayWebSocketClient', () => { await connectPromise; const inFlight = client.complete({ + id: 'req_inflight', model: 'anthropic/claude-sonnet-4-20250514', prompt: 'Finish me', }); await flushAsyncWork(); - const request = JSON.parse(firstWs?.sent[0] ?? '{}'); firstWs?.receive({ type: 'draining', message: 'server rolling restart', }); + expect(client.isDraining).toBe(true); expect(drainingMessage).toBe('server rolling restart'); - expect(client.state).toBe('draining'); + expect(client.state).toBe('reconnecting'); + expect(reconnectAttempt).toBe(1); + expect(MockWebSocket.instances).toHaveLength(2); - await expect( - client.complete({ - model: 'anthropic/claude-sonnet-4-20250514', - prompt: 'Should fail', - }) - ).rejects.toMatchObject({ - code: 'connection_draining', + const secondWs = MockWebSocket.instances[1]; + secondWs?.open(); + await flushAsyncWork(); + expect(client.state).toBe('connected'); + + const newRequestPromise = client.complete({ + id: 'req_new', + model: 'anthropic/claude-sonnet-4-20250514', + prompt: 'New request on fresh socket', + }); + await flushAsyncWork(); + + expect(JSON.parse(firstWs?.sent[0] ?? '{}')).toMatchObject({ id: 'req_inflight' }); + expect(JSON.parse(secondWs?.sent[0] ?? '{}')).toMatchObject({ + id: 'req_new', + prompt: 'New request on fresh socket', }); firstWs?.receive({ type: 'response', - id: request.id, + id: 'req_inflight', status: 'complete', content: 'done', }); + secondWs?.receive({ + type: 'response', + id: 'req_new', + status: 'complete', + content: 'fresh', + }); + await expect(inFlight).resolves.toMatchObject({ content: 'done' }); + await expect(newRequestPromise).resolves.toMatchObject({ content: 'fresh' }); firstWs?.close(1001, 'going away'); await flushAsyncWork(); - await new Promise((resolve) => setTimeout(resolve, 5)); - - expect(reconnectAttempt).toBe(1); - expect(MockWebSocket.instances.length).toBe(2); expect(client.isDraining).toBe(false); - - const secondWs = MockWebSocket.instances[1]; - secondWs?.open(); - await flushAsyncWork(); expect(client.state).toBe('connected'); + expect(MockWebSocket.instances).toHaveLength(2); }); it('multiplexes concurrent requests over one connection by request id', async () => { From 4806ebb3b632226e05ad45edc8500e2978d4b314 Mon Sep 17 00:00:00 2001 From: Jeff Haynie Date: Sat, 4 Jul 2026 10:16:52 -0500 Subject: [PATCH 3/7] Fix aigateway dist imports and require orgId only for CLI keys. Use .ts extension in websocket protocol import so Node can resolve dist/protocol.js, and gate WebSocket orgId on ck_* API keys. Co-authored-by: Cursor --- packages/aigateway/src/index.ts | 5 ++- packages/aigateway/src/protocol.ts | 4 ++ packages/aigateway/src/websocket.ts | 17 +++---- packages/aigateway/test/index.test.ts | 55 +++++++++++++++++++++++ packages/aigateway/test/websocket.test.ts | 16 +++++++ 5 files changed, 87 insertions(+), 10 deletions(-) diff --git a/packages/aigateway/src/index.ts b/packages/aigateway/src/index.ts index ec5562fa7..e2ef42e19 100644 --- a/packages/aigateway/src/index.ts +++ b/packages/aigateway/src/index.ts @@ -26,6 +26,7 @@ import { resolveServiceUrl, type Logger, } from '@agentuity/client'; +import { isCliApiKey } from './protocol.ts'; import { createAIGatewayWebSocketClient, type AIGatewayWebSocketClient, @@ -94,9 +95,9 @@ export class AIGatewayClient { 'API key is required for AI Gateway WebSocket connections. Provide apiKey when constructing AIGatewayClient or set AGENTUITY_AIGATEWAY_KEY / AGENTUITY_SDK_KEY.' ); } - if (!this.#orgId) { + if (isCliApiKey(this.#apiKey) && !this.#orgId) { throw new Error( - 'Organization ID is required for AI Gateway WebSocket connections. Provide orgId when constructing AIGatewayClient or set AGENTUITY_ORGID / AGENTUITY_ORG_ID / AGENTUITY_CLOUD_ORG_ID.' + '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.' ); } return createAIGatewayWebSocketClient({ diff --git a/packages/aigateway/src/protocol.ts b/packages/aigateway/src/protocol.ts index 2851ae424..a8ef0e151 100644 --- a/packages/aigateway/src/protocol.ts +++ b/packages/aigateway/src/protocol.ts @@ -1,5 +1,9 @@ import { z } from 'zod'; +export function isCliApiKey(apiKey: string): boolean { + return apiKey.startsWith('ck_'); +} + export const AIGatewayWSFrameType = { request: 'request', response: 'response', diff --git a/packages/aigateway/src/websocket.ts b/packages/aigateway/src/websocket.ts index 5e86ffcd8..b8bd65da0 100644 --- a/packages/aigateway/src/websocket.ts +++ b/packages/aigateway/src/websocket.ts @@ -6,11 +6,12 @@ import { AIGatewayWSFrameType, AIGatewayWSResponseStatus, buildAIGatewayWebSocketUrl, + isCliApiKey, parseAIGatewayWSServerFrame, type AIGatewayWSUsage, type AIGatewayWSServerError, type AIGatewayWSServerResponse, -} from './protocol'; +} from './protocol.ts'; export const AIGatewayWebSocketErrorCode = { connection_error: 'connection_error', @@ -120,16 +121,16 @@ function normalizeOrgId(orgId: string | undefined): string | undefined { return trimmed ? trimmed : undefined; } -function resolveOrgId(options: { orgId?: string }): string { +function resolveOrgId(options: { orgId?: string; apiKey: string }): string | undefined { const resolved = normalizeOrgId(options.orgId) ?? normalizeOrgId(getEnv('AGENTUITY_ORGID')) ?? normalizeOrgId(getEnv('AGENTUITY_ORG_ID')) ?? normalizeOrgId(getEnv('AGENTUITY_CLOUD_ORG_ID')); - if (!resolved) { + if (!resolved && isCliApiKey(options.apiKey)) { throw new AIGatewayWebSocketError({ message: - 'Organization ID is required. Provide orgId in options or set AGENTUITY_ORGID / AGENTUITY_ORG_ID / AGENTUITY_CLOUD_ORG_ID.', + 'Organization ID is required when using a CLI API key (ck_*). Provide orgId in options or set AGENTUITY_ORGID / AGENTUITY_ORG_ID / AGENTUITY_CLOUD_ORG_ID.', code: 'connection_error', }); } @@ -221,7 +222,7 @@ export class AIGatewayWebSocketClient { AIGatewayWebSocketOptions; readonly #apiKey: string; - readonly #orgId: string; + readonly #orgId: string | undefined; readonly #url: string; #activeLane: WebSocketLane | null = null; @@ -251,7 +252,7 @@ export class AIGatewayWebSocketClient { ...options, }; this.#apiKey = options.apiKey; - this.#orgId = resolveOrgId(options); + this.#orgId = resolveOrgId({ orgId: options.orgId, apiKey: options.apiKey }); this.#url = resolveWebSocketUrl(options); } @@ -267,7 +268,7 @@ export class AIGatewayWebSocketClient { return this.#url; } - get orgId(): string { + get orgId(): string | undefined { return this.#orgId; } @@ -440,7 +441,7 @@ export class AIGatewayWebSocketClient { const ws = new WebSocket(this.#url, { headers: { Authorization: `Bearer ${this.#apiKey}`, - 'x-agentuity-orgid': this.#orgId, + ...(this.#orgId ? { 'x-agentuity-orgid': this.#orgId } : {}), }, }); const lane: WebSocketLane = { diff --git a/packages/aigateway/test/index.test.ts b/packages/aigateway/test/index.test.ts index c24b8e237..cd1ed726b 100644 --- a/packages/aigateway/test/index.test.ts +++ b/packages/aigateway/test/index.test.ts @@ -274,3 +274,58 @@ describe('AIGatewayClient.completeStructured', () => { expect(result.data).toEqual({ name: 'ada', age: 42 }); }); }); + +describe('AIGatewayClient.createWebSocket orgId requirement', () => { + const ORIGINAL_ENV = { + AGENTUITY_ORGID: process.env.AGENTUITY_ORGID, + AGENTUITY_ORG_ID: process.env.AGENTUITY_ORG_ID, + AGENTUITY_CLOUD_ORG_ID: process.env.AGENTUITY_CLOUD_ORG_ID, + }; + + beforeEach(() => { + delete process.env.AGENTUITY_ORGID; + delete process.env.AGENTUITY_ORG_ID; + delete process.env.AGENTUITY_CLOUD_ORG_ID; + process.env.AGENTUITY_AIGATEWAY_URL = 'https://aigateway.test'; + }); + + afterEach(() => { + for (const [key, value] of Object.entries(ORIGINAL_ENV)) { + if (value === undefined) { + delete process.env[key]; + } else { + process.env[key] = value; + } + } + }); + + test('allows SDK keys without orgId', () => { + const client = new AIGatewayClient({ + apiKey: 'ag_test_sdk_key', + url: 'https://aigateway.test', + }); + + expect(() => client.createWebSocket()).not.toThrow(); + }); + + test('requires orgId for CLI keys', () => { + const client = new AIGatewayClient({ + apiKey: 'ck_test_cli_key', + url: 'https://aigateway.test', + }); + + expect(() => client.createWebSocket()).toThrow( + 'Organization ID is required for AI Gateway WebSocket connections when using a CLI API key (ck_*).' + ); + }); + + test('allows CLI keys when orgId is provided', () => { + const client = new AIGatewayClient({ + apiKey: 'ck_test_cli_key', + orgId: 'org_explicit', + url: 'https://aigateway.test', + }); + + expect(() => client.createWebSocket()).not.toThrow(); + }); +}); diff --git a/packages/aigateway/test/websocket.test.ts b/packages/aigateway/test/websocket.test.ts index 85bfb5d8e..421371b79 100644 --- a/packages/aigateway/test/websocket.test.ts +++ b/packages/aigateway/test/websocket.test.ts @@ -119,6 +119,22 @@ describe('AIGatewayWebSocketClient', () => { MockWebSocket.instances = []; }); + it('connects without org header for SDK keys', async () => { + const client = new AIGatewayWebSocketClient({ + apiKey: 'ag_test_sdk_key', + url: 'https://aigateway.example', + }); + + const connectPromise = client.connect(); + const ws = MockWebSocket.instances[0]; + expect(ws?.headers).toEqual({ + Authorization: 'Bearer ag_test_sdk_key', + }); + + ws?.open(); + await connectPromise; + }); + it('connects with auth headers and completes compact requests', async () => { const client = new AIGatewayWebSocketClient({ apiKey: 'ag_test', From ad2e4c2e84f7379adb203ce2e9593e1b1c64cef8 Mon Sep 17 00:00:00 2001 From: Jeff Haynie Date: Sun, 5 Jul 2026 21:18:09 -0500 Subject: [PATCH 4/7] fixes --- packages/aigateway/src/websocket.ts | 89 ++++++++++-- packages/aigateway/test/websocket.test.ts | 161 ++++++++++++++++++++++ 2 files changed, 242 insertions(+), 8 deletions(-) diff --git a/packages/aigateway/src/websocket.ts b/packages/aigateway/src/websocket.ts index b8bd65da0..1adbc546d 100644 --- a/packages/aigateway/src/websocket.ts +++ b/packages/aigateway/src/websocket.ts @@ -1,5 +1,5 @@ import { StructuredError } from '@agentuity/adapter'; -import { resolveRegion, resolveServiceUrl } from '@agentuity/client'; +import { resolveRegion, resolveServiceUrl, type Logger } from '@agentuity/client'; import { getEnv, getServiceUrls } from '@agentuity/config'; import { z } from 'zod'; import { @@ -12,6 +12,7 @@ import { type AIGatewayWSServerError, type AIGatewayWSServerResponse, } from './protocol.ts'; +import { getAIGatewayCompletionText, getAIGatewayStreamDeltaText } from './service.ts'; export const AIGatewayWebSocketErrorCode = { connection_error: 'connection_error', @@ -85,6 +86,9 @@ export interface AIGatewayWebSocketOptions { maxReconnectDelayMs?: number; maxReconnectAttempts?: number; defaultTimeoutMs?: number; + /** Log outbound request and inbound response frames. Defaults to AGENTUITY_AIGATEWAY_WS_DEBUG. */ + debug?: boolean; + logger?: Logger; onOpen?: () => void; onClose?: (code: number, reason: string) => void; onError?: (error: AIGatewayWebSocketErrorInstance) => void; @@ -154,6 +158,21 @@ function isTerminalCloseCode(code: number): boolean { return code >= 4000 && code < 5000; } +function isTruthyEnv(value: string | undefined): boolean { + if (!value) { + return false; + } + const normalized = value.trim().toLowerCase(); + return normalized === '1' || normalized === 'true' || normalized === 'yes'; +} + +function resolveDebugEnabled(options: AIGatewayWebSocketOptions): boolean { + if (options.debug !== undefined) { + return options.debug; + } + return isTruthyEnv(getEnv('AGENTUITY_AIGATEWAY_WS_DEBUG')); +} + function createRequestId(): string { if (typeof crypto !== 'undefined' && typeof crypto.randomUUID === 'function') { return crypto.randomUUID(); @@ -166,25 +185,42 @@ function buildRequestFrame( id: string ): Record { const parsed = AIGatewayWSRequestOptionsSchema.parse(options); + const compact = + options.compact ?? (parsed.data === undefined ? (parsed.compact ?? true) : false); const frame: Record = { type: AIGatewayWSFrameType.request, id, - compact: parsed.compact ?? true, + compact, }; - if (parsed.compact) { + + const applySharedFields = () => { if (parsed.model !== undefined) frame.model = parsed.model; - if (parsed.prompt !== undefined) frame.prompt = parsed.prompt; - if (parsed.system !== undefined) frame.system = parsed.system; if (parsed.thinking !== undefined) frame.thinking = parsed.thinking; if (parsed.temperature !== undefined) frame.temperature = parsed.temperature; if (parsed.max_tokens !== undefined) frame.max_tokens = parsed.max_tokens; if (parsed.stream !== undefined) frame.stream = parsed.stream; - } else if (parsed.data !== undefined) { - frame.data = parsed.data; + }; + + if (compact) { + applySharedFields(); + if (parsed.prompt !== undefined) frame.prompt = parsed.prompt; + if (parsed.system !== undefined) frame.system = parsed.system; + } else { + applySharedFields(); + if (parsed.data !== undefined) frame.data = parsed.data; } + return frame; } +/** @internal Exported for unit tests and SDK consumers building frames manually. */ +export function buildAIGatewayWebSocketRequestFrame( + options: AIGatewayWSRequestOptions, + id: string +): Record { + return buildRequestFrame(options, id); +} + function mapCompleteResponse(response: AIGatewayWSServerResponse): AIGatewayWSResult { return { id: response.id, @@ -224,6 +260,7 @@ export class AIGatewayWebSocketClient { readonly #apiKey: string; readonly #orgId: string | undefined; readonly #url: string; + readonly #debugEnabled: boolean; #activeLane: WebSocketLane | null = null; #retiringLane: WebSocketLane | null = null; @@ -254,6 +291,18 @@ export class AIGatewayWebSocketClient { this.#apiKey = options.apiKey; this.#orgId = resolveOrgId({ orgId: options.orgId, apiKey: options.apiKey }); this.#url = resolveWebSocketUrl(options); + this.#debugEnabled = resolveDebugEnabled(options); + } + + #logDebug(message: string, detail?: unknown): void { + if (!this.#debugEnabled && !resolveDebugEnabled(this.#options)) { + return; + } + if (detail !== undefined) { + console.log('[aigateway-ws]', message, detail); + } else { + console.log('[aigateway-ws]', message); + } } get state(): AIGatewayWebSocketState { @@ -463,6 +512,7 @@ export class AIGatewayWebSocketClient { if (lane.role === 'active') { this.#reconnectAttempts = 0; this.#setState('connected'); + this.#logDebug('connected', { url: this.#url, orgId: this.#orgId }); this.#options.onOpen?.(); this.#finishConnect(); } @@ -572,6 +622,7 @@ export class AIGatewayWebSocketClient { } #handleMessage(raw: unknown, lane: WebSocketLane): void { + this.#logDebug('recv raw', raw); let parsed: unknown; try { parsed = typeof raw === 'string' ? JSON.parse(raw) : raw; @@ -597,15 +648,18 @@ export class AIGatewayWebSocketClient { } if (frame.type === AIGatewayWSFrameType.draining) { + this.#logDebug('recv draining', frame); this.#handleDraining(frame.message ?? 'server is draining', lane); return; } if (frame.type === AIGatewayWSFrameType.error) { + this.#logDebug('recv error', frame); this.#handleErrorFrame(frame); return; } + this.#logDebug('recv', frame); this.#handleResponseFrame(frame); } @@ -625,6 +679,22 @@ export class AIGatewayWebSocketClient { } if (response.status === AIGatewayWSResponseStatus.delta) { + // Full-mode (compact: false) provider-native frames carry payload in `data`, + // not top-level `delta`. Passthrough the native chunk only — Hub rehydrates + // provider SSE from `event` frames; also emitting extracted `delta` duplicates text. + if (response.data !== undefined && response.delta === undefined) { + const extracted = getAIGatewayStreamDeltaText(response.data); + if (extracted) { + pending.accumulated += extracted; + } + pending.push?.({ + type: 'event', + event: response.event ?? 'data', + data: response.data, + }); + return; + } + const delta = response.delta ?? ''; pending.accumulated += delta; pending.push?.({ type: 'delta', delta }); @@ -647,9 +717,11 @@ export class AIGatewayWebSocketClient { } if (response.status === AIGatewayWSResponseStatus.complete) { + const content = + response.content ?? (pending.accumulated || getAIGatewayCompletionText(response.data)); const result = mapCompleteResponse({ ...response, - content: response.content ?? pending.accumulated, + content: content || undefined, }); clearTimeout(pending.timeout); this.#pending.delete(response.id); @@ -766,6 +838,7 @@ export class AIGatewayWebSocketClient { code: 'connection_error', }); } + this.#logDebug('send', payload); lane.ws.send(JSON.stringify(payload)); } diff --git a/packages/aigateway/test/websocket.test.ts b/packages/aigateway/test/websocket.test.ts index 421371b79..bb6660364 100644 --- a/packages/aigateway/test/websocket.test.ts +++ b/packages/aigateway/test/websocket.test.ts @@ -183,6 +183,62 @@ describe('AIGatewayWebSocketClient', () => { }); }); + it('logs outbound and inbound frames when debug is enabled', async () => { + const logs: unknown[][] = []; + const originalLog = console.log; + console.log = (...args: unknown[]) => { + logs.push(args); + }; + + try { + const client = new AIGatewayWebSocketClient({ + apiKey: 'ag_test', + orgId: 'org_test', + url: 'https://aigateway.example', + debug: true, + }); + + const connectPromise = client.connect(); + const ws = MockWebSocket.instances[0]; + ws?.open(); + await connectPromise; + + expect(logs.some((row) => row[1] === 'connected')).toBe(true); + + const streamPromise = (async () => { + for await (const _event of client.stream({ + model: 'anthropic/claude-sonnet-4-20250514', + prompt: 'Hello', + stream: true, + })) { + // drain + } + })(); + + await flushAsyncWork(); + const request = JSON.parse(MockWebSocket.instances[0]?.sent[0] ?? '{}'); + expect(logs.some((row) => row[1] === 'send' && row[2]?.type === 'request')).toBe(true); + + MockWebSocket.instances[0]?.receive({ + type: 'response', + id: request.id, + status: 'delta', + delta: 'Hi', + }); + MockWebSocket.instances[0]?.receive({ + type: 'response', + id: request.id, + status: 'complete', + content: 'Hi', + }); + + await streamPromise; + expect(logs.some((row) => row[1] === 'recv' && row[2]?.status === 'delta')).toBe(true); + } finally { + console.log = originalLog; + } + }); + it('streams delta and thinking frames before complete', async () => { const client = new AIGatewayWebSocketClient({ apiKey: 'ag_test', @@ -407,4 +463,109 @@ describe('AIGatewayWebSocketClient', () => { requestId: 'req_error', }); }); + + it('streams full-mode provider-native delta frames from response.data', async () => { + const client = new AIGatewayWebSocketClient({ + apiKey: 'ag_test', + orgId: 'org_test', + url: 'https://aigateway.example', + }); + + const connectPromise = client.connect(); + const ws = MockWebSocket.instances[0]; + ws?.open(); + await connectPromise; + + const events: Array<{ type: string; delta?: string; event?: string }> = []; + const streamPromise = (async () => { + for await (const event of client.stream({ + compact: false, + model: 'oss/glm-5p2', + data: { model: 'GLM-5.2', messages: [{ role: 'user', content: 'Hi' }] }, + stream: true, + })) { + if (event.type === 'delta') { + events.push({ type: event.type, delta: event.delta }); + } else if (event.type === 'event') { + events.push({ type: event.type, event: event.event }); + } else if (event.type === 'complete') { + events.push({ type: event.type, delta: event.result.content }); + } + } + })(); + + await flushAsyncWork(); + const request = JSON.parse(MockWebSocket.instances[0]?.sent[0] ?? '{}'); + + MockWebSocket.instances[0]?.receive({ + type: 'response', + id: request.id, + status: 'delta', + data: { + object: 'chat.completion.chunk', + choices: [{ index: 0, delta: { content: 'Hello' } }], + }, + }); + MockWebSocket.instances[0]?.receive({ + type: 'response', + id: request.id, + status: 'complete', + data: { + choices: [{ index: 0, message: { role: 'assistant', content: 'Hello' } }], + }, + }); + + await streamPromise; + expect(events).toEqual([ + { type: 'event', event: 'data' }, + { type: 'complete', delta: 'Hello' }, + ]); + }); + + it('includes model and stream on full-mode request frames', async () => { + const client = new AIGatewayWebSocketClient({ + apiKey: 'ag_test', + orgId: 'org_test', + url: 'https://aigateway.example', + }); + + const connectPromise = client.connect(); + const ws = MockWebSocket.instances[0]; + ws?.open(); + await connectPromise; + + const resultPromise = client.complete({ + compact: false, + model: 'openai/gpt-4.1-mini', + data: { + model: 'gpt-4.1-mini', + messages: [{ role: 'user', content: 'Hello' }], + }, + stream: false, + }); + await flushAsyncWork(); + + const request = JSON.parse(ws?.sent[0] ?? '{}'); + expect(request).toMatchObject({ + type: 'request', + compact: false, + model: 'openai/gpt-4.1-mini', + stream: false, + data: { + model: 'gpt-4.1-mini', + messages: [{ role: 'user', content: 'Hello' }], + }, + }); + + ws?.receive({ + type: 'response', + id: request.id, + status: 'complete', + content: 'Hi there', + }); + + await expect(resultPromise).resolves.toMatchObject({ + content: 'Hi there', + }); + }); }); From 53168afa57a4d10a14b4071d39bd9ab9a59b65d8 Mon Sep 17 00:00:00 2001 From: Jeff Haynie Date: Sun, 5 Jul 2026 21:31:52 -0500 Subject: [PATCH 5/7] fix(aigateway): resolve all 9 inline/nitpick findings - Use AIGatewayWebSocketError in createWebSocket guard clauses - Guard cancel() in timeout handler to always reject pending - Fix options spread order to preserve defaults - Import StructuredError from @agentuity/core - Check streamError after async generator loop exits - Close idle retiring lane instead of waiting for CLOSED - Remove leaking AGENTUITY_AIGATEWAY_URL from test setup - Replace deprecated .passthrough() with z.looseObject() - Route #logDebug through configured logger --- packages/aigateway/src/index.ts | 17 ++++++++----- packages/aigateway/src/protocol.ts | 36 +++++++++++++-------------- packages/aigateway/src/websocket.ts | 29 +++++++++++++++------ packages/aigateway/test/index.test.ts | 1 - 4 files changed, 50 insertions(+), 33 deletions(-) diff --git a/packages/aigateway/src/index.ts b/packages/aigateway/src/index.ts index e2ef42e19..9a0c2c6df 100644 --- a/packages/aigateway/src/index.ts +++ b/packages/aigateway/src/index.ts @@ -29,6 +29,7 @@ import { import { isCliApiKey } from './protocol.ts'; import { createAIGatewayWebSocketClient, + AIGatewayWebSocketError, type AIGatewayWebSocketClient, type AIGatewayWebSocketOptions, } from './websocket.ts'; @@ -91,14 +92,18 @@ export class AIGatewayClient { } = {} ): AIGatewayWebSocketClient { if (!this.#apiKey) { - throw new Error( - 'API key is required for AI Gateway WebSocket connections. Provide apiKey when constructing AIGatewayClient or set AGENTUITY_AIGATEWAY_KEY / AGENTUITY_SDK_KEY.' - ); + 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 Error( - '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.' - ); + 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, diff --git a/packages/aigateway/src/protocol.ts b/packages/aigateway/src/protocol.ts index a8ef0e151..4115e0e24 100644 --- a/packages/aigateway/src/protocol.ts +++ b/packages/aigateway/src/protocol.ts @@ -27,25 +27,23 @@ export const AIGatewayWSUsageSchema = z.object({ export type AIGatewayWSUsage = z.infer; -export const AIGatewayWSServerResponseSchema = z - .object({ - 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(), - }) - .passthrough(); +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; diff --git a/packages/aigateway/src/websocket.ts b/packages/aigateway/src/websocket.ts index 1adbc546d..a0f6a2cbf 100644 --- a/packages/aigateway/src/websocket.ts +++ b/packages/aigateway/src/websocket.ts @@ -1,4 +1,4 @@ -import { StructuredError } from '@agentuity/adapter'; +import { StructuredError } from '@agentuity/core'; import { resolveRegion, resolveServiceUrl, type Logger } from '@agentuity/client'; import { getEnv, getServiceUrls } from '@agentuity/config'; import { z } from 'zod'; @@ -281,12 +281,12 @@ export class AIGatewayWebSocketClient { }); } this.#options = { + ...options, autoReconnect: options.autoReconnect ?? true, reconnectDelayMs: options.reconnectDelayMs ?? DEFAULT_RECONNECT_DELAY_MS, maxReconnectDelayMs: options.maxReconnectDelayMs ?? DEFAULT_MAX_RECONNECT_DELAY_MS, maxReconnectAttempts: options.maxReconnectAttempts ?? DEFAULT_MAX_RECONNECT_ATTEMPTS, defaultTimeoutMs: options.defaultTimeoutMs ?? DEFAULT_TIMEOUT_MS, - ...options, }; this.#apiKey = options.apiKey; this.#orgId = resolveOrgId({ orgId: options.orgId, apiKey: options.apiKey }); @@ -295,10 +295,16 @@ export class AIGatewayWebSocketClient { } #logDebug(message: string, detail?: unknown): void { - if (!this.#debugEnabled && !resolveDebugEnabled(this.#options)) { + if (!this.#debugEnabled) { return; } - if (detail !== undefined) { + if (this.#options.logger) { + if (detail !== undefined) { + this.#options.logger.debug(`[aigateway-ws] ${message}`, detail); + } else { + this.#options.logger.debug(`[aigateway-ws] ${message}`); + } + } else if (detail !== undefined) { console.log('[aigateway-ws]', message, detail); } else { console.log('[aigateway-ws]', message); @@ -420,7 +426,11 @@ export class AIGatewayWebSocketClient { notify?.(); }, timeout: setTimeout(() => { - this.cancel(id); + try { + this.cancel(id); + } catch { + // cancel can fail if socket is not OPEN; always reject + } this.#rejectPending( id, new AIGatewayWebSocketError({ @@ -461,6 +471,9 @@ export class AIGatewayWebSocketClient { done = true; } } + if (streamError) { + throw streamError; + } } finally { clearTimeout(pending.timeout); this.#pending.delete(id); @@ -778,9 +791,11 @@ export class AIGatewayWebSocketClient { if (!this.#retiringLane || this.#retiringLane.requestIds.size > 0) { return; } - if (this.#retiringLane.ws.readyState === WebSocket.CLOSED) { - this.#retiringLane = null; + const ws = this.#retiringLane.ws; + if (ws.readyState !== WebSocket.CLOSED) { + ws.close(1000, 'lane idle'); } + this.#retiringLane = null; } #isTrackedLane(lane: WebSocketLane): boolean { diff --git a/packages/aigateway/test/index.test.ts b/packages/aigateway/test/index.test.ts index cd1ed726b..ff97d798d 100644 --- a/packages/aigateway/test/index.test.ts +++ b/packages/aigateway/test/index.test.ts @@ -286,7 +286,6 @@ describe('AIGatewayClient.createWebSocket orgId requirement', () => { delete process.env.AGENTUITY_ORGID; delete process.env.AGENTUITY_ORG_ID; delete process.env.AGENTUITY_CLOUD_ORG_ID; - process.env.AGENTUITY_AIGATEWAY_URL = 'https://aigateway.test'; }); afterEach(() => { From bf55858c66231efad522fd8133ec4472e200c87d Mon Sep 17 00:00:00 2001 From: Jeff Haynie Date: Sun, 5 Jul 2026 21:37:57 -0500 Subject: [PATCH 6/7] fix(aigateway): add @agentuity/core to declared dependencies --- packages/aigateway/package.json | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/aigateway/package.json b/packages/aigateway/package.json index 98375123f..2375bfa07 100644 --- a/packages/aigateway/package.json +++ b/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" From 53d1e0bb07d6b92d10bff43eddf080b508c97e2c Mon Sep 17 00:00:00 2001 From: Jeff Haynie Date: Sun, 5 Jul 2026 21:43:03 -0500 Subject: [PATCH 7/7] fix(aigateway): add ../core to tsconfig project references --- packages/aigateway/tsconfig.json | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/packages/aigateway/tsconfig.json b/packages/aigateway/tsconfig.json index 6a51b3923..49c0df624 100644 --- a/packages/aigateway/tsconfig.json +++ b/packages/aigateway/tsconfig.json @@ -7,5 +7,10 @@ }, "include": ["src/**/*"], "exclude": ["node_modules", "dist"], - "references": [{ "path": "../adapter" }, { "path": "../client" }, { "path": "../config" }] + "references": [ + { "path": "../adapter" }, + { "path": "../client" }, + { "path": "../config" }, + { "path": "../core" } + ] }