From f260a1817b9a3bb693408e298c56443f97823f80 Mon Sep 17 00:00:00 2001 From: B1nh M1nh <43268322+b1nhm1nh@users.noreply.github.com> Date: Sun, 2 Aug 2026 09:34:27 +0700 Subject: [PATCH] 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. --- custom-server.js | 41 +++- open-sse/providers/registry/ollama.js | 1 + open-sse/services/usage.js | 2 +- open-sse/services/usage/misc.js | 91 ++++++-- .../usage/components/ProviderLimits/utils.js | 16 ++ src/shared/services/initializeApp.js | 6 + src/sse/services/backgroundTokenRefresh.js | 195 +++++++++++++++++ src/sse/services/tokenRefresh.js | 8 +- tests/unit/background-token-refresh.test.js | 203 ++++++++++++++++++ tests/unit/ollama-usage.test.js | 165 ++++++++++++++ 10 files changed, 710 insertions(+), 18 deletions(-) create mode 100644 src/sse/services/backgroundTokenRefresh.js create mode 100644 tests/unit/background-token-refresh.test.js create mode 100644 tests/unit/ollama-usage.test.js diff --git a/custom-server.js b/custom-server.js index 6e39683f..f21d4366 100644 --- a/custom-server.js +++ b/custom-server.js @@ -1,7 +1,42 @@ const http = require("http"); +const path = require("path"); +const { pathToFileURL } = require("url"); const origCreate = http.createServer.bind(http); +let backgroundRefreshStarted = false; + +function startBackgroundTokenRefreshFromCustomServer() { + if (backgroundRefreshStarted) return; + backgroundRefreshStarted = true; + // Prefer source path (repo / standalone that still has src). Fail-open if missing + // — initializeApp also starts the same scheduler when the Next app boots. + const modPath = path.join(__dirname, "src", "sse", "services", "backgroundTokenRefresh.js"); + import(pathToFileURL(modPath).href) + .then((m) => { + try { + m.startBackgroundTokenRefresh(); + } catch (e) { + console.error("[BackgroundTokenRefresh] start failed:", e && e.message ? e.message : e); + } + const stop = () => { + try { + m.stopBackgroundTokenRefresh(); + } catch { + /* ignore */ + } + }; + process.once("SIGINT", stop); + process.once("SIGTERM", stop); + }) + .catch((e) => { + // Expected in published CLI standalone (src/ not on disk). App bootstrap covers it. + if (process.env.DEBUG_BACKGROUND_TOKEN_REFRESH) { + console.error("[BackgroundTokenRefresh] import failed:", e && e.message ? e.message : e); + } + }); +} + // Wrap Next standalone HTTP server: derive client IP from the TCP socket // (unspoofable) and strip client-supplied forwarding headers so downstream // rate-limiting keys on the real peer address instead of attacker-controlled XFF. @@ -26,7 +61,11 @@ http.createServer = (...args) => { if (viaProxy) req.headers["x-9r-via-proxy"] = "1"; return handler(req, res); }; - return origCreate(...rest, wrapped); + const server = origCreate(...rest, wrapped); + server.once("listening", () => { + startBackgroundTokenRefreshFromCustomServer(); + }); + return server; }; require("./server.js"); diff --git a/open-sse/providers/registry/ollama.js b/open-sse/providers/registry/ollama.js index 69923aa1..1e938c58 100644 --- a/open-sse/providers/registry/ollama.js +++ b/open-sse/providers/registry/ollama.js @@ -32,5 +32,6 @@ export default { serviceKinds: ["llm"], features: { usage: true, + usageApikey: true, }, }; diff --git a/open-sse/services/usage.js b/open-sse/services/usage.js index 65819e72..d1c361b9 100644 --- a/open-sse/services/usage.js +++ b/open-sse/services/usage.js @@ -39,7 +39,7 @@ const USAGE_HANDLERS = { qoder: (c) => getQoderUsage(c.accessToken, c.proxyOptions), qwen: (c) => getQwenUsage(c.accessToken, c.providerSpecificData), iflow: (c) => getIflowUsage(c.accessToken), - ollama: (c) => getOllamaUsage(c.accessToken), + ollama: (c) => getOllamaUsage(c.apiKey, c.providerSpecificData, c.proxyOptions), glm: (c) => getGlmUsage(c.apiKey, c.provider, c.proxyOptions), "glm-cn": (c) => getGlmUsage(c.apiKey, c.provider, c.proxyOptions), minimax: (c) => getMiniMaxUsage(c.apiKey, c.provider, c.proxyOptions), diff --git a/open-sse/services/usage/misc.js b/open-sse/services/usage/misc.js index 6ce012fa..6a4323c8 100644 --- a/open-sse/services/usage/misc.js +++ b/open-sse/services/usage/misc.js @@ -46,23 +46,86 @@ export async function getIflowUsage(accessToken) { /** * Ollama Cloud Usage - * Ollama Cloud uses an API key from ollama.com/settings/keys - * and has no public usage API — free tier has light usage limits (resets every 5h & 7d). - * This returns an informational message with the plan details. + * GET https://ollama.com/api/usage — session (5h) + weekly (7d) `usage` is a 0..1 + * ratio (1.0 = limit reached, e.g. weekly 100% used). No reset timestamp exposed. + * POST https://ollama.com/api/me — plan label (fail-open). + * Auth: Authorization: Bearer */ -export async function getOllamaUsage(accessToken, providerSpecificData) { +export async function getOllamaUsage(apiKey, providerSpecificData, proxyOptions = null) { + if (!apiKey) { + return { message: "Ollama Cloud API key not available." }; + } + try { - // Ollama Cloud does not expose a public quota/usage API. - // The provider is configured as noAuth with a notice explaining limits. - // We return a graceful message so the UI shows a friendly state instead of an error. - const plan = providerSpecificData?.plan || "Free"; - return { - plan, - message: "Ollama Cloud uses a free tier with light usage limits (resets every 5h & 7d). For detailed usage tracking, visit ollama.com/settings/keys.", - quotas: [], - }; + const response = await proxyAwareFetch("https://ollama.com/api/usage", { + headers: { + Authorization: `Bearer ${apiKey}`, + Accept: "application/json", + }, + }, proxyOptions); + + if (response.status === 401 || response.status === 403) { + return { message: "Ollama Cloud API key invalid or expired." }; + } + + if (!response.ok) { + return { message: `Ollama Cloud usage API error (${response.status}).` }; + } + + let data; + try { + data = await response.json(); + } catch { + return { message: "Ollama Cloud usage response was not JSON." }; + } + + // Best-effort plan label from /api/me + const me = await proxyAwareFetch("https://ollama.com/api/me", { + method: "POST", + headers: { + Authorization: `Bearer ${apiKey}`, + Accept: "application/json", + "Content-Length": "0", + }, + }, proxyOptions).then((r) => (r.ok ? r.json() : null)).catch(() => null); + + const planRaw = typeof me?.Plan === "string" ? me.Plan : ""; + const plan = planRaw + ? planRaw.charAt(0).toUpperCase() + planRaw.slice(1).toLowerCase() + : "Ollama Cloud"; + + const limits = data?.limits && typeof data.limits === "object" ? data.limits : {}; + + // Ollama `usage` is a 0..1 ratio (1.0 = limit reached). Convert to a 0..100 + // bar. Do NOT set absolute `remaining` — QuotaTable reads remainingPercentage. + function ratioQuota(usageRatio, resetAt = null) { + const ratio = Math.max(0, Math.min(1, Number(usageRatio) || 0)); + const usedPct = Math.round(ratio * 100); + return { used: usedPct, total: 100, remainingPercentage: 100 - usedPct, resetAt, unlimited: false }; + } + + const sessionRaw = limits.session?.usage; + const weeklyRaw = limits.weekly?.usage; + const sessionNum = Number(sessionRaw); + const weeklyNum = Number(weeklyRaw); + const hasSession = sessionRaw !== undefined && sessionRaw !== null && !Number.isNaN(sessionNum); + const hasWeekly = weeklyRaw !== undefined && weeklyRaw !== null && !Number.isNaN(weeklyNum); + + if (!hasSession && !hasWeekly) { + return { + plan, + message: "Ollama Cloud connected. No usage limits reported.", + quotas: {}, + }; + } + + const quotas = {}; + if (hasSession) quotas["Session (5h)"] = ratioQuota(sessionNum); + if (hasWeekly) quotas["Weekly (7d)"] = ratioQuota(weeklyNum); + + return { plan, quotas }; } catch (error) { - return { message: "Unable to fetch Ollama Cloud usage." }; + return { message: `Ollama Cloud error: ${error.message}` }; } } diff --git a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.js b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.js index b3a61c19..9f185b83 100644 --- a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.js +++ b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.js @@ -522,6 +522,22 @@ export function parseQuotaData(provider, data) { } break; + case "ollama": + // Session (5h) / Weekly (7d) usage % from ollama.com/api/usage. + // remainingPercentage only — no absolute remaining (UI treats remaining as %). + if (data.quotas) { + Object.entries(data.quotas).forEach(([name, quota]) => { + normalizedQuotas.push({ + name, + used: quota.used || 0, + total: quota.total || 0, + resetAt: quota.resetAt || null, + remainingPercentage: quota.remainingPercentage, + }); + }); + } + break; + default: // Generic fallback for unknown providers if (data.quotas) { diff --git a/src/shared/services/initializeApp.js b/src/shared/services/initializeApp.js index 536b6676..5b27f774 100644 --- a/src/shared/services/initializeApp.js +++ b/src/shared/services/initializeApp.js @@ -112,6 +112,12 @@ async function runHeavyStartup() { .then(({ startQuotaAutoPing }) => startQuotaAutoPing()) .catch((e) => console.log("[AutoPing] scheduler start failed:", e.message)); } + + // Proactive OAuth token refresh (e.g. grok-cli ~6h TTL). Module is idempotent + // and also started from custom-server.js when that entry is used. + import("@/sse/services/backgroundTokenRefresh.js") + .then(({ startBackgroundTokenRefresh }) => startBackgroundTokenRefresh()) + .catch((e) => console.log("[BackgroundTokenRefresh] scheduler start failed:", e.message)); } function hasQuotaAutoPingEnabled(settings) { diff --git a/src/sse/services/backgroundTokenRefresh.js b/src/sse/services/backgroundTokenRefresh.js new file mode 100644 index 00000000..b10a86ba --- /dev/null +++ b/src/sse/services/backgroundTokenRefresh.js @@ -0,0 +1,195 @@ +// Background proactive OAuth token refresh — independent of inbound requests. +// Fail-open everywhere: tick errors and per-connection failures never kill the interval. + +import * as log from "../utils/logger.js"; +import { getRefreshLeadMs } from "open-sse/services/tokenRefresh.js"; +import { getCredentialExpiryMs } from "open-sse/services/oauthCredentialManager.js"; + +/** Refresh when expiry is within 30 minutes (or the provider on-request lead, whichever larger). */ +export const BACKGROUND_REFRESH_LEAD_MS = 30 * 60 * 1000; +const DEFAULT_INTERVAL_MS = 5 * 60 * 1000; +const INITIAL_DELAY_MS = 10 * 1000; + +let started = false; +let intervalHandle = null; +let initialTimeoutHandle = null; +let tickRunning = false; + +function isTruthyEnv(value) { + if (value == null || value === "") return false; + const v = String(value).trim().toLowerCase(); + return v === "1" || v === "true" || v === "yes" || v === "on"; +} + +function isNonServerRuntime() { + if (typeof window !== "undefined") return true; + const phase = process.env.NEXT_PHASE || ""; + if ( + phase === "phase-production-build" || + phase === "phase-export" || + phase === "phase-static" + ) { + return true; + } + // Next.js build / static generation markers + if (process.env.NEXT_RUNTIME === "edge") return true; + return false; +} + +/** + * Pure selection: OAuth connections with a refreshToken whose access token + * expires within max(provider on-request lead, BACKGROUND_REFRESH_LEAD_MS). + * + * @param {Array} connections + * @param {number} [nowMs] + * @returns {Array} + */ +export function selectConnectionsNeedingRefresh(connections, nowMs = Date.now()) { + if (!Array.isArray(connections) || connections.length === 0) return []; + + const out = []; + for (const conn of connections) { + if (!conn) continue; + + const authType = String(conn.authType || "").toLowerCase().replace(/_/g, ""); + if (authType !== "oauth") continue; + if (!conn.refreshToken) continue; + + const expiresAtMs = getCredentialExpiryMs(conn); + if (expiresAtMs === null) continue; + + const providerLead = getRefreshLeadMs(conn.provider); + const leadMs = Math.max( + Number.isFinite(providerLead) ? providerLead : 0, + BACKGROUND_REFRESH_LEAD_MS + ); + + if (expiresAtMs - nowMs < leadMs) { + out.push(conn); + } + } + return out; +} + +async function loadActiveConnections() { + // Dynamic import avoids circular load with db / app graph at module eval time. + const { getProviderConnections } = await import("../../lib/db/repos/connectionsRepo.js"); + return getProviderConnections({ isActive: true }); +} + +async function refreshOne(connection) { + const { checkAndRefreshToken } = await import("./tokenRefresh.js"); + return checkAndRefreshToken(connection.provider, connection, { force: true }); +} + +/** + * One scheduler tick. Fail-open at top level and per connection. + * @param {{ loadConnections?: Function, refreshConnection?: Function }} [deps] + */ +export async function runBackgroundTokenRefreshTick(deps = {}) { + if (tickRunning) { + log.debug("BG_TOKEN_REFRESH", "Tick already running, skip"); + return; + } + tickRunning = true; + try { + const load = deps.loadConnections || loadActiveConnections; + const refresh = deps.refreshConnection || refreshOne; + + const connections = await load(); + const due = selectConnectionsNeedingRefresh(connections, Date.now()); + + if (due.length === 0) { + log.debug("BG_TOKEN_REFRESH", "No connections due for refresh", { + active: Array.isArray(connections) ? connections.length : 0, + }); + return; + } + + log.info("BG_TOKEN_REFRESH", "Refreshing due OAuth connections", { + due: due.length, + ids: due.map((c) => c.id).filter(Boolean), + }); + + await Promise.allSettled( + due.map(async (conn) => { + try { + await refresh(conn); + log.info("BG_TOKEN_REFRESH", "Connection refresh finished", { + id: conn.id, + provider: conn.provider, + }); + } catch (err) { + log.warn("BG_TOKEN_REFRESH", "Connection refresh failed (swallowed)", { + id: conn?.id, + provider: conn?.provider, + error: err?.message ?? String(err), + }); + } + }) + ); + } catch (err) { + log.warn("BG_TOKEN_REFRESH", "Tick failed (swallowed)", { + error: err?.message ?? String(err), + }); + } finally { + tickRunning = false; + } +} + +/** + * Start the background interval. Safe to call multiple times (no-op if already started). + * @param {{ intervalMs?: number }} [opts] + * @returns {boolean} true if started this call + */ +export function startBackgroundTokenRefresh({ intervalMs } = {}) { + if (started) return false; + if (isTruthyEnv(process.env.DISABLE_BACKGROUND_TOKEN_REFRESH)) { + log.info("BG_TOKEN_REFRESH", "Disabled via DISABLE_BACKGROUND_TOKEN_REFRESH"); + return false; + } + if (isNonServerRuntime()) { + log.debug("BG_TOKEN_REFRESH", "Skip start outside long-running server runtime"); + return false; + } + + started = true; + const period = Number.isFinite(intervalMs) && intervalMs > 0 ? intervalMs : DEFAULT_INTERVAL_MS; + + const safeTick = () => { + runBackgroundTokenRefreshTick().catch((err) => { + log.warn("BG_TOKEN_REFRESH", "Unhandled tick rejection (swallowed)", { + error: err?.message ?? String(err), + }); + }); + }; + + // First pass soon after boot so idle connections don't wait a full interval. + initialTimeoutHandle = setTimeout(safeTick, INITIAL_DELAY_MS); + if (initialTimeoutHandle.unref) initialTimeoutHandle.unref(); + + intervalHandle = setInterval(safeTick, period); + if (intervalHandle.unref) intervalHandle.unref(); + + log.info("BG_TOKEN_REFRESH", "Scheduler started", { + intervalMs: period, + initialDelayMs: INITIAL_DELAY_MS, + leadMs: BACKGROUND_REFRESH_LEAD_MS, + }); + return true; +} + +export function stopBackgroundTokenRefresh() { + if (initialTimeoutHandle) { + clearTimeout(initialTimeoutHandle); + initialTimeoutHandle = null; + } + if (intervalHandle) { + clearInterval(intervalHandle); + intervalHandle = null; + } + if (started) { + started = false; + log.info("BG_TOKEN_REFRESH", "Scheduler stopped"); + } +} diff --git a/src/sse/services/tokenRefresh.js b/src/sse/services/tokenRefresh.js index d6fe288e..4e75a8ad 100644 --- a/src/sse/services/tokenRefresh.js +++ b/src/sse/services/tokenRefresh.js @@ -216,16 +216,20 @@ export async function updateProviderCredentials(connectionId, newCredentials) { * * @param {string} provider * @param {object} credentials + * @param {{ force?: boolean }} [options] force=true skips the on-request lead check + * (used by background scheduler which applies a larger lead). Request path omits this. * @returns {Promise} updated credentials object */ -export async function checkAndRefreshToken(provider, credentials) { +export async function checkAndRefreshToken(provider, credentials, options = {}) { let creds = { ...credentials }; if (!creds.connectionId && creds.id) { creds.connectionId = creds.id; } + const force = options?.force === true; + // ── 1. Regular access-token expiry ──────────────────────────────────────── - if (_shouldRefreshCredentials(provider, creds)) { + if (force || _shouldRefreshCredentials(provider, creds)) { const expiresAt = creds.expiresAt ? new Date(creds.expiresAt).getTime() : null; const remaining = expiresAt ? expiresAt - Date.now() : null; const refreshLead = _getRefreshLeadMs(provider); diff --git a/tests/unit/background-token-refresh.test.js b/tests/unit/background-token-refresh.test.js new file mode 100644 index 00000000..62507545 --- /dev/null +++ b/tests/unit/background-token-refresh.test.js @@ -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(); + }); +}); diff --git a/tests/unit/ollama-usage.test.js b/tests/unit/ollama-usage.test.js new file mode 100644 index 00000000..fd8587e3 --- /dev/null +++ b/tests/unit/ollama-usage.test.js @@ -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, + }); + }); +});