Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
125 changes: 125 additions & 0 deletions src/quota.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,125 @@
// Provider-agnostic quota model (#725 Part 1). Pure and side-effect free:
// callers hand in window readings from whatever source they have — the live
// quota endpoints the menubar already fetches, or windows derived from local
// token history — and get back one merged, provenance-tagged view with pace.
//
// Provenance is load-bearing: a derived window is an estimate and every
// surface must be able to render it as one. Same discipline as
// costIsEstimated — a guess never renders as a measurement.

export type QuotaWindowKind = 'five_hour' | 'weekly' | 'monthly' | 'credits' | { custom: string }

export type QuotaSource = 'live' | 'derived'

export type QuotaPace = {
/** Elapsed fraction of the window, 0..1. */
expectedFraction: number
/** usedFraction − expectedFraction; positive = ahead of pace (deficit). */
deltaFraction: number
/** Linear projection of usedFraction at the reset boundary. */
projectedAtReset: number
/**
* When the window hits 100% at the current pace. Absent when the pace
* doesn't overflow, or on short windows where a whole-window linear ETA
* isn't defensible (one burst on a 5h window reads as "runs out in 40min",
* then recovers). Deficit still reports there.
*/
exhaustsAt?: Date
}

export type QuotaWindow = {
kind: QuotaWindowKind
/** Display label; provider-supplied when live. */
label: string
/** 0..1. Callers clamp out-of-range provider values before storing. */
usedFraction: number
resetsAt?: Date
windowSeconds?: number
source: QuotaSource
/** Computed via computePace/withPace, never persisted. */
pace?: QuotaPace
}

export type PlanQuota = {
provider: string
/** e.g. "pro", "max_20x" when known. */
plan?: string
windows: QuotaWindow[]
asOf: Date
}

/** No pace until this fraction of the window has elapsed — projecting a week
* off the first few minutes is noise, not signal. */
export const QUOTA_PACE_MIN_ELAPSED_FRACTION = 0.03

/** Windows at or under this length report deficit but no exhaustion ETA. */
export const QUOTA_PACE_ETA_MAX_WINDOW_SECONDS = 6 * 3600

/** Stable identity for merging: live and derived readings of the same window
* kind describe the same underlying quota. */
export function quotaWindowKey(kind: QuotaWindowKind): string {
return typeof kind === 'string' ? kind : `custom:${kind.custom}`
}

/**
* Whole-window linear pace with the guards that keep it honest. Returns
* undefined — never a fabricated value — when the window lacks the inputs,
* hasn't elapsed enough to say anything, is exhausted (the bar already says
* it), or its reset time is skewed (in the past, or further out than one
* window length).
*/
export function computePace(window: QuotaWindow, now: Date = new Date()): QuotaPace | undefined {
const { resetsAt, windowSeconds } = window
if (!resetsAt || !windowSeconds || windowSeconds <= 0) return undefined
const remainingSeconds = (resetsAt.getTime() - now.getTime()) / 1000
if (remainingSeconds <= 0 || remainingSeconds > windowSeconds) return undefined
const elapsedSeconds = windowSeconds - remainingSeconds
const expectedFraction = elapsedSeconds / windowSeconds
if (expectedFraction < QUOTA_PACE_MIN_ELAPSED_FRACTION) return undefined
const used = Math.min(Math.max(window.usedFraction, 0), 1)
if (used >= 1) return undefined

const projectedAtReset = used / expectedFraction
const pace: QuotaPace = {
expectedFraction,
deltaFraction: used - expectedFraction,
projectedAtReset,
}
if (projectedAtReset > 1 && windowSeconds > QUOTA_PACE_ETA_MAX_WINDOW_SECONDS) {
const usedPerSecond = used / elapsedSeconds
if (usedPerSecond > 0) {
pace.exhaustsAt = new Date(now.getTime() + ((1 - used) / usedPerSecond) * 1000)
}
}
return pace
}

/**
* Merge live and derived readings of a provider's windows. Live wins per
* window kind; derived fills the gaps. Provenance stays on each surviving
* window. Order: live windows first (in given order), then unmatched derived.
*/
export function mergeQuotaWindows(live: QuotaWindow[], derived: QuotaWindow[]): QuotaWindow[] {
const liveKeys = new Set(live.map(w => quotaWindowKey(w.kind)))
return [...live, ...derived.filter(w => !liveKeys.has(quotaWindowKey(w.kind)))]
}

