|
| 1 | +import { afterEach, describe, expect, test } from "bun:test"; |
| 2 | +import { |
| 3 | + clearGatherRoutedModelsInflight, |
| 4 | + gatherRoutedModels as gatherRoutedModelsDirect, |
| 5 | + resetCatalogRuntimeStateForTests, |
| 6 | + type ComboCatalogOmission, |
| 7 | +} from "../src/codex/catalog"; |
| 8 | +import { clearModelCache } from "../src/codex/model-cache"; |
| 9 | +import { withStubbedProviderFetch } from "./helpers/catalog-provider-fetch"; |
| 10 | +import type { OcxConfig } from "../src/types"; |
| 11 | + |
| 12 | +const originalFetch = globalThis.fetch; |
| 13 | + |
| 14 | +const gatherRoutedModels: typeof gatherRoutedModelsDirect = (config, options) => |
| 15 | + gatherRoutedModelsDirect(withStubbedProviderFetch(config), options); |
| 16 | + |
| 17 | +afterEach(() => { |
| 18 | + globalThis.fetch = originalFetch; |
| 19 | + clearModelCache(); |
| 20 | + clearGatherRoutedModelsInflight(); |
| 21 | + resetCatalogRuntimeStateForTests(); |
| 22 | +}); |
| 23 | + |
| 24 | +describe("gatherRoutedModels single-flight", () => { |
| 25 | + test("concurrent callers with the same provider set share one upstream discovery", async () => { |
| 26 | + let fetchCount = 0; |
| 27 | + let release!: () => void; |
| 28 | + const gate = new Promise<void>(resolve => { |
| 29 | + release = resolve; |
| 30 | + }); |
| 31 | + |
| 32 | + globalThis.fetch = (async () => { |
| 33 | + fetchCount += 1; |
| 34 | + await gate; |
| 35 | + return new Response(JSON.stringify({ data: [{ id: "model-a" }] }), { |
| 36 | + status: 200, |
| 37 | + headers: { "content-type": "application/json" }, |
| 38 | + }); |
| 39 | + }) as typeof fetch; |
| 40 | + |
| 41 | + const config: OcxConfig = { |
| 42 | + port: 10100, |
| 43 | + defaultProvider: "slow", |
| 44 | + providers: { |
| 45 | + slow: { |
| 46 | + adapter: "openai-chat", |
| 47 | + baseUrl: "https://api.example.test/v1", |
| 48 | + models: [], |
| 49 | + }, |
| 50 | + }, |
| 51 | + }; |
| 52 | + |
| 53 | + const first = gatherRoutedModels(config); |
| 54 | + const second = gatherRoutedModels(config); |
| 55 | + // Both must have joined before the live fetch resolves. |
| 56 | + await Promise.resolve(); |
| 57 | + expect(fetchCount).toBe(1); |
| 58 | + release(); |
| 59 | + const [a, b] = await Promise.all([first, second]); |
| 60 | + expect(fetchCount).toBe(1); |
| 61 | + expect(a.map(m => `${m.provider}/${m.id}`)).toEqual(["slow/model-a"]); |
| 62 | + expect(b).toEqual(a); |
| 63 | + }); |
| 64 | + |
| 65 | + test("joiners still receive comboOmissions from the shared flight", async () => { |
| 66 | + globalThis.fetch = (async () => |
| 67 | + new Response(JSON.stringify({ data: [{ id: "m1" }] }), { |
| 68 | + status: 200, |
| 69 | + headers: { "content-type": "application/json" }, |
| 70 | + })) as typeof fetch; |
| 71 | + |
| 72 | + const config: OcxConfig = { |
| 73 | + port: 10100, |
| 74 | + defaultProvider: "a", |
| 75 | + providers: { |
| 76 | + a: { |
| 77 | + adapter: "openai-chat", |
| 78 | + baseUrl: "https://api.example.test/v1", |
| 79 | + models: [], |
| 80 | + }, |
| 81 | + }, |
| 82 | + combos: { |
| 83 | + incomplete: { |
| 84 | + strategy: "failover", |
| 85 | + stickyLimit: 1, |
| 86 | + defaultEffort: "medium", |
| 87 | + alias: null, |
| 88 | + targets: [ |
| 89 | + { provider: "a", model: "m1", weight: 1 }, |
| 90 | + { provider: "missing", model: "x", weight: 1 }, |
| 91 | + ], |
| 92 | + }, |
| 93 | + }, |
| 94 | + }; |
| 95 | + |
| 96 | + const omissionsA: ComboCatalogOmission[] = []; |
| 97 | + const omissionsB: ComboCatalogOmission[] = []; |
| 98 | + await Promise.all([ |
| 99 | + gatherRoutedModels(config, { comboOmissions: omissionsA }), |
| 100 | + gatherRoutedModels(config, { comboOmissions: omissionsB }), |
| 101 | + ]); |
| 102 | + expect(omissionsA.some(item => item.id === "incomplete")).toBe(true); |
| 103 | + expect(omissionsB).toEqual(omissionsA); |
| 104 | + }); |
| 105 | + |
| 106 | + test("distinct provider sets keep separate in-flight gathers (no slot eviction)", async () => { |
| 107 | + let releaseA!: () => void; |
| 108 | + let releaseB!: () => void; |
| 109 | + const gateA = new Promise<void>(resolve => { releaseA = resolve; }); |
| 110 | + const gateB = new Promise<void>(resolve => { releaseB = resolve; }); |
| 111 | + const fetchByHost = new Map<string, number>(); |
| 112 | + |
| 113 | + globalThis.fetch = (async (input: RequestInfo | URL) => { |
| 114 | + const url = String(input); |
| 115 | + const host = url.includes("provider-a") ? "a" : url.includes("provider-b") ? "b" : "other"; |
| 116 | + fetchByHost.set(host, (fetchByHost.get(host) ?? 0) + 1); |
| 117 | + if (host === "a") await gateA; |
| 118 | + else await gateB; |
| 119 | + return new Response(JSON.stringify({ data: [{ id: `model-${host}` }] }), { |
| 120 | + status: 200, |
| 121 | + headers: { "content-type": "application/json" }, |
| 122 | + }); |
| 123 | + }) as typeof fetch; |
| 124 | + |
| 125 | + const configA: OcxConfig = { |
| 126 | + port: 10100, |
| 127 | + defaultProvider: "a", |
| 128 | + providers: { |
| 129 | + a: { |
| 130 | + adapter: "openai-chat", |
| 131 | + baseUrl: "https://provider-a.example.test/v1", |
| 132 | + models: [], |
| 133 | + }, |
| 134 | + }, |
| 135 | + }; |
| 136 | + const configB: OcxConfig = { |
| 137 | + port: 10100, |
| 138 | + defaultProvider: "b", |
| 139 | + providers: { |
| 140 | + b: { |
| 141 | + adapter: "openai-chat", |
| 142 | + baseUrl: "https://provider-b.example.test/v1", |
| 143 | + models: [], |
| 144 | + }, |
| 145 | + }, |
| 146 | + }; |
| 147 | + |
| 148 | + const firstA = gatherRoutedModels(configA); |
| 149 | + const firstB = gatherRoutedModels(configB); |
| 150 | + const secondA = gatherRoutedModels(configA); |
| 151 | + await Promise.resolve(); |
| 152 | + expect(fetchByHost.get("a")).toBe(1); |
| 153 | + expect(fetchByHost.get("b")).toBe(1); |
| 154 | + |
| 155 | + releaseA(); |
| 156 | + releaseB(); |
| 157 | + const [a1, b1, a2] = await Promise.all([firstA, firstB, secondA]); |
| 158 | + expect(fetchByHost.get("a")).toBe(1); |
| 159 | + expect(fetchByHost.get("b")).toBe(1); |
| 160 | + expect(a1.map(m => `${m.provider}/${m.id}`)).toEqual(["a/model-a"]); |
| 161 | + expect(b1.map(m => `${m.provider}/${m.id}`)).toEqual(["b/model-b"]); |
| 162 | + expect(a2).toEqual(a1); |
| 163 | + }); |
| 164 | +}); |
0 commit comments