opencode-antigravity-auth/src/plugin/model-specific-quota.test.ts at test/streaming-transformer · ghzhost/opencode-antigravity-auth · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import { describe, it, expect, beforeEach } from "vitest";
import { AccountManager } from "./accounts";
import type { OAuthAuthDetails } from "./types";

describe("Model-specific Gemini quota", () => {
  let manager: AccountManager;
  const auth: OAuthAuthDetails = {
    type: "oauth",
    refresh: "test-refresh",
    access: "test-access",
    expires: Date.now() + 3600000,
  };

  beforeEach(() => {
    manager = new AccountManager(auth);
  });

  it("blocks only the specific Gemini model when markRateLimited is called with a model", () => {
    const account = manager.getCurrentAccountForFamily("gemini")!;
    const modelPro = "gemini-1.5-pro";
    const modelFlash = "gemini-1.5-flash";

    // Mark gemini-1.5-pro as rate limited on antigravity
    manager.markRateLimited(account, 60000, "gemini", "antigravity", modelPro);

    // gemini-1.5-pro should be rate limited for antigravity
    expect(manager.isRateLimitedForHeaderStyle(account, "gemini", "antigravity", modelPro)).toBe(true);

    // gemini-1.5-flash should NOT be rate limited for antigravity
    expect(manager.isRateLimitedForHeaderStyle(account, "gemini", "antigravity", modelFlash)).toBe(false);

    // General gemini (no model) should NOT be rate limited
    expect(manager.isRateLimitedForHeaderStyle(account, "gemini", "antigravity")).toBe(false);
  });

  it("falls back to gemini-cli only for the specific model", () => {
    const account = manager.getCurrentAccountForFamily("gemini")!;
    const modelPro = "gemini-1.5-pro";
    const modelFlash = "gemini-1.5-flash";

    // Mark gemini-1.5-pro as rate limited on antigravity
    manager.markRateLimited(account, 60000, "gemini", "antigravity", modelPro);

    // Available header style for Pro should be gemini-cli
    expect(manager.getAvailableHeaderStyle(account, "gemini", modelPro)).toBe("gemini-cli");

    // Available header style for Flash should still be antigravity
    expect(manager.getAvailableHeaderStyle(account, "gemini", modelFlash)).toBe("antigravity");
  });

  it("returns null when all header styles are exhausted for the specific model on a single account", () => {
    const manager2 = new AccountManager(auth);

    const account = manager2.getCurrentAccountForFamily("gemini")!;
    const modelPro = "gemini-1.5-pro";
    const modelFlash = "gemini-1.5-flash";

    manager2.markRateLimited(account, 60000, "gemini", "antigravity", modelPro);
    manager2.markRateLimited(account, 60000, "gemini", "gemini-cli", modelPro);

    // No other account available, so returns null for the rate-limited model
    expect(manager2.getCurrentOrNextForFamily("gemini", modelPro)).toBeNull();

    // Flash should still return the same account since it's not rate-limited
    const flashAccount = manager2.getCurrentOrNextForFamily("gemini", modelFlash);
    expect(flashAccount).toBe(account);
  });

  it("base family rate limit blocks all models in that family", () => {
    const account = manager.getCurrentAccountForFamily("gemini")!;
    const modelPro = "gemini-1.5-pro";

    // Mark base gemini-antigravity as rate limited
    manager.markRateLimited(account, 60000, "gemini", "antigravity");

    // All Gemini models should now be blocked for antigravity on this account
    expect(manager.isRateLimitedForHeaderStyle(account, "gemini", "antigravity", modelPro)).toBe(true);
    expect(manager.isRateLimitedForHeaderStyle(account, "gemini", "antigravity", "gemini-1.5-flash")).toBe(true);
  });
});