@oh-my-pi/pi-ai 6.9.69 → 8.0.0

package/src/usage/openai-codex.ts

@@ -0,0 +1,393 @@
+import { Buffer } from "node:buffer";
+import { CODEX_BASE_URL } from "$ai/providers/openai-codex/constants";
+import type {
+	UsageAmount,
+	UsageCache,
+	UsageFetchContext,
+	UsageFetchParams,
+	UsageLimit,
+	UsageProvider,
+	UsageReport,
+	UsageWindow,
+} from "$ai/usage";
+
+const CODEX_USAGE_PATH = "wham/usage";
+const DEFAULT_CACHE_TTL_MS = 60_000;
+const JWT_AUTH_CLAIM = "https://api.openai.com/auth";
+const JWT_PROFILE_CLAIM = "https://api.openai.com/profile";
+
+interface CodexUsageWindowPayload {
+	used_percent?: number;
+	limit_window_seconds?: number;
+	reset_after_seconds?: number;
+	reset_at?: number;
+}
+
+interface CodexUsageRateLimitPayload {
+	allowed?: boolean;
+	limit_reached?: boolean;
+	primary_window?: CodexUsageWindowPayload | null;
+	secondary_window?: CodexUsageWindowPayload | null;
+}
+
+interface CodexUsagePayload {
+	plan_type?: string;
+	rate_limit?: CodexUsageRateLimitPayload | null;
+}
+
+interface ParsedUsageWindow {
+	usedPercent?: number;
+	limitWindowSeconds?: number;
+	resetAfterSeconds?: number;
+	resetAt?: number;
+}
+
+interface ParsedUsage {
+	planType?: string;
+	allowed?: boolean;
+	limitReached?: boolean;
+	primary?: ParsedUsageWindow;
+	secondary?: ParsedUsageWindow;
+	raw: CodexUsagePayload;
+}
+
+interface JwtPayload {
+	[JWT_AUTH_CLAIM]?: {
+		chatgpt_account_id?: string;
+	};
+	[JWT_PROFILE_CLAIM]?: {
+		email?: string;
+	};
+}
+
+const isRecord = (value: unknown): value is Record<string, unknown> =>
+	typeof value === "object" && value !== null && !Array.isArray(value);
+
+const toNumber = (value: unknown): number | undefined => {
+	if (typeof value === "number" && Number.isFinite(value)) return value;
+	if (typeof value === "string") {
+		const trimmed = value.trim();
+		if (!trimmed) return undefined;
+		const parsed = Number(trimmed);
+		if (Number.isFinite(parsed)) return parsed;
+	}
+	return undefined;
+};
+
+const toBoolean = (value: unknown): boolean | undefined => {
+	if (typeof value === "boolean") return value;
+	return undefined;
+};
+
+function base64UrlDecode(input: string): string {
+	const base64 = input.replace(/-/g, "+").replace(/_/g, "/");
+	const padLen = (4 - (base64.length % 4)) % 4;
+	const padded = base64 + "=".repeat(padLen);
+	return Buffer.from(padded, "base64").toString("utf8");
+}
+
+function parseJwt(token: string): JwtPayload | null {
+	const parts = token.split(".");
+	if (parts.length !== 3) return null;
+	try {
+		const payloadJson = base64UrlDecode(parts[1]);
+		return JSON.parse(payloadJson) as JwtPayload;
+	} catch {
+		return null;
+	}
+}
+
+function extractAccountId(token: string | undefined): string | undefined {
+	if (!token) return undefined;
+	const payload = parseJwt(token);
+	return payload?.[JWT_AUTH_CLAIM]?.chatgpt_account_id ?? undefined;
+}
+
+function extractEmail(token: string | undefined): string | undefined {
+	if (!token) return undefined;
+	const payload = parseJwt(token);
+	return payload?.[JWT_PROFILE_CLAIM]?.email ?? undefined;
+}
+
+function parseUsageWindow(payload: unknown): ParsedUsageWindow | undefined {
+	if (!isRecord(payload)) return undefined;
+	const usedPercent = toNumber(payload.used_percent);
+	const limitWindowSeconds = toNumber(payload.limit_window_seconds);
+	const resetAfterSeconds = toNumber(payload.reset_after_seconds);
+	const resetAt = toNumber(payload.reset_at);
+	if (
+		usedPercent === undefined &&
+		limitWindowSeconds === undefined &&
+		resetAfterSeconds === undefined &&
+		resetAt === undefined
+	) {
+		return undefined;
+	}
+	return {
+		usedPercent,
+		limitWindowSeconds,
+		resetAfterSeconds,
+		resetAt,
+	};
+}
+
+function parseUsagePayload(payload: unknown): ParsedUsage | null {
+	if (!isRecord(payload)) return null;
+	const planType = typeof payload.plan_type === "string" ? payload.plan_type : undefined;
+	const rateLimit = isRecord(payload.rate_limit) ? payload.rate_limit : undefined;
+	if (!rateLimit) return null;
+	const parsed: ParsedUsage = {
+		planType,
+		allowed: toBoolean(rateLimit.allowed),
+		limitReached: toBoolean(rateLimit.limit_reached),
+		primary: parseUsageWindow(rateLimit.primary_window),
+		secondary: parseUsageWindow(rateLimit.secondary_window),
+		raw: payload as CodexUsagePayload,
+	};
+	if (!parsed.primary && !parsed.secondary && parsed.allowed === undefined && parsed.limitReached === undefined) {
+		return null;
+	}
+	return parsed;
+}
+
+function normalizeCodexBaseUrl(baseUrl?: string): string {
+	const fallback = CODEX_BASE_URL;
+	const trimmed = baseUrl?.trim() ? baseUrl.trim() : fallback;
+	const base = trimmed.replace(/\/+$/, "");
+	const lower = base.toLowerCase();
+	if (
+		(lower.startsWith("https://chatgpt.com") || lower.startsWith("https://chat.openai.com")) &&
+		!lower.includes("/backend-api")
+	) {
+		return `${base}/backend-api`;
+	}
+	return base;
+}
+
+function buildCodexUsageUrl(baseUrl: string): string {
+	const normalized = baseUrl.endsWith("/") ? baseUrl : `${baseUrl}/`;
+	return `${normalized}${CODEX_USAGE_PATH}`;
+}
+
+function formatWindowLabel(value: number, unit: "hour" | "day"): string {
+	const rounded = Math.round(value);
+	const suffix = rounded === 1 ? unit : `${unit}s`;
+	return `${rounded} ${suffix}`;
+}
+
+function buildWindowLabel(seconds: number): { id: string; label: string } {
+	const daySeconds = 86_400;
+	if (seconds >= daySeconds) {
+		const days = Math.round(seconds / daySeconds);
+		return { id: `${days}d`, label: formatWindowLabel(days, "day") };
+	}
+	const hours = Math.max(1, Math.round(seconds / 3600));
+	return { id: `${hours}h`, label: formatWindowLabel(hours, "hour") };
+}
+
+function resolveResetTimes(window: ParsedUsageWindow, nowMs: number): Pick<UsageWindow, "resetsAt" | "resetInMs"> {
+	const resetAt = window.resetAt;
+	if (resetAt !== undefined) {
+		const resetAtMs = resetAt > 1_000_000_000_000 ? resetAt : resetAt * 1000;
+		if (Number.isFinite(resetAtMs)) {
+			return { resetsAt: resetAtMs, resetInMs: resetAtMs - nowMs };
+		}
+	}
+	if (window.resetAfterSeconds !== undefined) {
+		const resetInMs = window.resetAfterSeconds * 1000;
+		return { resetsAt: nowMs + resetInMs, resetInMs };
+	}
+	return {};
+}
+
+function buildUsageWindow(window: ParsedUsageWindow, key: string, nowMs: number): UsageWindow {
+	if (window.limitWindowSeconds !== undefined) {
+		const { id, label } = buildWindowLabel(window.limitWindowSeconds);
+		const durationMs = window.limitWindowSeconds * 1000;
+		return { id, label, durationMs, ...resolveResetTimes(window, nowMs) };
+	}
+	const fallbackLabel = key === "primary" ? "Primary window" : "Secondary window";
+	return { id: key, label: fallbackLabel, ...resolveResetTimes(window, nowMs) };
+}
+
+function buildUsageAmount(window: ParsedUsageWindow): UsageAmount {
+	const usedPercent = window.usedPercent;
+	if (usedPercent === undefined) {
+		return { unit: "percent" };
+	}
+	const clamped = Math.min(Math.max(usedPercent, 0), 100);
+	const usedFraction = clamped / 100;
+	return {
+		used: clamped,
+		limit: 100,
+		remaining: Math.max(0, 100 - clamped),
+		usedFraction,
+		remainingFraction: Math.max(0, 1 - usedFraction),
+		unit: "percent",
+	};
+}
+
+function buildUsageStatus(usedFraction?: number, limitReached?: boolean): UsageLimit["status"] {
+	if (limitReached) return "exhausted";
+	if (usedFraction === undefined) return "unknown";
+	if (usedFraction >= 1) return "exhausted";
+	if (usedFraction >= 0.9) return "warning";
+	return "ok";
+}
+
+function buildUsageLimit(args: {
+	key: "primary" | "secondary";
+	window: ParsedUsageWindow;
+	accountId?: string;
+	planType?: string;
+	limitReached?: boolean;
+	nowMs: number;
+}): UsageLimit {
+	const usageWindow = buildUsageWindow(args.window, args.key, args.nowMs);
+	const amount = buildUsageAmount(args.window);
+	return {
+		id: `openai-codex:${args.key}`,
+		label: usageWindow.label,
+		scope: {
+			provider: "openai-codex",
+			accountId: args.accountId,
+			tier: args.planType,
+			windowId: usageWindow.id,
+			shared: true,
+		},
+		window: usageWindow,
+		amount,
+		status: buildUsageStatus(amount.usedFraction, args.limitReached),
+	};
+}
+
+function resolveCacheExpiry(args: { report: UsageReport | null; nowMs: number }): number {
+	const { report, nowMs } = args;
+	if (!report) return nowMs + DEFAULT_CACHE_TTL_MS;
+	const exhausted = report.limits.some((limit) => limit.status === "exhausted");
+	const resetCandidates = report.limits
+		.map((limit) => limit.window?.resetsAt)
+		.filter((value): value is number => typeof value === "number" && Number.isFinite(value));
+	const earliestReset = resetCandidates.length > 0 ? Math.min(...resetCandidates) : undefined;
+	if (exhausted && earliestReset) return earliestReset;
+	if (earliestReset) return Math.min(nowMs + DEFAULT_CACHE_TTL_MS, earliestReset);
+	return nowMs + DEFAULT_CACHE_TTL_MS;
+}
+
+async function getCachedReport(
+	cache: UsageCache,
+	cacheKey: string,
+	nowMs: number,
+): Promise<UsageReport | null | undefined> {
+	const cached = await cache.get(cacheKey);
+	if (!cached) return undefined;
+	if (cached.expiresAt <= nowMs) return undefined;
+	return cached.value;
+}
+
+async function setCachedReport(
+	cache: UsageCache,
+	cacheKey: string,
+	report: UsageReport | null,
+	expiresAt: number,
+): Promise<void> {
+	await cache.set(cacheKey, { value: report, expiresAt });
+}
+
+export const openaiCodexUsageProvider: UsageProvider = {
+	id: "openai-codex",
+	supports(params: UsageFetchParams): boolean {
+		return params.provider === "openai-codex" && params.credential.type === "oauth";
+	},
+	async fetchUsage(params: UsageFetchParams, ctx: UsageFetchContext): Promise<UsageReport | null> {
+		if (params.provider !== "openai-codex") return null;
+		const { credential } = params;
+		if (credential.type !== "oauth") return null;
+
+		const accessToken = credential.accessToken;
+		if (!accessToken) return null;
+
+		const nowMs = ctx.now();
+		if (credential.expiresAt !== undefined && credential.expiresAt <= nowMs) {
+			ctx.logger?.warn("Codex usage token expired", { provider: params.provider });
+			return null;
+		}
+
+		const baseUrl = normalizeCodexBaseUrl(params.baseUrl);
+		const accountId = credential.accountId ?? extractAccountId(accessToken);
+		const cacheKey = `usage:openai-codex:${accountId ?? "unknown"}:${baseUrl}`;
+		const cached = await getCachedReport(ctx.cache, cacheKey, nowMs);
+		if (cached !== undefined) return cached;
+
+		const headers: Record<string, string> = {
+			Authorization: `Bearer ${accessToken}`,
+			"User-Agent": "OpenCode-Status-Plugin/1.0",
+		};
+		if (accountId) {
+			headers["ChatGPT-Account-Id"] = accountId;
+		}
+
+		const url = buildCodexUsageUrl(baseUrl);
+		let payload: unknown;
+		try {
+			const response = await ctx.fetch(url, { headers, signal: params.signal });
+			if (!response.ok) {
+				ctx.logger?.warn("Codex usage request failed", { status: response.status, provider: params.provider });
+				return null;
+			}
+			payload = await response.json();
+		} catch (error) {
+			ctx.logger?.warn("Codex usage request error", { provider: params.provider, error: String(error) });
+			return null;
+		}
+
+		const parsed = parseUsagePayload(payload);
+		if (!parsed) {
+			ctx.logger?.warn("Codex usage response invalid", { provider: params.provider });
+			return null;
+		}
+
+		const limits: UsageLimit[] = [];
+		if (parsed.primary) {
+			limits.push(
+				buildUsageLimit({
+					key: "primary",
+					window: parsed.primary,
+					accountId,
+					planType: parsed.planType,
+					limitReached: parsed.limitReached,
+					nowMs,
+				}),
+			);
+		}
+		if (parsed.secondary) {
+			limits.push(
+				buildUsageLimit({
+					key: "secondary",
+					window: parsed.secondary,
+					accountId,
+					planType: parsed.planType,
+					limitReached: parsed.limitReached,
+					nowMs,
+				}),
+			);
+		}
+
+		const report: UsageReport = {
+			provider: "openai-codex",
+			fetchedAt: nowMs,
+			limits,
+			metadata: {
+				planType: parsed.planType,
+				allowed: parsed.allowed,
+				limitReached: parsed.limitReached,
+				email: credential.email ?? extractEmail(accessToken),
+			},
+			raw: parsed.raw,
+		};
+
+		const expiresAt = resolveCacheExpiry({ report, nowMs });
+		await setCachedReport(ctx.cache, cacheKey, report, expiresAt);
+		return report;
+	},
+};

