Skip to content

Commit 240db4f

Browse files
JamesRobert20JamesRobert20
committed
fix(zoo-gateway): address edelauna review feedback
1 parent 9feabae commit 240db4f

4 files changed

Lines changed: 106 additions & 59 deletions

File tree

‎src/api/providers/fetchers/__tests__/modelCache.spec.ts‎

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1019,18 +1019,17 @@ describe("MODEL_CACHE_EMPTY_RESPONSE throttling", () => {
10191019
expect(resultB).toEqual(accountBModels)
10201020
})
10211021

1022-
it("never deduplicates concurrent zoo-gateway fetches, even for the same token", async () => {
1023-
// Auth-scoped providers skip dedupedFetch unconditionally, so even two calls carrying
1024-
// an identical token each fire their own provider fetch -- there is no in-flight sharing
1025-
// to key correctly or incorrectly for these providers.
1022+
it("deduplicates concurrent zoo-gateway fetches for the same session token", async () => {
1023+
// Auth-scoped providers use inFlightAuthScopedFetch to coalesce concurrent calls
1024+
// for the same cache key so two panels opening simultaneously share one fetch.
10261025
freshMockGetZooGatewayModels.mockResolvedValue(zooGatewayOk({}))
10271026

10281027
await Promise.all([
10291028
freshGetModels({ provider: providerIdentifiers.zooGateway, apiKey: "same-token" }),
10301029
freshGetModels({ provider: providerIdentifiers.zooGateway, apiKey: "same-token" }),
10311030
])
10321031

1033-
expect(freshMockGetZooGatewayModels).toHaveBeenCalledTimes(2)
1032+
expect(freshMockGetZooGatewayModels).toHaveBeenCalledTimes(1)
10341033
})
10351034
})
10361035

@@ -1292,6 +1291,17 @@ describe("auth session cache", () => {
12921291
expect(mockSet).not.toHaveBeenCalled()
12931292
})
12941293

