diff --git a/open-sse/utils/diagnostics.ts b/open-sse/utils/diagnostics.ts index 545203c3ad..b323da4496 100644 --- a/open-sse/utils/diagnostics.ts +++ b/open-sse/utils/diagnostics.ts @@ -208,6 +208,10 @@ export function detectMalformedNonStream(resp: unknown): MalformedReason | null // falls through to `empty_choices` → a false 502 (#5108, regression from #4942). if (body.type === "message" && Array.isArray(body.content)) { const hasOutput = (body.content as unknown[]).some((block) => { + // A malformed/partial provider response could carry a null (or non-object) + // entry in `content`; guard before type-asserting so the detector never + // throws on `null.type` (that would crash the whole non-stream classifier). + if (block === null || typeof block !== "object") return false; const b = block as Record; // Text block with visible text. if (b.type === "text" && typeof b.text === "string" && (b.text as string).length > 0) { diff --git a/src/domain/quotaCache.ts b/src/domain/quotaCache.ts index 5d1242923d..cb03598f90 100644 --- a/src/domain/quotaCache.ts +++ b/src/domain/quotaCache.ts @@ -22,6 +22,7 @@ import { cleanupOldSnapshots, getLatestQuotaSnapshotsForConnection, } from "@/lib/db/quotaSnapshots"; +import { recordProviderQuotaResetEventIfChanged } from "@/lib/db/quotaResetEvents"; // ─── Types ────────────────────────────────────────────────────────────────── @@ -236,6 +237,19 @@ export function setQuotaCache( (quotaInfo.total > 0 ? Math.round(((quotaInfo.total - (quotaInfo.used || 0)) / quotaInfo.total) * 100) : 0); + recordProviderQuotaResetEventIfChanged({ + provider, + connectionId, + windowKey, + currentResetAt: quotaInfo.resetAt ?? null, + currentRemainingPercentage: remainingPercentage, + previousObservation: prior?.quotas?.[windowKey] + ? { + resetAt: prior.quotas[windowKey].resetAt, + remainingPercentage: prior.quotas[windowKey].remainingPercentage, + } + : null, + }); // #4438 — only persist on the first observation or a real change. if (!quotaSnapshotChanged(prior, windowKey, remainingPercentage, entry.exhausted)) continue; try { diff --git a/src/lib/db/migrations/108_provider_quota_reset_events.sql b/src/lib/db/migrations/108_provider_quota_reset_events.sql new file mode 100644 index 0000000000..2565cf074a --- /dev/null +++ b/src/lib/db/migrations/108_provider_quota_reset_events.sql @@ -0,0 +1,26 @@ +-- Migration: observed provider quota reset windows +-- Records real upstream quota window transitions so API-key USD weekly caps can +-- align with the provider-observed reset instead of assuming resetAt - 7 days. + +CREATE TABLE IF NOT EXISTS provider_quota_reset_events ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + provider TEXT NOT NULL, + connection_id TEXT NOT NULL, + window_key TEXT NOT NULL, + window_started_at TEXT NOT NULL, + window_resets_at TEXT NOT NULL, + observed_at TEXT NOT NULL, + previous_remaining_percentage REAL, + new_remaining_percentage REAL, + previous_used_percentage REAL, + new_used_percentage REAL, + raw_data TEXT, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + UNIQUE(connection_id, window_key, window_started_at, window_resets_at) +); + +CREATE INDEX IF NOT EXISTS idx_provider_quota_reset_events_connection_window + ON provider_quota_reset_events(connection_id, window_key, window_resets_at); + +CREATE INDEX IF NOT EXISTS idx_provider_quota_reset_events_provider_observed + ON provider_quota_reset_events(provider, observed_at); diff --git a/src/lib/db/quotaResetEvents.ts b/src/lib/db/quotaResetEvents.ts new file mode 100644 index 0000000000..86701e65c8 --- /dev/null +++ b/src/lib/db/quotaResetEvents.ts @@ -0,0 +1,206 @@ +import { getDbInstance } from "./core"; + +interface StatementLike { + all: (...params: unknown[]) => TRow[]; + get: (...params: unknown[]) => TRow | undefined; + run: (...params: unknown[]) => { changes?: number }; +} + +interface DbLike { + prepare: (sql: string) => StatementLike; +} + +interface QuotaObservation { + resetAt: string | null; + remainingPercentage: number | null; +} + +interface ResetEventInput { + provider: string; + connectionId: string; + windowKey: string; + currentResetAt: string | null; + currentRemainingPercentage: number | null; + previousObservation?: QuotaObservation | null; + observedAt?: string; +} + +interface ResetEventWindowRow { + windowStartedAt: string; + windowResetsAt: string; + observedAt: string; +} + +interface QuotaSnapshotObservationRow { + nextResetAt: string | null; + remainingPercentage: number | null; +} + +function toNumberOrNull(value: unknown): number | null { + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string" && value.trim()) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : null; + } + return null; +} + +function clampPercent(value: number | null): number | null { + if (value === null || !Number.isFinite(value)) return null; + return Math.max(0, Math.min(100, value)); +} + +function usedPercent(remainingPercentage: number | null): number | null { + const remaining = clampPercent(remainingPercentage); + return remaining === null ? null : Math.max(0, Math.min(100, 100 - remaining)); +} + +function parseResetIso(value: string | null): string | null { + if (!value) return null; + const parsed = Date.parse(value); + if (!Number.isFinite(parsed)) return null; + return new Date(parsed).toISOString(); +} + +function resetDay(value: string | null): string | null { + const iso = parseResetIso(value); + return iso ? iso.slice(0, 10) : null; +} + +function normalizeWindowKey(value: string): string { + return value + .toLowerCase() + .replace(/[^a-z0-9]+/g, " ") + .trim(); +} + +function isPrimaryWeeklyWindow(windowKey: string): boolean { + const normalized = normalizeWindowKey(windowKey); + return ( + (normalized.includes("weekly") || normalized.includes("7d")) && !normalized.includes("sonnet") + ); +} + +function getLatestSnapshotObservation( + connectionId: string, + windowKey: string +): QuotaObservation | null { + const db = getDbInstance() as unknown as DbLike; + try { + const row = db + .prepare( + ` + SELECT + next_reset_at as nextResetAt, + remaining_percentage as remainingPercentage + FROM quota_snapshots + WHERE connection_id = ? + AND LOWER(window_key) = LOWER(?) + AND next_reset_at IS NOT NULL + ORDER BY created_at DESC, id DESC + LIMIT 1 + ` + ) + .get(connectionId, windowKey); + if (!row) return null; + return { + resetAt: row.nextResetAt, + remainingPercentage: toNumberOrNull(row.remainingPercentage), + }; + } catch (error: unknown) { + if (error instanceof Error && error.message.includes("no such table")) return null; + throw error; + } +} + +export function recordProviderQuotaResetEventIfChanged(input: ResetEventInput): void { + if (!input.connectionId || !input.windowKey || !isPrimaryWeeklyWindow(input.windowKey)) return; + + const currentResetIso = parseResetIso(input.currentResetAt); + if (!currentResetIso) return; + + const previous = + input.previousObservation ?? getLatestSnapshotObservation(input.connectionId, input.windowKey); + const previousResetIso = parseResetIso(previous?.resetAt ?? null); + if (!previousResetIso) return; + + const previousResetMs = Date.parse(previousResetIso); + const currentResetMs = Date.parse(currentResetIso); + if (!Number.isFinite(previousResetMs) || !Number.isFinite(currentResetMs)) return; + if (currentResetMs <= previousResetMs) return; + if (resetDay(previousResetIso) === resetDay(currentResetIso)) return; + + const previousRemaining = clampPercent(toNumberOrNull(previous?.remainingPercentage)); + const currentRemaining = clampPercent(toNumberOrNull(input.currentRemainingPercentage)); + const observedAt = parseResetIso(input.observedAt ?? null) ?? new Date().toISOString(); + + try { + const db = getDbInstance() as unknown as DbLike; + db.prepare( + ` + INSERT OR IGNORE INTO provider_quota_reset_events + (provider, connection_id, window_key, window_started_at, window_resets_at, + observed_at, previous_remaining_percentage, new_remaining_percentage, + previous_used_percentage, new_used_percentage, raw_data) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ` + ).run( + input.provider, + input.connectionId, + input.windowKey, + previousResetIso, + currentResetIso, + observedAt, + previousRemaining, + currentRemaining, + usedPercent(previousRemaining), + usedPercent(currentRemaining), + null + ); + } catch (error: unknown) { + if (error instanceof Error && error.message.includes("no such table")) return; + throw error; + } +} + +export function getProviderQuotaWindowStartIso( + connectionId: string, + targetResetAtIso: string, + nowMs = Date.now() +): string | null { + if (!connectionId || !targetResetAtIso) return null; + const targetDay = resetDay(targetResetAtIso); + if (!targetDay) return null; + + const db = getDbInstance() as unknown as DbLike; + const nowIso = new Date(nowMs).toISOString(); + + try { + const rows = db + .prepare( + ` + SELECT + window_started_at as windowStartedAt, + window_resets_at as windowResetsAt, + observed_at as observedAt + FROM provider_quota_reset_events + WHERE connection_id = @connectionId + AND LOWER(window_key) LIKE '%weekly%' + AND LOWER(window_key) NOT LIKE '%sonnet%' + AND observed_at <= @nowIso + ORDER BY observed_at DESC, id DESC + ` + ) + .all({ connectionId, nowIso }); + + for (const row of rows) { + if (resetDay(row.windowResetsAt) === targetDay) { + return parseResetIso(row.windowStartedAt); + } + } + return null; + } catch (error: unknown) { + if (error instanceof Error && error.message.includes("no such table")) return null; + throw error; + } +} diff --git a/src/lib/localDb.ts b/src/lib/localDb.ts index 1fcadda8f7..98a24aecb7 100755 --- a/src/lib/localDb.ts +++ b/src/lib/localDb.ts @@ -324,6 +324,7 @@ export { } from "./db/quotaSnapshots"; export * from "./db/sessionAccountAffinity"; +export * from "./db/quotaResetEvents"; export type { QuotaSnapshotRow, ProviderUtilizationPoint } from "@/shared/types/utilization"; diff --git a/src/lib/usage/apiKeyUsageLimits.ts b/src/lib/usage/apiKeyUsageLimits.ts index c9d2165df6..92bd2fadcf 100644 --- a/src/lib/usage/apiKeyUsageLimits.ts +++ b/src/lib/usage/apiKeyUsageLimits.ts @@ -1,5 +1,6 @@ import { getDbInstance } from "@/lib/db/core"; import type { ProviderLimitsCacheEntry } from "@/lib/db/providerLimits"; +import { getProviderQuotaWindowStartIso } from "@/lib/db/quotaResetEvents"; import { calculateCost } from "./costCalculator"; import { buildErrorBody, sanitizeErrorMessage } from "@omniroute/open-sse/utils/error.ts"; @@ -258,6 +259,20 @@ function getObservedWeeklyWindowStartIso( } } +// Prefer the persisted, provider-observed window start (recorded by +// quotaResetEvents on real reset transitions); fall back to inferring it from +// historical snapshots when no observed event is available yet. +function getWeeklyWindowStartIso( + connectionId: string, + targetResetAtIso: string, + nowMs: number +): string | null { + return ( + getProviderQuotaWindowStartIso(connectionId, targetResetAtIso, nowMs) ?? + getObservedWeeklyWindowStartIso(connectionId, targetResetAtIso, nowMs) + ); +} + async function resolveDeps(deps: ApiKeyUsageLimitDeps): Promise> { const providers = deps.getProviderConnectionById && deps.getProviderConnections @@ -301,7 +316,7 @@ async function getProviderWeeklyWindow( resetCandidates.push({ connectionId: connection.id, resetAtIso: resetAt, - observedWindowStartIso: getObservedWeeklyWindowStartIso(connection.id, resetAt, nowMs), + observedWindowStartIso: getWeeklyWindowStartIso(connection.id, resetAt, nowMs), }); } } @@ -316,7 +331,7 @@ async function getProviderWeeklyWindow( resetCandidates.push({ connectionId: connection.id, resetAtIso: resetAt, - observedWindowStartIso: getObservedWeeklyWindowStartIso(connection.id, resetAt, nowMs), + observedWindowStartIso: getWeeklyWindowStartIso(connection.id, resetAt, nowMs), }); } } diff --git a/tests/unit/diagnostics.test.ts b/tests/unit/diagnostics.test.ts index bda2e9e394..9c7df879f0 100644 --- a/tests/unit/diagnostics.test.ts +++ b/tests/unit/diagnostics.test.ts @@ -126,6 +126,23 @@ test("detectMalformedNonStream returns null for a Claude-native message carrying assert.equal(detectMalformedNonStream(body), null); }); +test("detectMalformedNonStream tolerates a null block in a Claude-native content array", () => { + // A malformed/partial provider response could carry a null entry in `content`. + // The detector must not throw (TypeError on null.type) — it skips the null + // block and classifies by the remaining valid blocks. + const body = { + type: "message", + role: "assistant", + content: [null, { type: "text", text: "still here" }], + }; + assert.equal(detectMalformedNonStream(body), null); +}); + +test("detectMalformedNonStream returns 'empty_choices' for a Claude-native message of only null blocks", () => { + const body = { type: "message", role: "assistant", content: [null] }; + assert.equal(detectMalformedNonStream(body), "empty_choices"); +}); + test("detectMalformedNonStream returns null when tool_calls present", () => { const body = { choices: [ diff --git a/tests/unit/lib/quota-reset-events.test.ts b/tests/unit/lib/quota-reset-events.test.ts new file mode 100644 index 0000000000..cb91316fb7 --- /dev/null +++ b/tests/unit/lib/quota-reset-events.test.ts @@ -0,0 +1,86 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-quota-reset-events-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../../src/lib/db/core.ts"); +const { recordProviderQuotaResetEventIfChanged, getProviderQuotaWindowStartIso } = await import( + "../../../src/lib/db/quotaResetEvents.ts" +); + +// Force migrations (incl. 108_provider_quota_reset_events) to run. +core.getDbInstance(); + +const CONN = "conn-1"; +const PROVIDER = "antigravity"; +const PREV_RESET = "2026-01-08T00:00:00.000Z"; +const CUR_RESET = "2026-01-15T00:00:00.000Z"; // +7d, different day +const OBSERVED = "2026-01-15T00:05:00.000Z"; + +test.after(() => { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("records a weekly window transition and getWindowStart returns the prior window start", () => { + recordProviderQuotaResetEventIfChanged({ + provider: PROVIDER, + connectionId: CONN, + windowKey: "weekly", + currentResetAt: CUR_RESET, + currentRemainingPercentage: 100, + previousObservation: { resetAt: PREV_RESET, remainingPercentage: 5 }, + observedAt: OBSERVED, + }); + + // For the new window (resets at CUR_RESET) the observed start is PREV_RESET. + const start = getProviderQuotaWindowStartIso(CONN, CUR_RESET, Date.parse(OBSERVED) + 1000); + assert.equal(start, PREV_RESET); +}); + +test("getWindowStart returns null for a reset day with no recorded event", () => { + assert.equal( + getProviderQuotaWindowStartIso(CONN, "2026-02-01T00:00:00.000Z", Date.parse(OBSERVED) + 1000), + null + ); +}); + +test("does not record when previous and current reset fall on the same day (no transition)", () => { + recordProviderQuotaResetEventIfChanged({ + provider: PROVIDER, + connectionId: "conn-sameday", + windowKey: "weekly", + currentResetAt: "2026-03-10T23:00:00.000Z", + currentRemainingPercentage: 80, + previousObservation: { resetAt: "2026-03-10T01:00:00.000Z", remainingPercentage: 70 }, + observedAt: "2026-03-10T23:30:00.000Z", + }); + assert.equal( + getProviderQuotaWindowStartIso( + "conn-sameday", + "2026-03-10T23:00:00.000Z", + Date.parse("2026-03-11T00:00:00.000Z") + ), + null + ); +}); + +test("does not record for a non-weekly (e.g. daily) window", () => { + recordProviderQuotaResetEventIfChanged({ + provider: PROVIDER, + connectionId: "conn-daily", + windowKey: "daily", + currentResetAt: CUR_RESET, + currentRemainingPercentage: 100, + previousObservation: { resetAt: PREV_RESET, remainingPercentage: 5 }, + observedAt: OBSERVED, + }); + assert.equal( + getProviderQuotaWindowStartIso("conn-daily", CUR_RESET, Date.parse(OBSERVED) + 1000), + null + ); +});