@aryee337/aery-ai 0.2.28 → 0.2.29

package/src/usage/gemini.ts

@@ -0,0 +1,250 @@
+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"],
+	},
+	{
+		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-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;
+	},
+};

package/src/usage/github-copilot.ts

@@ -0,0 +1,421 @@
+/**
+ * GitHub Copilot usage provider.
+ *
+ * Normalizes Copilot quota usage into the shared UsageReport schema.
+ */
+import type {
+	UsageAmount,
+	UsageFetchContext,
+	UsageFetchParams,
+	UsageLimit,
+	UsageProvider,
+	UsageReport,
+	UsageStatus,
+	UsageWindow,
+} from "../usage";
+import { isRecord, toBoolean, toNumber } from "../utils";
+import { OPENCODE_HEADERS } from "../utils/oauth/github-copilot";
+
+type CopilotQuotaDetail = {
+	entitlement: number;
+	overage_count: number;
+	overage_permitted: boolean;
+	percent_remaining: number;
+	quota_id: string;
+	quota_remaining: number;
+	remaining: number;
+	unlimited: boolean;
+};
+
+type CopilotQuotaSnapshots = {
+	chat?: CopilotQuotaDetail;
+	completions?: CopilotQuotaDetail;
+	premium_interactions?: CopilotQuotaDetail;
+};
+
+type CopilotUsageResponse = {
+	copilot_plan: string;
+	quota_reset_date: string;
+	quota_snapshots: CopilotQuotaSnapshots;
+};
+
+type BillingUsageItem = {
+	product: string;
+	sku: string;
+	model?: string;
+	unitType: string;
+	grossQuantity: number;
+	netQuantity: number;
+	limit?: number;
+};
+
+type BillingUsageResponse = {
+	timePeriod: { year: number; month?: number };
+	user: string;
+	usageItems: BillingUsageItem[];
+};
+
+function resolveGitHubApiBaseUrl(params: UsageFetchParams): string {
+	const baseUrl = params.baseUrl?.replace(/\/$/, "");
+	if (baseUrl && !baseUrl.includes("githubcopilot.com")) return baseUrl;
+	const enterpriseUrl = params.credential.enterpriseUrl?.trim();
+	if (!enterpriseUrl) return "https://api.github.com";
+	if (enterpriseUrl.startsWith("http://") || enterpriseUrl.startsWith("https://")) {
+		return enterpriseUrl.replace(/\/$/, "");
+	}
+	if (enterpriseUrl.startsWith("api.")) {
+		return `https://${enterpriseUrl}`;
+	}
+	return `https://api.${enterpriseUrl}`;
+}
+
+function buildWindow(resetDate: string | undefined): UsageWindow | undefined {
+	if (!resetDate) return undefined;
+	const resetAt = Date.parse(resetDate);
+	if (!Number.isFinite(resetAt)) return undefined;
+	return {
+		id: "monthly",
+		label: "Monthly",
+		resetsAt: resetAt,
+	};
+}
+
+function buildAmount(used: number | undefined, limit: number | undefined, unit: UsageAmount["unit"]): UsageAmount {
+	const safeLimit = limit !== undefined && Number.isFinite(limit) ? limit : undefined;
+	const safeUsed = used !== undefined && Number.isFinite(used) ? used : undefined;
+	const remaining = safeLimit !== undefined && safeUsed !== undefined ? Math.max(0, safeLimit - safeUsed) : undefined;
+	const usedFraction =
+		safeLimit !== undefined && safeUsed !== undefined && safeLimit > 0 ? safeUsed / safeLimit : undefined;
+	const remainingFraction =
+		safeLimit !== undefined && remaining !== undefined && safeLimit > 0 ? remaining / safeLimit : undefined;
+	return {
+		used: safeUsed,
+		limit: safeLimit,
+		remaining,
+		usedFraction,
+		remainingFraction,
+		unit,
+	};
+}
+
+function deriveStatus(amount: UsageAmount, unlimited: boolean): UsageStatus {
+	if (unlimited) return "ok";
+	if (amount.remainingFraction === undefined) return "unknown";
+	if (amount.remainingFraction <= 0) return "exhausted";
+	if (amount.remainingFraction <= 0.1) return "warning";
+	return "ok";
+}
+
+function parseQuotaDetail(value: unknown): CopilotQuotaDetail | null {
+	if (!isRecord(value)) return null;
+	const entitlement = toNumber(value.entitlement);
+	const remaining = toNumber(value.remaining);
+	const percentRemaining = toNumber(value.percent_remaining);
+	const unlimited = toBoolean(value.unlimited);
+	if (
+		entitlement === undefined ||
+		remaining === undefined ||
+		percentRemaining === undefined ||
+		unlimited === undefined
+	) {
+		return null;
+	}
+	const overageCount = toNumber(value.overage_count) ?? 0;
+	const overagePermitted = toBoolean(value.overage_permitted) ?? false;
+	const quotaId = typeof value.quota_id === "string" ? value.quota_id : "";
+	const quotaRemaining = toNumber(value.quota_remaining) ?? remaining;
+	return {
+		entitlement,
+		overage_count: overageCount,
+		overage_permitted: overagePermitted,
+		percent_remaining: percentRemaining,
+		quota_id: quotaId,
+		quota_remaining: quotaRemaining,
+		remaining,
+		unlimited,
+	};
+}
+
+async function fetchJson(ctx: UsageFetchContext, url: string, init: RequestInit): Promise<unknown> {
+	const response = await ctx.fetch(url, init);
+	if (!response.ok) {
+		const text = await response.text();
+		throw new Error(`${response.status} ${response.statusText}: ${text}`);
+	}
+	return response.json();
+}
+
+async function resolveGitHubUsername(
+	ctx: UsageFetchContext,
+	baseUrl: string,
+	token: string,
+	signal?: AbortSignal,
+): Promise<string | undefined> {
+	try {
+		const data = await fetchJson(ctx, `${baseUrl}/user`, {
+			headers: {
+				Accept: "application/vnd.github+json",
+				Authorization: `Bearer ${token}`,
+				"X-GitHub-Api-Version": "2022-11-28",
+			},
+			signal,
+		});
+		if (!isRecord(data)) return undefined;
+		return typeof data.login === "string" ? data.login : undefined;
+	} catch {
+		return undefined;
+	}
+}
+
+async function fetchInternalUsage(
+	ctx: UsageFetchContext,
+	githubApiBaseUrl: string,
+	token: string,
+	signal?: AbortSignal,
+): Promise<CopilotUsageResponse> {
+	const headers: Record<string, string> = {
+		"Content-Type": "application/json",
+		Accept: "application/json",
+		Authorization: `Bearer ${token}`,
+		...OPENCODE_HEADERS,
+	};
+	const data = await fetchJson(ctx, `${githubApiBaseUrl}/copilot_internal/user`, { headers, signal });
+	if (!isRecord(data)) throw new Error("Invalid Copilot usage response");
+	return data as CopilotUsageResponse;
+}
+
+async function fetchBillingUsage(
+	ctx: UsageFetchContext,
+	baseUrl: string,
+	username: string,
+	token: string,
+	signal?: AbortSignal,
+): Promise<BillingUsageResponse> {
+	const data = await fetchJson(
+		ctx,
+		`${baseUrl}/users/${encodeURIComponent(username)}/settings/billing/premium_request/usage`,
+		{
+			headers: {
+				Accept: "application/vnd.github+json",
+				Authorization: `Bearer ${token}`,
+				"X-GitHub-Api-Version": "2022-11-28",
+			},
+			signal,
+		},
+	);
+
+	if (!isRecord(data)) throw new Error("Invalid Copilot billing usage response");
+	return data as BillingUsageResponse;
+}
+
+function buildLimitFromQuota(
+	key: string,
+	label: string,
+	quota: CopilotQuotaDetail,
+	plan: string,
+	window: UsageWindow | undefined,
+	accountId?: string,
+): UsageLimit {
+	const used = quota.unlimited ? undefined : Math.max(0, quota.entitlement - quota.remaining);
+	const limit = quota.unlimited ? undefined : quota.entitlement;
+	const amount = buildAmount(used, limit, "requests");
+	const notes: string[] = [];
+	if (quota.unlimited) notes.push("Unlimited");
+	if (quota.overage_count > 0) {
+		notes.push(`Overage requests: ${quota.overage_count}`);
+	}
+	return {
+		id: `copilot:${key}`,
+		label,
+		scope: {
+			provider: "github-copilot",
+			accountId,
+			tier: plan,
+			windowId: window?.id,
+		},
+		window,
+		amount,
+		status: deriveStatus(amount, quota.unlimited),
+		notes: notes.length > 0 ? notes : undefined,
+	};
+}
+
+function normalizeQuotaSnapshots(
+	data: CopilotUsageResponse,
+	accountId?: string,
+): { limits: UsageLimit[]; window?: UsageWindow } {
+	const window = buildWindow(data.quota_reset_date);
+	const snapshots = data.quota_snapshots ?? {};
+	const limits: UsageLimit[] = [];
+	const premium = parseQuotaDetail(snapshots.premium_interactions);
+	if (premium) {
+		limits.push(buildLimitFromQuota("premium", "Premium Requests", premium, data.copilot_plan, window, accountId));
+	}
+	const chat = parseQuotaDetail(snapshots.chat);
+	if (chat && !chat.unlimited) {
+		limits.push(buildLimitFromQuota("chat", "Chat Requests", chat, data.copilot_plan, window, accountId));
+	}
+	const completions = parseQuotaDetail(snapshots.completions);
+	if (completions && !completions.unlimited) {
+		limits.push(buildLimitFromQuota("completions", "Completions", completions, data.copilot_plan, window, accountId));
+	}
+	return { limits, window };
+}
+
+function normalizeBillingUsage(data: BillingUsageResponse): UsageLimit[] {
+	const limits: UsageLimit[] = [];
+	const periodLabel = data.timePeriod.month
+		? `${data.timePeriod.year}-${String(data.timePeriod.month).padStart(2, "0")}`
+		: `${data.timePeriod.year}`;
+	const window: UsageWindow = {
+		id: "billing-period",
+		label: periodLabel,
+	};
+
+	const premiumItems = data.usageItems.filter(
+		item => item.sku === "Copilot Premium Request" || item.sku.includes("Premium"),
+	);
+	const totalUsed = premiumItems.reduce((sum, item) => sum + item.grossQuantity, 0);
+	const totalLimit = premiumItems.reduce((sum, item) => sum + (item.limit ?? 0), 0) || undefined;
+	const totalAmount = buildAmount(totalUsed, totalLimit, "requests");
+	limits.push({
+		id: "copilot:premium",
+		label: "Premium Requests",
+		scope: {
+			provider: "github-copilot",
+			accountId: data.user,
+			windowId: window.id,
+		},
+		window,
+		amount: totalAmount,
+		status: deriveStatus(totalAmount, false),
+	});
+
+	for (const item of data.usageItems) {
+		if (!item.model) continue;
+		if (item.grossQuantity <= 0) continue;
+		const amount = buildAmount(item.grossQuantity, item.limit, "requests");
+		limits.push({
+			id: `copilot:model:${item.model}`,
+			label: `Model ${item.model}`,
+			scope: {
+				provider: "github-copilot",
+				accountId: data.user,
+				modelId: item.model,
+				windowId: window.id,
+			},
+			window,
+			amount,
+			status: deriveStatus(amount, false),
+		});
+	}
+
+	return limits;
+}
+
+export const githubCopilotUsageProvider: UsageProvider = {
+	id: "github-copilot",
+	supports: ({ provider, credential }) => {
+		if (provider !== "github-copilot") return false;
+		if (credential.type === "oauth") {
+			return Boolean(credential.refreshToken || credential.accessToken);
+		}
+		return Boolean(credential.apiKey);
+	},
+	fetchUsage: async (params, ctx) => {
+		if (!githubCopilotUsageProvider.supports?.(params)) return null;
+
+		const githubApiBaseUrl = resolveGitHubApiBaseUrl(params);
+		let report: UsageReport | null = null;
+
+		if (params.credential.type === "api_key") {
+			let username: string | undefined;
+			const candidate =
+				params.credential.accountId || params.credential.metadata?.username || params.credential.metadata?.user;
+			if (typeof candidate === "string" && candidate.trim()) {
+				username = candidate.trim();
+			}
+			if (!username && params.credential.apiKey) {
+				username = await resolveGitHubUsername(ctx, githubApiBaseUrl, params.credential.apiKey, params.signal);
+			}
+			if (!username) {
+				ctx.logger?.warn("Copilot usage requires username for billing API", { provider: params.provider });
+			} else if (params.credential.apiKey) {
+				try {
+					const billing = await fetchBillingUsage(
+						ctx,
+						githubApiBaseUrl,
+						username,
+						params.credential.apiKey,
+						params.signal,
+					);
+					report = {
+						provider: "github-copilot",
+						fetchedAt: Date.now(),
+						limits: normalizeBillingUsage(billing),
+						metadata: {
+							accountId: billing.user,
+							account: billing.user,
+							period: billing.timePeriod,
+						},
+					};
+				} catch (error) {
+					ctx.logger?.warn("Copilot usage fetch failed", { error: String(error) });
+				}
+			}
+			if (!report && params.credential.apiKey) {
+				try {
+					const usage = await fetchInternalUsage(ctx, githubApiBaseUrl, params.credential.apiKey, params.signal);
+					const normalized = normalizeQuotaSnapshots(usage, username);
+					report = {
+						provider: "github-copilot",
+						fetchedAt: Date.now(),
+						limits: normalized.limits,
+						metadata: {
+							accountId: username,
+							plan: usage.copilot_plan,
+							quotaResetDate: usage.quota_reset_date,
+						},
+						raw: usage,
+					};
+				} catch (error) {
+					ctx.logger?.warn("Copilot usage fetch failed", { error: String(error) });
+				}
+			}
+		} else {
+			const { refreshToken, accessToken } = params.credential;
+			if (!refreshToken && !accessToken) return null;
+			const oauthToken = refreshToken || accessToken;
+			if (!oauthToken) return null;
+			const githubToken = refreshToken ?? accessToken;
+			if (!githubToken) return null;
+			try {
+				const usage = await fetchInternalUsage(ctx, githubApiBaseUrl, githubToken, params.signal);
+				let accountId = params.credential.accountId;
+				if (!accountId && refreshToken) {
+					accountId = await resolveGitHubUsername(ctx, githubApiBaseUrl, refreshToken, params.signal);
+				}
+				if (!accountId && accessToken) {
+					accountId = await resolveGitHubUsername(ctx, githubApiBaseUrl, accessToken, params.signal);
+				}
+				const normalized = normalizeQuotaSnapshots(usage, accountId);
+				report = {
+					provider: "github-copilot",
+					fetchedAt: Date.now(),
+					limits: normalized.limits,
+					metadata: {
+						accountId,
+						email: params.credential.email,
+						plan: usage.copilot_plan,
+						quotaResetDate: usage.quota_reset_date,
+					},
+					raw: usage,
+				};
+			} catch (error) {
+				ctx.logger?.warn("Copilot usage fetch failed", { error: String(error) });
+			}
+		}
+
+		return report;
+	},
+};