1294+
it("concurrent getModels calls for the same session key share one fetch", async () => {
1295+
freshMockGetZooGatewayModels.mockResolvedValue(zooGatewayOk(zooModels))
1296+
const options = { provider: providerIdentifiers.zooGateway, apiKey: "session-token" }
1297+
1298+
const [first, second] = await Promise.all([freshGetModels(options), freshGetModels(options)])
1299+
1300+
expect(first).toEqual(zooModels)
1301+
expect(second).toEqual(zooModels)
1302+
expect(freshMockGetZooGatewayModels).toHaveBeenCalledTimes(1)
1303+
})
1304+
12951305
it("revalidates with ETag after TTL expires and keeps the prior catalog on 304", async () => {
12961306
vi.useFakeTimers()
12971307
try {
@@ -1311,6 +1321,12 @@ describe("auth session cache", () => {
13111321
expect(freshMockGetZooGatewayModels).toHaveBeenLastCalledWith(
13121322
expect.objectContaining({ ifNoneMatch: '"v1"' }),
13131323
)
1324+
1325+
// A third call within the refreshed TTL window must not trigger another fetch,
1326+
// proving touchAuthSessionEntry actually updated fetchedAt.
1327+
const third = await freshGetModels(options)
1328+
expect(third).toEqual(zooModels)
1329+
expect(freshMockGetZooGatewayModels).toHaveBeenCalledTimes(2)
13141330
} finally {
13151331
vi.useRealTimers()
13161332
}
@@ -1346,6 +1362,13 @@ describe("auth session cache", () => {
13461362
}
13471363
})
13481364

1365+
it("getModels throws on the first call when the provider fetch fails and there is no prior cache", async () => {
1366+
freshMockGetZooGatewayModels.mockRejectedValue(new Error("network down"))
1367+
const options = { provider: providerIdentifiers.zooGateway, apiKey: "session-token" }
1368+
1369+
await expect(freshGetModels(options)).rejects.toThrow("network down")
1370+
})
1371+
13491372
it("refreshModels returns an empty catalog when refresh throws and no session cache exists", async () => {
13501373
freshMockGetZooGatewayModels.mockRejectedValue(new Error("network down"))
13511374
const options = { provider: providerIdentifiers.zooGateway, apiKey: "session-token" }

‎src/api/providers/fetchers/__tests__/zoo-gateway.spec.ts‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,22 @@ describe("Zoo Gateway Fetchers", () => {
108108
expect(result.etag).toBe('"catalog-abc"')
109109
})
110110

111-
it("reads ETag from the capitalized response header", async () => {
111+
it("reads etag from the lowercase response header (Node.js normalises headers)", async () => {
112+
mockAxiosGet.mockResolvedValueOnce({
113+
...mockResponse,
114+
headers: { etag: '"catalog-lower"' },
115+
})
116+
117+
const result = await getZooGatewayModels(gatewayOptions())
118+
119+
expect(result.kind).toBe("ok")
120+
if (result.kind !== "ok") return
121+
expect(result.etag).toBe('"catalog-lower"')
122+
})
123+
124+
it("returns undefined etag when only the capitalized ETag header is present (unreachable in real HTTP)", async () => {
125+
// Node.js lowercases all HTTP response headers before they reach userland,
126+
// so response.headers.ETag is always undefined in production.
112127
mockAxiosGet.mockResolvedValueOnce({
113128
...mockResponse,
114129
headers: { ETag: '"catalog-capital"' },
@@ -118,7 +133,7 @@ describe("Zoo Gateway Fetchers", () => {
118133

119134
expect(result.kind).toBe("ok")
120135
if (result.kind !== "ok") return
121-
expect(result.etag).toBe('"catalog-capital"')
136+
expect(result.etag).toBeUndefined()
122137
})
123138

124139
it("sends If-None-Match when provided and returns not_modified on 304", async () => {

‎src/api/providers/fetchers/modelCache.ts‎

Lines changed: 60 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,11 @@ const modelRecordSchema = z.record(z.string(), modelInfoSchema)
4444
// deduplicate each other's in-flight refreshes.
4545
const inFlightRefresh = new Map<string, Promise<ModelRecord>>()
4646

47+
// Same single-flight guard for auth-scoped providers. Without this, two concurrent getModels()
48+
// calls (e.g. two panels opening on startup) both miss the session cache and each fire their own
49+
// provider fetch. Keyed on the same compound cache key as authSessionCache.
50+
const inFlightAuthScopedFetch = new Map<string, Promise<ModelRecord>>()
51+
4752
// Cache keys (see getCacheKey) for which we've already reported an empty model response this
4853
// session. A persistently-empty endpoint (e.g. misconfigured server) would otherwise re-fire this
4954
// event on every cache refresh; gate it to at most once per distinct provider+server+key identity
@@ -237,44 +242,62 @@ async function resolveAuthScopedModels(
237242
return existing.models
238243
}
239244

240-
try {
241-
const fetched = await fetchAuthScopedModelsFromProvider(
242-
assertAuthScopedGetModelsOptions(options),
243-
existing?.etag,
244-
)
245+
// Coalesce concurrent fetches for the same session identity so two panels
246+
// opening simultaneously don't each fire their own provider request.
247+
const existingFlight = inFlightAuthScopedFetch.get(cacheKey)
248+
if (existingFlight && !forceRefresh) {
249+
return existingFlight
250+
}
251+
252+
const fetchPromise = (async () => {
253+
try {
254+
const fetched = await fetchAuthScopedModelsFromProvider(
255+
assertAuthScopedGetModelsOptions(options),
256+
existing?.etag,
257+
)
258+
259+
if (fetched.notModified) {
260+
// Re-read from the Map: a concurrent sign-out could have cleared
261+
// the entry between when we captured `existing` and now.
262+
const current = getAuthSessionEntry(cacheKey)
263+
if (current && Object.keys(current.models).length > 0) {
264+
touchAuthSessionEntry(cacheKey, current)
265+
reportedEmptyModelResponse.delete(cacheKey)
266+
return current.models
267+
}
268+
} else {
269+
const modelCount = Object.keys(fetched.models).length
270+
if (modelCount > 0) {
271+
setAuthSessionEntry(cacheKey, fetched.models, fetched.etag)
272+
reportedEmptyModelResponse.delete(cacheKey)
273+
return fetched.models
274+
}
275+
276+
captureModelCacheEmptyResponseOnce(provider, cacheKey, {
277+
context: forceRefresh ? "refreshModels" : "getModels",
278+
hasExistingCache: Boolean(existing && Object.keys(existing.models).length > 0),
279+
...(existing ? { existingCacheSize: Object.keys(existing.models).length } : {}),
280+
})
281+
}
245282

246-
if (fetched.notModified) {
247283
if (existing && Object.keys(existing.models).length > 0) {
248-
touchAuthSessionEntry(cacheKey, existing)
249-
reportedEmptyModelResponse.delete(cacheKey)
250284
return existing.models
251285
}
252-
} else {
253-
const modelCount = Object.keys(fetched.models).length
254-
if (modelCount > 0) {
255-
setAuthSessionEntry(cacheKey, fetched.models, fetched.etag)
256-
reportedEmptyModelResponse.delete(cacheKey)
257-
return fetched.models
258-
}
259-
260-
captureModelCacheEmptyResponseOnce(provider, cacheKey, {
261-
context: forceRefresh ? "refreshModels" : "getModels",
262-
hasExistingCache: Boolean(existing && Object.keys(existing.models).length > 0),
263-
...(existing ? { existingCacheSize: Object.keys(existing.models).length } : {}),
264-
})
265-
}
266286

267-
if (existing && Object.keys(existing.models).length > 0) {
268-
return existing.models
287+
return fetched.models
288+
} catch (error) {
289+
if (existing && Object.keys(existing.models).length > 0) {
290+
return existing.models
291+
}
292+
throw error
269293
}
294+
})()
270295

271-
return fetched.models
272-
} catch (error) {
273-
if (existing && Object.keys(existing.models).length > 0) {
274-
return existing.models
275-
}
276-
throw error
277-
}
296+
// Wrap cleanup in the promise itself so the stored value IS what callers await —
297+
// a detached .finally() would produce an unhandled rejection when the fetch fails.
298+
const fetchWithCleanup = fetchPromise.finally(() => inFlightAuthScopedFetch.delete(cacheKey))
299+
inFlightAuthScopedFetch.set(cacheKey, fetchWithCleanup)
300+
return fetchWithCleanup
278301
}
279302

280303
// Memoize derived digests so the deliberately-structureless KDF runs at most once per
@@ -389,6 +412,12 @@ async function readModels(cacheKey: string): Promise<ModelRecord | undefined> {
389412
async function fetchModelsFromProvider(options: GetModelsOptions): Promise<ModelRecord> {
390413
const { provider } = options
391414

415+
if (isAuthScopedProvider(provider)) {
416+
throw new Error(
417+
`fetchModelsFromProvider must not be called for auth-scoped provider "${provider}" — use resolveAuthScopedModels instead`,
418+
)
419+
}
420+
392421
let models: ModelRecord
393422

394423
switch (provider) {
@@ -432,21 +461,6 @@ async function fetchModelsFromProvider(options: GetModelsOptions): Promise<Model
432461
case providerIdentifiers.moonshot:
433462
models = await getMoonshotModels(options.baseUrl, options.apiKey)
434463
break
435-
case providerIdentifiers.zooGateway: {
436-
const result = await getZooGatewayModels({
437-
zooSessionToken: options.apiKey,
438-
zooGatewayBaseUrl: options.baseUrl,
439-
})
440-
if (result.kind === "not_modified") {
441-
models = {}
442-
break
443-
}
444-
models = result.models
445-
break
446-
}
447-
case providerIdentifiers.kimiCode:
448-
models = await getKimiCodeModels(options.apiKey)
449-
break
450464
default: {
451465
// Ensures router is exhaustively checked if RouterName is a strict union.
452466
const exhaustiveCheck: never = provider

‎src/api/providers/fetchers/zoo-gateway.ts‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -53,12 +53,7 @@ export async function getZooGatewayModels(
5353
return { kind: "not_modified" }
5454
}
5555

56-
const etag =
57-
typeof response.headers.etag === "string"
58-
? response.headers.etag
59-
: typeof response.headers.ETag === "string"
60-
? response.headers.ETag
61-
: undefined
56+
const etag = typeof response.headers.etag === "string" ? response.headers.etag : undefined
6257

6358
const result = vercelAiGatewayModelsResponseSchema.safeParse(response.data)
6459

0 commit comments

Comments
 (0)