gateway.ts · Source

gateway.tspackages/ai/src/gateway.tsraw ← src20,440 B← Back to the app
1/**
2 * @factory/ai — AiGateway (ai-architecture §2.3–§2.4, TASK-0121).
3 *
4 * The only place provider types are imported. Capability surface:
5 *
6 *   - `chat(params)` → `AsyncIterable<ChatChunk>` — tool calling is a
7 *     parameter of chat, not a separate method
8 *   - `chatCollected(params)` → `Promise<ChatResult>` (collects the stream)
9 *   - `embed(params)` → `EmbedResult` (Anthropic is routed to the fallback
10 *     provider — it has no embedding API)
11 *   - `resolve(workspaceId)` → the resolved config (for `POST /ai/ping`)
12 *
13 * Reliability (arch §2.4):
14 *   - retries: 429/5xx/network, exponential backoff 1 s → 4 s → 16 s, max 3 —
15 *     **only if no output has been yielded yet** (re-streaming partial output
16 *     would duplicate content)
17 *   - fallback: after retries exhausted → one attempt on the fallback
18 *     provider/model (usage metered to the fallback provider)
19 *   - circuit breaker: per provider, 5 consecutive failures → open 30 s
20 *     (half-open probe)
21 *   - budgets: submit-time check via `BudgetService` (fast fail, 402/429)
22 *   - usage: every completed or failed-after-consume call appends an
23 *     `ai_usage` record via the injected ledger (arch §7.1)
24 *
25 * **Recorded decision (spec conflict, rule 29):** TASK-0121 names the codes
26 * `AI_BUDGET_EXCEEDED`/`SERVICE_UNAVAILABLE`, which the frozen
27 * `@factory/shared` `ErrorCodes` does not define. The gateway uses the shared
28 * envelope (`AI_ERROR`, and `RATE_LIMITED` where the shared code exists) and
29 * carries the task's code in `details.code` (see `usage.ts` module doc); the
30 * API layer maps `details.code` onto the canonical table.
31 */
32 
33import { ApiError, ErrorCodes } from "@factory/shared";
34 
35import type { ResolvedAiConfig } from "./config.js";
36import { ProviderHttpError } from "./providers/http.js";
37import type {
38  ChatChunk,
39  ChatMessage,
40  ChatRequest,
41  ChatResult,
42  EmbedResult,
43  ProviderAdapter,
44  ProviderId,
45  ToolSpec,
46} from "./types.js";
47import { validateTools } from "./tools.js";
48import type { Budgets, UsageLedger } from "./usage.js";import { BudgetService } from "./usage.js";
49 
50export interface ChatParams {
51  workspaceId: string;
52  userId?: string;
53  capability?: string;
54  messages: ChatMessage[];
55  model?: string;
56  tools?: ToolSpec[];
57  toolChoice?: ChatRequest["toolChoice"];
58  temperature?: number;
59  maxTokens?: number;
60  stream?: boolean;
61  metadata?: Record<string, string>;
62  /**
63   * TASK-0122: caller cancellation (the SSE endpoint passes the request's
64   * AbortSignal). Aborting mid-stream stops the provider call and ends the
65   * stream with a non-retryable `ai_aborted` error (never retried or
66   * re-routed to the fallback — the caller already went away).
67   */
68  signal?: AbortSignal;
69}
70 
71export interface EmbedParams {
72  workspaceId: string;
73  userId?: string;
74  texts: string[];
75  model?: string;
76  metadata?: Record<string, string>;
77}
78 
79/** App layer provides adapters from config + decrypted provider secrets. */
80export interface ProviderRegistry {
81  get(providerId: ProviderId, config: ResolvedAiConfig): ProviderAdapter;
82}
83 
84export interface AiGatewayOptions {
85  resolveConfig: (workspaceId: string) => ResolvedAiConfig;
86  registry: ProviderRegistry;
87  budgets: BudgetService;
88  ledger: UsageLedger;
89  now?: () => number;
90  /** Injectable sleep for backoff (tests pass a no-op). */
91  sleep?: (ms: number) => Promise<void>;
92  /** Max retries after the first attempt (arch §2.4: 3). */
93  maxRetries?: number;
94}
95 
96const BACKOFFS_MS: [number, number, number] = [1_000, 4_000, 16_000];
97 
98/** Per-provider circuit breaker: N consecutive failures → open for openMs. */
99export class CircuitBreaker {
100  private failures = 0;
101  private openUntil = 0;
102  private readonly now: () => number;
103  private readonly threshold: number;
104  private readonly openMs: number;
105  constructor(now: () => number, threshold = 5, openMs = 30_000) {
106    this.now = now;
107    this.threshold = threshold;
108    this.openMs = openMs;
109  }
110  get state(): "closed" | "open" | "half-open" {
111    if (this.failures >= this.threshold) {
112      return this.now() < this.openUntil ? "open" : "half-open";
113    }
114    return "closed";
115  }
116  failure(): void {
117    this.failures += 1;
118    if (this.failures >= this.threshold) this.openUntil = this.now() + this.openMs;
119  }
120  success(): void {
121    this.failures = 0;
122    this.openUntil = 0;
123  }
124}
125 
126/**
127 * Retryable at the gateway level per arch §2.4: provider 429 / 5xx /
128 * network. The final (post-retries) failure is then normalized — a provider
129 * 429 surfaces as `RATE_LIMITED`, 5xx/transport as `AI_ERROR` +
130 * `SERVICE_UNAVAILABLE` (module doc).
131 */
132/** TASK-0122: true when the error is a caller-initiated stream abort. */
133function isAborted(err: unknown, signal?: AbortSignal): boolean {
134  if (signal?.aborted) return true;
135  if (err instanceof Error) return err.name === "AbortError" || err.message === "The operation was aborted";
136  return false;
137}
138 
139function isRetryable(err: unknown, signal?: AbortSignal): boolean {
140  // A caller-initiated abort is never retryable (the caller went away —
141  // retrying would re-spend tokens for no one, and a fallback re-route
142  // would stream a duplicate answer into the dead socket).
143  if (isAborted(err, signal)) return false;
144  if (err instanceof ProviderHttpError) return err.retryable;
145  if (err instanceof ApiError) return err.code === ErrorCodes.TIMEOUT;
146  return true; // transport-level failures (e.g. "fetch failed") are retryable
147}
148 
149/**
150 * Normalize a provider failure to the canonical error envelope (module doc
151 * explains the recorded code decision): 401/403 → AI_ERROR +
152 * `details.code = SERVICE_UNAVAILABLE {service:"ai", reason:"auth"}`;
153 * 429 → RATE_LIMITED; timeout/transport → AI_ERROR +
154 * `{reason:"ai_timeout", retryable:true}`. Raw provider bodies never leak
155 * (only a truncated status + message is kept server-side).
156 */
157export function normalizeProviderError(err: unknown, provider: ProviderId): ApiError {
158  if (err instanceof ApiError) return err;
159  if (err instanceof ProviderHttpError) {
160    if (err.status === 401 || err.status === 403) {
161      return new ApiError(ErrorCodes.AI_ERROR, {
162        message: "AI provider authentication failed. Check the provider key configuration.",
163        status: 502,
164        details: { code: "SERVICE_UNAVAILABLE", service: "ai", reason: "auth", provider },
165      });
166    }
167    if (err.status === 429) {
168      return new ApiError(ErrorCodes.RATE_LIMITED, {
169        message: "AI provider rate limit reached. Try again shortly.",
170        details: { provider },
171      });
172    }
173    return new ApiError(ErrorCodes.AI_ERROR, {
174      message: "AI provider request failed.",
175      status: 502,
176      details: {
177        code: "SERVICE_UNAVAILABLE",
178        service: "ai",
179        reason: err.status === 504 ? "ai_timeout" : "provider_error",
180        retryable: err.retryable,
181        provider,
182        httpStatus: err.status,
183      },
184    });
185  }
186  return new ApiError(ErrorCodes.AI_ERROR, {
187    message: "AI provider request failed.",
188    status: 502,
189    details: {
190      code: "SERVICE_UNAVAILABLE",
191      service: "ai",
192      reason: "ai_timeout",
193      retryable: true,
194      provider,
195    },
196  });
197}
198 
199function metadataFor(params: {
200  workspaceId: string;
201  userId?: string;
202  capability?: string;
203  metadata?: Record<string, string>;
204}): Record<string, string> {
205  return {
206    capability: params.capability ?? "chat",
207    workspaceId: params.workspaceId,
208    ...(params.userId ? { userId: params.userId } : {}),
209    ...params.metadata,
210  };
211}
212 
213export class AiGateway {
214  private readonly options: AiGatewayOptions;
215  private readonly now: () => number;
216  private readonly sleep: (ms: number) => Promise<void>;
217  private readonly maxRetries: number;
218  private readonly breakers = new Map<ProviderId, CircuitBreaker>();
219 
220  constructor(options: AiGatewayOptions) {
221    this.options = options;
222    this.now = options.now ?? Date.now;
223    this.sleep = options.sleep ?? ((ms) => new Promise((r) => setTimeout(r, ms)));
224    this.maxRetries = options.maxRetries ?? 3;
225  }
226 
227  /** Resolved config for a workspace (drives `POST /ai/ping`). */
228  resolve(workspaceId: string): ResolvedAiConfig {
229    const config = this.options.resolveConfig(workspaceId);
230    if (!config.enabled) {
231      throw new ApiError(ErrorCodes.NOT_IMPLEMENTED, {
232        message: "AI is disabled for this workspace.",
233        details: { code: "AI_DISABLED" },
234      });
235    }
236    return config;
237  }
238 
239  breaker(provider: ProviderId): CircuitBreaker {
240    let b = this.breakers.get(provider);
241    if (!b) {
242      b = new CircuitBreaker(this.now);
243      this.breakers.set(provider, b);
244    }
245    return b;
246  }
247 
248  private async adapterFor(config: ResolvedAiConfig, providerId: ProviderId): Promise<ProviderAdapter> {
249    const state = this.breaker(providerId).state;
250    if (state === "open") {
251      throw new ApiError(ErrorCodes.AI_ERROR, {
252        message: "AI provider is temporarily unavailable (circuit open).",
253        status: 503,
254        details: { code: "SERVICE_UNAVAILABLE", service: "ai", reason: "circuit_open", provider: providerId },
255      });
256    }
257    const adapter = this.options.registry.get(providerId, config);
258    if (adapter.id !== providerId) {
259      throw new ApiError(ErrorCodes.INTERNAL, {
260        message: `registry returned adapter "${adapter.id}", expected "${providerId}"`,
261      });
262    }
263    return adapter;
264  }
265 
266  /**
267   * Streaming chat (the base primitive). Retries + fallback happen
268   * pre-output only; once a chunk is yielded the caller owns the stream.
269   */
270  async *chat(params: ChatParams): AsyncIterable<ChatChunk> {
271    const config = this.resolve(params.workspaceId);
272    const model = params.model ?? config.chatModel;
273 
274    // tool schema validation (strict subset) — reject before any spend
275    if (params.tools && params.tools.length > 0) {
276      const result = validateTools(params.tools);
277      if (!result.ok) {
278        throw new ApiError(ErrorCodes.VALIDATION_ERROR, { message: `invalid tool schema: ${result.error}` });
279      }
280    }
281 
282    // maxOutputTokens clamp (silent, arch §7.2)
283    const maxTokens = this.options.budgets.clampMaxTokens(params.maxTokens, config.budgets);
284 
285    // submit-time budget check (fast fail; estimate = 2× maxTokens)
286    this.options.budgets.check(params.workspaceId, config.budgets, 2 * maxTokens);
287 
288    const baseRequest: ChatRequest = {
289      model,
290      messages: params.messages,
291      stream: true,
292      metadata: metadataFor(params),
293      ...(params.signal ? { signal: params.signal } : {}),
294    };
295    if (params.temperature !== undefined) baseRequest.temperature = params.temperature;
296    if (maxTokens !== undefined) baseRequest.maxTokens = maxTokens;
297    if (params.tools && params.tools.length > 0) {
298      baseRequest.tools = params.tools;
299      if (params.toolChoice !== undefined) baseRequest.toolChoice = params.toolChoice;
300    }
301 
302    let adapter = await this.adapterFor(config, config.provider);
303    let request = baseRequest;
304    let attempts = 0;
305    let usedFallback = false;
306 
307    for (;;) {
308      let yielded = false;
309      let usage = { in: 0, out: 0, total: 0 };
310      let finishReason: ChatResult["finishReason"] = "stop";
311      const started = this.now();
312      try {
313        // A pre-stream abort (the caller cancelled before any provider
314        // output) maps to the same non-retryable `ai_aborted` path.
315        if (params.signal?.aborted) {
316          const preAbort = new Error("The operation was aborted");
317          preAbort.name = "AbortError";
318          throw preAbort;
319        }
320        for await (const chunk of adapter.chat(request)) {
321          // Per-chunk abort check (TASK-0122): adapters backed by real
322          // fetch streams honor the signal themselves, but plain-JS
323          // generators (the mock) yield back-to-back — without this check a
324          // mid-stream abort would only surface at the adapter's end.
325          if (params.signal?.aborted) {
326            const abortErr = new Error("The operation was aborted");
327            abortErr.name = "AbortError";
328            throw abortErr;
329          }
330          yielded = true;
331          if (chunk.done) {
332            if (chunk.usage) usage = chunk.usage;
333            if (chunk.finishReason) finishReason = chunk.finishReason;
334          }
335          yield chunk;
336        }
337        this.breaker(adapter.id).success();
338        void this.recordChat(params, adapter.id, request.model, usage, this.now() - started, "ok");
339        return;
340      } catch (err) {
341        if (isAborted(err, params.signal)) {
342          // Caller-initiated cancel (TASK-0122): record what was consumed,
343          // never retry/re-route (the caller is gone), and surface a
344          // non-retryable `ai_aborted` error the service distinguishes from
345          // a provider failure. Deliberately NOT counted as a breaker
346          // failure — the provider did nothing wrong (a "Stop" click must
347          // not trip the circuit after enough cancellations).
348          void this.recordChat(params, adapter.id, request.model, usage, this.now() - started, "error");
349          throw new ApiError(ErrorCodes.AI_ERROR, {
350            message: "AI stream aborted by the caller.",
351            status: 502,
352            details: { code: "SERVICE_UNAVAILABLE", service: "ai", reason: "ai_aborted", provider: adapter.id },
353          });
354        }
355        this.breaker(adapter.id).failure();
356        if (yielded) {
357          // mid-stream failure: never re-stream (would duplicate content)
358          void this.recordChat(params, adapter.id, request.model, usage, this.now() - started, "error");
359          throw normalizeProviderError(err, adapter.id);
360        }
361        const retryable = isRetryable(err, params.signal);
362        if (retryable && attempts < this.maxRetries) {
363          attempts += 1;
364          await this.sleep(BACKOFFS_MS[Math.min(attempts - 1, BACKOFFS_MS.length - 1)]!);
365          continue;
366        }
367        if (!usedFallback && config.fallback && config.fallback.provider !== adapter.id) {
368          usedFallback = true;
369          try {
370            adapter = await this.adapterFor(config, config.fallback.provider);
371            request = { ...baseRequest, model: config.fallback.chatModel };
372            continue;
373          } catch {
374            // fall through to normalizing the primary failure
375          }
376        }
377        void this.recordChat(params, adapter.id, request.model, usage, this.now() - started, "error");
378        throw normalizeProviderError(err, adapter.id);
379      }
380    }
381  }
382 
383  /** Non-streaming chat: collects the chunk stream into a ChatResult. */
384  async chatCollected(params: ChatParams): Promise<ChatResult> {
385    const config = this.resolve(params.workspaceId);
386    const model = params.model ?? config.chatModel;
387    const started = this.now();
388    let content = "";
389    const toolCalls: ChatResult["toolCalls"] = [];
390    let usage = { in: 0, out: 0, total: 0 };
391    let finishReason: ChatResult["finishReason"] = "stop";
392    for await (const chunk of this.chat(params)) {
393      if (chunk.delta) content += chunk.delta;
394      if (chunk.toolCallDelta) toolCalls.push(chunk.toolCallDelta);
395      if (chunk.done) {
396        if (chunk.usage) usage = chunk.usage;
397        if (chunk.finishReason) finishReason = chunk.finishReason;
398      }
399    }
400    return {
401      content,
402      toolCalls,
403      usage,
404      model,
405      latencyMs: this.now() - started,
406      provider: config.provider,
407      finishReason,
408    };
409  }
410 
411  /**
412   * Embeddings. If the primary provider has no embedding API (Anthropic) the
413   * gateway routes to `embedFallbackProvider` (default `openai`, arch §2.2).
414   * Retries on 429/5xx with the same backoff schedule.
415   */
416  async embed(params: EmbedParams): Promise<EmbedResult> {
417    const config = this.resolve(params.workspaceId);
418    const model = params.model ?? config.embedModel;
419    const started = this.now();
420 
421    // Anthropic has no embed API: route straight to the fallback provider.
422    let providerId: ProviderId =
423      config.provider === "anthropic" && config.embedFallbackProvider ? config.embedFallbackProvider : config.provider;
424    let adapter = await this.adapterFor(config, providerId);
425    let attempts = 0;
426    let usedEmbedFallback = false;
427 
428    const record = (status: "ok" | "error", result?: EmbedResult): Promise<unknown> =>
429      this.options.ledger.record({
430        workspaceId: params.workspaceId,
431        ...(params.userId ? { userId: params.userId } : {}),
432        ...(params.metadata?.capability ? { capability: params.metadata.capability } : {}),
433        kind: "embed",
434        model,
435        provider: adapter.id,
436        tokensIn: 0,
437        tokensOut: result?.usage.tokens ?? estimateTokens(params.texts),
438        latencyMs: result?.latencyMs ?? this.now() - started,
439        costUsd: 0, // embed prices are tracked per-model in the app layer
440        requestKind: "batch",
441        status,
442      });
443 
444    for (;;) {
445      try {
446        const result = await adapter.embed({
447          model,
448          texts: params.texts,
449          metadata: { workspaceId: params.workspaceId, ...(params.userId ? { userId: params.userId } : {}), ...params.metadata },
450        });
451        this.breaker(adapter.id).success();
452        void record("ok", result);
453        return result;
454      } catch (err) {
455        this.breaker(adapter.id).failure();
456        const retryable = isRetryable(err);
457        if (retryable && attempts < this.maxRetries) {
458          attempts += 1;
459          await this.sleep(BACKOFFS_MS[Math.min(attempts - 1, BACKOFFS_MS.length - 1)]!);
460          continue;
461        }
462        if (
463          !usedEmbedFallback &&
464          config.embedFallbackProvider &&
465          config.embedFallbackProvider !== adapter.id
466        ) {
467          usedEmbedFallback = true;
468          try {
469            adapter = await this.adapterFor(config, config.embedFallbackProvider);
470            continue;
471          } catch {
472            // fall through to normalizing the primary failure
473          }
474        }
475        void record("error");
476        throw normalizeProviderError(err, adapter.id);
477      }
478    }
479  }
480 
481  /** Ping/diagnostic: resolved provider/model + budget state (no tokens). */
482  async ping(workspaceId: string): Promise<{
483    provider: ProviderId;
484    chatModel: string;
485    embedModel: string;
486    budget: { limit: number; used: number };
487  }> {
488    const config = this.resolve(workspaceId);
489    const monthStart = new Date(this.now());
490    const used = this.options.ledger.tokensSince(workspaceId, new Date(monthStart.getFullYear(), monthStart.getMonth(), 1).getTime());
491    return {
492      provider: config.provider,
493      chatModel: config.chatModel,
494      embedModel: config.embedModel,
495      budget: { limit: config.budgets.tokensPerMonth, used },
496    };
497  }
498 
499  private async recordChat(
500    params: ChatParams,
501    provider: ProviderId,
502    model: string,
503    usage: { in: number; out: number; total: number },
504    latencyMs: number,
505    status: "ok" | "error" | "timeout",
506  ): Promise<void> {
507    await this.options.ledger.record({
508      workspaceId: params.workspaceId,
509      ...(params.userId ? { userId: params.userId } : {}),
510      ...(params.capability ? { capability: params.capability } : {}),
511      kind: "chat",
512      model,
513      provider,
514      tokensIn: usage.in,
515      tokensOut: usage.out,
516      latencyMs,
517      costUsd: this.costFor(model, usage.in, usage.out),
518      requestKind: "stream",
519      status,
520    });
521  }
522 
523  /** Cost in USD from the static price table (0 for OSS/unknown models). */
524  private costFor(model: string, tokensIn: number, tokensOut: number): number {
525    const price = MODEL_PRICES[model];
526    if (!price) return 0;
527    return (tokensIn / 1000) * price.in + (tokensOut / 1000) * price.out;
528  }
529}
530 
531/** Static price table (USD per 1k tokens). The app layer can extend this. */
532export const MODEL_PRICES: Record<string, { in: number; out: number }> = {
533  "gpt-4o": { in: 2.5, out: 10 },
534  "gpt-4o-mini": { in: 0.15, out: 0.6 },
535  "text-embedding-3-small": { in: 0.02, out: 0 },
536  "text-embedding-3-large": { in: 0.13, out: 0 },
537  "claude-sonnet-4-20250514": { in: 3, out: 15 },
538  "claude-opus-4-20250514": { in: 15, out: 75 },
539  "claude-3-5-haiku-20241022": { in: 0.8, out: 4 },
540  "gemini-2.0-flash": { in: 0.1, out: 0.4 },
541  "text-embedding-004": { in: 0.025, out: 0 },
542};
543 
544function estimateTokens(texts: string[]): number {
545  return Math.ceil(texts.join(" ").length / 4);
546}
547