feat: add quota tracking
This commit is contained in:
@@ -1,9 +1,12 @@
|
||||
import { Agent, Model, Plugin, Provider } from "@opencode/plugin";
|
||||
import type { RpcRegistration } from "@opencode/plugin/promise/rpc";
|
||||
import type { ModelInfo } from "@opencode/client";
|
||||
import { z } from "zod";
|
||||
import { gptAgents, settings, workerNames } from "./config.js";
|
||||
import { discover } from "./provider.js";
|
||||
import { descriptions, routingGuide } from "./routing.js";
|
||||
import { QuotaTracker } from "./quota.js";
|
||||
import { QuotaRpc } from "./quota-rpc.js";
|
||||
|
||||
// Permission actions use wildcards, not filesystem globs.
|
||||
function matches(value: string, pattern: string) {
|
||||
@@ -50,6 +53,14 @@ export default Plugin.define({
|
||||
let codegraphTools: string[] = [];
|
||||
let agentTransformReady = false;
|
||||
const controller = new AbortController();
|
||||
let quotaRegistration: RpcRegistration<typeof QuotaRpc> | undefined;
|
||||
const quota = new QuotaTracker(context, async () => {
|
||||
await quotaRegistration?.events.emit("updated", {});
|
||||
});
|
||||
quotaRegistration = await context.rpc.register(QuotaRpc, {
|
||||
status: async ({ sessionID }) => (await quota.status(sessionID)) ?? null,
|
||||
});
|
||||
await quota.init();
|
||||
// Subscribe before the first read: configuration can publish its catalog while
|
||||
// asynchronous setup is running.
|
||||
const updates = (async () => {
|
||||
@@ -69,6 +80,20 @@ export default Plugin.define({
|
||||
inherited = permissions;
|
||||
permissionKey = key;
|
||||
await context.agent.reload();
|
||||
} else if (
|
||||
event.type === "session.execution.succeeded" ||
|
||||
event.type === "session.execution.failed" ||
|
||||
event.type === "session.execution.interrupted" ||
|
||||
event.type === "session.idle"
|
||||
) {
|
||||
await quota.complete(event.data.sessionID);
|
||||
} else if (event.type === "credential.updated") {
|
||||
await quota.authChanged();
|
||||
} else if (
|
||||
event.type === "credential.switched" &&
|
||||
String(event.data.integrationID) === "openai"
|
||||
) {
|
||||
await quota.authChanged(true);
|
||||
}
|
||||
}
|
||||
})().catch((error) => {
|
||||
@@ -251,6 +276,7 @@ export default Plugin.define({
|
||||
properties: { ...properties, agent: { type: "string", enum: workers } },
|
||||
};
|
||||
});
|
||||
await context.session.hook("model.request", async (event) => quota.start(event.sessionID));
|
||||
await context.tool.hook("execute.before", (event) => {
|
||||
if (event.tool !== "subagent") return;
|
||||
const { agent } = delegation.parse(event.input);
|
||||
@@ -278,6 +304,8 @@ export default Plugin.define({
|
||||
return async () => {
|
||||
controller.abort();
|
||||
await updates;
|
||||
await quota.dispose();
|
||||
await quotaRegistration?.dispose();
|
||||
};
|
||||
},
|
||||
});
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
import { Rpc } from "@opencode/plugin/rpc";
|
||||
import { z } from "zod";
|
||||
|
||||
const window = z.object({
|
||||
id: z.string(),
|
||||
label: z.string().nullable(),
|
||||
seconds: z.number(),
|
||||
used: z.number(),
|
||||
resetsAt: z.number(),
|
||||
remaining: z.number(),
|
||||
sessionConsumed: z.number(),
|
||||
});
|
||||
|
||||
export const QuotaRpc = Rpc.define({
|
||||
id: "stanp-opencode.quota",
|
||||
methods: {
|
||||
status: {
|
||||
input: z.object({ sessionID: z.string() }),
|
||||
output: z
|
||||
.object({ plan: z.string(), fetchedAt: z.number(), windows: z.array(window) })
|
||||
.nullable(),
|
||||
},
|
||||
},
|
||||
events: { updated: { schema: z.object({}) } },
|
||||
});
|
||||
+429
@@ -0,0 +1,429 @@
|
||||
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";
|
||||
|
||||
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; sessionConsumed: 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(),
|
||||
sessions: z.record(
|
||||
z.string(),
|
||||
z.record(
|
||||
z.string(),
|
||||
z.object({ resetsAt: z.number(), lastUsed: z.number(), consumed: z.number() }),
|
||||
),
|
||||
),
|
||||
});
|
||||
type StoredState = z.infer<typeof storedState>;
|
||||
|
||||
type QuotaContext = Pick<Plugin.Context, "integration" | "session" | "storage">;
|
||||
type QuotaFetch = (input: string | URL | Request, init?: RequestInit) => Promise<Response>;
|
||||
|
||||
const quotaWindow = z.object({
|
||||
limit_window_seconds: z.number().finite().positive(),
|
||||
used_percent: z.number().finite(),
|
||||
reset_at: z.number().finite(),
|
||||
});
|
||||
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
|
||||
.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 Map<string, 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() {
|
||||
try {
|
||||
await this.resolveCredential();
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
? await fingerprint(`account:${id}`)
|
||||
: (this.state?.account ?? (await fingerprint(`credential:${credential.refresh}`)));
|
||||
if (generation !== this.generation) return;
|
||||
if (this.state?.account !== account) {
|
||||
const parsed = storedState.safeParse(await this.context.storage.get(STORAGE_KEY));
|
||||
if (generation !== this.generation) return;
|
||||
const stored = parsed.success ? parsed.data : undefined;
|
||||
if (generation !== this.generation) return;
|
||||
this.state = stored?.account === account ? stored : { account, sessions: {} };
|
||||
}
|
||||
if (generation !== this.generation) return;
|
||||
if (!this.timer) this.timer = setInterval(() => void this.refresh(), this.pollMs);
|
||||
return credential;
|
||||
}
|
||||
|
||||
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 controller = new AbortController();
|
||||
this.controller = controller;
|
||||
const timeout = setTimeout(() => controller.abort(), TIMEOUT_MS);
|
||||
let response: Response;
|
||||
try {
|
||||
const headers = new Headers({
|
||||
Authorization: `Bearer ${credential.access}`,
|
||||
Accept: "application/json",
|
||||
});
|
||||
const id = accountID(credential);
|
||||
if (id) headers.set("ChatGPT-Account-Id", id);
|
||||
response = await this.fetcher(USAGE_URL, { headers, signal: controller.signal });
|
||||
} finally {
|
||||
clearTimeout(timeout);
|
||||
if (this.controller === controller) this.controller = undefined;
|
||||
}
|
||||
if (!response.ok || generation !== this.generation) return;
|
||||
// SAFETY: JSON.parse semantics from Response.json produce only JSON values;
|
||||
// quotaPayload validates all fields used below.
|
||||
const payload = (await response.json()) as Schema.Json;
|
||||
if (generation !== this.generation || !this.state) return;
|
||||
const parsed = quotaPayload.safeParse(payload);
|
||||
if (parsed.success && parsed.data.account_id) {
|
||||
const account = await fingerprint(`account:${parsed.data.account_id}`);
|
||||
if (generation !== this.generation || !this.state) return;
|
||||
if (account !== this.state.account) this.state = { account, sessions: {} };
|
||||
}
|
||||
const snapshot = normalizeQuota(payload);
|
||||
if (!snapshot) return;
|
||||
this.apply(snapshot);
|
||||
await this.persist();
|
||||
await this.changed();
|
||||
return snapshot;
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
private apply(snapshot: QuotaSnapshot) {
|
||||
if (!this.state) return;
|
||||
for (const root of this.active.keys()) {
|
||||
const tracked = (this.state.sessions[root] ??= {});
|
||||
for (const window of snapshot.windows) {
|
||||
const previous = tracked[window.id];
|
||||
if (!previous) {
|
||||
tracked[window.id] = {
|
||||
resetsAt: window.resetsAt,
|
||||
lastUsed: window.used,
|
||||
consumed: 0,
|
||||
};
|
||||
continue;
|
||||
}
|
||||
if (previous.resetsAt !== window.resetsAt) {
|
||||
previous.resetsAt = window.resetsAt;
|
||||
previous.lastUsed = window.used;
|
||||
previous.consumed += window.used;
|
||||
} else if (window.used >= previous.lastUsed) {
|
||||
previous.consumed += window.used - previous.lastUsed;
|
||||
previous.lastUsed = window.used;
|
||||
}
|
||||
}
|
||||
}
|
||||
this.state.latest = snapshot;
|
||||
}
|
||||
|
||||
private async root(sessionID: string) {
|
||||
let current = sessionID;
|
||||
const seen = new Set<string>();
|
||||
while (!seen.has(current)) {
|
||||
seen.add(current);
|
||||
const session = await this.context.session.get({ sessionID: current });
|
||||
if (!session.parentID) return current;
|
||||
current = session.parentID;
|
||||
}
|
||||
return sessionID;
|
||||
}
|
||||
|
||||
async start(sessionID: string) {
|
||||
const generation = this.generation;
|
||||
try {
|
||||
const root = await this.root(sessionID);
|
||||
if (generation !== this.generation) return;
|
||||
const family = this.active.get(root);
|
||||
if (family?.has(sessionID)) return;
|
||||
const snapshot = await this.refresh();
|
||||
if (!snapshot || !this.state || generation !== this.generation) return;
|
||||
if (!family) {
|
||||
const previous = this.state.sessions[root] ?? {};
|
||||
this.state.sessions[root] = Object.fromEntries(
|
||||
snapshot.windows.map((window) => [
|
||||
window.id,
|
||||
{
|
||||
resetsAt: window.resetsAt,
|
||||
lastUsed: window.used,
|
||||
consumed: previous[window.id]?.consumed ?? 0,
|
||||
},
|
||||
]),
|
||||
);
|
||||
}
|
||||
await this.persist();
|
||||
if (generation !== this.generation) return;
|
||||
await this.changed();
|
||||
if (generation !== this.generation) return;
|
||||
const active = family ?? new Set<string>();
|
||||
active.add(sessionID);
|
||||
this.active.set(root, active);
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
async complete(sessionID: string) {
|
||||
try {
|
||||
const root = await this.root(sessionID);
|
||||
const family = this.active.get(root);
|
||||
if (!family?.has(sessionID)) return;
|
||||
await this.refresh();
|
||||
family.delete(sessionID);
|
||||
if (family.size === 0) this.active.delete(root);
|
||||
await this.persist();
|
||||
await this.changed();
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
async status(sessionID: string): Promise<QuotaStatus | undefined> {
|
||||
if (!this.state?.latest) return;
|
||||
const tracked = this.state.sessions[await this.root(sessionID)];
|
||||
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),
|
||||
sessionConsumed: tracked?.[window.id]?.consumed ?? 0,
|
||||
})),
|
||||
};
|
||||
}
|
||||
|
||||
private async persist() {
|
||||
if (!this.state) return;
|
||||
const value: Schema.Json = this.state.latest
|
||||
? {
|
||||
account: this.state.account,
|
||||
sessions: this.state.sessions,
|
||||
latest: this.state.latest,
|
||||
}
|
||||
: { account: this.state.account, sessions: this.state.sessions };
|
||||
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;
|
||||
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);
|
||||
}
|
||||
}
|
||||
+78
@@ -0,0 +1,78 @@
|
||||
import { Plugin } from "@opencode/plugin/tui";
|
||||
import { createResource, createSignal, For, Show } from "solid-js";
|
||||
import type { QuotaStatus, QuotaWindow } from "./quota.js";
|
||||
import { QuotaRpc } from "./quota-rpc.js";
|
||||
|
||||
function duration(seconds: number) {
|
||||
if (seconds % 86_400 === 0) return `${seconds / 86_400}d`;
|
||||
if (seconds % 3_600 === 0) return `${seconds / 3_600}h`;
|
||||
return `${Math.round(seconds / 60)}m`;
|
||||
}
|
||||
|
||||
function percent(value: number) {
|
||||
return value < 0.1 && value > 0 ? "<0.1" : value.toFixed(value < 10 ? 1 : 0);
|
||||
}
|
||||
|
||||
export function formatQuotaWindow(
|
||||
window: QuotaWindow & { remaining: number; sessionConsumed: number },
|
||||
) {
|
||||
const label = window.label
|
||||
? `${window.label} ${duration(window.seconds)}`
|
||||
: duration(window.seconds);
|
||||
return `${label}: ${percent(window.remaining)}% left · ~${percent(window.sessionConsumed)}pp this session`;
|
||||
}
|
||||
|
||||
export async function loadQuotaStatus(load: () => Promise<QuotaStatus | null>) {
|
||||
try {
|
||||
return await load();
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
export default Plugin.define({
|
||||
id: "stanp-opencode.quota-tui",
|
||||
setup(context) {
|
||||
const quota = context.client.rpc(QuotaRpc);
|
||||
const [revision, setRevision] = createSignal(0);
|
||||
const retry = setInterval(() => {
|
||||
setRevision((value) => value + 1);
|
||||
}, 60_000);
|
||||
const unsubscribe = quota.events.on("updated", () => {
|
||||
setRevision((value) => value + 1);
|
||||
});
|
||||
|
||||
const Slot = (props: { sessionID: string }) => {
|
||||
const [status] = createResource(
|
||||
() => `${props.sessionID}:${revision()}`,
|
||||
() => loadQuotaStatus(() => quota.status({ sessionID: props.sessionID })),
|
||||
);
|
||||
return (
|
||||
<Show when={status()}>
|
||||
<box flexDirection="column">
|
||||
<text fg={context.theme.text.default}>
|
||||
<b>ChatGPT quota ({status()!.plan})</b>
|
||||
</text>
|
||||
<For each={status()!.windows}>
|
||||
{(window) => (
|
||||
<text fg={context.theme.text.muted}>
|
||||
{formatQuotaWindow(window)}
|
||||
</text>
|
||||
)}
|
||||
</For>
|
||||
</box>
|
||||
</Show>
|
||||
);
|
||||
};
|
||||
|
||||
const remove = context.ui.slot({
|
||||
append: "sidebar.footer",
|
||||
render: ({ sessionID }) => <Slot sessionID={sessionID} />,
|
||||
});
|
||||
return () => {
|
||||
remove();
|
||||
unsubscribe();
|
||||
clearInterval(retry);
|
||||
};
|
||||
},
|
||||
});
|
||||
Reference in New Issue
Block a user