Skip to content

Commit fa44261

Browse files
fix(quotas): compare the API key directly instead of hashing it for the cache
1 parent 96f41c6 commit fa44261

2 files changed

Lines changed: 43 additions & 22 deletions

File tree

‎src/services/language-models/__tests__/providerQuotas.spec.ts‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -710,7 +710,7 @@ describe("getProviderQuotas", () => {
710710
expect(fetchSpy).toHaveBeenCalledTimes(2)
711711
})
712712

713-
it("refetches when the API key changes, without keeping the key in the cache key", async () => {
713+
it("refetches when the API key changes", async () => {
714714
const fetchSpy = vi.fn(async () => new Response(JSON.stringify({ data: { limit: 10, limit_remaining: 5 } })))
715715
await getProviderQuotas(
716716
[{ apiProvider: providerIdentifiers.openrouter, openRouterApiKey: "sk-one" } as never],
@@ -725,6 +725,29 @@ describe("getProviderQuotas", () => {
725725
expect(fetchSpy).toHaveBeenCalledTimes(2)
726726
})
727727

728+
it("does not hand a result fetched with one key to a caller using another, even while the first is in flight", async () => {
729+
const releases: Array<(r: Response) => void> = []
730+
const fetchSpy = vi.fn(() => new Promise<Response>((resolve) => releases.push(resolve)))
731+
const ok = () => new Response(JSON.stringify({ data: { limit: 10, limit_remaining: 5 } }))
732+
733+
const first = getProviderQuotas(
734+
[{ apiProvider: providerIdentifiers.openrouter, openRouterApiKey: "sk-one" } as never],
735+
{},
736+
fetchSpy,
737+
)
738+
const second = getProviderQuotas(
739+
[{ apiProvider: providerIdentifiers.openrouter, openRouterApiKey: "sk-two" } as never],
740+
{},
741+
fetchSpy,
742+
)
743+
await vi.waitFor(() => expect(fetchSpy).toHaveBeenCalledTimes(2))
744+
releases.forEach((release) => release(ok()))
745+
await Promise.all([first, second])
746+
747+
// Each key was fetched on its own; neither piggy-backed on the other's request.
748+
expect(fetchSpy).toHaveBeenCalledTimes(2)
749+
})
750+
728751
it("shares one request between concurrent callers for the same credentials", async () => {
729752
let release: (r: Response) => void = () => {}
730753
const fetchSpy = vi.fn(() => new Promise<Response>((resolve) => (release = resolve)))

‎src/services/language-models/providerQuotas.ts‎

Lines changed: 19 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,6 @@ import {
55
type ProviderSettings,
66
} from "@roo-code/types"
77

8-
import { createHash } from "node:crypto"
9-
108
import { getGitHubCopilotQuota } from "./copilotQuota"
119
import {
1210
QuotaHttpError,
@@ -554,19 +552,16 @@ export async function loadProfiles(source: ProfileSource): Promise<ProviderSetti
554552
return profiles.filter((p): p is ProviderSettings => !!p)
555553
}
556554

557-
// Successful results only, per provider and per credential set. A missing key, a
558-
// sign-out or a transient error is never pinned, so fixing it shows up on the
559-
// next open instead of after the TTL.
560-
const cache = new Map<string, { at: number; result: ProviderQuota }>()
561-
const inflight = new Map<string, Promise<ProviderQuota>>()
562-
563-
/** Identifies the key and URLs a result was fetched with, without keeping the key itself. */
564-
function fingerprint(provider: KeyedProvider, creds: KeyedCredentials): string {
565-
const digest = createHash("sha256")
566-
.update(`${creds.apiKey}\n${(creds.urls ?? []).join("\n")}`)
567-
.digest("hex")
568-
.slice(0, 16)
569-
return `${provider.providerId}:${digest}`
555+
// Successful results only, per provider and URL set. A missing key, a sign-out or a
556+
// transient error is never pinned, so fixing it shows up on the next open instead of
557+
// after the TTL. An entry is reused only for the exact key it was fetched with, so
558+
// switching keys refetches. The key is compared, never hashed or used in a map key.
559+
const cache = new Map<string, { at: number; apiKey: string; result: ProviderQuota }>()
560+
const inflight = new Map<string, { apiKey: string; request: Promise<ProviderQuota> }>()
561+
562+
/** Identifies a provider and the URLs it is queried at (the credential itself is not part of it). */
563+
function cacheKey(provider: KeyedProvider, creds: KeyedCredentials): string {
564+
return `${provider.providerId}|${(creds.urls ?? []).join(",")}`
570565
}
571566

572567
function quotaForKeyedProvider(
@@ -580,21 +575,24 @@ function quotaForKeyedProvider(
580575
if (skipped) return Promise.resolve(skipped)
581576

582577
const resolved = creds as KeyedCredentials
583-
const key = fingerprint(provider, resolved)
578+
const apiKey = resolved.apiKey as string
579+
const key = cacheKey(provider, resolved)
584580
const hit = cache.get(key)
585-
if (!force && hit && Date.now() - hit.at < CACHE_TTL_MS) return Promise.resolve(hit.result)
581+
if (!force && hit && hit.apiKey === apiKey && Date.now() - hit.at < CACHE_TTL_MS) return Promise.resolve(hit.result)
586582

587583
// Concurrent callers (e.g. a refresh while the first load is still running) share one request.
588584
const pending = inflight.get(key)
589-
if (pending) return pending
585+
if (pending && pending.apiKey === apiKey) return pending.request
590586

591587
const request = fetchKeyed(provider, resolved, fetchImpl)
592588
.then((result) => {
593-
if (result.ok) cache.set(key, { at: Date.now(), result })
589+
if (result.ok) cache.set(key, { at: Date.now(), apiKey, result })
594590
return result
595591
})
596-
.finally(() => inflight.delete(key))
597-
inflight.set(key, request)
592+
.finally(() => {
593+
if (inflight.get(key)?.request === request) inflight.delete(key)
594+
})
595+
inflight.set(key, { apiKey, request })
598596
return request
599597
}
600598

0 commit comments

Comments
 (0)