Compare commits
2 Commits
2031ab0506
...
615762b0ac
| Author | SHA1 | Date | |
|---|---|---|---|
| 615762b0ac | |||
| 4b06f45b2c |
82
index.ts
82
index.ts
@@ -1,7 +1,4 @@
|
|||||||
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
import { AuthStorage, type AuthCredential, type ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||||
import { readFileSync } from "node:fs";
|
|
||||||
import { homedir } from "node:os";
|
|
||||||
import { join } from "node:path";
|
|
||||||
import { statusbarConfig } from "./config";
|
import { statusbarConfig } from "./config";
|
||||||
import { contextModule, costModule, directoryModule, modelModule, thinkingModule } from "./modules/basic";
|
import { contextModule, costModule, directoryModule, modelModule, thinkingModule } from "./modules/basic";
|
||||||
import { usageModule } from "./modules/usage";
|
import { usageModule } from "./modules/usage";
|
||||||
@@ -12,7 +9,6 @@ import { zaiUsageProvider } from "./usage/zai";
|
|||||||
import type { ModuleContext, ModuleSpec, RenderedModule, StatusbarState } from "./types";
|
import type { ModuleContext, ModuleSpec, RenderedModule, StatusbarState } from "./types";
|
||||||
import type { Provider, UsageCredential, UsageProvider, UsageReport } from "./usage";
|
import type { Provider, UsageCredential, UsageProvider, UsageReport } from "./usage";
|
||||||
|
|
||||||
const AUTH_PATH = join(homedir(), ".pi", "agent", "auth.json");
|
|
||||||
const REFRESH_MS = 60_000;
|
const REFRESH_MS = 60_000;
|
||||||
const ANTHROPIC_REFRESH_MS = 15 * 60_000;
|
const ANTHROPIC_REFRESH_MS = 15 * 60_000;
|
||||||
|
|
||||||
@@ -32,36 +28,49 @@ function activeProvider(ctx: any): Provider | undefined {
|
|||||||
return undefined;
|
return undefined;
|
||||||
}
|
}
|
||||||
|
|
||||||
function readPiCredential(provider: Provider): UsageCredential | undefined {
|
function credentialString(raw: unknown, key: string): string | undefined {
|
||||||
try {
|
return isRecord(raw) && typeof raw[key] === "string" ? raw[key] : undefined;
|
||||||
const auth = JSON.parse(readFileSync(AUTH_PATH, "utf8"));
|
}
|
||||||
if (!isRecord(auth) || !isRecord(auth[provider])) return undefined;
|
|
||||||
|
|
||||||
const raw = auth[provider];
|
function buildUsageCredential(raw: AuthCredential | undefined, apiKey: string): UsageCredential {
|
||||||
if (raw.type === "oauth") {
|
if (raw?.type === "oauth") {
|
||||||
return {
|
return {
|
||||||
type: "oauth",
|
type: "oauth",
|
||||||
accessToken: typeof raw.access === "string" ? raw.access : undefined,
|
accessToken: apiKey,
|
||||||
refreshToken: typeof raw.refresh === "string" ? raw.refresh : undefined,
|
refreshToken: raw.refresh,
|
||||||
expiresAt: typeof raw.expires === "number" ? raw.expires : undefined,
|
expiresAt: raw.expires,
|
||||||
accountId: typeof raw.accountId === "string" ? raw.accountId : undefined,
|
accountId: credentialString(raw, "accountId"),
|
||||||
email: typeof raw.email === "string" ? raw.email : undefined,
|
email: credentialString(raw, "email"),
|
||||||
metadata: raw,
|
metadata: raw,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
if (raw.type === "api_key") {
|
|
||||||
return {
|
return {
|
||||||
type: "api_key",
|
type: "api_key",
|
||||||
apiKey: typeof raw.key === "string" ? raw.key : undefined,
|
apiKey,
|
||||||
accountId: typeof raw.accountId === "string" ? raw.accountId : undefined,
|
accountId: credentialString(raw, "accountId"),
|
||||||
email: typeof raw.email === "string" ? raw.email : undefined,
|
email: credentialString(raw, "email"),
|
||||||
metadata: raw,
|
metadata: raw,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
} catch {}
|
|
||||||
|
|
||||||
return undefined;
|
async function readPiCredential(authStorage: AuthStorage, provider: Provider): Promise<UsageCredential | undefined> {
|
||||||
|
authStorage.reload();
|
||||||
|
const apiKey = await authStorage.getApiKey(provider);
|
||||||
|
if (!apiKey) return undefined;
|
||||||
|
return buildUsageCredential(authStorage.get(provider), apiKey);
|
||||||
|
}
|
||||||
|
|
||||||
|
async function forceRefreshPiCredential(authStorage: AuthStorage, provider: Provider): Promise<UsageCredential> {
|
||||||
|
authStorage.reload();
|
||||||
|
const raw = authStorage.get(provider);
|
||||||
|
const oauthProvider = authStorage.getOAuthProviders().find(candidate => candidate.id === provider);
|
||||||
|
if (raw?.type !== "oauth" || !oauthProvider) throw new Error("login expired");
|
||||||
|
|
||||||
|
// Refresh Provider OAuth Token
|
||||||
|
const refreshed = await oauthProvider.refreshToken(raw);
|
||||||
|
authStorage.set(provider, { type: "oauth", ...refreshed });
|
||||||
|
return buildUsageCredential(authStorage.get(provider), oauthProvider.getApiKey(refreshed));
|
||||||
}
|
}
|
||||||
|
|
||||||
function renderModule(moduleCtx: ModuleContext, spec: ModuleSpec): RenderedModule {
|
function renderModule(moduleCtx: ModuleContext, spec: ModuleSpec): RenderedModule {
|
||||||
@@ -105,6 +114,7 @@ export default function piStatusbarExtension(pi: ExtensionAPI) {
|
|||||||
let latestCtx: any;
|
let latestCtx: any;
|
||||||
let requestRender: (() => void) | undefined;
|
let requestRender: (() => void) | undefined;
|
||||||
const statusbarState: StatusbarState = {};
|
const statusbarState: StatusbarState = {};
|
||||||
|
const authStorage = AuthStorage.create();
|
||||||
|
|
||||||
function updateThinkingLevel() {
|
function updateThinkingLevel() {
|
||||||
try {
|
try {
|
||||||
@@ -154,7 +164,7 @@ export default function piStatusbarExtension(pi: ExtensionAPI) {
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
const credential = readPiCredential(provider);
|
const credential = await readPiCredential(authStorage, provider);
|
||||||
if (!credential) {
|
if (!credential) {
|
||||||
statusbarState.report = { provider, fetchedAt: Date.now(), limits: [] };
|
statusbarState.report = { provider, fetchedAt: Date.now(), limits: [] };
|
||||||
statusbarState.error = "not logged in";
|
statusbarState.error = "not logged in";
|
||||||
@@ -163,25 +173,37 @@ export default function piStatusbarExtension(pi: ExtensionAPI) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
inFlight?.abort();
|
inFlight?.abort();
|
||||||
inFlight = new AbortController();
|
const controller = new AbortController();
|
||||||
|
inFlight = controller;
|
||||||
statusbarState.report = undefined;
|
statusbarState.report = undefined;
|
||||||
statusbarState.error = undefined;
|
statusbarState.error = undefined;
|
||||||
rerender(ctx);
|
rerender(ctx);
|
||||||
|
|
||||||
try {
|
try {
|
||||||
const report = await usageProviders[provider].fetchUsage({
|
let activeCredential = credential;
|
||||||
provider,
|
const usageCtx = {
|
||||||
credential,
|
|
||||||
baseUrl: ctx.model?.baseUrl,
|
|
||||||
signal: inFlight.signal,
|
|
||||||
}, {
|
|
||||||
fetch: globalThis.fetch.bind(globalThis),
|
fetch: globalThis.fetch.bind(globalThis),
|
||||||
logger: {
|
logger: {
|
||||||
debug: () => undefined,
|
debug: () => undefined,
|
||||||
warn: () => undefined,
|
warn: () => undefined,
|
||||||
},
|
},
|
||||||
|
};
|
||||||
|
const fetchParams = () => ({
|
||||||
|
provider,
|
||||||
|
credential: activeCredential,
|
||||||
|
baseUrl: ctx.model?.baseUrl,
|
||||||
|
signal: controller.signal,
|
||||||
});
|
});
|
||||||
|
|
||||||
|
let report: UsageReport | null;
|
||||||
|
try {
|
||||||
|
report = await usageProviders[provider].fetchUsage(fetchParams(), usageCtx);
|
||||||
|
} catch (error) {
|
||||||
|
if (!(error instanceof Error) || error.message !== "unauthorized") throw error;
|
||||||
|
activeCredential = await forceRefreshPiCredential(authStorage, provider);
|
||||||
|
report = await usageProviders[provider].fetchUsage(fetchParams(), usageCtx);
|
||||||
|
}
|
||||||
|
|
||||||
statusbarState.report = report ?? { provider, fetchedAt: Date.now(), limits: [] };
|
statusbarState.report = report ?? { provider, fetchedAt: Date.now(), limits: [] };
|
||||||
statusbarState.error = report ? undefined : "unavailable";
|
statusbarState.error = report ? undefined : "unavailable";
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
|
|||||||
@@ -78,19 +78,27 @@ function parseIsoTime(value: string | undefined): number | undefined {
|
|||||||
function parseBucket(bucket: unknown): ParsedUsageBucket | undefined {
|
function parseBucket(bucket: unknown): ParsedUsageBucket | undefined {
|
||||||
if (!isRecord(bucket)) return undefined;
|
if (!isRecord(bucket)) return undefined;
|
||||||
const utilization = toNumber(bucket.utilization);
|
const utilization = toNumber(bucket.utilization);
|
||||||
const resetsAt = parseIsoTime(typeof bucket.resets_at === "string" ? bucket.resets_at : undefined);
|
const resetsAt = parseIsoTime(
|
||||||
|
typeof bucket.resets_at === "string" ? bucket.resets_at : undefined,
|
||||||
|
);
|
||||||
if (utilization === undefined && resetsAt === undefined) {
|
if (utilization === undefined && resetsAt === undefined) {
|
||||||
return undefined;
|
return undefined;
|
||||||
}
|
}
|
||||||
return { utilization, resetsAt };
|
return { utilization, resetsAt };
|
||||||
}
|
}
|
||||||
|
|
||||||
function getPayloadString(payload: Record<string, unknown>, key: string): string | undefined {
|
function getPayloadString(
|
||||||
|
payload: Record<string, unknown>,
|
||||||
|
key: string,
|
||||||
|
): string | undefined {
|
||||||
const value = payload[key];
|
const value = payload[key];
|
||||||
return typeof value === "string" && value.trim() ? value.trim() : undefined;
|
return typeof value === "string" && value.trim() ? value.trim() : undefined;
|
||||||
}
|
}
|
||||||
|
|
||||||
function extractUsageIdentity(payload: ClaudeUsageResponse, orgId?: string): { accountId?: string; email?: string } {
|
function extractUsageIdentity(
|
||||||
|
payload: ClaudeUsageResponse,
|
||||||
|
orgId?: string,
|
||||||
|
): { accountId?: string; email?: string } {
|
||||||
if (!isRecord(payload)) return { accountId: orgId };
|
if (!isRecord(payload)) return { accountId: orgId };
|
||||||
const accountId =
|
const accountId =
|
||||||
getPayloadString(payload, "account_id") ??
|
getPayloadString(payload, "account_id") ??
|
||||||
@@ -108,7 +116,12 @@ function extractUsageIdentity(payload: ClaudeUsageResponse, orgId?: string): { a
|
|||||||
}
|
}
|
||||||
|
|
||||||
function hasUsageData(payload: ClaudeUsageResponse): boolean {
|
function hasUsageData(payload: ClaudeUsageResponse): boolean {
|
||||||
return Boolean(payload.five_hour || payload.seven_day || payload.seven_day_opus || payload.seven_day_sonnet);
|
return Boolean(
|
||||||
|
payload.five_hour ||
|
||||||
|
payload.seven_day ||
|
||||||
|
payload.seven_day_opus ||
|
||||||
|
payload.seven_day_sonnet,
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
async function fetchUsagePayload(
|
async function fetchUsagePayload(
|
||||||
@@ -123,12 +136,16 @@ async function fetchUsagePayload(
|
|||||||
try {
|
try {
|
||||||
const response = await ctx.fetch(url, { headers, signal });
|
const response = await ctx.fetch(url, { headers, signal });
|
||||||
if (!response.ok) {
|
if (!response.ok) {
|
||||||
ctx.logger?.warn("Claude usage fetch failed", { status: response.status, statusText: response.statusText });
|
ctx.logger?.warn("Claude usage fetch failed", {
|
||||||
|
status: response.status,
|
||||||
|
statusText: response.statusText,
|
||||||
|
});
|
||||||
return null;
|
return null;
|
||||||
}
|
}
|
||||||
const payload = (await response.json()) as ClaudeUsageResponse;
|
const payload = (await response.json()) as ClaudeUsageResponse;
|
||||||
lastPayload = payload;
|
lastPayload = payload;
|
||||||
const orgId = response.headers.get("anthropic-organization-id")?.trim() || undefined;
|
const orgId =
|
||||||
|
response.headers.get("anthropic-organization-id")?.trim() || undefined;
|
||||||
lastOrgId = orgId ?? lastOrgId;
|
lastOrgId = orgId ?? lastOrgId;
|
||||||
if (payload && isRecord(payload) && hasUsageData(payload)) {
|
if (payload && isRecord(payload) && hasUsageData(payload)) {
|
||||||
return { payload, orgId };
|
return { payload, orgId };
|
||||||
@@ -139,7 +156,9 @@ async function fetchUsagePayload(
|
|||||||
}
|
}
|
||||||
|
|
||||||
if (attempt < MAX_RETRIES - 1) {
|
if (attempt < MAX_RETRIES - 1) {
|
||||||
await new Promise(resolve => setTimeout(resolve, BASE_RETRY_DELAY_MS * 2 ** attempt));
|
await new Promise((resolve) =>
|
||||||
|
setTimeout(resolve, BASE_RETRY_DELAY_MS * 2 ** attempt),
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -181,7 +200,9 @@ async function resolveEmail(
|
|||||||
return profile?.account?.email;
|
return profile?.account?.email;
|
||||||
}
|
}
|
||||||
|
|
||||||
function buildUsageAmount(utilization: number | undefined): UsageAmount | undefined {
|
function buildUsageAmount(
|
||||||
|
utilization: number | undefined,
|
||||||
|
): UsageAmount | undefined {
|
||||||
if (utilization === undefined) return undefined;
|
if (utilization === undefined) return undefined;
|
||||||
const clamped = Math.min(Math.max(utilization, 0), 100);
|
const clamped = Math.min(Math.max(utilization, 0), 100);
|
||||||
const usedFraction = clamped / 100;
|
const usedFraction = clamped / 100;
|
||||||
@@ -195,7 +216,9 @@ function buildUsageAmount(utilization: number | undefined): UsageAmount | undefi
|
|||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
function buildUsageStatus(usedFraction: number | undefined): UsageStatus | undefined {
|
function buildUsageStatus(
|
||||||
|
usedFraction: number | undefined,
|
||||||
|
): UsageStatus | undefined {
|
||||||
if (usedFraction === undefined) return undefined;
|
if (usedFraction === undefined) return undefined;
|
||||||
if (usedFraction >= 1) return "exhausted";
|
if (usedFraction >= 1) return "exhausted";
|
||||||
if (usedFraction >= 0.9) return "warning";
|
if (usedFraction >= 0.9) return "warning";
|
||||||
@@ -220,7 +243,9 @@ function buildUsageLimit(args: {
|
|||||||
id: args.windowId,
|
id: args.windowId,
|
||||||
label: args.windowLabel,
|
label: args.windowLabel,
|
||||||
durationMs: args.durationMs,
|
durationMs: args.durationMs,
|
||||||
...(args.bucket.resetsAt !== undefined ? { resetsAt: args.bucket.resetsAt } : {}),
|
...(args.bucket.resetsAt !== undefined
|
||||||
|
? { resetsAt: args.bucket.resetsAt }
|
||||||
|
: {}),
|
||||||
};
|
};
|
||||||
return {
|
return {
|
||||||
id: args.id,
|
id: args.id,
|
||||||
@@ -237,7 +262,10 @@ function buildUsageLimit(args: {
|
|||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
async function fetchClaudeUsage(params: UsageFetchParams, ctx: UsageFetchContext): Promise<UsageReport | null> {
|
async function fetchClaudeUsage(
|
||||||
|
params: UsageFetchParams,
|
||||||
|
ctx: UsageFetchContext,
|
||||||
|
): Promise<UsageReport | null> {
|
||||||
if (params.provider !== "anthropic") return null;
|
if (params.provider !== "anthropic") return null;
|
||||||
const credential = params.credential;
|
const credential = params.credential;
|
||||||
if (credential.type !== "oauth" || !credential.accessToken) return null;
|
if (credential.type !== "oauth" || !credential.accessToken) return null;
|
||||||
@@ -249,7 +277,12 @@ async function fetchClaudeUsage(params: UsageFetchParams, ctx: UsageFetchContext
|
|||||||
authorization: `Bearer ${credential.accessToken}`,
|
authorization: `Bearer ${credential.accessToken}`,
|
||||||
};
|
};
|
||||||
|
|
||||||
const payloadResult = await fetchUsagePayload(url, headers, ctx, params.signal);
|
const payloadResult = await fetchUsagePayload(
|
||||||
|
url,
|
||||||
|
headers,
|
||||||
|
ctx,
|
||||||
|
params.signal,
|
||||||
|
);
|
||||||
if (!payloadResult || !isRecord(payloadResult.payload)) return null;
|
if (!payloadResult || !isRecord(payloadResult.payload)) return null;
|
||||||
const { payload, orgId } = payloadResult;
|
const { payload, orgId } = payloadResult;
|
||||||
|
|
||||||
@@ -304,7 +337,8 @@ async function fetchClaudeUsage(params: UsageFetchParams, ctx: UsageFetchContext
|
|||||||
if (limits.length === 0) return null;
|
if (limits.length === 0) return null;
|
||||||
const identity = extractUsageIdentity(payload, orgId);
|
const identity = extractUsageIdentity(payload, orgId);
|
||||||
const accountId = identity.accountId ?? credential.accountId;
|
const accountId = identity.accountId ?? credential.accountId;
|
||||||
const email = identity.email ?? (await resolveEmail(params, ctx, baseUrl, headers));
|
const email =
|
||||||
|
identity.email ?? (await resolveEmail(params, ctx, baseUrl, headers));
|
||||||
|
|
||||||
const report: UsageReport = {
|
const report: UsageReport = {
|
||||||
provider: params.provider,
|
provider: params.provider,
|
||||||
@@ -324,14 +358,18 @@ async function fetchClaudeUsage(params: UsageFetchParams, ctx: UsageFetchContext
|
|||||||
export const claudeUsageProvider: UsageProvider = {
|
export const claudeUsageProvider: UsageProvider = {
|
||||||
id: "anthropic",
|
id: "anthropic",
|
||||||
fetchUsage: fetchClaudeUsage,
|
fetchUsage: fetchClaudeUsage,
|
||||||
supports: params => params.provider === "anthropic" && params.credential.type === "oauth",
|
supports: (params) =>
|
||||||
|
params.provider === "anthropic" && params.credential.type === "oauth",
|
||||||
};
|
};
|
||||||
|
|
||||||
export const claudeRankingStrategy: CredentialRankingStrategy = {
|
export const claudeRankingStrategy: CredentialRankingStrategy = {
|
||||||
findWindowLimits(report) {
|
findWindowLimits(report) {
|
||||||
const primary = report.limits.find(l => l.id === "anthropic:5h");
|
const primary = report.limits.find((l) => l.id === "anthropic:5h");
|
||||||
const secondary = report.limits.find(l => l.id === "anthropic:7d");
|
const secondary = report.limits.find((l) => l.id === "anthropic:7d");
|
||||||
return { primary, secondary };
|
return { primary, secondary };
|
||||||
},
|
},
|
||||||
windowDefaults: { primaryMs: 5 * 60 * 60 * 1000, secondaryMs: 7 * 24 * 60 * 60 * 1000 },
|
windowDefaults: {
|
||||||
|
primaryMs: 5 * 60 * 60 * 1000,
|
||||||
|
secondaryMs: 7 * 24 * 60 * 60 * 1000,
|
||||||
|
},
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -126,8 +126,11 @@ function parseUsageWindow(payload: unknown): ParsedUsageWindow | undefined {
|
|||||||
|
|
||||||
function parseUsagePayload(payload: unknown): ParsedUsage | null {
|
function parseUsagePayload(payload: unknown): ParsedUsage | null {
|
||||||
if (!isRecord(payload)) return null;
|
if (!isRecord(payload)) return null;
|
||||||
const planType = typeof payload.plan_type === "string" ? payload.plan_type : undefined;
|
const planType =
|
||||||
const rateLimit = isRecord(payload.rate_limit) ? payload.rate_limit : undefined;
|
typeof payload.plan_type === "string" ? payload.plan_type : undefined;
|
||||||
|
const rateLimit = isRecord(payload.rate_limit)
|
||||||
|
? payload.rate_limit
|
||||||
|
: undefined;
|
||||||
if (!rateLimit) return null;
|
if (!rateLimit) return null;
|
||||||
const parsed: ParsedUsage = {
|
const parsed: ParsedUsage = {
|
||||||
planType,
|
planType,
|
||||||
@@ -137,7 +140,12 @@ function parseUsagePayload(payload: unknown): ParsedUsage | null {
|
|||||||
secondary: parseUsageWindow(rateLimit.secondary_window),
|
secondary: parseUsageWindow(rateLimit.secondary_window),
|
||||||
raw: payload as CodexUsagePayload,
|
raw: payload as CodexUsagePayload,
|
||||||
};
|
};
|
||||||
if (!parsed.primary && !parsed.secondary && parsed.allowed === undefined && parsed.limitReached === undefined) {
|
if (
|
||||||
|
!parsed.primary &&
|
||||||
|
!parsed.secondary &&
|
||||||
|
parsed.allowed === undefined &&
|
||||||
|
parsed.limitReached === undefined
|
||||||
|
) {
|
||||||
return null;
|
return null;
|
||||||
}
|
}
|
||||||
return parsed;
|
return parsed;
|
||||||
@@ -149,7 +157,8 @@ function normalizeCodexBaseUrl(baseUrl?: string): string {
|
|||||||
const base = trimmed.replace(/\/+$/, "");
|
const base = trimmed.replace(/\/+$/, "");
|
||||||
const lower = base.toLowerCase();
|
const lower = base.toLowerCase();
|
||||||
if (
|
if (
|
||||||
(lower.startsWith("https://chatgpt.com") || lower.startsWith("https://chat.openai.com")) &&
|
(lower.startsWith("https://chatgpt.com") ||
|
||||||
|
lower.startsWith("https://chat.openai.com")) &&
|
||||||
!lower.includes("/backend-api")
|
!lower.includes("/backend-api")
|
||||||
) {
|
) {
|
||||||
return `${base}/backend-api`;
|
return `${base}/backend-api`;
|
||||||
@@ -178,7 +187,10 @@ function buildWindowLabel(seconds: number): { id: string; label: string } {
|
|||||||
return { id: `${hours}h`, label: formatWindowLabel(hours, "hour") };
|
return { id: `${hours}h`, label: formatWindowLabel(hours, "hour") };
|
||||||
}
|
}
|
||||||
|
|
||||||
function resolveResetTime(window: ParsedUsageWindow, nowMs: number): number | undefined {
|
function resolveResetTime(
|
||||||
|
window: ParsedUsageWindow,
|
||||||
|
nowMs: number,
|
||||||
|
): number | undefined {
|
||||||
const resetAt = window.resetAt;
|
const resetAt = window.resetAt;
|
||||||
if (resetAt !== undefined) {
|
if (resetAt !== undefined) {
|
||||||
const resetAtMs = resetAt > 1_000_000_000_000 ? resetAt : resetAt * 1000;
|
const resetAtMs = resetAt > 1_000_000_000_000 ? resetAt : resetAt * 1000;
|
||||||
@@ -190,15 +202,29 @@ function resolveResetTime(window: ParsedUsageWindow, nowMs: number): number | un
|
|||||||
return undefined;
|
return undefined;
|
||||||
}
|
}
|
||||||
|
|
||||||
function buildUsageWindow(window: ParsedUsageWindow, key: string, nowMs: number): UsageWindow {
|
function buildUsageWindow(
|
||||||
|
window: ParsedUsageWindow,
|
||||||
|
key: string,
|
||||||
|
nowMs: number,
|
||||||
|
): UsageWindow {
|
||||||
const resetsAt = resolveResetTime(window, nowMs);
|
const resetsAt = resolveResetTime(window, nowMs);
|
||||||
if (window.limitWindowSeconds !== undefined) {
|
if (window.limitWindowSeconds !== undefined) {
|
||||||
const { id, label } = buildWindowLabel(window.limitWindowSeconds);
|
const { id, label } = buildWindowLabel(window.limitWindowSeconds);
|
||||||
const durationMs = window.limitWindowSeconds * 1000;
|
const durationMs = window.limitWindowSeconds * 1000;
|
||||||
return { id, label, durationMs, ...(resetsAt !== undefined ? { resetsAt } : {}) };
|
return {
|
||||||
|
id,
|
||||||
|
label,
|
||||||
|
durationMs,
|
||||||
|
...(resetsAt !== undefined ? { resetsAt } : {}),
|
||||||
|
};
|
||||||
}
|
}
|
||||||
const fallbackLabel = key === "primary" ? "Primary window" : "Secondary window";
|
const fallbackLabel =
|
||||||
return { id: key, label: fallbackLabel, ...(resetsAt !== undefined ? { resetsAt } : {}) };
|
key === "primary" ? "Primary window" : "Secondary window";
|
||||||
|
return {
|
||||||
|
id: key,
|
||||||
|
label: fallbackLabel,
|
||||||
|
...(resetsAt !== undefined ? { resetsAt } : {}),
|
||||||
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
function buildUsageAmount(window: ParsedUsageWindow): UsageAmount {
|
function buildUsageAmount(window: ParsedUsageWindow): UsageAmount {
|
||||||
@@ -218,7 +244,10 @@ function buildUsageAmount(window: ParsedUsageWindow): UsageAmount {
|
|||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
function buildUsageStatus(usedFraction?: number, limitReached?: boolean): UsageLimit["status"] {
|
function buildUsageStatus(
|
||||||
|
usedFraction?: number,
|
||||||
|
limitReached?: boolean,
|
||||||
|
): UsageLimit["status"] {
|
||||||
if (limitReached) return "exhausted";
|
if (limitReached) return "exhausted";
|
||||||
if (usedFraction === undefined) return "unknown";
|
if (usedFraction === undefined) return "unknown";
|
||||||
if (usedFraction >= 1) return "exhausted";
|
if (usedFraction >= 1) return "exhausted";
|
||||||
@@ -255,9 +284,14 @@ function buildUsageLimit(args: {
|
|||||||
export const openaiCodexUsageProvider: UsageProvider = {
|
export const openaiCodexUsageProvider: UsageProvider = {
|
||||||
id: "openai-codex",
|
id: "openai-codex",
|
||||||
supports(params: UsageFetchParams): boolean {
|
supports(params: UsageFetchParams): boolean {
|
||||||
return params.provider === "openai-codex" && params.credential.type === "oauth";
|
return (
|
||||||
|
params.provider === "openai-codex" && params.credential.type === "oauth"
|
||||||
|
);
|
||||||
},
|
},
|
||||||
async fetchUsage(params: UsageFetchParams, ctx: UsageFetchContext): Promise<UsageReport | null> {
|
async fetchUsage(
|
||||||
|
params: UsageFetchParams,
|
||||||
|
ctx: UsageFetchContext,
|
||||||
|
): Promise<UsageReport | null> {
|
||||||
if (params.provider !== "openai-codex") return null;
|
if (params.provider !== "openai-codex") return null;
|
||||||
const { credential } = params;
|
const { credential } = params;
|
||||||
if (credential.type !== "oauth") return null;
|
if (credential.type !== "oauth") return null;
|
||||||
@@ -267,7 +301,9 @@ export const openaiCodexUsageProvider: UsageProvider = {
|
|||||||
|
|
||||||
const nowMs = Date.now();
|
const nowMs = Date.now();
|
||||||
if (credential.expiresAt !== undefined && credential.expiresAt <= nowMs) {
|
if (credential.expiresAt !== undefined && credential.expiresAt <= nowMs) {
|
||||||
ctx.logger?.warn("Codex usage token expired", { provider: params.provider });
|
ctx.logger?.warn("Codex usage token expired", {
|
||||||
|
provider: params.provider,
|
||||||
|
});
|
||||||
return null;
|
return null;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -288,19 +324,27 @@ export const openaiCodexUsageProvider: UsageProvider = {
|
|||||||
try {
|
try {
|
||||||
const response = await ctx.fetch(url, { headers, signal: params.signal });
|
const response = await ctx.fetch(url, { headers, signal: params.signal });
|
||||||
if (!response.ok) {
|
if (!response.ok) {
|
||||||
ctx.logger?.warn("Codex usage request failed", { status: response.status, provider: params.provider });
|
ctx.logger?.warn("Codex usage request failed", {
|
||||||
|
status: response.status,
|
||||||
|
provider: params.provider,
|
||||||
|
});
|
||||||
return null;
|
return null;
|
||||||
}
|
}
|
||||||
payload = await response.json();
|
payload = await response.json();
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
ctx.logger?.warn("Codex usage request error", { provider: params.provider, error: String(error) });
|
ctx.logger?.warn("Codex usage request error", {
|
||||||
|
provider: params.provider,
|
||||||
|
error: String(error),
|
||||||
|
});
|
||||||
return null;
|
return null;
|
||||||
}
|
}
|
||||||
|
|
||||||
const parsed = parseUsagePayload(payload);
|
const parsed = parseUsagePayload(payload);
|
||||||
const planType =
|
const planType =
|
||||||
parsed?.planType ??
|
parsed?.planType ??
|
||||||
(isRecord(payload) && typeof payload.plan_type === "string" ? payload.plan_type : undefined);
|
(isRecord(payload) && typeof payload.plan_type === "string"
|
||||||
|
? payload.plan_type
|
||||||
|
: undefined);
|
||||||
|
|
||||||
const limits: UsageLimit[] = [];
|
const limits: UsageLimit[] = [];
|
||||||
if (parsed?.primary) {
|
if (parsed?.primary) {
|
||||||
@@ -350,17 +394,24 @@ const FIVE_HOUR_MS = 5 * 60 * 60 * 1000;
|
|||||||
|
|
||||||
export const codexRankingStrategy: CredentialRankingStrategy = {
|
export const codexRankingStrategy: CredentialRankingStrategy = {
|
||||||
findWindowLimits(report) {
|
findWindowLimits(report) {
|
||||||
const findLimit = (key: "primary" | "secondary"): UsageLimit | undefined => {
|
const findLimit = (
|
||||||
const direct = report.limits.find(l => l.id === `openai-codex:${key}`);
|
key: "primary" | "secondary",
|
||||||
|
): UsageLimit | undefined => {
|
||||||
|
const direct = report.limits.find((l) => l.id === `openai-codex:${key}`);
|
||||||
if (direct) return direct;
|
if (direct) return direct;
|
||||||
const byId = report.limits.find(l => l.id.toLowerCase().includes(key));
|
const byId = report.limits.find((l) => l.id.toLowerCase().includes(key));
|
||||||
if (byId) return byId;
|
if (byId) return byId;
|
||||||
const windowId = key === "secondary" ? "7d" : "1h";
|
const windowId = key === "secondary" ? "7d" : "1h";
|
||||||
return report.limits.find(l => l.scope.windowId?.toLowerCase() === windowId);
|
return report.limits.find(
|
||||||
|
(l) => l.scope.windowId?.toLowerCase() === windowId,
|
||||||
|
);
|
||||||
};
|
};
|
||||||
return { primary: findLimit("primary"), secondary: findLimit("secondary") };
|
return { primary: findLimit("primary"), secondary: findLimit("secondary") };
|
||||||
},
|
},
|
||||||
windowDefaults: { primaryMs: 60 * 60 * 1000, secondaryMs: 7 * 24 * 60 * 60 * 1000 },
|
windowDefaults: {
|
||||||
|
primaryMs: 60 * 60 * 1000,
|
||||||
|
secondaryMs: 7 * 24 * 60 * 60 * 1000,
|
||||||
|
},
|
||||||
hasPriorityBoost(primary) {
|
hasPriorityBoost(primary) {
|
||||||
if (!primary) return false;
|
if (!primary) return false;
|
||||||
const windowId = primary.scope.windowId?.toLowerCase();
|
const windowId = primary.scope.windowId?.toLowerCase();
|
||||||
@@ -372,6 +423,10 @@ export const codexRankingStrategy: CredentialRankingStrategy = {
|
|||||||
Math.abs(durationMs - FIVE_HOUR_MS) <= 60_000);
|
Math.abs(durationMs - FIVE_HOUR_MS) <= 60_000);
|
||||||
if (!isFiveHourWindow) return false;
|
if (!isFiveHourWindow) return false;
|
||||||
const usedFraction = primary.amount.usedFraction;
|
const usedFraction = primary.amount.usedFraction;
|
||||||
return typeof usedFraction === "number" && Number.isFinite(usedFraction) && usedFraction === 0;
|
return (
|
||||||
|
typeof usedFraction === "number" &&
|
||||||
|
Number.isFinite(usedFraction) &&
|
||||||
|
usedFraction === 0
|
||||||
|
);
|
||||||
},
|
},
|
||||||
};
|
};
|
||||||
|
|||||||
29
usage/zai.ts
29
usage/zai.ts
@@ -75,7 +75,8 @@ function buildUsageAmount(args: {
|
|||||||
: args.used !== undefined && args.limit !== undefined && args.limit > 0
|
: args.used !== undefined && args.limit !== undefined && args.limit > 0
|
||||||
? Math.min(args.used / args.limit, 1)
|
? Math.min(args.used / args.limit, 1)
|
||||||
: undefined;
|
: undefined;
|
||||||
const remainingFraction = usedFraction !== undefined ? Math.max(1 - usedFraction, 0) : undefined;
|
const remainingFraction =
|
||||||
|
usedFraction !== undefined ? Math.max(1 - usedFraction, 0) : undefined;
|
||||||
return {
|
return {
|
||||||
used: args.used,
|
used: args.used,
|
||||||
limit: args.limit,
|
limit: args.limit,
|
||||||
@@ -86,7 +87,9 @@ function buildUsageAmount(args: {
|
|||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
function getUsageStatus(usedFraction: number | undefined): UsageStatus | undefined {
|
function getUsageStatus(
|
||||||
|
usedFraction: number | undefined,
|
||||||
|
): UsageStatus | undefined {
|
||||||
if (usedFraction === undefined) return undefined;
|
if (usedFraction === undefined) return undefined;
|
||||||
if (usedFraction >= 1) return "exhausted";
|
if (usedFraction >= 1) return "exhausted";
|
||||||
if (usedFraction >= 0.9) return "warning";
|
if (usedFraction >= 0.9) return "warning";
|
||||||
@@ -107,7 +110,10 @@ function buildModelUsageUrl(baseUrl: string, now: Date): string {
|
|||||||
return `${baseUrl}${MODEL_USAGE_PATH}?startTime=${encodeURIComponent(startTime)}&endTime=${encodeURIComponent(endTime)}`;
|
return `${baseUrl}${MODEL_USAGE_PATH}?startTime=${encodeURIComponent(startTime)}&endTime=${encodeURIComponent(endTime)}`;
|
||||||
}
|
}
|
||||||
|
|
||||||
async function fetchZaiUsage(params: UsageFetchParams, ctx: UsageFetchContext): Promise<UsageReport | null> {
|
async function fetchZaiUsage(
|
||||||
|
params: UsageFetchParams,
|
||||||
|
ctx: UsageFetchContext,
|
||||||
|
): Promise<UsageReport | null> {
|
||||||
if (params.provider !== "zai") return null;
|
if (params.provider !== "zai") return null;
|
||||||
const credential = params.credential;
|
const credential = params.credential;
|
||||||
if (credential.type !== "api_key" || !credential.apiKey) return null;
|
if (credential.type !== "api_key" || !credential.apiKey) return null;
|
||||||
@@ -127,7 +133,10 @@ async function fetchZaiUsage(params: UsageFetchParams, ctx: UsageFetchContext):
|
|||||||
signal: params.signal,
|
signal: params.signal,
|
||||||
});
|
});
|
||||||
if (!response.ok) {
|
if (!response.ok) {
|
||||||
ctx.logger?.warn("ZAI usage fetch failed", { status: response.status, statusText: response.statusText });
|
ctx.logger?.warn("ZAI usage fetch failed", {
|
||||||
|
status: response.status,
|
||||||
|
statusText: response.statusText,
|
||||||
|
});
|
||||||
return null;
|
return null;
|
||||||
}
|
}
|
||||||
payload = (await response.json()) as ZaiQuotaPayload;
|
payload = (await response.json()) as ZaiQuotaPayload;
|
||||||
@@ -138,11 +147,16 @@ async function fetchZaiUsage(params: UsageFetchParams, ctx: UsageFetchContext):
|
|||||||
|
|
||||||
if (!payload) return null;
|
if (!payload) return null;
|
||||||
if (payload.success !== true) {
|
if (payload.success !== true) {
|
||||||
ctx.logger?.warn("ZAI usage response invalid", { code: payload.code, message: payload.msg });
|
ctx.logger?.warn("ZAI usage response invalid", {
|
||||||
|
code: payload.code,
|
||||||
|
message: payload.msg,
|
||||||
|
});
|
||||||
return null;
|
return null;
|
||||||
}
|
}
|
||||||
|
|
||||||
const limitsPayload = Array.isArray(payload.data?.limits) ? payload.data?.limits : [];
|
const limitsPayload = Array.isArray(payload.data?.limits)
|
||||||
|
? payload.data?.limits
|
||||||
|
: [];
|
||||||
const limits: UsageLimit[] = [];
|
const limits: UsageLimit[] = [];
|
||||||
|
|
||||||
for (const rawLimit of limitsPayload) {
|
for (const rawLimit of limitsPayload) {
|
||||||
@@ -243,5 +257,6 @@ async function fetchZaiUsage(params: UsageFetchParams, ctx: UsageFetchContext):
|
|||||||
export const zaiUsageProvider: UsageProvider = {
|
export const zaiUsageProvider: UsageProvider = {
|
||||||
id: "zai",
|
id: "zai",
|
||||||
fetchUsage: fetchZaiUsage,
|
fetchUsage: fetchZaiUsage,
|
||||||
supports: params => params.provider === "zai" && params.credential.type === "api_key",
|
supports: (params) =>
|
||||||
|
params.provider === "zai" && params.credential.type === "api_key",
|
||||||
};
|
};
|
||||||
|
|||||||
Reference in New Issue
Block a user