package/src/usage/zai.ts

@@ -0,0 +1,292 @@
+import type {
+	UsageAmount,
+	UsageFetchContext,
+	UsageFetchParams,
+	UsageLimit,
+	UsageProvider,
+	UsageReport,
+	UsageStatus,
+	UsageWindow,
+} from "$ai/usage";
+
+const DEFAULT_ENDPOINT = "https://api.z.ai";
+const QUOTA_PATH = "/api/monitor/usage/quota/limit";
+const MODEL_USAGE_PATH = "/api/monitor/usage/model-usage";
+const DEFAULT_CACHE_TTL_MS = 60_000;
+const SEVEN_DAYS_MS = 7 * 24 * 60 * 60 * 1000;
+
+function normalizeZaiBaseUrl(baseUrl?: string): string {
+	if (!baseUrl || !baseUrl.trim()) return DEFAULT_ENDPOINT;
+	try {
+		return new URL(baseUrl.trim()).origin;
+	} catch {
+		return DEFAULT_ENDPOINT;
+	}
+}
+
+interface ZaiUsageLimitItem {
+	type?: string;
+	usage?: number;
+	currentValue?: number;
+	percentage?: number;
+	remaining?: number;
+	nextResetTime?: number;
+}
+
+interface ZaiQuotaPayload {
+	success?: boolean;
+	code?: number;
+	msg?: string;
+	data?: {
+		limits?: ZaiUsageLimitItem[];
+	};
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function toNumber(value: unknown): number | undefined {
+	if (typeof value === "number" && Number.isFinite(value)) return value;
+	if (typeof value === "string" && value.trim()) {
+		const parsed = Number(value);
+		return Number.isFinite(parsed) ? parsed : undefined;
+	}
+	return undefined;
+}
+
+function parseMillis(value: unknown): number | undefined {
+	const parsed = toNumber(value);
+	if (parsed === undefined) return undefined;
+	return parsed > 1_000_000_000_000 ? parsed : parsed * 1000;
+}
+
+function parseLimitItem(value: unknown): ZaiUsageLimitItem | null {
+	if (!isRecord(value)) return null;
+	const type = typeof value.type === "string" ? value.type : undefined;
+	if (!type) return null;
+	return {
+		type,
+		usage: toNumber(value.usage),
+		currentValue: toNumber(value.currentValue),
+		percentage: toNumber(value.percentage),
+		remaining: toNumber(value.remaining),
+		nextResetTime: parseMillis(value.nextResetTime),
+	};
+}
+
+function buildUsageAmount(args: {
+	used: number | undefined;
+	limit: number | undefined;
+	remaining: number | undefined;
+	unit: UsageAmount["unit"];
+	percentage?: number;
+}): UsageAmount {
+	const usedFraction =
+		args.percentage !== undefined
+			? Math.min(Math.max(args.percentage / 100, 0), 1)
+			: args.used !== undefined && args.limit !== undefined && args.limit > 0
+				? Math.min(args.used / args.limit, 1)
+				: undefined;
+	const remainingFraction = usedFraction !== undefined ? Math.max(1 - usedFraction, 0) : undefined;
+	return {
+		used: args.used,
+		limit: args.limit,
+		remaining: args.remaining,
+		usedFraction,
+		remainingFraction,
+		unit: args.unit,
+	};
+}
+
+function buildUsageWindow(
+	id: string,
+	label: string,
+	resetsAt: number | undefined,
+	now: number,
+): UsageWindow | undefined {
+	if (!resetsAt) return { id, label };
+	const resetInMs = Math.max(0, resetsAt - now);
+	return {
+		id,
+		label,
+		resetsAt,
+		resetInMs,
+	};
+}
+
+function getUsageStatus(usedFraction: number | undefined): UsageStatus | undefined {
+	if (usedFraction === undefined) return undefined;
+	if (usedFraction >= 1) return "exhausted";
+	if (usedFraction >= 0.9) return "warning";
+	return "ok";
+}
+
+function buildCacheKey(params: UsageFetchParams): string {
+	const credential = params.credential;
+	const account = credential.accountId ?? credential.email ?? "unknown";
+	const token = credential.apiKey ?? credential.accessToken;
+	const fingerprint = token && typeof token === "string" ? Bun.hash(token).toString(16) : "anonymous";
+	const baseUrl = params.baseUrl ?? DEFAULT_ENDPOINT;
+	return `usage:${params.provider}:${account}:${fingerprint}:${baseUrl}`;
+}
+
+function resolveCacheExpiry(now: number, limits: UsageLimit[]): number {
+	const earliestReset = limits
+		.map((limit) => limit.window?.resetsAt)
+		.filter((value): value is number => typeof value === "number" && Number.isFinite(value))
+		.reduce((min, value) => (min === undefined ? value : Math.min(min, value)), undefined as number | undefined);
+	if (!earliestReset) return now + DEFAULT_CACHE_TTL_MS;
+	return Math.min(earliestReset, now + DEFAULT_CACHE_TTL_MS);
+}
+
+function formatDate(value: Date): string {
+	const pad = (input: number) => String(input).padStart(2, "0");
+	return `${value.getFullYear()}-${pad(value.getMonth() + 1)}-${pad(value.getDate())}+${pad(value.getHours())}:${pad(
+		value.getMinutes(),
+	)}:${pad(value.getSeconds())}`;
+}
+
+function buildModelUsageUrl(baseUrl: string, now: Date): string {
+	const start = new Date(now.getTime() - SEVEN_DAYS_MS);
+	const startTime = formatDate(start);
+	const endTime = formatDate(now);
+	return `${baseUrl}${MODEL_USAGE_PATH}?startTime=${encodeURIComponent(startTime)}&endTime=${encodeURIComponent(endTime)}`;
+}
+
+async function fetchZaiUsage(params: UsageFetchParams, ctx: UsageFetchContext): Promise<UsageReport | null> {
+	if (params.provider !== "zai") return null;
+	const credential = params.credential;
+	if (credential.type !== "api_key" || !credential.apiKey) return null;
+
+	const cacheKey = buildCacheKey(params);
+	const cachedEntry = await ctx.cache.get(cacheKey);
+	const now = ctx.now();
+	if (cachedEntry && cachedEntry.expiresAt > now) return cachedEntry.value;
+
+	const baseUrl = normalizeZaiBaseUrl(params.baseUrl);
+	const url = `${baseUrl}${QUOTA_PATH}`;
+	const headers: Record<string, string> = {
+		Authorization: credential.apiKey,
+		"Content-Type": "application/json",
+		"User-Agent": "OpenCode-Status-Plugin/1.0",
+	};
+
+	let payload: ZaiQuotaPayload | null = null;
+	try {
+		const response = await ctx.fetch(url, {
+			headers,
+			signal: params.signal,
+		});
+		if (!response.ok) {
+			ctx.logger?.warn("ZAI usage fetch failed", { status: response.status, statusText: response.statusText });
+			return null;
+		}
+		payload = (await response.json()) as ZaiQuotaPayload;
+	} catch (error) {
+		ctx.logger?.warn("ZAI usage fetch error", { error: String(error) });
+		return null;
+	}
+
+	if (!payload) return null;
+	if (payload.success !== true) {
+		ctx.logger?.warn("ZAI usage response invalid", { code: payload.code, message: payload.msg });
+		return null;
+	}
+
+	const limitsPayload = Array.isArray(payload.data?.limits) ? payload.data?.limits : [];
+	const limits: UsageLimit[] = [];
+
+	for (const rawLimit of limitsPayload) {
+		const parsed = parseLimitItem(rawLimit);
+		if (!parsed) continue;
+		if (parsed.type === "TOKENS_LIMIT") {
+			const amount = buildUsageAmount({
+				used: parsed.currentValue,
+				limit: parsed.usage,
+				remaining: parsed.remaining,
+				percentage: parsed.percentage,
+				unit: "tokens",
+			});
+			const window = buildUsageWindow("quota", "Quota", parsed.nextResetTime, now);
+			limits.push({
+				id: "zai:tokens",
+				label: "ZAI Token Quota",
+				scope: {
+					provider: params.provider,
+					windowId: window?.id ?? "quota",
+					shared: true,
+				},
+				window,
+				amount,
+				status: getUsageStatus(amount.usedFraction),
+			});
+		}
+		if (parsed.type === "TIME_LIMIT") {
+			const window = buildUsageWindow("quota", "Quota", undefined, now);
+			const amount = buildUsageAmount({
+				used: parsed.currentValue,
+				limit: parsed.usage,
+				remaining: parsed.remaining,
+				percentage: parsed.percentage,
+				unit: "requests",
+			});
+			limits.push({
+				id: "zai:requests",
+				label: "ZAI Request Quota",
+				scope: {
+					provider: params.provider,
+					windowId: "quota",
+					shared: true,
+				},
+				window,
+				amount,
+				status: getUsageStatus(amount.usedFraction),
+			});
+		}
+	}
+
+	if (limits.length === 0) return null;
+
+	const report: UsageReport = {
+		provider: params.provider,
+		fetchedAt: now,
+		limits,
+		metadata: {
+			endpoint: url,
+			accountId: credential.accountId,
+			email: credential.email,
+		},
+		raw: payload,
+	};
+
+	const expiresAt = resolveCacheExpiry(now, limits);
+	await ctx.cache.set(cacheKey, { value: report, expiresAt });
+
+	const modelUsageUrl = buildModelUsageUrl(baseUrl, new Date(now));
+	try {
+		const response = await ctx.fetch(modelUsageUrl, {
+			headers,
+			signal: params.signal,
+		});
+		if (response.ok) {
+			const modelUsagePayload = (await response.json()) as unknown;
+			if (isRecord(modelUsagePayload)) {
+				report.metadata = {
+					...report.metadata,
+					modelUsage: modelUsagePayload,
+				};
+			}
+		}
+	} catch (error) {
+		ctx.logger?.debug("ZAI model usage fetch failed", { error: String(error) });
+	}
+
+	return report;
+}
+
+export const zaiUsageProvider: UsageProvider = {
+	id: "zai",
+	fetchUsage: fetchZaiUsage,
+	supports: (params) => params.provider === "zai" && params.credential.type === "api_key",
+};