@linxiraos/pi-ai 1.0.0

package/src/usage/cursor.ts

@@ -0,0 +1,335 @@
+import { extractCursorAccessTokenUserId } from "../registry/oauth/cursor";
+import type {
+	UsageAmount,
+	UsageFetchContext,
+	UsageFetchParams,
+	UsageLimit,
+	UsageProvider,
+	UsageReport,
+	UsageStatus,
+	UsageWindow,
+} from "../usage";
+import { toNumber } from "./shared";
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function parseTimestamp(value: unknown): number | undefined {
+	const numeric = toNumber(value);
+	if (numeric !== undefined) return numeric < 1_000_000_000_000 ? numeric * 1000 : numeric;
+	if (typeof value !== "string" || !value.trim()) return undefined;
+	const parsed = Date.parse(value);
+	return Number.isFinite(parsed) ? parsed : undefined;
+}
+
+const DEFAULT_CURSOR_BASE_URL = "https://api2.cursor.sh";
+
+function normalizeCursorBaseUrl(baseUrl?: string): string {
+	if (!baseUrl) return DEFAULT_CURSOR_BASE_URL;
+	return baseUrl.replace(/\/+$/, "");
+}
+
+type CursorUsageSource = "auth-usage" | "usage-summary" | "auth-me";
+
+async function fetchCursorJson(
+	ctx: UsageFetchContext,
+	url: string,
+	init: RequestInit,
+	source: CursorUsageSource,
+): Promise<unknown | undefined> {
+	try {
+		const response = await ctx.fetch(url, init);
+		if (!response.ok) {
+			ctx.logger?.warn("Cursor usage request failed", {
+				status: response.status,
+				provider: "cursor",
+				source,
+			});
+			return undefined;
+		}
+		return await response.json();
+	} catch (error) {
+		ctx.logger?.warn("Cursor usage request error", {
+			provider: "cursor",
+			source,
+			error: String(error),
+		});
+		return undefined;
+	}
+}
+
+function deriveResetsAt(payload: Record<string, unknown>): number | undefined {
+	const endKeys = ["billingCycleEnd", "endOfMonth", "resetsAt", "nextReset"];
+	for (const key of endKeys) {
+		const parsed = parseTimestamp(payload[key]);
+		if (parsed !== undefined) return parsed;
+	}
+
+	const startKeys = ["startOfMonth", "billingCycleStart", "startOfBillingCycle"];
+	for (const key of startKeys) {
+		const parsed = parseTimestamp(payload[key]);
+		if (parsed !== undefined) {
+			const date = new Date(parsed);
+			date.setUTCMonth(date.getUTCMonth() + 1);
+			return date.getTime();
+		}
+	}
+	return undefined;
+}
+
+function resolveCursorStatus(usedFraction: number | undefined): UsageStatus {
+	if (usedFraction === undefined) return "unknown";
+	if (usedFraction >= 1) return "exhausted";
+	if (usedFraction >= 0.9) return "warning";
+	return "ok";
+}
+
+export function parseCursorIndividualUsage(payload: unknown, fetchedAt = Date.now()): UsageReport | null {
+	if (!isRecord(payload) || !isRecord(payload.individualUsage) || !isRecord(payload.individualUsage.overall)) {
+		return null;
+	}
+	const individual = payload.individualUsage.overall;
+	if (individual.enabled === false) return null;
+
+	const reportedUsed = toNumber(individual.used);
+	const reportedRemaining = toNumber(individual.remaining);
+	const hasValidUsed = reportedUsed !== undefined && reportedUsed >= 0;
+	const hasValidRemaining = reportedRemaining !== undefined && reportedRemaining >= 0;
+	const limit = toNumber(individual.limit);
+
+	let amount: UsageAmount;
+	if (individual.limit === null || individual.limit === undefined) {
+		if (!hasValidUsed) return null;
+		amount = { used: reportedUsed / 100, unit: "usd" };
+	} else {
+		if (limit === undefined || limit <= 0) return null;
+		let used: number;
+		if (reportedUsed !== undefined && reportedUsed > 0) {
+			used = reportedUsed;
+		} else if (hasValidRemaining && reportedRemaining < limit) {
+			used = Math.max(0, limit - reportedRemaining);
+		} else if (hasValidUsed) {
+			used = reportedUsed;
+		} else {
+			return null;
+		}
+		const remaining = Math.max(0, limit - used);
+		amount = {
+			used: used / 100,
+			limit: limit / 100,
+			remaining: remaining / 100,
+			usedFraction: used / limit,
+			remainingFraction: remaining / limit,
+			unit: "usd",
+		};
+	}
+
+	const resetsAt = deriveResetsAt(payload);
+	const window: UsageWindow = {
+		id: "monthly",
+		label: "Monthly",
+		...(resetsAt !== undefined ? { resetsAt } : {}),
+	};
+	const limitEntry: UsageLimit = {
+		id: "cursor:usd:individual-overall",
+		label: "Personal Usage",
+		scope: {
+			provider: "cursor",
+			windowId: window.id,
+		},
+		window,
+		amount,
+		...(amount.usedFraction !== undefined ? { status: resolveCursorStatus(amount.usedFraction) } : {}),
+	};
+	return {
+		provider: "cursor",
+		fetchedAt,
+		limits: [limitEntry],
+		raw: payload,
+	};
+}
+
+export function parseCursorUsage(payload: unknown, fetchedAt = Date.now()): UsageReport | null {
+	if (!isRecord(payload)) return null;
+	const limits: UsageLimit[] = [];
+	const resetsAt = deriveResetsAt(payload);
+
+	const window: UsageWindow = {
+		id: "monthly",
+		label: "Monthly",
+		...(resetsAt !== undefined ? { resetsAt } : {}),
+	};
+
+	for (const [key, value] of Object.entries(payload)) {
+		if (!isRecord(value)) continue;
+
+		// used can be: numRequests, used, amountUsed, usdUsed
+		const usedVal =
+			toNumber(value.numRequests) ?? toNumber(value.used) ?? toNumber(value.amountUsed) ?? toNumber(value.usdUsed);
+
+		// limit can be: maxRequestUsage, limit, amountLimit, usdLimit
+		const limitVal =
+			toNumber(value.maxRequestUsage) ??
+			toNumber(value.limit) ??
+			toNumber(value.amountLimit) ??
+			toNumber(value.usdLimit);
+
+		if (usedVal !== undefined && limitVal !== undefined) {
+			const isUsd =
+				key === "planUsage" ||
+				key.toLowerCase().includes("usd") ||
+				key.toLowerCase().includes("billing") ||
+				key.toLowerCase().includes("stripe");
+
+			const unit = isUsd ? "usd" : "requests";
+			const cleanBucket = key.toLowerCase().trim();
+			const limitId = isUsd ? `cursor:usd:${cleanBucket}` : `cursor:requests:${cleanBucket}`;
+
+			const label = isUsd ? `${key} spend` : `${key} requests`;
+
+			const amount: UsageAmount = {
+				used: usedVal,
+				limit: limitVal,
+				remaining: Math.max(0, limitVal - usedVal),
+				usedFraction: limitVal > 0 ? usedVal / limitVal : 0,
+				remainingFraction: limitVal > 0 ? Math.max(0, limitVal - usedVal) / limitVal : 0,
+				unit,
+			};
+
+			const status = resolveCursorStatus(amount.usedFraction);
+
+			limits.push({
+				id: limitId,
+				label,
+				scope: {
+					provider: "cursor",
+					...(window ? { windowId: window.id } : {}),
+				},
+				...(window ? { window } : {}),
+				amount,
+				status,
+			});
+		}
+	}
+
+	if (limits.length === 0) {
+		return null;
+	}
+
+	return {
+		provider: "cursor",
+		fetchedAt,
+		limits,
+		raw: payload,
+	};
+}
+
+export const cursorUsageProvider: UsageProvider = {
+	id: "cursor",
+	supports(params: UsageFetchParams): boolean {
+		if (params.provider !== "cursor") return false;
+		const { credential } = params;
+		if (credential.type === "oauth") {
+			return Boolean(credential.accessToken);
+		}
+		if (credential.type === "api_key") {
+			return Boolean(credential.apiKey);
+		}
+		return false;
+	},
+	async fetchUsage(params: UsageFetchParams, ctx: UsageFetchContext): Promise<UsageReport | null> {
+		if (params.provider !== "cursor") return null;
+		const { credential } = params;
+		const token = credential.type === "oauth" ? credential.accessToken : credential.apiKey;
+		if (!token) return null;
+
+		const baseUrl = normalizeCursorBaseUrl(params.baseUrl ?? credential.apiEndpoint);
+		const url = `${baseUrl}/auth/usage`;
+		const headers: Record<string, string> = {
+			Accept: "application/json",
+			Authorization: `Bearer ${token}`,
+		};
+		const fetchedAt = Date.now();
+
+		const legacyReportPromise = fetchCursorJson(
+			ctx,
+			url,
+			{
+				headers,
+				signal: params.signal,
+			},
+			"auth-usage",
+		).then(payload => parseCursorUsage(payload, fetchedAt));
+
+		let summaryReportPromise = Promise.resolve<UsageReport | null>(null);
+		let profileEmailPromise = Promise.resolve<string | undefined>(undefined);
+		if (credential.type === "oauth" && baseUrl === DEFAULT_CURSOR_BASE_URL) {
+			const userId = extractCursorAccessTokenUserId(token);
+			if (userId) {
+				const sessionHeaders: Record<string, string> = {
+					Accept: "application/json",
+					Cookie: `WorkosCursorSessionToken=${encodeURIComponent(`${userId}::${token}`)}`,
+				};
+				summaryReportPromise = fetchCursorJson(
+					ctx,
+					"https://cursor.com/api/usage-summary",
+					{
+						headers: sessionHeaders,
+						signal: params.signal,
+					},
+					"usage-summary",
+				).then(payload => parseCursorIndividualUsage(payload, fetchedAt));
+				profileEmailPromise = fetchCursorJson(
+					ctx,
+					"https://cursor.com/api/auth/me",
+					{
+						headers: sessionHeaders,
+						signal: params.signal,
+					},
+					"auth-me",
+				).then(payload => {
+					if (
+						!isRecord(payload) ||
+						payload.sub !== userId ||
+						typeof payload.email !== "string" ||
+						!payload.email.trim()
+					) {
+						return undefined;
+					}
+					return payload.email.trim();
+				});
+			}
+		}
+
+		const [legacyReport, summaryReport, profileEmail] = await Promise.all([
+			legacyReportPromise,
+			summaryReportPromise,
+			profileEmailPromise,
+		]);
+		let report: UsageReport | null;
+		if (legacyReport && summaryReport) {
+			report = {
+				provider: "cursor",
+				fetchedAt,
+				limits: [...legacyReport.limits, ...summaryReport.limits],
+				raw: {
+					authUsage: legacyReport.raw,
+					usageSummary: summaryReport.raw,
+				},
+			};
+		} else {
+			report = legacyReport ?? summaryReport;
+		}
+		if (!report) return null;
+
+		const email = profileEmail ?? credential.email?.trim();
+		const metadata = {
+			...(email ? { email } : {}),
+			...(credential.accountId ? { accountId: credential.accountId } : {}),
+			...(credential.projectId ? { projectId: credential.projectId } : {}),
+		};
+		if (Object.keys(metadata).length > 0) report.metadata = metadata;
+		return report;
+	},
+};