/**
* Assemble a plan's merged quota view with pace attached. A provider with no
* readings from either source yields an empty windows list — absent data is
* absent, never zeros.
*/
export function buildPlanQuota(input: {
provider: string
plan?: string
live?: QuotaWindow[]
derived?: QuotaWindow[]
now?: Date
}): PlanQuota {
const now = input.now ?? new Date()
const windows = mergeQuotaWindows(input.live ?? [], input.derived ?? []).map(w => {
const pace = computePace(w, now)
return pace ? { ...w, pace } : { ...w, pace: undefined }
})
return { provider: input.provider, plan: input.plan, windows, asOf: now }
}
181 changes: 181 additions & 0 deletions tests/quota.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,181 @@
import { describe, expect, it } from 'vitest'

import {
QUOTA_PACE_ETA_MAX_WINDOW_SECONDS,
QUOTA_PACE_MIN_ELAPSED_FRACTION,
buildPlanQuota,
computePace,
mergeQuotaWindows,
quotaWindowKey,
type QuotaWindow,
} from '../src/quota.js'

const NOW = new Date('2026-07-18T12:00:00Z')
const WEEK = 7 * 24 * 3600
const FIVE_HOURS = 5 * 3600

function window(overrides: Partial<QuotaWindow> = {}): QuotaWindow {
return {
kind: 'weekly',
label: 'Weekly',
usedFraction: 0.5,
windowSeconds: WEEK,
resetsAt: resetsAfterElapsedFraction(0.5, WEEK),
source: 'live',
...overrides,
}
}

/** resetsAt such that `fraction` of the window has elapsed at NOW. */
function resetsAfterElapsedFraction(fraction: number, windowSeconds: number): Date {
return new Date(NOW.getTime() + windowSeconds * (1 - fraction) * 1000)
}

describe('quotaWindowKey', () => {
it('distinguishes named kinds and custom kinds', () => {
expect(quotaWindowKey('weekly')).toBe('weekly')
expect(quotaWindowKey({ custom: 'Codex-Spark' })).toBe('custom:Codex-Spark')
expect(quotaWindowKey({ custom: 'weekly' })).not.toBe(quotaWindowKey('weekly'))
})
})

describe('computePace', () => {
it('reports on-pace at the window midpoint', () => {
const pace = computePace(window(), NOW)
expect(pace?.expectedFraction).toBeCloseTo(0.5, 6)
expect(pace?.deltaFraction).toBeCloseTo(0, 6)
expect(pace?.projectedAtReset).toBeCloseTo(1, 6)
expect(pace?.exhaustsAt).toBeUndefined()
})

it('computes deficit with an exhaustion ETA strictly before the reset', () => {
const resetsAt = resetsAfterElapsedFraction(0.5, WEEK)
const pace = computePace(window({ usedFraction: 0.8, resetsAt }), NOW)
expect(pace?.deltaFraction).toBeCloseTo(0.3, 6)
expect(pace?.projectedAtReset).toBeCloseTo(1.6, 6)
// 80% in 3.5 days → 100% at 4.375 days elapsed.
const expectedHit = NOW.getTime() + WEEK * (0.5 * (1 / 0.8) - 0.5) * 1000
expect(pace?.exhaustsAt?.getTime()).toBeCloseTo(expectedHit, -3)
expect(pace!.exhaustsAt!.getTime()).toBeLessThan(resetsAt.getTime())
})

it('reports reserve without an ETA', () => {
const pace = computePace(window({ usedFraction: 0.2 }), NOW)
expect(pace?.deltaFraction).toBeCloseTo(-0.3, 6)
expect(pace?.projectedAtReset).toBeCloseTo(0.4, 6)
expect(pace?.exhaustsAt).toBeUndefined()
})

it('suppresses the ETA on short windows but keeps the deficit', () => {
const pace = computePace(
window({
usedFraction: 0.9,
windowSeconds: FIVE_HOURS,
resetsAt: resetsAfterElapsedFraction(0.5, FIVE_HOURS),
}),
NOW
)
expect(pace?.deltaFraction).toBeCloseTo(0.4, 6)
expect(pace?.projectedAtReset).toBeGreaterThan(1)
expect(pace?.exhaustsAt).toBeUndefined()
})

it('says nothing early in the window', () => {
const early = window({ resetsAt: resetsAfterElapsedFraction(0.02, WEEK) })
expect(computePace(early, NOW)).toBeUndefined()
expect(0.02).toBeLessThan(QUOTA_PACE_MIN_ELAPSED_FRACTION)
})

it('says nothing on skewed or missing inputs', () => {
expect(computePace(window({ resetsAt: new Date(NOW.getTime() - 60_000) }), NOW)).toBeUndefined()
expect(
computePace(window({ resetsAt: new Date(NOW.getTime() + (WEEK + 3600) * 1000) }), NOW)
).toBeUndefined()
expect(computePace(window({ resetsAt: undefined }), NOW)).toBeUndefined()
expect(computePace(window({ windowSeconds: undefined }), NOW)).toBeUndefined()
expect(computePace(window({ windowSeconds: 0 }), NOW)).toBeUndefined()
})

it('says nothing on an exhausted window, clamping over-range input', () => {
expect(computePace(window({ usedFraction: 1 }), NOW)).toBeUndefined()
expect(computePace(window({ usedFraction: 1.3 }), NOW)).toBeUndefined()
})

it('treats zero usage mid-window as pure reserve', () => {
const pace = computePace(window({ usedFraction: 0 }), NOW)
expect(pace?.deltaFraction).toBeCloseTo(-0.5, 6)
expect(pace?.projectedAtReset).toBe(0)
expect(pace?.exhaustsAt).toBeUndefined()
})

it('keeps the ETA boundary aligned with the exported constant', () => {
const boundary = window({
usedFraction: 0.9,
windowSeconds: QUOTA_PACE_ETA_MAX_WINDOW_SECONDS,
resetsAt: resetsAfterElapsedFraction(0.5, QUOTA_PACE_ETA_MAX_WINDOW_SECONDS),
})
expect(computePace(boundary, NOW)?.exhaustsAt).toBeUndefined()
const above = window({
usedFraction: 0.9,
windowSeconds: QUOTA_PACE_ETA_MAX_WINDOW_SECONDS + 3600,
resetsAt: resetsAfterElapsedFraction(0.5, QUOTA_PACE_ETA_MAX_WINDOW_SECONDS + 3600),
})
expect(computePace(above, NOW)?.exhaustsAt).toBeInstanceOf(Date)
})
})

