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" diff --git a/packages/aigateway/src/index.ts b/packages/aigateway/src/index.ts index dbb5091bb..9a0c2c6df 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,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 & { + 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 { return this.#service.listModels(); } diff --git a/packages/aigateway/src/protocol.ts b/packages/aigateway/src/protocol.ts new file mode 100644 index 000000000..4115e0e24 --- /dev/null +++ b/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; + +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; + +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..a0f6a2cbf --- /dev/null +++ b/packages/aigateway/src/websocket.ts @@ -0,0 +1,887 @@ +import { StructuredError } from '@agentuity/core'; +import { resolveRegion, resolveServiceUrl, type Logger } from '@agentuity/client'; +import { getEnv, getServiceUrls } from '@agentuity/config'; +import { z } from 'zod'; +import { + AIGatewayWSFrameType, + AIGatewayWSResponseStatus, + buildAIGatewayWebSocketUrl, + isCliApiKey, + parseAIGatewayWSServerFrame, + type AIGatewayWSUsage, + type AIGatewayWSServerError, + type AIGatewayWSServerResponse, +} from './protocol.ts'; +import { getAIGatewayCompletionText, getAIGatewayStreamDeltaText } from './service.ts'; + +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; + /** 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; + 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; +} + +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; +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; 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 && isCliApiKey(options.apiKey)) { + throw new AIGatewayWebSocketError({ + message: + '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', + }); + } + 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 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(); + } + return `req_${Date.now()}_${Math.random().toString(36).slice(2, 10)}`; +} + +function buildRequestFrame( + options: AIGatewayWSRequestOptions, + 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, + }; + + const applySharedFields = () => { + if (parsed.model !== undefined) frame.model = parsed.model; + 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; + }; + + 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, + 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 | undefined; + readonly #url: string; + readonly #debugEnabled: boolean; + + #activeLane: WebSocketLane | null = null; + #retiringLane: WebSocketLane | null = null; + #state: AIGatewayWebSocketState = 'closed'; + #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(); + + constructor(options: AIGatewayWebSocketOptions) { + if (!options.apiKey) { + throw new AIGatewayWebSocketError({ + message: 'API key is required', + code: 'connection_error', + }); + } + 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, + }; + 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) { + return; + } + 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); + } + } + + get state(): AIGatewayWebSocketState { + return this.#state; + } + + get isDraining(): boolean { + return this.#retiringLane !== null; + } + + get url(): string { + return this.#url; + } + + get orgId(): string | undefined { + return this.#orgId; + } + + connect(): Promise { + if (this.#activeLane?.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; + if (!this.#activeLane) { + this.#beginActiveConnection(); + } + }); + 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, + }) + ); + 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({ + message: reason || 'WebSocket closed by client', + code: 'connection_closed', + closeCode: code, + closeReason: reason, + }) + ); + } + + cancel(requestId: string): void { + const lane = this.#findLaneForRequest(requestId); + if (!lane) { + return; + } + this.#sendJsonOnLane(lane, { + 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(() => { + try { + this.cancel(id); + } catch { + // cancel can fail if socket is not OPEN; always reject + } + 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), 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; + } + } + if (streamError) { + throw streamError; + } + } finally { + clearTimeout(pending.timeout); + this.#pending.delete(id); + this.#clearRetiringLaneIfIdle(); + } + } + + #assertCanSendRequest(): void { + if (this.#intentionallyClosed) { + throw new AIGatewayWebSocketError({ + message: 'WebSocket client is closed', + code: 'connection_closed', + }); + } + } + + #beginActiveConnection(): void { + if (this.#activeLane) { + return; + } + + this.#clearReconnectTimer(); + this.#setState( + this.#reconnectAttempts > 0 || this.#retiringLane ? 'reconnecting' : 'connecting' + ); + + const ws = new WebSocket(this.#url, { + headers: { + Authorization: `Bearer ${this.#apiKey}`, + ...(this.#orgId ? { 'x-agentuity-orgid': this.#orgId } : {}), + }, + }); + 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 (!this.#isTrackedLane(lane)) { + return; + } + if (lane.role === 'active') { + this.#reconnectAttempts = 0; + this.#setState('connected'); + this.#logDebug('connected', { url: this.#url, orgId: this.#orgId }); + this.#options.onOpen?.(); + this.#finishConnect(); + } + }; + + ws.onmessage = (event) => { + if (!this.#isTrackedLane(lane)) { + return; + } + this.#handleMessage(event.data, lane); + }; + + ws.onerror = () => { + if (!this.#isTrackedLane(lane)) { + return; + } + if (lane.role === 'active') { + this.#options.onError?.( + new AIGatewayWebSocketError({ + message: 'WebSocket connection error', + code: 'connection_error', + }) + ); + } + }; + + ws.onclose = (event: CloseEvent) => { + if (!this.#isTrackedLane(lane)) { + return; + } + this.#handleLaneClose(lane, event); + }; + } + + #handleLaneClose(lane: WebSocketLane, event: CloseEvent): void { + const terminalClose = isTerminalCloseCode(event.code); + if (terminalClose) { + this.#intentionallyClosed = true; + } + + this.#options.onClose?.(event.code, event.reason); + + if (lane.role === 'retiring') { + this.#retiringLane = null; + this.#rejectLanePending( + lane, + new AIGatewayWebSocketError({ + 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 (this.#activeLane?.ws.readyState === WebSocket.OPEN) { + this.#setState('connected'); + } + return; + } + + 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'); + } + } + + #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 { + this.#logDebug('recv raw', raw); + 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.#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); + } + + #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) { + // 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 }); + 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 content = + response.content ?? (pending.accumulated || getAIGatewayCompletionText(response.data)); + const result = mapCompleteResponse({ + ...response, + content: content || undefined, + }); + clearTimeout(pending.timeout); + this.#pending.delete(response.id); + this.#removeRequestFromLanes(response.id); + pending.push?.({ type: 'complete', result }); + pending.resolve(result); + this.#clearRetiringLaneIfIdle(); + } + } + + #rejectPending(id: string, error: Error): void { + const pending = this.#pending.get(id); + if (!pending) { + return; + } + clearTimeout(pending.timeout); + this.#pending.delete(id); + this.#removeRequestFromLanes(id); + pending.reject(error); + 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 of [...this.#pending.keys()]) { + this.#rejectPending(id, error); + } + } + + #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; + } + return null; + } + + #clearRetiringLaneIfIdle(): void { + if (!this.#retiringLane || this.#retiringLane.requestIds.size > 0) { + return; + } + const ws = this.#retiringLane.ws; + if (ws.readyState !== WebSocket.CLOSED) { + ws.close(1000, 'lane idle'); + } + this.#retiringLane = null; + } + + #isTrackedLane(lane: WebSocketLane): boolean { + return this.#activeLane === lane || this.#retiringLane === lane; + } + + #scheduleReconnect(): void { + if (this.#intentionallyClosed || !this.#options.autoReconnect || this.#activeLane) { + 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.#beginActiveConnection(); + }, delay); + } + + #clearReconnectTimer(): void { + if (this.#reconnectTimer) { + clearTimeout(this.#reconnectTimer); + this.#reconnectTimer = null; + } + } + + #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.#logDebug('send', payload); + lane.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/index.test.ts b/packages/aigateway/test/index.test.ts index c24b8e237..ff97d798d 100644 --- a/packages/aigateway/test/index.test.ts +++ b/packages/aigateway/test/index.test.ts @@ -274,3 +274,57 @@ 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; + }); + + 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 new file mode 100644 index 000000000..bb6660364 --- /dev/null +++ b/packages/aigateway/test/websocket.test.ts @@ -0,0 +1,571 @@ +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 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', + 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('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', + 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('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', + 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({ + id: 'req_inflight', + model: 'anthropic/claude-sonnet-4-20250514', + prompt: 'Finish me', + }); + await flushAsyncWork(); + + firstWs?.receive({ + type: 'draining', + message: 'server rolling restart', + }); + + expect(client.isDraining).toBe(true); + expect(drainingMessage).toBe('server rolling restart'); + expect(client.state).toBe('reconnecting'); + expect(reconnectAttempt).toBe(1); + expect(MockWebSocket.instances).toHaveLength(2); + + 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: '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(); + expect(client.isDraining).toBe(false); + expect(client.state).toBe('connected'); + expect(MockWebSocket.instances).toHaveLength(2); + }); + + 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', + }); + }); + + 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', + }); + }); +}); 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" } + ] }