Skip to content

Commit 4583f97

Browse files
committed
feat(oauth): harden model family classification against gemma prefix collision
- Tighten classifyModelFamilyForQuota to require explicit gemini keyword, preventing gemma from matching Gem family. - Add regression test in tests/oauth/oauth-account-quota-rank.test.ts.
1 parent b091f28 commit 4583f97

2 files changed

Lines changed: 27 additions & 2 deletions

File tree

src/oauth/account-quota-rank.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,9 @@ const PASSIVE_HEADROOM_MAX_AGE_MS = 60 * 60_000;
5050
* Returns "Gem" for Gemini models, "Cla" for Claude/Opus/Sonnet, or undefined
5151
* for unknown models (which falls back to all-window ranking).
5252
*/
53-
function classifyModelFamilyForQuota(modelId: string): string | undefined {
53+
function classifyModelFamilyForQuota(modelId: string): "Gem" | "Cla" | undefined {
5454
const lower = modelId.toLowerCase();
55-
if (lower.includes("gemini") || lower.startsWith("gem")) return "Gem";
55+
if (lower.includes("gemini") || lower === "gemini" || lower.startsWith("gemini-")) return "Gem";
5656
if (lower.includes("claude") || lower.includes("opus") || lower.includes("sonnet") || lower.includes("gpt-oss") || lower.includes("gpt_oss")) return "Cla";
5757
return undefined;
5858
}

tests/oauth/oauth-account-quota-rank.test.ts

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,5 +354,30 @@ describe("model-aware failover and pre-dispatch integration", () => {
354354
const ranked = rankAccountsByHeadroom("google-antigravity", ring, "gemini-3.8-flash");
355355
expect(ranked).toEqual(ring);
356356
});
357+
358+
test("does not misclassify gemma as Gemini (no Gem prefix pollution)", () => {
359+
const idGem = "antigravity-gem";
360+
const idCla = "antigravity-cla";
361+
const ring = [idGem, idCla];
362+
363+
setCachedProviderAccountQuotaForTests("google-antigravity", idGem, {
364+
customWindows: [
365+
{ label: "Gem", percent: 10 },
366+
{ label: "Cla", percent: 90 },
367+
],
368+
updatedAt: Date.now(),
369+
});
370+
setCachedProviderAccountQuotaForTests("google-antigravity", idCla, {
371+
customWindows: [
372+
{ label: "Gem", percent: 90 },
373+
{ label: "Cla", percent: 10 },
374+
],
375+
updatedAt: Date.now(),
376+
});
377+
378+
// When requestedModelId is gemma, it should return undefined family and compare all windows (Math.max of 10 and 90 = 90 for both)
379+
const ranked = rankAccountsByHeadroom("google-antigravity", ring, "gemma-2-9b-it");
380+
expect(ranked).toEqual(ring);
381+
});
357382
});
358383

0 commit comments

Comments
 (0)