stanp-opencode/src/quota/quota.ts

import type { Plugin } from "@opencode/plugin";
import type { Credential } from "@opencode/schema/credential";
import type { Schema } from "effect";
import { z } from "zod";

const INTEGRATION_ID = "openai";
const USAGE_URL = "https://chatgpt.com/backend-api/wham/usage";
const POLL_MS = 60_000;
const TIMEOUT_MS = 5_000;
const STORAGE_KEY = "chatgpt-quota";
const jsonValue = z.json();

function compactLimitIdentity(value: string) {
    return value.toLowerCase().replaceAll(/[^a-z0-9]+/gu, "");
}

function isSparkLimit(meteredFeature?: string, limitName?: string | null) {
    return (
        (meteredFeature !== undefined &&
            compactLimitIdentity(meteredFeature) === "codexbengalfox") ||
        (limitName !== null &&
            limitName !== undefined &&
            compactLimitIdentity(limitName) === "gpt53codexspark")
    );
}

function hasReserveToken(value: string) {
    return /(^|[^a-z0-9])reserve($|[^a-z0-9])/iu.test(value);
}

function isReserveLimit(meteredFeature?: string, limitName?: string | null) {
    return (
        (meteredFeature !== undefined && hasReserveToken(meteredFeature)) ||
        (limitName !== null && limitName !== undefined && hasReserveToken(limitName))
    );
}

function isExcludedLimit(meteredFeature?: string, limitName?: string | null) {
    return isSparkLimit(meteredFeature, limitName) || isReserveLimit(meteredFeature, limitName);
}

function isExcludedWindow(window: Pick<QuotaWindow, "id" | "label">) {
    return isExcludedLimit(window.id.split(":", 1)[0], window.label);
}

export interface QuotaWindow {
    id: string;
    label: string | null;
    seconds: number;
    used: number;
    resetsAt: number;
}

export interface QuotaSnapshot {
    plan: string;
    fetchedAt: number;
    windows: QuotaWindow[];
}

export interface QuotaStatus {
    plan: string;
    fetchedAt: number;
    windows: Array<QuotaWindow & { remaining: number }>;
}

const storedState = z.object({
    account: z.string(),
    latest: z
        .object({
            plan: z.string(),
            fetchedAt: z.number(),
            windows: z.array(
                z.object({
                    id: z.string(),
                    label: z.string().nullable(),
                    seconds: z.number(),
                    used: z.number(),
                    resetsAt: z.number(),
                }),
            ),
        })
        .optional(),
});
type StoredState = z.infer<typeof storedState>;

function withoutExcludedState(state: StoredState): StoredState {
    const windows = state.latest?.windows.filter((window) => !isExcludedWindow(window));
    return {
        ...state,
        latest:
            state.latest && windows !== undefined && windows.length > 0
                ? { ...state.latest, windows }
                : undefined,
    };
}

type QuotaContext = Pick<Plugin.Context, "integration" | "storage">;
type QuotaFetch = (input: string | URL | Request, init?: RequestInit) => Promise<Response>;

const quotaWindow = z.object({
    limit_window_seconds: z.number().positive(),
    used_percent: z.number(),
    reset_at: z.number(),
});
const additionalRateLimitIdentity = z.object({
    limit_name: z.string().optional(),
    metered_feature: z.string().optional(),
});
const quotaPayload = z.object({
    plan_type: z.string(),
    account_id: z.string().optional(),
    rate_limit: z
        .object({
            allowed: z.boolean().optional(),
            primary_window: quotaWindow.nullish(),
            secondary_window: quotaWindow.nullish(),
        })
        .nullish(),
    additional_rate_limits: z.preprocess(
        (value) =>
            Array.isArray(value)
                ? value.filter((item) => {
                      const identity = additionalRateLimitIdentity.safeParse(item);
                      return (
                          !identity.success ||
                          !isExcludedLimit(identity.data.metered_feature, identity.data.limit_name)
                      );
                  })
                : value,
        z
            .array(
                z.object({
                    limit_name: z.string(),
                    metered_feature: z.string(),
                    rate_limit: z
                        .object({
                            allowed: z.boolean().optional(),
                            primary_window: quotaWindow.nullish(),
                            secondary_window: quotaWindow.nullish(),
                        })
                        .nullish(),
                }),
            )
            .nullish(),
    ),
});

