embedding-openai.test.ts 47 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337
  1. /**
  2. * embedding-openai.test.ts - Tests for OpenAIEmbeddingsProvider (HTTP backend).
  3. *
  4. * Uses a mock fetch — no network required. Covers:
  5. * - 200 happy path
  6. * - 429 → retry → success
  7. * - 503 persistent → exhausted retries → null
  8. * - 4xx (non-429) → no retry, immediate failure
  9. * - batch chunking (>64 items → multiple HTTP calls)
  10. * - timeout / abort
  11. * - malformed JSON / missing data array
  12. * - circuit breaker open + half-open recovery
  13. * - dimension probing
  14. * - healthcheck endpoint
  15. */
  16. import { describe, test, expect, vi } from "vitest";
  17. import {
  18. OpenAIEmbeddingsProvider,
  19. CircuitBreaker,
  20. CircuitOpenError,
  21. HttpError,
  22. isRetryableStatus,
  23. chunkArray,
  24. RETRY_BACKOFFS_MS,
  25. } from "../src/embedding/openai.js";
  26. // ─────────────────────────── Helpers ─────────────────────────────────────────
  27. function mockResponse(status: number, body: unknown, opts?: { delayMs?: number }): Response {
  28. const text = typeof body === "string" ? body : JSON.stringify(body);
  29. const init: ResponseInit = {
  30. status,
  31. headers: { "content-type": "application/json" },
  32. };
  33. if (opts?.delayMs) {
  34. // Synchronous Response — test code awaits it directly, so delayMs would
  35. // need to be implemented in the fetch wrapper, not here.
  36. }
  37. return new Response(text, init);
  38. }
  39. function makeFetchSequence(responses: Array<() => Promise<Response> | Response>): {
  40. fetchImpl: typeof fetch;
  41. calls: { url: string; init?: RequestInit }[];
  42. } {
  43. const calls: { url: string; init?: RequestInit }[] = [];
  44. let i = 0;
  45. const fetchImpl = (async (input: RequestInfo | URL, init?: RequestInit) => {
  46. const url = typeof input === "string" ? input : input.toString();
  47. calls.push({ url, init });
  48. if (i >= responses.length) throw new Error(`Mock fetch exhausted at call ${i + 1}`);
  49. const r = responses[i++]!();
  50. return r instanceof Promise ? r : r;
  51. }) as typeof fetch;
  52. return { fetchImpl, calls };
  53. }
  54. function fakeEmbedding(dim: number, seed = 0): number[] {
  55. return Array.from({ length: dim }, (_, i) => Math.sin(seed + i) * 0.5);
  56. }
  57. function embeddingsResponse(texts: string[], dim = 4): Response {
  58. return mockResponse(200, {
  59. object: "list",
  60. model: "embeddinggemma:300m",
  61. data: texts.map((_, i) => ({
  62. object: "embedding",
  63. index: i,
  64. embedding: fakeEmbedding(dim, i * 7),
  65. })),
  66. });
  67. }
  68. // ─────────────────────────── Pure helpers ────────────────────────────────────
  69. describe("isRetryableStatus", () => {
  70. test("429 retryable", () => expect(isRetryableStatus(429)).toBe(true));
  71. test("503 retryable", () => expect(isRetryableStatus(503)).toBe(true));
  72. test("400 NOT retryable", () => expect(isRetryableStatus(400)).toBe(false));
  73. test("401 NOT retryable", () => expect(isRetryableStatus(401)).toBe(false));
  74. test("404 NOT retryable", () => expect(isRetryableStatus(404)).toBe(false));
  75. test("500 NOT retryable", () => expect(isRetryableStatus(500)).toBe(false));
  76. test("502 NOT retryable", () => expect(isRetryableStatus(502)).toBe(false));
  77. test("200 NOT retryable", () => expect(isRetryableStatus(200)).toBe(false));
  78. });
  79. describe("chunkArray", () => {
  80. test("empty input → empty output", () => {
  81. expect(chunkArray([], 5)).toEqual([]);
  82. });
  83. test("input ≤ size → single chunk", () => {
  84. expect(chunkArray([1, 2, 3], 5)).toEqual([[1, 2, 3]]);
  85. });
  86. test("input = size → single chunk", () => {
  87. expect(chunkArray([1, 2, 3, 4, 5], 5)).toEqual([[1, 2, 3, 4, 5]]);
  88. });
  89. test("input > size → multiple chunks", () => {
  90. expect(chunkArray([1, 2, 3, 4, 5, 6, 7], 3)).toEqual([
  91. [1, 2, 3],
  92. [4, 5, 6],
  93. [7],
  94. ]);
  95. });
  96. test("65 items at size 64 → 64 + 1", () => {
  97. const items = Array.from({ length: 65 }, (_, i) => i);
  98. const chunks = chunkArray(items, 64);
  99. expect(chunks.length).toBe(2);
  100. expect(chunks[0]!.length).toBe(64);
  101. expect(chunks[1]!.length).toBe(1);
  102. });
  103. test("size < 1 throws", () => {
  104. expect(() => chunkArray([1, 2, 3], 0)).toThrow();
  105. expect(() => chunkArray([1, 2, 3], -1)).toThrow();
  106. });
  107. });
  108. // ─────────────────────────── Circuit Breaker ─────────────────────────────────
  109. describe("CircuitBreaker", () => {
  110. test("starts closed", () => {
  111. const cb = new CircuitBreaker();
  112. expect(cb.getState()).toBe("closed");
  113. expect(cb.shouldFailFast()).toBe(false);
  114. });
  115. test("stays closed below minSamples even with all-failures", () => {
  116. const cb = new CircuitBreaker({ minSamples: 4, threshold: 0.5 });
  117. cb.recordFailure();
  118. cb.recordFailure();
  119. cb.recordFailure();
  120. expect(cb.getState()).toBe("closed");
  121. });
  122. test("opens when failure rate exceeds threshold", () => {
  123. const cb = new CircuitBreaker({ minSamples: 4, threshold: 0.5 });
  124. cb.recordFailure();
  125. cb.recordFailure();
  126. cb.recordFailure();
  127. cb.recordFailure();
  128. expect(cb.getState()).toBe("open");
  129. expect(cb.shouldFailFast()).toBe(true);
  130. });
  131. test("transitions OPEN → HALF-OPEN after openDurationMs", () => {
  132. let now = 1_000_000;
  133. const cb = new CircuitBreaker({
  134. minSamples: 4,
  135. threshold: 0.5,
  136. openDurationMs: 5000,
  137. now: () => now,
  138. });
  139. for (let i = 0; i < 4; i++) cb.recordFailure();
  140. expect(cb.getState()).toBe("open");
  141. now += 5001;
  142. expect(cb.getState()).toBe("half-open");
  143. });
  144. test("HALF-OPEN + success → CLOSED", () => {
  145. let now = 1_000_000;
  146. const cb = new CircuitBreaker({
  147. minSamples: 4,
  148. threshold: 0.5,
  149. openDurationMs: 5000,
  150. now: () => now,
  151. });
  152. for (let i = 0; i < 4; i++) cb.recordFailure();
  153. now += 5001;
  154. cb.recordSuccess(); // half-open probe
  155. expect(cb.getState()).toBe("closed");
  156. });
  157. test("HALF-OPEN + failure → re-OPEN", () => {
  158. let now = 1_000_000;
  159. const cb = new CircuitBreaker({
  160. minSamples: 4,
  161. threshold: 0.5,
  162. openDurationMs: 5000,
  163. now: () => now,
  164. });
  165. for (let i = 0; i < 4; i++) cb.recordFailure();
  166. now += 5001;
  167. expect(cb.getState()).toBe("half-open");
  168. cb.recordFailure();
  169. expect(cb.getState()).toBe("open");
  170. });
  171. test("samples outside window are dropped", () => {
  172. let now = 1_000_000;
  173. const cb = new CircuitBreaker({
  174. minSamples: 4,
  175. threshold: 0.5,
  176. windowMs: 1000,
  177. now: () => now,
  178. });
  179. cb.recordFailure();
  180. cb.recordFailure();
  181. now += 1500; // window expired
  182. cb.recordSuccess();
  183. cb.recordSuccess();
  184. cb.recordSuccess();
  185. cb.recordSuccess();
  186. // Old failures should be discarded; rate = 0/4 < threshold
  187. expect(cb.getState()).toBe("closed");
  188. });
  189. test("reset() clears state", () => {
  190. const cb = new CircuitBreaker({ minSamples: 2, threshold: 0.5 });
  191. cb.recordFailure();
  192. cb.recordFailure();
  193. expect(cb.getState()).toBe("open");
  194. cb.reset();
  195. expect(cb.getState()).toBe("closed");
  196. });
  197. });
  198. // ─────────────────────────── HappyPath ───────────────────────────────────────
  199. describe("OpenAIEmbeddingsProvider — happy path", () => {
  200. test("single embed call → 200 success", async () => {
  201. const { fetchImpl, calls } = makeFetchSequence([
  202. () => embeddingsResponse(["hello"], 4),
  203. ]);
  204. const p = new OpenAIEmbeddingsProvider({
  205. endpoint: "https://ai.example.com",
  206. fetchImpl,
  207. });
  208. const r = await p.embed("hello");
  209. expect(r).not.toBeNull();
  210. expect(r!.embedding.length).toBe(4);
  211. expect(r!.model).toBe("embeddinggemma");
  212. expect(p.getDimensions()).toBe(4);
  213. expect(calls.length).toBe(1);
  214. expect(calls[0]!.url).toBe("https://ai.example.com/v1/embeddings");
  215. });
  216. test("strips trailing slashes from endpoint", async () => {
  217. const { fetchImpl, calls } = makeFetchSequence([
  218. () => embeddingsResponse(["x"], 2),
  219. ]);
  220. const p = new OpenAIEmbeddingsProvider({
  221. endpoint: "https://ai.example.com////",
  222. fetchImpl,
  223. });
  224. await p.embed("x");
  225. expect(calls[0]!.url).toBe("https://ai.example.com/v1/embeddings");
  226. });
  227. test("batch of 3 → 1 HTTP call", async () => {
  228. const { fetchImpl, calls } = makeFetchSequence([
  229. () => embeddingsResponse(["a", "b", "c"], 3),
  230. ]);
  231. const p = new OpenAIEmbeddingsProvider({
  232. endpoint: "https://ai.example.com",
  233. fetchImpl,
  234. });
  235. const result = await p.embedBatch(["a", "b", "c"]);
  236. expect(result.length).toBe(3);
  237. expect(result.every((r) => r !== null)).toBe(true);
  238. expect(calls.length).toBe(1);
  239. });
  240. test("respects custom modelId / upstreamModel in request body", async () => {
  241. const { fetchImpl, calls } = makeFetchSequence([
  242. () => embeddingsResponse(["x"], 2),
  243. ]);
  244. const p = new OpenAIEmbeddingsProvider({
  245. endpoint: "https://ai.example.com",
  246. modelId: "embeddinggemma",
  247. upstreamModel: "embeddinggemma:300m",
  248. fetchImpl,
  249. });
  250. const r = await p.embed("x");
  251. expect(r!.model).toBe("embeddinggemma");
  252. const body = JSON.parse(calls[0]!.init!.body as string);
  253. expect(body.model).toBe("embeddinggemma:300m");
  254. });
  255. test("Authorization header set when apiKey provided", async () => {
  256. const { fetchImpl, calls } = makeFetchSequence([
  257. () => embeddingsResponse(["x"], 2),
  258. ]);
  259. const p = new OpenAIEmbeddingsProvider({
  260. endpoint: "https://ai.example.com",
  261. apiKey: "sk-test-123",
  262. fetchImpl,
  263. });
  264. await p.embed("x");
  265. const headers = calls[0]!.init!.headers as Record<string, string>;
  266. expect(headers["Authorization"]).toBe("Bearer sk-test-123");
  267. });
  268. test("Authorization header omitted when apiKey not provided", async () => {
  269. const { fetchImpl, calls } = makeFetchSequence([
  270. () => embeddingsResponse(["x"], 2),
  271. ]);
  272. const p = new OpenAIEmbeddingsProvider({
  273. endpoint: "https://ai.example.com",
  274. fetchImpl,
  275. });
  276. await p.embed("x");
  277. const headers = calls[0]!.init!.headers as Record<string, string>;
  278. expect(headers["Authorization"]).toBeUndefined();
  279. });
  280. });
  281. // ─────────────────────────── Batch chunking ──────────────────────────────────
  282. describe("OpenAIEmbeddingsProvider — batch chunking", () => {
  283. test("100 items at batchSize=64 → 2 HTTP calls (64 + 36)", async () => {
  284. const { fetchImpl, calls } = makeFetchSequence([
  285. () => embeddingsResponse(Array.from({ length: 64 }, () => "x"), 4),
  286. () => embeddingsResponse(Array.from({ length: 36 }, () => "x"), 4),
  287. ]);
  288. const p = new OpenAIEmbeddingsProvider({
  289. endpoint: "https://ai.example.com",
  290. fetchImpl,
  291. batchSize: 64,
  292. });
  293. const texts = Array.from({ length: 100 }, (_, i) => `text-${i}`);
  294. const result = await p.embedBatch(texts);
  295. expect(result.length).toBe(100);
  296. expect(result.every((r) => r !== null)).toBe(true);
  297. expect(calls.length).toBe(2);
  298. const body0 = JSON.parse(calls[0]!.init!.body as string);
  299. const body1 = JSON.parse(calls[1]!.init!.body as string);
  300. expect(body0.input.length).toBe(64);
  301. expect(body1.input.length).toBe(36);
  302. });
  303. test("custom batchSize=10 → multiple smaller calls", async () => {
  304. const { fetchImpl, calls } = makeFetchSequence([
  305. () => embeddingsResponse(Array.from({ length: 10 }, () => "x"), 2),
  306. () => embeddingsResponse(Array.from({ length: 10 }, () => "x"), 2),
  307. () => embeddingsResponse(Array.from({ length: 5 }, () => "x"), 2),
  308. ]);
  309. const p = new OpenAIEmbeddingsProvider({
  310. endpoint: "https://ai.example.com",
  311. fetchImpl,
  312. batchSize: 10,
  313. });
  314. const texts = Array.from({ length: 25 }, (_, i) => `t${i}`);
  315. const result = await p.embedBatch(texts);
  316. expect(result.length).toBe(25);
  317. expect(result.every((r) => r !== null)).toBe(true);
  318. expect(calls.length).toBe(3);
  319. });
  320. test("empty input → no HTTP calls", async () => {
  321. const { fetchImpl, calls } = makeFetchSequence([]);
  322. const p = new OpenAIEmbeddingsProvider({
  323. endpoint: "https://ai.example.com",
  324. fetchImpl,
  325. });
  326. const result = await p.embedBatch([]);
  327. expect(result).toEqual([]);
  328. expect(calls.length).toBe(0);
  329. });
  330. });
  331. // ─────────────────────────── Retry behavior ──────────────────────────────────
  332. describe("OpenAIEmbeddingsProvider — retry on 429/503", () => {
  333. test("429 → retry → success", async () => {
  334. const sleepCalls: number[] = [];
  335. const { fetchImpl, calls } = makeFetchSequence([
  336. () => mockResponse(429, { error: "rate limit" }),
  337. () => embeddingsResponse(["x"], 4),
  338. ]);
  339. const p = new OpenAIEmbeddingsProvider({
  340. endpoint: "https://ai.example.com",
  341. fetchImpl,
  342. retryBackoffsMs: [10, 20, 40],
  343. sleep: async (ms) => {
  344. sleepCalls.push(ms);
  345. },
  346. });
  347. const r = await p.embed("x");
  348. expect(r).not.toBeNull();
  349. expect(calls.length).toBe(2);
  350. expect(sleepCalls).toEqual([10]);
  351. });
  352. test("503 → retry → success", async () => {
  353. const sleepCalls: number[] = [];
  354. const { fetchImpl, calls } = makeFetchSequence([
  355. () => mockResponse(503, { error: "service unavailable" }),
  356. () => embeddingsResponse(["x"], 4),
  357. ]);
  358. const p = new OpenAIEmbeddingsProvider({
  359. endpoint: "https://ai.example.com",
  360. fetchImpl,
  361. retryBackoffsMs: [5, 10, 20],
  362. sleep: async (ms) => {
  363. sleepCalls.push(ms);
  364. },
  365. });
  366. const r = await p.embed("x");
  367. expect(r).not.toBeNull();
  368. expect(calls.length).toBe(2);
  369. expect(sleepCalls).toEqual([5]);
  370. });
  371. test("503 persistent → exhausted retries → null result", async () => {
  372. const sleepCalls: number[] = [];
  373. const { fetchImpl, calls } = makeFetchSequence([
  374. () => mockResponse(503, "down"),
  375. () => mockResponse(503, "down"),
  376. () => mockResponse(503, "down"),
  377. () => mockResponse(503, "down"),
  378. ]);
  379. const p = new OpenAIEmbeddingsProvider({
  380. endpoint: "https://ai.example.com",
  381. fetchImpl,
  382. retryBackoffsMs: [1, 2, 4],
  383. sleep: async (ms) => {
  384. sleepCalls.push(ms);
  385. },
  386. });
  387. const r = await p.embed("x");
  388. expect(r).toBeNull();
  389. expect(calls.length).toBe(4); // initial + 3 retries
  390. expect(sleepCalls).toEqual([1, 2, 4]);
  391. });
  392. test("default backoff schedule is 1s/4s/16s", () => {
  393. expect(RETRY_BACKOFFS_MS).toEqual([1000, 4000, 16000]);
  394. });
  395. test("4xx (non-429) → immediate failure, no retry", async () => {
  396. const sleepCalls: number[] = [];
  397. const { fetchImpl, calls } = makeFetchSequence([
  398. () => mockResponse(401, { error: "unauthorized" }),
  399. ]);
  400. const p = new OpenAIEmbeddingsProvider({
  401. endpoint: "https://ai.example.com",
  402. fetchImpl,
  403. retryBackoffsMs: [10, 20, 40],
  404. sleep: async (ms) => {
  405. sleepCalls.push(ms);
  406. },
  407. });
  408. const r = await p.embed("x");
  409. expect(r).toBeNull();
  410. expect(calls.length).toBe(1); // no retries
  411. expect(sleepCalls).toEqual([]);
  412. });
  413. test("404 → immediate failure, no retry", async () => {
  414. const { fetchImpl, calls } = makeFetchSequence([
  415. () => mockResponse(404, "not found"),
  416. ]);
  417. const p = new OpenAIEmbeddingsProvider({
  418. endpoint: "https://ai.example.com",
  419. fetchImpl,
  420. retryBackoffsMs: [1],
  421. sleep: async () => {},
  422. });
  423. await p.embed("x");
  424. expect(calls.length).toBe(1);
  425. });
  426. });
  427. // ─────────────────────────── Malformed responses ─────────────────────────────
  428. describe("OpenAIEmbeddingsProvider — malformed responses", () => {
  429. test("malformed JSON → null result", async () => {
  430. const { fetchImpl } = makeFetchSequence([
  431. () => new Response("not-json{}", { status: 200 }),
  432. ]);
  433. const p = new OpenAIEmbeddingsProvider({
  434. endpoint: "https://ai.example.com",
  435. fetchImpl,
  436. retryBackoffsMs: [],
  437. sleep: async () => {},
  438. });
  439. const r = await p.embed("x");
  440. expect(r).toBeNull();
  441. });
  442. test("missing data array → null result", async () => {
  443. const { fetchImpl } = makeFetchSequence([
  444. () => mockResponse(200, { object: "list", model: "x" }),
  445. ]);
  446. const p = new OpenAIEmbeddingsProvider({
  447. endpoint: "https://ai.example.com",
  448. fetchImpl,
  449. retryBackoffsMs: [],
  450. sleep: async () => {},
  451. });
  452. const r = await p.embed("x");
  453. expect(r).toBeNull();
  454. });
  455. test("data item index out of range → null result", async () => {
  456. const { fetchImpl } = makeFetchSequence([
  457. () =>
  458. mockResponse(200, {
  459. object: "list",
  460. data: [
  461. { index: 5, embedding: [0.1, 0.2] }, // out of range for 1 input
  462. ],
  463. }),
  464. ]);
  465. const p = new OpenAIEmbeddingsProvider({
  466. endpoint: "https://ai.example.com",
  467. fetchImpl,
  468. retryBackoffsMs: [],
  469. sleep: async () => {},
  470. });
  471. const r = await p.embed("x");
  472. expect(r).toBeNull();
  473. });
  474. test("missing index uses row order (Gemini OpenAI-compat)", async () => {
  475. const { fetchImpl } = makeFetchSequence([
  476. () =>
  477. mockResponse(200, {
  478. object: "list",
  479. model: "gemini-embedding-001",
  480. data: [
  481. { object: "embedding", embedding: [0.1, 0.2] },
  482. { object: "embedding", embedding: [0.7, 0.8] },
  483. ],
  484. }),
  485. ]);
  486. const p = new OpenAIEmbeddingsProvider({
  487. endpoint: "https://generativelanguage.googleapis.com/v1beta/openai",
  488. fetchImpl,
  489. });
  490. const result = await p.embedBatch(["zero", "one"]);
  491. expect(result.length).toBe(2);
  492. expect(result[0]!.embedding).toEqual([0.1, 0.2]);
  493. expect(result[1]!.embedding).toEqual([0.7, 0.8]);
  494. });
  495. test("response handles out-of-order data array (sorts by index)", async () => {
  496. const { fetchImpl } = makeFetchSequence([
  497. () =>
  498. mockResponse(200, {
  499. object: "list",
  500. data: [
  501. { index: 1, embedding: [0.7, 0.8] },
  502. { index: 0, embedding: [0.1, 0.2] },
  503. ],
  504. }),
  505. ]);
  506. const p = new OpenAIEmbeddingsProvider({
  507. endpoint: "https://ai.example.com",
  508. fetchImpl,
  509. });
  510. const result = await p.embedBatch(["zero", "one"]);
  511. expect(result.length).toBe(2);
  512. expect(result[0]!.embedding).toEqual([0.1, 0.2]);
  513. expect(result[1]!.embedding).toEqual([0.7, 0.8]);
  514. });
  515. });
  516. // ─────────────────────────── Timeout / abort ─────────────────────────────────
  517. describe("OpenAIEmbeddingsProvider — timeout and abort", () => {
  518. test("user abort signal → null result + no further calls", async () => {
  519. const { fetchImpl, calls } = makeFetchSequence([
  520. () => embeddingsResponse(["a", "b"], 2),
  521. ]);
  522. const p = new OpenAIEmbeddingsProvider({
  523. endpoint: "https://ai.example.com",
  524. fetchImpl,
  525. batchSize: 1,
  526. });
  527. const ctrl = new AbortController();
  528. ctrl.abort(new Error("user cancelled"));
  529. const result = await p.embedBatch(["a", "b"], { signal: ctrl.signal });
  530. expect(result).toEqual([null, null]);
  531. expect(calls.length).toBe(0); // signal aborted before first call
  532. });
  533. test("per-attempt timeout aborts a slow request", async () => {
  534. let aborted = false;
  535. const fetchImpl = (async (_url: any, init?: RequestInit) => {
  536. return await new Promise<Response>((_resolve, reject) => {
  537. const sig = init?.signal;
  538. sig?.addEventListener("abort", () => {
  539. aborted = true;
  540. reject(new DOMException("aborted", "AbortError"));
  541. });
  542. });
  543. }) as typeof fetch;
  544. const p = new OpenAIEmbeddingsProvider({
  545. endpoint: "https://ai.example.com",
  546. fetchImpl,
  547. timeoutMs: 50,
  548. retryBackoffsMs: [],
  549. sleep: async () => {},
  550. });
  551. const r = await p.embed("hello");
  552. expect(r).toBeNull();
  553. expect(aborted).toBe(true);
  554. });
  555. });
  556. // ─────────────────────────── Circuit breaker integration ─────────────────────
  557. describe("OpenAIEmbeddingsProvider — circuit breaker integration", () => {
  558. test("repeated failures eventually trip breaker → CircuitOpenError", async () => {
  559. // 4 chunks of size 1 = 4 sample slots. All fail with 401 → 4 failures → breaker opens.
  560. const { fetchImpl } = makeFetchSequence([
  561. () => mockResponse(401, "fail"),
  562. () => mockResponse(401, "fail"),
  563. () => mockResponse(401, "fail"),
  564. () => mockResponse(401, "fail"),
  565. () => mockResponse(401, "fail"), // shouldn't be reached
  566. ]);
  567. const p = new OpenAIEmbeddingsProvider({
  568. endpoint: "https://ai.example.com",
  569. fetchImpl,
  570. batchSize: 1,
  571. retryBackoffsMs: [],
  572. sleep: async () => {},
  573. });
  574. // First call: 4 sub-chunks, all fail, breaker opens during 4th
  575. const result1 = await p.embedBatch(["a", "b", "c", "d"]);
  576. expect(result1.every((x) => x === null)).toBe(true);
  577. expect(p.breaker.getState()).toBe("open");
  578. // Second call: breaker fails fast
  579. await expect(p.embedBatch(["e"])).rejects.toBeInstanceOf(CircuitOpenError);
  580. });
  581. test("breaker recovers after openDuration → success closes it", async () => {
  582. let now = 1_000_000;
  583. const { fetchImpl } = makeFetchSequence([
  584. () => mockResponse(401, "fail"),
  585. () => mockResponse(401, "fail"),
  586. () => mockResponse(401, "fail"),
  587. () => mockResponse(401, "fail"),
  588. () => embeddingsResponse(["recovered"], 2),
  589. ]);
  590. const p = new OpenAIEmbeddingsProvider({
  591. endpoint: "https://ai.example.com",
  592. fetchImpl,
  593. batchSize: 1,
  594. retryBackoffsMs: [],
  595. sleep: async () => {},
  596. now: () => now,
  597. });
  598. // Override breaker with a shorter open duration
  599. (p as any).breaker = new CircuitBreaker({
  600. minSamples: 4,
  601. threshold: 0.5,
  602. openDurationMs: 1000,
  603. now: () => now,
  604. });
  605. await p.embedBatch(["a", "b", "c", "d"]);
  606. expect((p as any).breaker.getState()).toBe("open");
  607. now += 1500;
  608. expect((p as any).breaker.getState()).toBe("half-open");
  609. const r = await p.embed("recovered");
  610. expect(r).not.toBeNull();
  611. expect((p as any).breaker.getState()).toBe("closed");
  612. });
  613. });
  614. // ─────────────────────────── Healthcheck ─────────────────────────────────────
  615. describe("OpenAIEmbeddingsProvider — healthcheck", () => {
  616. test("healthcheck pings GET /health when available", async () => {
  617. const { fetchImpl, calls } = makeFetchSequence([
  618. () =>
  619. mockResponse(200, {
  620. status: "ok",
  621. model: "embeddinggemma:300m",
  622. }),
  623. ]);
  624. const p = new OpenAIEmbeddingsProvider({
  625. endpoint: "https://ai.example.com",
  626. fetchImpl,
  627. });
  628. const h = await p.healthcheck();
  629. expect(h.ok).toBe(true);
  630. expect(calls.length).toBe(1);
  631. expect(calls[0]!.url).toBe("https://ai.example.com/health");
  632. expect(calls[0]!.init!.method).toBe("GET");
  633. });
  634. test("healthcheck failure → falls through to embed probe", async () => {
  635. const { fetchImpl } = makeFetchSequence([
  636. () => mockResponse(404, "no /health"),
  637. // Then fall back to /v1/embeddings probe
  638. () => embeddingsResponse(["healthcheck"], 4),
  639. ]);
  640. const p = new OpenAIEmbeddingsProvider({
  641. endpoint: "https://ai.example.com",
  642. fetchImpl,
  643. });
  644. const h = await p.healthcheck();
  645. // 404 isn't an exception, it returns ok:false from the /health branch
  646. // The fallback probe is only triggered on actual exceptions.
  647. expect(h.ok).toBe(false);
  648. expect(h.detail).toContain("404");
  649. });
  650. });
  651. // ─────────────────────────── HttpError ───────────────────────────────────────
  652. describe("HttpError", () => {
  653. test("preserves status and body preview", () => {
  654. const err = new HttpError(429, "rate limit exceeded");
  655. expect(err.status).toBe(429);
  656. expect(err.bodyPreview).toBe("rate limit exceeded");
  657. expect(err.message).toContain("HTTP 429");
  658. });
  659. test("truncates long bodies in message", () => {
  660. const longBody = "x".repeat(500);
  661. const err = new HttpError(500, longBody);
  662. expect(err.message.length).toBeLessThan(longBody.length + 200);
  663. });
  664. });
  665. // ─────────────────────────── lastError tracking (i-vm1lxwry) ────────────────
  666. describe("OpenAIEmbeddingsProvider — getLastError (i-vm1lxwry)", () => {
  667. test("returns undefined before first call", () => {
  668. const { fetchImpl } = makeFetchSequence([]);
  669. const p = new OpenAIEmbeddingsProvider({
  670. endpoint: "https://ai.example.com",
  671. fetchImpl,
  672. });
  673. expect(p.getLastError()).toBeUndefined();
  674. });
  675. test("captures HTTP status + endpoint after non-retryable failure", async () => {
  676. const { fetchImpl } = makeFetchSequence([
  677. () => mockResponse(500, "internal error: GPU OOM"),
  678. ]);
  679. const p = new OpenAIEmbeddingsProvider({
  680. endpoint: "https://ai.example.com",
  681. fetchImpl,
  682. retryBackoffsMs: [],
  683. sleep: async () => {},
  684. });
  685. const r = await p.embed("hello");
  686. expect(r).toBeNull();
  687. const lastErr = p.getLastError();
  688. expect(lastErr).toBeDefined();
  689. expect(lastErr).toContain("https://ai.example.com/v1/embeddings");
  690. expect(lastErr).toContain("status=500");
  691. expect(lastErr).toContain("internal error: GPU OOM");
  692. });
  693. test("captures malformed-JSON error message", async () => {
  694. const { fetchImpl } = makeFetchSequence([
  695. () => new Response("not json at all", { status: 200, headers: { "content-type": "application/json" } }),
  696. ]);
  697. const p = new OpenAIEmbeddingsProvider({
  698. endpoint: "https://ai.example.com",
  699. fetchImpl,
  700. retryBackoffsMs: [],
  701. sleep: async () => {},
  702. });
  703. const r = await p.embed("hello");
  704. expect(r).toBeNull();
  705. const lastErr = p.getLastError();
  706. expect(lastErr).toBeDefined();
  707. expect(lastErr).toContain("https://ai.example.com/v1/embeddings");
  708. expect(lastErr).toMatch(/error="/);
  709. });
  710. test("clears lastError after a fully-successful sweep", async () => {
  711. const { fetchImpl } = makeFetchSequence([
  712. () => mockResponse(500, "fail"),
  713. () => embeddingsResponse(["recovered"], 4),
  714. ]);
  715. const p = new OpenAIEmbeddingsProvider({
  716. endpoint: "https://ai.example.com",
  717. fetchImpl,
  718. retryBackoffsMs: [],
  719. sleep: async () => {},
  720. });
  721. // First call fails — lastError set
  722. const r1 = await p.embed("first");
  723. expect(r1).toBeNull();
  724. expect(p.getLastError()).toBeDefined();
  725. // Second call succeeds — lastError cleared
  726. const r2 = await p.embed("recovered");
  727. expect(r2).not.toBeNull();
  728. expect(p.getLastError()).toBeUndefined();
  729. });
  730. test("getEndpoint() exposes configured endpoint (no trailing slash)", () => {
  731. const { fetchImpl } = makeFetchSequence([]);
  732. const p = new OpenAIEmbeddingsProvider({
  733. endpoint: "https://ai.example.com//",
  734. fetchImpl,
  735. });
  736. expect(p.getEndpoint()).toBe("https://ai.example.com");
  737. });
  738. });
  739. // ─────────────────────────── dispose ─────────────────────────────────────────
  740. // ─────────────────────────── Concurrent dispatch (i-fkpnar9i Phase 1 #1) ─────
  741. /**
  742. * Fetch helper that lets a test:
  743. * - count how many requests are in-flight at any moment
  744. * - control resolution order via per-call deferred promises
  745. * - inspect the start order vs resolution order
  746. *
  747. * Each `responses[i]` returns a Response (or Promise<Response>); the helper
  748. * wraps each call so it awaits a `gate[i]` deferred BEFORE responding. Tests
  749. * call `release(i)` to let the i-th request settle. Useful for testing that
  750. * concurrent dispatch actually overlaps requests.
  751. */
  752. function makeGatedFetchSequence(count: number): {
  753. fetchImpl: typeof fetch;
  754. inFlight: () => number;
  755. startOrder: number[];
  756. release: (idx: number, response: Response) => void;
  757. releaseAll: (responseFor: (idx: number) => Response) => void;
  758. } {
  759. const gates: Array<{ resolve: (r: Response) => void }> = [];
  760. const startOrder: number[] = [];
  761. let inFlight = 0;
  762. let nextStart = 0;
  763. for (let i = 0; i < count; i++) {
  764. let resolveFn: (r: Response) => void = () => {};
  765. new Promise<Response>((resolve) => {
  766. resolveFn = resolve;
  767. });
  768. // re-create properly:
  769. let r2: (x: Response) => void = () => {};
  770. const p = new Promise<Response>((resolve) => {
  771. r2 = resolve;
  772. });
  773. gates.push({ resolve: r2 });
  774. // attach the unresolved promise back to the slot via a closure (below)
  775. (gates[i] as any).promise = p;
  776. }
  777. const fetchImpl = (async (_input: RequestInfo | URL, _init?: RequestInit) => {
  778. const idx = nextStart++;
  779. if (idx >= count) throw new Error(`gated fetch exhausted at ${idx + 1}`);
  780. startOrder.push(idx);
  781. inFlight++;
  782. try {
  783. const r = await (gates[idx] as any).promise;
  784. return r as Response;
  785. } finally {
  786. inFlight--;
  787. }
  788. }) as typeof fetch;
  789. return {
  790. fetchImpl,
  791. inFlight: () => inFlight,
  792. startOrder,
  793. release: (idx: number, response: Response) => gates[idx]!.resolve(response),
  794. releaseAll: (responseFor: (idx: number) => Response) => {
  795. for (let i = 0; i < count; i++) gates[i]!.resolve(responseFor(i));
  796. },
  797. };
  798. }
  799. describe("OpenAIEmbeddingsProvider — concurrent dispatch (i-fkpnar9i)", () => {
  800. test("default concurrency is 4 — 8 chunks of size 1, max in-flight = 4", async () => {
  801. const N = 8;
  802. const gated = makeGatedFetchSequence(N);
  803. const p = new OpenAIEmbeddingsProvider({
  804. endpoint: "https://ai.example.com",
  805. fetchImpl: gated.fetchImpl,
  806. batchSize: 1, // each text becomes its own chunk
  807. });
  808. // Expected concurrency=4 default
  809. expect((p as any).concurrency).toBe(4);
  810. const texts = Array.from({ length: N }, (_, i) => `t${i}`);
  811. const promise = p.embedBatch(texts);
  812. // Yield to the microtask queue so workers can start their first dispatch.
  813. // Multiple yields needed for chained `await this.requestWithRetry → await fetch`.
  814. for (let i = 0; i < 5; i++) await Promise.resolve();
  815. expect(gated.inFlight()).toBe(4);
  816. expect(gated.startOrder).toEqual([0, 1, 2, 3]);
  817. // Release first 4 in reverse order — concurrent dispatch should still
  818. // preserve input order in `results` because each worker writes to its
  819. // pre-computed slot.
  820. for (let i = 3; i >= 0; i--) {
  821. gated.release(i, embeddingsResponse([`t${i}`], 4));
  822. }
  823. // Yield to let workers pick up next chunks
  824. for (let i = 0; i < 10; i++) await Promise.resolve();
  825. expect(gated.inFlight()).toBeGreaterThan(0); // 4 more workers should be in flight
  826. expect(gated.startOrder.length).toBe(8); // all 8 dispatched
  827. // Release the rest
  828. for (let i = 4; i < N; i++) {
  829. gated.release(i, embeddingsResponse([`t${i}`], 4));
  830. }
  831. const result = await promise;
  832. expect(result.length).toBe(N);
  833. // Critical: input order preserved despite out-of-order resolution
  834. for (let i = 0; i < N; i++) {
  835. expect(result[i]).not.toBeNull();
  836. expect(result[i]!.model).toBe("embeddinggemma");
  837. }
  838. expect(gated.inFlight()).toBe(0);
  839. });
  840. test("explicit concurrency=2 — only 2 in-flight at any moment", async () => {
  841. const N = 6;
  842. const gated = makeGatedFetchSequence(N);
  843. const p = new OpenAIEmbeddingsProvider({
  844. endpoint: "https://ai.example.com",
  845. fetchImpl: gated.fetchImpl,
  846. batchSize: 1,
  847. concurrency: 2,
  848. });
  849. const texts = Array.from({ length: N }, (_, i) => `t${i}`);
  850. const promise = p.embedBatch(texts);
  851. for (let i = 0; i < 5; i++) await Promise.resolve();
  852. expect(gated.inFlight()).toBe(2);
  853. // Cycle: release one, wait, expect new one started
  854. gated.release(0, embeddingsResponse(["t0"], 4));
  855. for (let i = 0; i < 10; i++) await Promise.resolve();
  856. expect(gated.inFlight()).toBe(2); // still 2 — slot filled by t2
  857. // Drain the rest
  858. for (let i = 1; i < N; i++) {
  859. gated.release(i, embeddingsResponse([`t${i}`], 4));
  860. for (let j = 0; j < 5; j++) await Promise.resolve();
  861. }
  862. const result = await promise;
  863. expect(result.every((r) => r !== null)).toBe(true);
  864. });
  865. test("concurrency=1 reproduces legacy sequential behavior", async () => {
  866. const N = 4;
  867. const gated = makeGatedFetchSequence(N);
  868. const p = new OpenAIEmbeddingsProvider({
  869. endpoint: "https://ai.example.com",
  870. fetchImpl: gated.fetchImpl,
  871. batchSize: 1,
  872. concurrency: 1,
  873. });
  874. const texts = Array.from({ length: N }, (_, i) => `t${i}`);
  875. const promise = p.embedBatch(texts);
  876. for (let i = 0; i < 5; i++) await Promise.resolve();
  877. expect(gated.inFlight()).toBe(1);
  878. expect(gated.startOrder).toEqual([0]);
  879. // Release one at a time, confirm the next starts only after.
  880. for (let i = 0; i < N; i++) {
  881. gated.release(i, embeddingsResponse([`t${i}`], 4));
  882. for (let j = 0; j < 5; j++) await Promise.resolve();
  883. }
  884. await promise;
  885. // All started in order (sequential)
  886. expect(gated.startOrder).toEqual([0, 1, 2, 3]);
  887. });
  888. test("results in input order even when the LAST chunk resolves first", async () => {
  889. const N = 4;
  890. const gated = makeGatedFetchSequence(N);
  891. const p = new OpenAIEmbeddingsProvider({
  892. endpoint: "https://ai.example.com",
  893. fetchImpl: gated.fetchImpl,
  894. batchSize: 1,
  895. concurrency: 4,
  896. });
  897. const texts = ["alpha", "beta", "gamma", "delta"];
  898. const promise = p.embedBatch(texts);
  899. // Wait for all 4 to be in flight, then resolve LAST first
  900. for (let i = 0; i < 5; i++) await Promise.resolve();
  901. expect(gated.inFlight()).toBe(4);
  902. gated.release(3, embeddingsResponse(["delta"], 4));
  903. gated.release(2, embeddingsResponse(["gamma"], 4));
  904. gated.release(1, embeddingsResponse(["beta"], 4));
  905. gated.release(0, embeddingsResponse(["alpha"], 4));
  906. const result = await promise;
  907. expect(result.length).toBe(N);
  908. // Each input slot got its own embedding — input order preserved
  909. expect(result[0]).not.toBeNull();
  910. expect(result[1]).not.toBeNull();
  911. expect(result[2]).not.toBeNull();
  912. expect(result[3]).not.toBeNull();
  913. });
  914. test("dimensions recorded correctly even if the first-resolving chunk is not chunk 0", async () => {
  915. const gated = makeGatedFetchSequence(2);
  916. const p = new OpenAIEmbeddingsProvider({
  917. endpoint: "https://ai.example.com",
  918. fetchImpl: gated.fetchImpl,
  919. batchSize: 1,
  920. concurrency: 2,
  921. });
  922. const promise = p.embedBatch(["a", "b"]);
  923. for (let i = 0; i < 5; i++) await Promise.resolve();
  924. // Resolve chunk 1 first with 7-dim, then chunk 0 with 7-dim
  925. gated.release(1, embeddingsResponse(["b"], 7));
  926. for (let i = 0; i < 5; i++) await Promise.resolve();
  927. gated.release(0, embeddingsResponse(["a"], 7));
  928. await promise;
  929. expect(p.getDimensions()).toBe(7);
  930. });
  931. test("abort signal during concurrent run stops new dispatches; in-flight settle", async () => {
  932. const gated = makeGatedFetchSequence(8);
  933. const p = new OpenAIEmbeddingsProvider({
  934. endpoint: "https://ai.example.com",
  935. fetchImpl: gated.fetchImpl,
  936. batchSize: 1,
  937. concurrency: 4,
  938. });
  939. const ctrl = new AbortController();
  940. const texts = Array.from({ length: 8 }, (_, i) => `t${i}`);
  941. const promise = p.embedBatch(texts, { signal: ctrl.signal });
  942. for (let i = 0; i < 5; i++) await Promise.resolve();
  943. expect(gated.startOrder.length).toBe(4);
  944. // Resolve in-flight, then abort — remaining 4 should NOT dispatch
  945. gated.release(0, embeddingsResponse(["t0"], 4));
  946. gated.release(1, embeddingsResponse(["t1"], 4));
  947. gated.release(2, embeddingsResponse(["t2"], 4));
  948. gated.release(3, embeddingsResponse(["t3"], 4));
  949. ctrl.abort(new Error("operator cancelled"));
  950. for (let i = 0; i < 20; i++) await Promise.resolve();
  951. const result = await promise;
  952. // First 4 succeeded, last 4 are null (never dispatched after abort)
  953. expect(result.slice(0, 4).every((r) => r !== null)).toBe(true);
  954. expect(result.slice(4).every((r) => r === null)).toBe(true);
  955. // Total dispatched MUST be ≤ 5 (the abort can race with one extra
  956. // worker pulling the next idx before the abort flag is set; we cap
  957. // at first 4 + at-most-1 grace).
  958. expect(gated.startOrder.length).toBeLessThanOrEqual(5);
  959. // Audit string captured
  960. expect(p.getLastError()).toMatch(/aborted by caller/);
  961. });
  962. test("ctor rejects concurrency < 1", () => {
  963. expect(() => new OpenAIEmbeddingsProvider({
  964. endpoint: "https://ai.example.com",
  965. fetchImpl: (async () => mockResponse(200, {})) as typeof fetch,
  966. concurrency: 0,
  967. })).toThrow(/concurrency must be ≥ 1/);
  968. expect(() => new OpenAIEmbeddingsProvider({
  969. endpoint: "https://ai.example.com",
  970. fetchImpl: (async () => mockResponse(200, {})) as typeof fetch,
  971. concurrency: -3,
  972. })).toThrow(/concurrency must be ≥ 1/);
  973. });
  974. test("circuit-open observed mid-run is thrown after in-flight settle", async () => {
  975. // 8 chunks, all fail → breaker opens after the first 4 (minSamples=4).
  976. // Workers 0-3 dispatch in parallel, all fail, recordFailure × 4 → breaker
  977. // OPEN. Remaining workers see shouldFailFast() and set circuitTrippedDuringRun.
  978. const N = 8;
  979. const { fetchImpl } = makeFetchSequence(
  980. Array.from({ length: N }, () => () => mockResponse(401, "fail"))
  981. );
  982. const p = new OpenAIEmbeddingsProvider({
  983. endpoint: "https://ai.example.com",
  984. fetchImpl,
  985. batchSize: 1,
  986. concurrency: 4,
  987. retryBackoffsMs: [],
  988. sleep: async () => {},
  989. });
  990. // First 4 land before breaker opens (concurrent dispatch); after they all
  991. // fail the breaker tips OPEN. The next pull observes shouldFailFast().
  992. // Either the result resolves with all-null (legacy semantics — breaker
  993. // tripped AFTER all workers grabbed their chunk) OR throws CircuitOpenError
  994. // (breaker observed before next pull). Both are valid post-condition;
  995. // we just assert the state ends OPEN and the call completes.
  996. let res: Awaited<ReturnType<typeof p.embedBatch>> | undefined;
  997. let err: unknown;
  998. try {
  999. res = await p.embedBatch(Array.from({ length: N }, (_, i) => `t${i}`));
  1000. } catch (e) {
  1001. err = e;
  1002. }
  1003. expect(p.breaker.getState()).toBe("open");
  1004. if (err) {
  1005. expect(err).toBeInstanceOf(CircuitOpenError);
  1006. } else {
  1007. expect(res!.every((r) => r === null)).toBe(true);
  1008. }
  1009. });
  1010. });
  1011. describe("OpenAIEmbeddingsProvider — dispose", () => {
  1012. test("dispose resets the breaker", async () => {
  1013. const { fetchImpl } = makeFetchSequence([
  1014. () => mockResponse(401, "fail"),
  1015. () => mockResponse(401, "fail"),
  1016. () => mockResponse(401, "fail"),
  1017. () => mockResponse(401, "fail"),
  1018. ]);
  1019. const p = new OpenAIEmbeddingsProvider({
  1020. endpoint: "https://ai.example.com",
  1021. fetchImpl,
  1022. batchSize: 1,
  1023. retryBackoffsMs: [],
  1024. sleep: async () => {},
  1025. });
  1026. await p.embedBatch(["a", "b", "c", "d"]);
  1027. expect(p.breaker.getState()).toBe("open");
  1028. await p.dispose();
  1029. expect(p.breaker.getState()).toBe("closed");
  1030. });
  1031. });
  1032. // ────────── rate-limit budget + bulk lane (i-yghj098h) ───────────────────────
  1033. import {
  1034. parseRetryAfterMs,
  1035. BulkLaneGate,
  1036. RATE_LIMIT_MAX_BACKOFF_MS,
  1037. DEFAULT_RATE_LIMIT_RETRIES,
  1038. LANE_RECOVERY_STREAK,
  1039. } from "../src/embedding/openai.js";
  1040. /** The shape ai.mm.mk actually returns when the token bucket is empty. */
  1041. function rateLimitedResponse(seconds: number, opts?: { header?: boolean }): Response {
  1042. const headers: Record<string, string> = { "content-type": "application/json" };
  1043. if (opts?.header) headers["retry-after"] = String(seconds);
  1044. return new Response(
  1045. JSON.stringify({
  1046. detail: `Rate limit exceeded (tokens). Retry after ${seconds}s.`,
  1047. error: "rate_limited",
  1048. }),
  1049. { status: 429, headers },
  1050. );
  1051. }
  1052. describe("parseRetryAfterMs", () => {
  1053. test("delta-seconds header", () => {
  1054. expect(parseRetryAfterMs("29")).toBe(29_000);
  1055. });
  1056. test("HTTP-date header", () => {
  1057. const now = () => 1_000_000;
  1058. const when = new Date(now() + 12_000).toUTCString();
  1059. // toUTCString truncates to whole seconds — allow the rounding slack.
  1060. expect(parseRetryAfterMs(when, undefined, now)).toBeGreaterThan(11_000);
  1061. });
  1062. test("falls back to the ai.mm.mk body prose", () => {
  1063. expect(
  1064. parseRetryAfterMs(null, '{"detail":"Rate limit exceeded (tokens). Retry after 29s."}'),
  1065. ).toBe(29_000);
  1066. });
  1067. test("clamps absurd values", () => {
  1068. expect(parseRetryAfterMs("99999")).toBe(RATE_LIMIT_MAX_BACKOFF_MS);
  1069. });
  1070. test("undefined when nothing parseable", () => {
  1071. expect(parseRetryAfterMs(null, "no numbers here")).toBeUndefined();
  1072. expect(parseRetryAfterMs("0")).toBeUndefined();
  1073. });
  1074. });
  1075. describe("OpenAIEmbeddingsProvider — 429 honours Retry-After (i-yghj098h)", () => {
  1076. test("waits the server-advertised cooldown from the body, not the 1s schedule", async () => {
  1077. const sleepCalls: number[] = [];
  1078. const { fetchImpl } = makeFetchSequence([
  1079. () => rateLimitedResponse(29),
  1080. () => embeddingsResponse(["x"], 4),
  1081. ]);
  1082. const p = new OpenAIEmbeddingsProvider({
  1083. endpoint: "https://ai.mm.mk",
  1084. fetchImpl,
  1085. retryBackoffsMs: [1_000, 4_000, 16_000],
  1086. sleep: async (ms) => { sleepCalls.push(ms); },
  1087. });
  1088. const r = await p.embed("x");
  1089. expect(r).not.toBeNull();
  1090. expect(sleepCalls).toEqual([29_000]);
  1091. });
  1092. test("Retry-After header wins over the body", async () => {
  1093. const sleepCalls: number[] = [];
  1094. const { fetchImpl } = makeFetchSequence([
  1095. () => rateLimitedResponse(7, { header: true }),
  1096. () => embeddingsResponse(["x"], 4),
  1097. ]);
  1098. const p = new OpenAIEmbeddingsProvider({
  1099. endpoint: "https://ai.mm.mk",
  1100. fetchImpl,
  1101. retryBackoffsMs: [],
  1102. sleep: async (ms) => { sleepCalls.push(ms); },
  1103. });
  1104. expect(await p.embed("x")).not.toBeNull();
  1105. expect(sleepCalls).toEqual([7_000]);
  1106. });
  1107. test("429 gets its own budget — survives more rate limits than the generic schedule", async () => {
  1108. const sleepCalls: number[] = [];
  1109. const { fetchImpl, calls } = makeFetchSequence([
  1110. () => rateLimitedResponse(1),
  1111. () => rateLimitedResponse(1),
  1112. () => rateLimitedResponse(1),
  1113. () => rateLimitedResponse(1),
  1114. () => embeddingsResponse(["x"], 4),
  1115. ]);
  1116. const p = new OpenAIEmbeddingsProvider({
  1117. endpoint: "https://ai.mm.mk",
  1118. fetchImpl,
  1119. // A single generic retry — before the fix this run gave up on call 2.
  1120. retryBackoffsMs: [10],
  1121. sleep: async (ms) => { sleepCalls.push(ms); },
  1122. });
  1123. expect(await p.embed("x")).not.toBeNull();
  1124. expect(calls.length).toBe(5);
  1125. expect(sleepCalls.length).toBe(4);
  1126. expect(DEFAULT_RATE_LIMIT_RETRIES).toBeGreaterThanOrEqual(4);
  1127. });
  1128. test("429 exhaustion still fails, and does not trip the circuit breaker", async () => {
  1129. const { fetchImpl } = makeFetchSequence(
  1130. Array.from({ length: 12 }, () => () => rateLimitedResponse(1)),
  1131. );
  1132. const p = new OpenAIEmbeddingsProvider({
  1133. endpoint: "https://ai.mm.mk",
  1134. fetchImpl,
  1135. retryBackoffsMs: [],
  1136. rateLimitRetries: 1,
  1137. sleep: async () => {},
  1138. });
  1139. for (let i = 0; i < 6; i++) expect(await p.embed("x")).toBeNull();
  1140. // Backpressure must not become a 5-minute hard OPEN.
  1141. expect(p.breaker.getState()).toBe("closed");
  1142. });
  1143. });
  1144. describe("BulkLaneGate", () => {
  1145. test("halves permits per cooldown and coalesces concurrent penalties", async () => {
  1146. const slept: number[] = [];
  1147. const gate = new BulkLaneGate(4, async (ms) => { slept.push(ms); });
  1148. await Promise.all([gate.penalize(5_000), gate.penalize(5_000)]);
  1149. expect(slept).toEqual([5_000]);
  1150. expect(gate.permitCount).toBe(2);
  1151. await gate.penalize(5_000);
  1152. expect(gate.permitCount).toBe(1);
  1153. await gate.penalize(5_000);
  1154. expect(gate.permitCount).toBe(1); // floor
  1155. });
  1156. test("recovers one permit per success streak", async () => {
  1157. const gate = new BulkLaneGate(4, async () => {});
  1158. await gate.penalize(1_000);
  1159. expect(gate.permitCount).toBe(2);
  1160. for (let i = 0; i < LANE_RECOVERY_STREAK; i++) gate.noteSuccess();
  1161. expect(gate.permitCount).toBe(3);
  1162. });
  1163. test("acquire/release respects the current cap", async () => {
  1164. const gate = new BulkLaneGate(1, async () => {});
  1165. await gate.acquire();
  1166. let secondEntered = false;
  1167. const second = gate.acquire().then(() => { secondEntered = true; });
  1168. await Promise.resolve();
  1169. expect(secondEntered).toBe(false);
  1170. gate.release();
  1171. await second;
  1172. expect(secondEntered).toBe(true);
  1173. });
  1174. });
  1175. describe("bulk lane — X-AI-Caller attribution + shared cooldown", () => {
  1176. function callerHeader(init?: RequestInit): string {
  1177. const headers = (init?.headers ?? {}) as Record<string, string>;
  1178. return headers["X-AI-Caller"] ?? "";
  1179. }
  1180. test("multi-input batches are attributed as embeddings-bulk", async () => {
  1181. const { fetchImpl, calls } = makeFetchSequence([
  1182. () => embeddingsResponse(["a", "b"], 4),
  1183. ]);
  1184. const p = new OpenAIEmbeddingsProvider({
  1185. endpoint: "https://ai.mm.mk",
  1186. fetchImpl,
  1187. sleep: async () => {},
  1188. });
  1189. await p.embedBatch(["a", "b"]);
  1190. expect(callerHeader(calls[0]!.init)).toContain("via=embeddings-bulk");
  1191. expect(callerHeader(calls[0]!.init)).toContain("comp=qmd");
  1192. });
  1193. test("single query-time embeds stay on the interactive label", async () => {
  1194. const { fetchImpl, calls } = makeFetchSequence([
  1195. () => embeddingsResponse(["a"], 4),
  1196. ]);
  1197. const p = new OpenAIEmbeddingsProvider({
  1198. endpoint: "https://ai.mm.mk",
  1199. fetchImpl,
  1200. sleep: async () => {},
  1201. });
  1202. await p.embed("a");
  1203. expect(callerHeader(calls[0]!.init)).toContain("via=embeddings");
  1204. expect(callerHeader(calls[0]!.init)).not.toContain("via=embeddings-bulk");
  1205. });
  1206. test("one 429 pauses the whole bulk pool once and narrows concurrency", async () => {
  1207. const sleepCalls: number[] = [];
  1208. let call = 0;
  1209. const fetchImpl = (async (_input: unknown, _init?: RequestInit) => {
  1210. call++;
  1211. if (call === 1) return rateLimitedResponse(30);
  1212. return embeddingsResponse(["x"], 4);
  1213. }) as unknown as typeof fetch;
  1214. const p = new OpenAIEmbeddingsProvider({
  1215. endpoint: "https://ai.mm.mk",
  1216. fetchImpl,
  1217. batchSize: 1,
  1218. concurrency: 2,
  1219. retryBackoffsMs: [],
  1220. sleep: async (ms) => { sleepCalls.push(ms); },
  1221. });
  1222. const out = await p.embedBatch(["a", "b", "c", "d"]);
  1223. expect(out.every((r) => r !== null)).toBe(true);
  1224. // Exactly one cooldown, at the length the gateway asked for.
  1225. expect(sleepCalls).toEqual([30_000]);
  1226. // AIMD: the lane gave capacity back to interactive traffic.
  1227. expect(p.lane.permitCount).toBe(1);
  1228. });
  1229. });