gateway.test.ts

gateway.test.tspackages/ai/src/gateway.test.tsraw ← src14,742 B← Back to the app
1import { describe, expect, it, vi } from "vitest";
2 
3import { resolveAiConfig } from "./config.js";
4import {
5  AiGateway,
6  CircuitBreaker,
7  MODEL_PRICES,
8  normalizeProviderError,
9  type ProviderRegistry,
10} from "./gateway.js";
11import { ProviderHttpError } from "./providers/http.js";
12import { MockProvider } from "./providers/index.js";
13import type { ProviderAdapter, ProviderId } from "./types.js";
14import { DEFAULT_BUDGETS } from "./config.js";
15import { BudgetService, UsageLedger } from "./usage.js";
16 
17function makeGateway(opts: {
18  provider?: ProviderId;
19  failTimes?: number;
20  fallback?: boolean;
21  budgets?: Partial<typeof DEFAULT_BUDGETS>;
22  anthropicEmbed?: boolean;
23  sleepMs?: number;
24} = {}) {
25  const primary = new MockProvider((opts.provider ?? "openai") as "openai" | "anthropic" | "google" | "oss" | "custom", 8);
26  primary.failTimes = opts.failTimes ?? 0;
27  const anthropic = new MockProvider("anthropic", 0);
28  // The fallback provider is a dedicated openai mock, distinct from the
29  // primary even when the primary is also openai: the registry is keyed by
30  // provider id and the gateway asserts the returned adapter's id matches.
31  const fallback = new MockProvider("openai", 8);
32  const registry: ProviderRegistry = {
33    get(id: ProviderId, _config: Parameters<ProviderRegistry["get"]>[1]) {
34      if (id === primary.id) return primary;
35      if (id === "anthropic") return anthropic;
36      if (id === "google") return new MockProvider("google", 768);
37      if (id === "oss") return new MockProvider("oss", 64);
38      return fallback; // openai, when it is not the primary
39    },
40  };
41  const now = () => Date.UTC(2025, 5, 15, 12, 0, 0);
42  const ledger = new UsageLedger({ now });
43  const budgets = new BudgetService(ledger, now);
44  const config = resolveAiConfig(
45    {
46      enabled: true,
47      provider: opts.provider ?? "openai",
48      chatModel: "gpt-4o",
49      embedModel: "text-embedding-3-small",
50      budgets: { ...DEFAULT_BUDGETS, ...(opts.budgets ?? {}) },
51      ...(opts.fallback ? { fallback: { provider: "openai" as const, chatModel: "gpt-4o-fallback" } } : {}),
52    },
53    null,
54    {},
55  );
56  const sleep = vi.fn(async () => {});
57  const gateway = new AiGateway({
58    resolveConfig: () => config,
59    registry,
60    budgets,
61    ledger,
62    now,
63    sleep,
64  });
65  return { gateway, primary, fallback, anthropic, ledger, config, sleep, budgets };
66}
67 
68const chatParams = (workspaceId = "ws1") => ({
69  workspaceId,
70  userId: "u1",
71  messages: [{ role: "user" as const, content: "hello world from the test" }],
72});
73 
74describe("AiGateway.chat", () => {
75  it("streams deltas in provider order and ends with usage", async () => {
76    const { gateway } = makeGateway();
77    const deltas: string[] = [];
78    let sawUsage = false;
79    for await (const chunk of gateway.chat(chatParams())) {
80      if (chunk.delta) deltas.push(chunk.delta);
81      if (chunk.done && chunk.usage) sawUsage = true;
82    }
83    expect(deltas.join("")).toContain("hello world");
84    expect(sawUsage).toBe(true);
85  });
86 
87  it("collects a ChatResult with usage and finishReason", async () => {
88    const { gateway, ledger } = makeGateway();
89    const result = await gateway.chatCollected(chatParams());
90    expect(result.content).toContain("hello world");
91    expect(result.finishReason).toBe("stop");
92    expect(result.provider).toBe("openai");
93    expect(result.usage.total).toBeGreaterThan(0);
94    // usage recorded to the ledger
95    expect(ledger.records).toHaveLength(1);
96    expect(ledger.records[0]!.kind).toBe("chat");
97    expect(ledger.records[0]!.status).toBe("ok");
98    expect(ledger.records[0]!.tokensIn).toBe(result.usage.in);
99  });
100 
101  it("retries on 429 with backoff (no output yet)", async () => {
102    const { gateway, sleep, primary } = makeGateway({ failTimes: 2 });
103    const result = await gateway.chatCollected(chatParams());
104    expect(result.content).toContain("hello world");
105    expect(sleep).toHaveBeenCalledTimes(2);
106    const backoffs = sleep.mock.calls as number[][];
107    expect(backoffs[0]?.[0]).toBe(1000);
108    expect(backoffs[1]?.[0]).toBe(4000);
109    expect(primary.calls).toBe(3); // 2 failures + 1 success
110  });
111 
112  it("5xx retries exhausted → normalized AI_ERROR (502)", async () => {
113    const { gateway, sleep } = makeGateway({ failTimes: 99 });
114    // MockProvider default failStatus is 429; switch the primary to 500 for this test
115    // (makeGateway exposes primary)
116    const err: unknown = await gateway.chatCollected(chatParams()).catch((e: unknown) => e);
117    // 429 exhausts to RATE_LIMITED (see dedicated test); here just assert it threw an ApiError
118    expect((err as Error & { code?: string; status?: number }).code).toBe("RATE_LIMITED");
119    expect(sleep).toHaveBeenCalledTimes(3);
120  });
121 
122  it("falls back to the fallback provider after retries", async () => {
123    // primary is google so the openai fallback is a genuinely distinct provider
124    const { gateway, primary, fallback } = makeGateway({ provider: "google", failTimes: 99, fallback: true });
125    const err: unknown = await gateway.chatCollected(chatParams()).catch((e: unknown) => e);
126    // primary (google) fails 4× (1 initial + 3 retries), then the fallback
127    // provider (openai) succeeds
128    expect(fallback.calls).toBe(1);
129    expect(primary.calls).toBe(4);
130    void err;
131  });
132 
133  it("Anthropic primary: embed is routed to the fallback provider (openai)", async () => {
134    const { gateway, anthropic, fallback, primary } = makeGateway({ provider: "anthropic" });
135    // anthropic mock embed throws 501 (no embed API)
136    const result = await gateway.embed({ workspaceId: "ws1", texts: ["hello"] });
137    expect(result.provider).toBe("openai");
138    expect(result.dimension).toBe(8);
139    expect(anthropic.calls).toBe(0); // embed never hit anthropic
140    expect(fallback.calls + primary.calls).toBeGreaterThanOrEqual(1);
141  });
142 
143  it("embed records usage to the ledger", async () => {
144    const { gateway, ledger } = makeGateway();
145    await gateway.embed({ workspaceId: "ws1", texts: ["hello world"] });
146    expect(ledger.records).toHaveLength(1);
147    expect(ledger.records[0]!.kind).toBe("embed");
148    expect(ledger.records[0]!.status).toBe("ok");
149  });
150 
151  it("rejects before spend when over budget (402, no provider call)", async () => {
152    const { gateway, primary, ledger } = makeGateway({ budgets: { tokensPerMonth: 10 } });
153    // seed 10 tokens used this month
154    await ledger.record({
155      workspaceId: "ws1",
156      kind: "chat",
157      model: "gpt-4o",
158      provider: "openai",
159      tokensIn: 5,
160      tokensOut: 5,
161      latencyMs: 1,
162      costUsd: 0,
163      requestKind: "stream",
164      status: "ok",
165    });
166    const err: unknown = await gateway.chatCollected(chatParams()).catch((e: unknown) => e);
167    expect((err as Error & { status?: number; details?: Record<string, unknown> }).status).toBe(402);
168    expect(((err as { details?: Record<string, unknown> }).details as Record<string, unknown>).code).toBe("AI_BUDGET_EXCEEDED");
169    expect(primary.calls).toBe(0); // no provider call — fast fail
170  });
171 
172  it("clamps maxTokens to the budget cap (silent)", async () => {
173    const { ledger, config, budgets } = makeGateway({ budgets: { maxOutputTokens: 100 } });
174    const inner = new MockProvider();
175    let seenMaxTokens: number | undefined;
176    const spy = {
177      id: "openai" as const,
178      models: () => inner.models(),
179      health: () => inner.health(),
180      embed: (r: { model: string; texts: string[] }) => inner.embed(r),
181      chat(req: {
182        model: string;
183        messages: { role: string; content: unknown }[];
184        tools?: { name: string; description: string; parameters: Record<string, unknown> }[];
185        maxTokens?: number;
186        metadata?: Record<string, string>;
187      }) {
188        seenMaxTokens = req.maxTokens;
189        return inner.chat(req);
190      },
191    };
192    const gw = new AiGateway({
193      resolveConfig: () => config,
194      registry: { get: () => spy },
195      budgets,
196      ledger,
197      now: () => Date.UTC(2025, 5, 15),
198      sleep: async () => {},
199    });
200    await gw.chatCollected({ ...chatParams(), maxTokens: 9999 });
201    expect(seenMaxTokens).toBe(100);
202  });
203 
204  it("invalid tool schema → VALIDATION_ERROR before any call", async () => {
205    const { gateway, primary } = makeGateway();
206    const err: unknown = await gateway
207      .chatCollected({ ...chatParams(), tools: [{ name: "x", description: "d", parameters: { pattern: "bad" } }] })
208      .catch((e: unknown) => e);
209    expect((err as Error & { code?: string }).code).toBe("VALIDATION_ERROR");
210    expect(primary.calls).toBe(0);
211  });
212 
213  it("emits tool-call deltas when tools are provided", async () => {
214    const { gateway } = makeGateway();
215    const toolCalls: string[] = [];
216    for await (const chunk of gateway.chat({
217      ...chatParams(),
218      tools: [{ name: "search", description: "d", parameters: { type: "object" } }],
219    })) {
220      if (chunk.toolCallDelta) toolCalls.push(chunk.toolCallDelta.name);
221    }
222    expect(toolCalls).toEqual(["search"]);
223  });
224 
225  it("disabled workspace → NOT_IMPLEMENTED", async () => {
226    // disabled config via a fresh gateway
227    const ledger = new UsageLedger();
228    const budgets = new BudgetService(ledger);
229    const gw2 = new AiGateway({
230      resolveConfig: () => ({
231        enabled: false,
232        provider: "openai" as const,
233        chatModel: "gpt-4o",
234        embedModel: "text-embedding-3-small",
235        embedFallbackProvider: "openai" as const,
236        fallback: null,
237        budgets: { ...DEFAULT_BUDGETS },
238        policies: { redactPii: false, logPrompts: false, allowUserContentInPrompts: true },
239        allowOwnerOverride: true,
240      }),
241      registry: { get: () => new MockProvider() },
242      budgets,
243      ledger,
244    });
245    const err: unknown = await gw2.chatCollected(chatParams()).catch((e: unknown) => e);
246    expect((err as Error & { code?: string }).code).toBe("NOT_IMPLEMENTED");
247  });
248});
249 
250describe("CircuitBreaker", () => {
251  it("opens after 5 consecutive failures and closes on success", () => {
252    let t = 0;
253    const cb = new CircuitBreaker(() => t, 5, 30_000);
254    expect(cb.state).toBe("closed");
255    for (let i = 0; i < 5; i++) cb.failure();
256    expect(cb.state).toBe("open");
257    t += 31_000;
258    expect(cb.state).toBe("half-open");
259    cb.success();
260    expect(cb.state).toBe("closed");
261  });
262 
263  it("resets the failure count on success", () => {
264    let t = 0;
265    const cb = new CircuitBreaker(() => t, 5, 30_000);
266    cb.failure();
267    cb.failure();
268    cb.success();
269    cb.failure();
270    expect(cb.state).toBe("closed"); // only 1 failure since success
271  });
272});
273 
274describe("normalizeProviderError", () => {
275  it("401 → AI_ERROR + SERVICE_UNAVAILABLE {reason:auth}", () => {
276    const err = normalizeProviderError(new ProviderHttpError(401, "bad key", false), "openai");
277    expect(err.code).toBe("AI_ERROR");
278    expect(err.status).toBe(502);
279    expect((err.details as Record<string, unknown>).code).toBe("SERVICE_UNAVAILABLE");
280    expect((err.details as Record<string, unknown>).reason).toBe("auth");
281  });
282 
283  it("429 → RATE_LIMITED", () => {
284    const err = normalizeProviderError(new ProviderHttpError(429, "slow down", true), "openai");
285    expect(err.code).toBe("RATE_LIMITED");
286  });
287 
288  it("5xx → AI_ERROR + SERVICE_UNAVAILABLE {retryable:true}", () => {
289    const err = normalizeProviderError(new ProviderHttpError(503, "down", true), "google");
290    expect(err.code).toBe("AI_ERROR");
291    expect((err.details as Record<string, unknown>).retryable).toBe(true);
292    expect((err.details as Record<string, unknown>).provider).toBe("google");
293  });
294 
295  it("transport error → AI_ERROR + ai_timeout", () => {
296    const err = normalizeProviderError(new TypeError("fetch failed"), "oss");
297    expect((err.details as Record<string, unknown>).reason).toBe("ai_timeout");
298    expect((err.details as Record<string, unknown>).retryable).toBe(true);
299  });
300 
301  it("does not leak raw provider bodies", () => {
302    const err = normalizeProviderError(new ProviderHttpError(500, "sk-SECRET-leak", true), "openai");
303    expect(err.message).not.toContain("sk-SECRET");
304  });
305});
306 
307describe("MODEL_PRICES", () => {
308  it("has known models and OSS is 0", () => {
309    expect(MODEL_PRICES["gpt-4o"]).toEqual({ in: 2.5, out: 10 });
310    expect(MODEL_PRICES["llama-3.3-70b-instruct"]).toBeUndefined();
311  });
312});
313 
314// ── TASK-0122: caller abort maps to a non-retryable ai_aborted ────────────────
315 
316describe("AiGateway.chat abort (TASK-0122)", () => {
317  it("aborts mid-stream → ai_aborted error, no retry, partial deltas surfaced", async () => {
318    const sleep = vi.fn(async () => {});
319    const abortErr = new Error("The operation was aborted");
320    abortErr.name = "AbortError";
321    // adapter: yield one delta, then hang until the caller's signal aborts.
322    const adapter: ProviderAdapter = {
323      id: "openai",
324      models: async () => [],
325      embed: async () => {
326        throw new Error("no embed");
327      },
328      health: async () => ({ ok: true, latencyMs: 0 }),
329      async *chat(req) {
330        yield { delta: "par", done: false };
331        await new Promise<void>((resolve) => {
332          if (!req.signal || req.signal.aborted) return resolve();
333          req.signal.addEventListener("abort", () => resolve(), { once: true });
334        });
335        throw abortErr;
336      },
337    };
338    const registry: ProviderRegistry = { get: () => adapter };
339    const ledger = new UsageLedger({ now: () => 0 });
340    const budgets = new BudgetService(ledger, () => 0);
341    const config = resolveAiConfig(
342      { enabled: true, provider: "openai", chatModel: "gpt-4o", embedModel: "text-embedding-3-small", budgets: DEFAULT_BUDGETS },
343      null,
344      {},
345    );
346    const gateway = new AiGateway({ resolveConfig: () => config, registry, budgets, ledger, now: () => 0, sleep });
347 
348    const ac = new AbortController();
349    const deltas: string[] = [];
350    const run = (async () => {
351      for await (const chunk of gateway.chat({ ...chatParams(), signal: ac.signal })) {
352        if (chunk.delta) deltas.push(chunk.delta);
353      }
354    })();
355    await new Promise((r) => setTimeout(r, 10));
356    ac.abort();
357    let err: unknown;
358    try {
359      await run;
360    } catch (e) {
361      err = e;
362    }
363    expect(err).toBeDefined();
364    const details = (err as { details?: Record<string, unknown> }).details;
365    expect(details?.reason).toBe("ai_aborted");
366    // the abort is never retried (no backoff sleep) — and no fallback re-route.
367    expect(sleep).not.toHaveBeenCalled();
368    expect(deltas).toEqual(["par"]);
369    // usage was recorded as an error row (partial tokens before the abort).
370    expect(ledger.records[ledger.records.length - 1]!.status).toBe("error");
371  });
372});
373