function normalizeWindow(
    id: string,
    item: z.infer<typeof quotaWindow>,
    label?: string,
): QuotaWindow {
    return {
        id,
        label: label ?? null,
        seconds: item.limit_window_seconds,
        used: Math.max(0, Math.min(100, item.used_percent)),
        resetsAt: item.reset_at < 10_000_000_000 ? item.reset_at * 1000 : item.reset_at,
    };
}

/** Normalize valid windows from the documented Codex usage payload. */
export function normalizeQuota(
    input: Schema.Json,
    fetchedAt = Date.now(),
): QuotaSnapshot | undefined {
    const parsed = quotaPayload.safeParse(input);
    if (!parsed.success) return;
    const plan = parsed.data.plan_type.toLowerCase();
    const windows: QuotaWindow[] = [];
    if (parsed.data.rate_limit)
        for (const [position, window] of [
            ["primary", parsed.data.rate_limit.primary_window],
            ["secondary", parsed.data.rate_limit.secondary_window],
        ] as const)
            if (window) windows.push(normalizeWindow(`codex:${position}`, window));
    for (const additional of parsed.data.additional_rate_limits ?? []) {
        if (!additional.rate_limit) continue;
        for (const [position, window] of [
            ["primary", additional.rate_limit.primary_window],
            ["secondary", additional.rate_limit.secondary_window],
        ] as const)
            if (window)
                windows.push(
                    normalizeWindow(
                        `${additional.metered_feature}:${position}`,
                        window,
                        additional.limit_name,
                    ),
                );
    }
    if (windows.length === 0) return;
    return { plan, fetchedAt, windows };
}

function accountID(credential: Credential.OAuth) {
    const metadataID = z
        .string()
        .safeParse(credential.metadata?.["accountID"] ?? credential.metadata?.["accountId"]);
    if (metadataID.success) return metadataID.data;
    const parts = credential.access.split(".");
    if (parts.length !== 3) return;
    try {
        const claims = z
            .object({
                chatgpt_account_id: z.string().optional(),
                "https://api.openai.com/auth": z
                    .object({ chatgpt_account_id: z.string().optional() })
                    .optional(),
                organizations: z.array(z.object({ id: z.string() })).optional(),
            })
            .safeParse(JSON.parse(Buffer.from(parts[1] ?? "", "base64url").toString()));
        if (!claims.success) return;
        return (
            claims.data.chatgpt_account_id ??
            claims.data["https://api.openai.com/auth"]?.chatgpt_account_id ??
            claims.data.organizations?.[0]?.id
        );
    } catch {
        return;
    }
}

async function fingerprint(value: string) {
    const bytes = await crypto.subtle.digest("SHA-256", new TextEncoder().encode(value));
    return Buffer.from(bytes).toString("hex");
}

export class QuotaTracker {
    private state: StoredState | undefined;
    private timer: ReturnType<typeof setInterval> | undefined;
    private request: Promise<QuotaSnapshot | undefined> | undefined;
    private readonly requests = new Set<Promise<QuotaSnapshot | undefined>>();
    private readonly active = new Set<string>();
    private controller: AbortController | undefined;
    private generation = 0;
    private disposed = false;

    constructor(
        private readonly context: QuotaContext,
        private readonly changed: () => Promise<void>,
        private readonly fetcher: QuotaFetch = fetch,
        private readonly pollMs = POLL_MS,
    ) {}

    async init() {
        await this.refresh();
    }

    private async resolveCredential(generation = this.generation) {
        const connection = await this.context.integration.connection.active(INTEGRATION_ID);
        if (generation !== this.generation) return;
        const credential = connection
            ? await this.context.integration.connection.resolve(connection)
            : undefined;
        if (generation !== this.generation) return;
        if (credential?.type !== "oauth") {
            if (this.state) await this.clear();
            return;
        }
        const id = accountID(credential);
        const account =
            id !== undefined && id !== ""
                ? await fingerprint(`account:${id}`)
                : (this.state?.account ?? (await fingerprint(`credential:${credential.refresh}`)));
        if (generation !== this.generation) return;
        if (this.state?.account !== account) await this.load(account, generation);
        if (generation !== this.generation) return;
        this.timer ??= setInterval(() => void this.refresh(), this.pollMs);
        return credential;
    }

