feat: Ollama Cloud quota tracker + proactive background OAuth refresh
Ollama: replace informational stub with real quota tracker hitting ollama.com/api/usage (session 5h + weekly 7d, 0..1 ratio) and /api/me plan label; bind handler to apiKey + add features.usageApikey so apikey connections work.
Token refresh: add backgroundTokenRefresh scheduler that refreshes OAuth connections within max(provider lead, 30min) of expiry, independent of inbound traffic (10s after boot, then every 5min, unref'd timers, DISABLE_BACKGROUND_TOKEN_REFRESH kill-switch, fail-open per tick/connection). Registered from custom-server.js (listening) and initializeApp.js. checkAndRefreshToken gains opt-in {force} for the scheduler; request path unchanged.
This commit is contained in:
B1nh M1nh
203
tests/unit/background-token-refresh.test.js
Normal file
203
tests/unit/background-token-refresh.test.js
Normal file
@@ -0,0 +1,203 @@
|
||||
/**
|
||||
* Background OAuth token-refresh scheduler.
|
||||
*
|
||||
* Covers pure selection (selectConnectionsNeedingRefresh) and a fake tick that
|
||||
* exercises checkAndRefreshToken dispatch + fail-open per connection.
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
|
||||
|
||||
const NOW = Date.parse("2026-08-01T12:00:00.000Z");
|
||||
|
||||
function conn(overrides = {}) {
|
||||
return {
|
||||
id: "c1",
|
||||
provider: "grok-cli",
|
||||
authType: "oauth",
|
||||
refreshToken: "rt-1",
|
||||
expiresAt: new Date(NOW + 10 * 60 * 1000).toISOString(),
|
||||
isActive: true,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
describe("selectConnectionsNeedingRefresh", () => {
|
||||
beforeEach(() => {
|
||||
vi.useFakeTimers();
|
||||
vi.setSystemTime(NOW);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
vi.resetModules();
|
||||
});
|
||||
|
||||
it("selects oauth grok-cli connection expiring in 10 minutes", async () => {
|
||||
const { selectConnectionsNeedingRefresh } = await import(
|
||||
"../../src/sse/services/backgroundTokenRefresh.js"
|
||||
);
|
||||
const list = selectConnectionsNeedingRefresh(
|
||||
[conn({ expiresAt: new Date(NOW + 10 * 60 * 1000).toISOString() })],
|
||||
NOW
|
||||
);
|
||||
expect(list).toHaveLength(1);
|
||||
expect(list[0].id).toBe("c1");
|
||||
});
|
||||
|
||||
it("skips connection expiring in 2 hours", async () => {
|
||||
const { selectConnectionsNeedingRefresh } = await import(
|
||||
"../../src/sse/services/backgroundTokenRefresh.js"
|
||||
);
|
||||
const list = selectConnectionsNeedingRefresh(
|
||||
[conn({ expiresAt: new Date(NOW + 2 * 60 * 60 * 1000).toISOString() })],
|
||||
NOW
|
||||
);
|
||||
expect(list).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("never selects apikey connections", async () => {
|
||||
const { selectConnectionsNeedingRefresh } = await import(
|
||||
"../../src/sse/services/backgroundTokenRefresh.js"
|
||||
);
|
||||
const list = selectConnectionsNeedingRefresh(
|
||||
[
|
||||
conn({ authType: "apikey", refreshToken: "rt" }),
|
||||
conn({ id: "c2", authType: "api_key", refreshToken: "rt" }),
|
||||
],
|
||||
NOW
|
||||
);
|
||||
expect(list).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("skips oauth connection without refreshToken", async () => {
|
||||
const { selectConnectionsNeedingRefresh } = await import(
|
||||
"../../src/sse/services/backgroundTokenRefresh.js"
|
||||
);
|
||||
const list = selectConnectionsNeedingRefresh(
|
||||
[conn({ refreshToken: null }), conn({ id: "c2", refreshToken: undefined })],
|
||||
NOW
|
||||
);
|
||||
expect(list).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("selects already-expired oauth connection", async () => {
|
||||
const { selectConnectionsNeedingRefresh } = await import(
|
||||
"../../src/sse/services/backgroundTokenRefresh.js"
|
||||
);
|
||||
const list = selectConnectionsNeedingRefresh(
|
||||
[conn({ expiresAt: new Date(NOW - 60 * 1000).toISOString() })],
|
||||
NOW
|
||||
);
|
||||
expect(list).toHaveLength(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe("runBackgroundTokenRefreshTick", () => {
|
||||
beforeEach(() => {
|
||||
vi.useFakeTimers();
|
||||
vi.setSystemTime(NOW);
|
||||
vi.resetModules();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
it("calls refresh only for due connections and swallows per-connection errors", async () => {
|
||||
const due = conn({
|
||||
id: "due",
|
||||
expiresAt: new Date(NOW + 10 * 60 * 1000).toISOString(),
|
||||
});
|
||||
const notDue = conn({
|
||||
id: "not-due",
|
||||
expiresAt: new Date(NOW + 2 * 60 * 60 * 1000).toISOString(),
|
||||
});
|
||||
const apikey = conn({
|
||||
id: "key",
|
||||
authType: "apikey",
|
||||
expiresAt: new Date(NOW + 60 * 1000).toISOString(),
|
||||
});
|
||||
|
||||
const refreshConnection = vi.fn(async (c) => {
|
||||
if (c.id === "due") throw new Error("boom");
|
||||
return c;
|
||||
});
|
||||
const loadConnections = vi.fn(async () => [due, notDue, apikey]);
|
||||
|
||||
const { runBackgroundTokenRefreshTick } = await import(
|
||||
"../../src/sse/services/backgroundTokenRefresh.js"
|
||||
);
|
||||
|
||||
await expect(
|
||||
runBackgroundTokenRefreshTick({ loadConnections, refreshConnection })
|
||||
).resolves.toBeUndefined();
|
||||
|
||||
expect(loadConnections).toHaveBeenCalledTimes(1);
|
||||
expect(refreshConnection).toHaveBeenCalledTimes(1);
|
||||
expect(refreshConnection.mock.calls[0][0].id).toBe("due");
|
||||
});
|
||||
|
||||
it("does not call refresh when nothing is due", async () => {
|
||||
const refreshConnection = vi.fn();
|
||||
const loadConnections = vi.fn(async () => [
|
||||
conn({
|
||||
expiresAt: new Date(NOW + 3 * 60 * 60 * 1000).toISOString(),
|
||||
}),
|
||||
]);
|
||||
|
||||
const { runBackgroundTokenRefreshTick } = await import(
|
||||
"../../src/sse/services/backgroundTokenRefresh.js"
|
||||
);
|
||||
|
||||
await runBackgroundTokenRefreshTick({ loadConnections, refreshConnection });
|
||||
|
||||
expect(refreshConnection).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("swallows top-level load errors", async () => {
|
||||
const refreshConnection = vi.fn();
|
||||
const loadConnections = vi.fn(async () => {
|
||||
throw new Error("db down");
|
||||
});
|
||||
|
||||
const { runBackgroundTokenRefreshTick } = await import(
|
||||
"../../src/sse/services/backgroundTokenRefresh.js"
|
||||
);
|
||||
|
||||
await expect(
|
||||
runBackgroundTokenRefreshTick({ loadConnections, refreshConnection })
|
||||
).resolves.toBeUndefined();
|
||||
expect(refreshConnection).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe("start/stop guards", () => {
|
||||
afterEach(async () => {
|
||||
vi.unstubAllEnvs();
|
||||
const mod = await import("../../src/sse/services/backgroundTokenRefresh.js");
|
||||
mod.stopBackgroundTokenRefresh();
|
||||
vi.resetModules();
|
||||
});
|
||||
|
||||
it("honors DISABLE_BACKGROUND_TOKEN_REFRESH kill-switch", async () => {
|
||||
vi.stubEnv("DISABLE_BACKGROUND_TOKEN_REFRESH", "1");
|
||||
const { startBackgroundTokenRefresh, stopBackgroundTokenRefresh } = await import(
|
||||
"../../src/sse/services/backgroundTokenRefresh.js"
|
||||
);
|
||||
expect(startBackgroundTokenRefresh()).toBe(false);
|
||||
stopBackgroundTokenRefresh();
|
||||
});
|
||||
|
||||
it("is idempotent: second start is no-op", async () => {
|
||||
vi.stubEnv("DISABLE_BACKGROUND_TOKEN_REFRESH", "");
|
||||
const { startBackgroundTokenRefresh, stopBackgroundTokenRefresh } = await import(
|
||||
"../../src/sse/services/backgroundTokenRefresh.js"
|
||||
);
|
||||
const first = startBackgroundTokenRefresh({ intervalMs: 60_000 });
|
||||
const second = startBackgroundTokenRefresh({ intervalMs: 60_000 });
|
||||
expect(first).toBe(true);
|
||||
expect(second).toBe(false);
|
||||
stopBackgroundTokenRefresh();
|
||||
});
|
||||
});
|
||||
165
tests/unit/ollama-usage.test.js
Normal file
165
tests/unit/ollama-usage.test.js
Normal file
@@ -0,0 +1,165 @@
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
|
||||
vi.mock("../../open-sse/utils/proxyFetch.js", () => ({
|
||||
proxyAwareFetch: vi.fn(),
|
||||
}));
|
||||
|
||||
import { proxyAwareFetch } from "../../open-sse/utils/proxyFetch.js";
|
||||
import { getUsageForProvider } from "../../open-sse/services/usage.js";
|
||||
import {
|
||||
USAGE_SUPPORTED_PROVIDERS,
|
||||
USAGE_APIKEY_PROVIDERS,
|
||||
} from "../../src/shared/constants/providers.js";
|
||||
import { parseQuotaData } from "../../src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.js";
|
||||
|
||||
const USAGE_URL = "https://ollama.com/api/usage";
|
||||
const ME_URL = "https://ollama.com/api/me";
|
||||
|
||||
function jsonResponse(body, status = 200) {
|
||||
return new Response(JSON.stringify(body), {
|
||||
status,
|
||||
headers: { "Content-Type": "application/json" },
|
||||
});
|
||||
}
|
||||
|
||||
const SAMPLE_USAGE = {
|
||||
activity: {
|
||||
cost: "0.00000",
|
||||
period: {
|
||||
type: "last_4_weeks",
|
||||
starting_at: "2026-07-01T00:00:00Z",
|
||||
ending_at: "2026-07-29T00:00:00Z",
|
||||
},
|
||||
models: [],
|
||||
},
|
||||
limits: {
|
||||
session: { usage: 0, models: [] },
|
||||
weekly: {
|
||||
usage: 1,
|
||||
models: [
|
||||
{ name: "glm-5.2", request_count: 5967 },
|
||||
{ name: "kimi-k2.5", request_count: 2 },
|
||||
],
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const SAMPLE_ME = {
|
||||
Plan: "max",
|
||||
};
|
||||
|
||||
describe("ollama registry usage flags", () => {
|
||||
it("is listed for apikey quota dashboard", () => {
|
||||
expect(USAGE_SUPPORTED_PROVIDERS).toContain("ollama");
|
||||
expect(USAGE_APIKEY_PROVIDERS).toContain("ollama");
|
||||
});
|
||||
});
|
||||
|
||||
describe("getUsageForProvider(ollama)", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("GETs /api/usage with Bearer apiKey and POSTs /api/me for plan", async () => {
|
||||
proxyAwareFetch
|
||||
.mockResolvedValueOnce(jsonResponse(SAMPLE_USAGE))
|
||||
.mockResolvedValueOnce(jsonResponse(SAMPLE_ME));
|
||||
|
||||
const usage = await getUsageForProvider({
|
||||
provider: "ollama",
|
||||
apiKey: "k",
|
||||
providerSpecificData: {},
|
||||
});
|
||||
|
||||
expect(usage.message).toBeUndefined();
|
||||
expect(usage.plan).toBe("Max");
|
||||
expect(usage.quotas["Session (5h)"]).toMatchObject({
|
||||
used: 0,
|
||||
total: 100,
|
||||
remainingPercentage: 100,
|
||||
unlimited: false,
|
||||
});
|
||||
expect(usage.quotas["Weekly (7d)"]).toMatchObject({
|
||||
used: 100,
|
||||
total: 100,
|
||||
remainingPercentage: 0,
|
||||
unlimited: false,
|
||||
});
|
||||
// Must not set absolute remaining — UI treats remaining as %
|
||||
expect(usage.quotas["Session (5h)"].remaining).toBeUndefined();
|
||||
expect(usage.quotas["Weekly (7d)"].remaining).toBeUndefined();
|
||||
|
||||
expect(proxyAwareFetch).toHaveBeenCalledTimes(2);
|
||||
|
||||
const [usageUrl, usageOpts] = proxyAwareFetch.mock.calls[0];
|
||||
expect(usageUrl).toBe(USAGE_URL);
|
||||
expect(usageOpts.headers.Authorization).toBe("Bearer k");
|
||||
expect(usageOpts.headers.Accept).toBe("application/json");
|
||||
|
||||
const [meUrl, meOpts] = proxyAwareFetch.mock.calls[1];
|
||||
expect(meUrl).toBe(ME_URL);
|
||||
expect(meOpts.method).toBe("POST");
|
||||
expect(meOpts.headers.Authorization).toBe("Bearer k");
|
||||
expect(meOpts.headers["Content-Length"]).toBe("0");
|
||||
});
|
||||
|
||||
it("surfaces invalid key message on 401", async () => {
|
||||
proxyAwareFetch.mockResolvedValueOnce(
|
||||
jsonResponse({ error: "unauthorized" }, 401),
|
||||
);
|
||||
|
||||
const usage = await getUsageForProvider({
|
||||
provider: "ollama",
|
||||
apiKey: "bad",
|
||||
});
|
||||
|
||||
expect(usage.message).toMatch(/invalid/i);
|
||||
expect(proxyAwareFetch).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("returns message when apiKey missing", async () => {
|
||||
const usage = await getUsageForProvider({
|
||||
provider: "ollama",
|
||||
providerSpecificData: {},
|
||||
});
|
||||
|
||||
expect(usage.message).toMatch(/api key/i);
|
||||
expect(proxyAwareFetch).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe("parseQuotaData(ollama)", () => {
|
||||
it("forwards remainingPercentage for dashboard bars", () => {
|
||||
const rows = parseQuotaData("ollama", {
|
||||
plan: "Max",
|
||||
quotas: {
|
||||
"Session (5h)": {
|
||||
used: 0,
|
||||
total: 100,
|
||||
remainingPercentage: 100,
|
||||
resetAt: null,
|
||||
},
|
||||
"Weekly (7d)": {
|
||||
used: 100,
|
||||
total: 100,
|
||||
remainingPercentage: 0,
|
||||
resetAt: null,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
expect(rows).toHaveLength(2);
|
||||
expect(rows[0]).toMatchObject({
|
||||
name: "Session (5h)",
|
||||
used: 0,
|
||||
total: 100,
|
||||
remainingPercentage: 100,
|
||||
});
|
||||
expect(rows[1]).toMatchObject({
|
||||
name: "Weekly (7d)",
|
||||
used: 100,
|
||||
total: 100,
|
||||
remainingPercentage: 0,
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user