package/src/usage/gemini.ts

@@ -0,0 +1,258 @@
+import { getGeminiCliHeaders } from "../providers/google-gemini-cli";
+import type {
+	UsageAmount,
+	UsageFetchContext,
+	UsageFetchParams,
+	UsageLimit,
+	UsageProvider,
+	UsageReport,
+	UsageWindow,
+} from "../usage";
+
+// (Refresh is the sole responsibility of AuthStorage; no provider-direct refresh here.)
+
+const DEFAULT_ENDPOINT = "https://cloudcode-pa.googleapis.com";
+
+const GEMINI_TIER_MAP: Array<{ tier: string; models: string[] }> = [
+	{
+		tier: "3-Flash",
+		models: ["gemini-3-flash-preview", "gemini-3-flash", "gemini-3.5-flash"],
+	},
+	{
+		tier: "Flash",
+		models: ["gemini-2.5-flash", "gemini-2.5-flash-lite", "gemini-2.0-flash", "gemini-1.5-flash"],
+	},
+	{
+		tier: "Pro",
+		models: [
+			"gemini-2.5-pro",
+			"gemini-3-pro-preview",
+			"gemini-3.1-pro-preview",
+			"gemini-3-pro",
+			"gemini-3.1-pro",
+			"gemini-pro-agent",
+			"gemini-1.5-pro",
+		],
+	},
+];
+
+interface LoadCodeAssistResponse {
+	cloudaicompanionProject?: string | { id?: string };
+	currentTier?: { id?: string; name?: string };
+}
+
+interface RetrieveUserQuotaResponse {
+	buckets?: Array<{
+		modelId?: string;
+		remainingFraction?: number;
+		resetTime?: string;
+	}>;
+}
+
+function getProjectId(payload: LoadCodeAssistResponse | undefined): string | undefined {
+	if (!payload) return undefined;
+	if (typeof payload.cloudaicompanionProject === "string") {
+		return payload.cloudaicompanionProject;
+	}
+	if (payload.cloudaicompanionProject && typeof payload.cloudaicompanionProject === "object") {
+		return payload.cloudaicompanionProject.id;
+	}
+	return undefined;
+}
+
+function getModelTier(modelId: string): string | undefined {
+	for (const entry of GEMINI_TIER_MAP) {
+		if (entry.models.includes(modelId)) {
+			return entry.tier;
+		}
+	}
+	const normalized = modelId.toLowerCase();
+	if (normalized.includes("flash")) return "Flash";
+	if (normalized.includes("pro")) return "Pro";
+	return undefined;
+}
+
+function parseWindow(resetTime: string | undefined): UsageWindow {
+	if (!resetTime) {
+		return {
+			id: "quota",
+			label: "Quota window",
+		};
+	}
+	const resetsAt = Date.parse(resetTime);
+	if (Number.isNaN(resetsAt)) {
+		return {
+			id: "quota",
+			label: "Quota window",
+		};
+	}
+	return {
+		id: `reset-${resetsAt}`,
+		label: "Quota window",
+		resetsAt,
+	};
+}
+
+function buildAmount(remainingFraction: number | undefined): UsageAmount {
+	if (remainingFraction === undefined || !Number.isFinite(remainingFraction)) {
+		return { unit: "percent" };
+	}
+	const remaining = Math.min(Math.max(remainingFraction, 0), 1);
+	const used = Math.min(Math.max(1 - remaining, 0), 1);
+	return {
+		unit: "percent",
+		used: Math.round(used * 1000) / 10,
+		remaining: Math.round(remaining * 1000) / 10,
+		limit: 100,
+		usedFraction: used,
+		remainingFraction: remaining,
+	};
+}
+
+/**
+ * Return the OAuth access token to use against `/v1internal:*`. AuthStorage is
+ * the sole refresh authority (broker-aware, single-flighted, rotation-safe);
+ * if the token landed here expired or near-expired, the next usage cycle will
+ * carry a freshly-refreshed credential. Returning `undefined` short-circuits
+ * the probe rather than POSTing a stale token to Google.
+ */
+function resolveAccessToken(params: UsageFetchParams): string | undefined {
+	const { credential } = params;
+	if (credential.type !== "oauth") return undefined;
+	if (!credential.accessToken) return undefined;
+	if (credential.expiresAt !== undefined && credential.expiresAt <= Date.now()) {
+		return undefined;
+	}
+	return credential.accessToken;
+}
+
+async function loadCodeAssist(
+	params: UsageFetchParams,
+	ctx: UsageFetchContext,
+	accessToken: string,
+	baseUrl: string,
+	projectId?: string,
+): Promise<LoadCodeAssistResponse | undefined> {
+	const response = await ctx.fetch(`${baseUrl}/v1internal:loadCodeAssist`, {
+		method: "POST",
+		headers: {
+			Authorization: `Bearer ${accessToken}`,
+			"Content-Type": "application/json",
+			...getGeminiCliHeaders(),
+		},
+		body: JSON.stringify({
+			...(projectId ? { cloudaicompanionProject: projectId } : {}),
+			metadata: {
+				ideType: "IDE_UNSPECIFIED",
+				platform: "PLATFORM_UNSPECIFIED",
+				pluginType: "GEMINI",
+			},
+		}),
+		signal: params.signal,
+	});
+
+	if (!response.ok) {
+		const errorText = await response.text();
+		ctx.logger?.warn("Gemini CLI loadCodeAssist failed", {
+			status: response.status,
+			error: errorText,
+		});
+		return undefined;
+	}
+
+	return (await response.json()) as LoadCodeAssistResponse;
+}
+
+async function fetchQuota(
+	params: UsageFetchParams,
+	ctx: UsageFetchContext,
+	accessToken: string,
+	baseUrl: string,
+	projectId?: string,
+): Promise<RetrieveUserQuotaResponse | undefined> {
+	const response = await ctx.fetch(`${baseUrl}/v1internal:retrieveUserQuota`, {
+		method: "POST",
+		headers: {
+			Authorization: `Bearer ${accessToken}`,
+			"Content-Type": "application/json",
+			...getGeminiCliHeaders(),
+		},
+		body: JSON.stringify(projectId ? { project: projectId } : {}),
+		signal: params.signal,
+	});
+
+	if (!response.ok) {
+		const errorText = await response.text();
+		ctx.logger?.warn("Gemini CLI retrieveUserQuota failed", {
+			status: response.status,
+			error: errorText,
+		});
+		return undefined;
+	}
+
+	return (await response.json()) as RetrieveUserQuotaResponse;
+}
+
+export const googleGeminiCliUsageProvider: UsageProvider = {
+	id: "google-gemini-cli",
+	supports: ({ credential }) => credential.type === "oauth" && !!credential.accessToken,
+	async fetchUsage(params, ctx) {
+		const { credential } = params;
+		if (credential.type !== "oauth") {
+			return null;
+		}
+		const accessToken = resolveAccessToken(params);
+		if (!accessToken) {
+			return null;
+		}
+
+		const baseUrl = (params.baseUrl?.trim() || DEFAULT_ENDPOINT).replace(/\/$/, "");
+
+		const loadResponse = await loadCodeAssist(params, ctx, accessToken, baseUrl, credential.projectId);
+		const projectId = credential.projectId ?? getProjectId(loadResponse);
+		const quotaResponse = await fetchQuota(params, ctx, accessToken, baseUrl, projectId);
+		if (!quotaResponse) {
+			return null;
+		}
+
+		const limits: UsageLimit[] = [];
+		const buckets = quotaResponse.buckets ?? [];
+
+		buckets.forEach((bucket, index) => {
+			const modelId = bucket.modelId;
+			const window = parseWindow(bucket.resetTime);
+			const amount = buildAmount(bucket.remainingFraction);
+			const tier = modelId ? getModelTier(modelId) : undefined;
+			const label = modelId ? `Gemini ${modelId}` : "Gemini quota";
+			const id = `${modelId ?? "unknown"}:${window?.id ?? index}`;
+
+			limits.push({
+				id,
+				label,
+				scope: {
+					provider: params.provider,
+					accountId: credential.accountId,
+					projectId,
+					modelId,
+					tier,
+					windowId: window?.id,
+				},
+				window,
+				amount,
+			});
+		});
+
+		const report: UsageReport = {
+			provider: params.provider,
+			fetchedAt: Date.now(),
+			limits,
+			metadata: {
+				currentTierId: loadResponse?.currentTier?.id,
+				currentTierName: loadResponse?.currentTier?.name,
+			},
+			raw: quotaResponse,
+		};
+
+		return report;
+	},
+};