    private async load(account: string, generation: number) {
        const value = await this.context.storage.get(STORAGE_KEY);
        const parsed = storedState.safeParse(value);
        if (generation !== this.generation) return;
        const stored = parsed.success ? parsed.data : undefined;
        this.state = stored?.account === account ? withoutExcludedState(stored) : { account };
    }

    /** Reads state through a call so checks after `await` are not narrowed away. */
    private current(): StoredState | undefined {
        return this.state;
    }

    private async clear() {
        this.state = undefined;
        if (this.timer) clearInterval(this.timer);
        this.timer = undefined;
        await this.context.storage.remove(STORAGE_KEY);
        await this.changed();
    }

    async refresh(): Promise<QuotaSnapshot | undefined> {
        const generation = this.generation;
        if (this.disposed) return;
        if (this.request) return this.request;
        const request = this.query();
        this.request = request;
        this.requests.add(request);
        try {
            const result = await request;
            if (generation !== this.generation) return;
            return result ?? this.state?.latest;
        } finally {
            this.requests.delete(request);
            if (this.request === request) this.request = undefined;
        }
    }

    private async query() {
        const generation = this.generation;
        try {
            const credential = await this.resolveCredential(generation);
            if (!credential || !this.state || generation !== this.generation) return;
            const response = await this.fetchUsage(credential);
            if (!response.ok || generation !== this.generation) return;
            const payload = jsonValue.parse(await response.json());
            if (generation !== this.generation || !this.current()) return;
            const parsed = quotaPayload.safeParse(payload);
            if (
                parsed.success &&
                parsed.data.account_id !== undefined &&
                parsed.data.account_id !== ""
            ) {
                const account = await fingerprint(`account:${parsed.data.account_id}`);
                const state = this.current();
                if (generation !== this.generation || !state) return;
                if (account !== state.account) this.state = { account };
            }
            const snapshot = normalizeQuota(payload);
            if (!snapshot) return;
            this.apply(snapshot);
            await this.persist();
            await this.changed();
            return snapshot;
        } catch {
            return;
        }
    }

    private async fetchUsage(credential: Credential.OAuth) {
        const controller = new AbortController();
        this.controller = controller;
        const timeout = setTimeout(() => {
            controller.abort();
        }, TIMEOUT_MS);
        try {
            const headers = new Headers({
                Authorization: `Bearer ${credential.access}`,
                Accept: "application/json",
            });
            const id = accountID(credential);
            if (id !== undefined && id !== "") headers.set("ChatGPT-Account-Id", id);
            return await this.fetcher(USAGE_URL, { headers, signal: controller.signal });
        } finally {
            clearTimeout(timeout);
            if (this.controller === controller) this.controller = undefined;
        }
    }

    private apply(snapshot: QuotaSnapshot) {
        if (!this.state) return;
        this.state.latest = snapshot;
    }

    async start(sessionID: string) {
        if (this.active.has(sessionID)) return;
        const generation = this.generation;
        const snapshot = await this.refresh();
        if (!snapshot || generation !== this.generation) return;
        this.active.add(sessionID);
    }

    async complete(sessionID: string) {
        if (!this.active.delete(sessionID)) return;
        await this.refresh();
    }

    status(): QuotaStatus | undefined {
        if (!this.state?.latest) return;
        return {
            plan: this.state.latest.plan,
            fetchedAt: this.state.latest.fetchedAt,
            windows: this.state.latest.windows.map((window) => ({
                ...window,
                remaining: Math.max(0, 100 - window.used),
            })),
        };
    }

    private async persist() {
        if (!this.state) return;
        const value: Schema.Json = this.state.latest
            ? {
                  account: this.state.account,
                  latest: this.state.latest,
              }
            : { account: this.state.account };
        await this.context.storage.set(STORAGE_KEY, value);
    }

    async authChanged(switched = false) {
        if (switched) {
            this.generation++;
            this.controller?.abort();
            this.controller = undefined;
            this.request = undefined;
            this.active.clear();
            if (this.state) await this.clear();
        }
        await this.refresh();
    }

    async dispose() {
        this.disposed = true;
        this.generation++;
        this.controller?.abort();
        if (this.timer) clearInterval(this.timer);
        this.timer = undefined;
        await Promise.all(this.requests);
    }
}