describe('mergeQuotaWindows', () => {
const liveWeekly = window({ source: 'live', label: 'Weekly (live)' })
const derivedWeekly = window({ source: 'derived', label: 'Weekly (derived)' })
const derivedMonthly = window({ kind: 'monthly', source: 'derived', label: 'Monthly' })

it('prefers live per window kind and lets derived fill gaps', () => {
const merged = mergeQuotaWindows([liveWeekly], [derivedWeekly, derivedMonthly])
expect(merged).toHaveLength(2)
expect(merged[0]).toBe(liveWeekly)
expect(merged[1]).toBe(derivedMonthly)
})

it('keeps distinct custom kinds separate', () => {
const spark = window({ kind: { custom: 'Spark' }, source: 'live' })
const review = window({ kind: { custom: 'Review' }, source: 'derived' })
expect(mergeQuotaWindows([spark], [review])).toHaveLength(2)
})

it('is all-derived when nothing live exists, and empty when nothing exists', () => {
expect(mergeQuotaWindows([], [derivedWeekly])).toEqual([derivedWeekly])
expect(mergeQuotaWindows([], [])).toEqual([])
})
})

describe('buildPlanQuota', () => {
it('merges, attaches pace, and stamps provenance per window', () => {
const quota = buildPlanQuota({
provider: 'codex',
plan: 'pro',
live: [window({ usedFraction: 0.8 })],
derived: [window({ kind: 'monthly', source: 'derived', usedFraction: 0.2 })],
now: NOW,
})
expect(quota.windows).toHaveLength(2)
expect(quota.windows[0].source).toBe('live')
expect(quota.windows[0].pace?.deltaFraction).toBeCloseTo(0.3, 6)
expect(quota.windows[1].source).toBe('derived')
expect(quota.windows[1].pace?.deltaFraction).toBeCloseTo(-0.3, 6)
expect(quota.asOf).toBe(NOW)
})

it('yields an empty windows list when no source has readings — never zeros', () => {
const quota = buildPlanQuota({ provider: 'claude', now: NOW })
expect(quota.windows).toEqual([])
})

it('leaves pace undefined when the guards say nothing', () => {
const quota = buildPlanQuota({
provider: 'codex',
live: [window({ resetsAt: resetsAfterElapsedFraction(0.01, WEEK) })],
now: NOW,
})
expect(quota.windows[0].pace).toBeUndefined()
})
})