Merge pull request #2260 from insodimension/fix/usage-active-oauth-account-upstream

Show active OAuth account in usage

packages/coding-agent/src/modes/controllers/command-controller.ts

@@ -36,9 +36,10 @@ import { computeContextBreakdown, renderContextUsage } from "../../modes/utils/c
 import { buildHotkeysMarkdown } from "../../modes/utils/hotkeys-markdown";
 import { buildToolsMarkdown } from "../../modes/utils/tools-markdown";
 import type { AsyncJobSnapshotItem } from "../../session/agent-session";
-import type { AuthStorage } from "../../session/auth-storage";
+import type { AuthStorage, OAuthAccountIdentity } from "../../session/auth-storage";
 import type { NewSessionOptions } from "../../session/session-manager";
 import { formatShakeSummary, type ShakeMode, type ShakeResult } from "../../session/shake-types";
+import { limitMatchesActiveAccount } from "../../slash-commands/helpers/active-oauth-account";
 import { outputMeta } from "../../tools/output-meta";
 import { resolveToCwd, stripOuterDoubleQuotes } from "../../tools/path-utils";
 import { replaceTabs } from "../../tools/render-utils";
@@ -404,7 +405,16 @@ export class CommandController {
 		}
 
 		const availableWidth = Math.max(40, (this.ctx.ui.terminal.columns ?? 100) - 2);
-		const output = renderUsageReports(usageReports, theme, Date.now(), availableWidth);
+		const currentProvider = this.ctx.session.model?.provider;
+		const activeAccount = currentProvider
+			? this.ctx.session.modelRegistry.authStorage.getOAuthAccountIdentity(
+					currentProvider,
+					this.ctx.session.sessionId,
+				)
+			: undefined;
+		const output = renderUsageReports(usageReports, theme, Date.now(), availableWidth, provider =>
+			provider === currentProvider ? activeAccount : undefined,
+		);
 		this.ctx.present([new Spacer(1), new Text(output, 1, 0)]);
 	}
 
@@ -1311,12 +1321,17 @@ function formatAccountHeaderRow(
 	nowMs: number,
 	columnWidth: number,
 	uiTheme: typeof theme,
+	activeAccount?: OAuthAccountIdentity,
 ): string[] {
 	const parts = limits.map((limit, index) => {
 		const reset = formatResetShort(limit, nowMs);
+		const report = reports[index];
+		const active = report !== undefined && limitMatchesActiveAccount(report, limit, activeAccount);
+		const label = formatAccountLabel(limit, report, index);
 		return {
-			label: formatAccountLabel(limit, reports[index], index),
+			label: active ? `● ${label}` : label,
 			suffix: reset ? `(${reset})` : "",
+			active,
 		};
 	});
 	const maxSuffixWidth = parts.reduce((max, p) => Math.max(max, visibleWidth(p.suffix)), 0);
@@ -1327,16 +1342,18 @@ function formatAccountHeaderRow(
 	if (prefixBudget < 2) {
 		return parts.map(p => {
 			const full = p.suffix ? `${p.label} ${p.suffix}` : p.label;
-			return padColumn(truncateJobLabel(full, columnWidth), columnWidth);
+			const cell = padColumn(truncateJobLabel(full, columnWidth), columnWidth);
+			return p.active ? uiTheme.fg("accent", cell) : cell;
 		});
 	}
 
 	return parts.map(p => {
 		const prefix = truncateJobLabel(p.label, prefixBudget);
 		const prefixCell = prefix + " ".repeat(prefixBudget - visibleWidth(prefix));
-		if (!p.suffix) return prefixCell + " ".repeat(maxSuffixWidth + gap);
+		const styledPrefix = p.active ? uiTheme.fg("accent", prefixCell) : prefixCell;
+		if (!p.suffix) return styledPrefix + " ".repeat(maxSuffixWidth + gap);
 		const suffixPad = " ".repeat(maxSuffixWidth - visibleWidth(p.suffix));
-		return `${prefixCell} ${suffixPad}${uiTheme.fg("dim", p.suffix)}`;
+		return `${styledPrefix} ${suffixPad}${uiTheme.fg("dim", p.suffix)}`;
 	});
 }
 
@@ -1456,6 +1473,7 @@ function renderUsageReports(
 	uiTheme: typeof theme,
 	nowMs: number,
 	availableWidth: number,
+	resolveActiveAccount?: (provider: string) => OAuthAccountIdentity | undefined,
 ): string {
 	const lines: string[] = [];
 	const latestFetchedAt = Math.max(...reports.map(report => report.fetchedAt ?? 0));
@@ -1481,6 +1499,7 @@ function renderUsageReports(
 	for (const { provider, providerReports } of providerEntries) {
 		lines.push("");
 		const providerName = formatProviderName(provider);
+		const activeAccount = resolveActiveAccount?.(provider);
 
 		const limitGroups = new Map<
 			string,
@@ -1504,6 +1523,10 @@ function renderUsageReports(
 		}
 
 		lines.push(uiTheme.bold(uiTheme.fg("accent", providerName)));
+		const activeAccountLabel = activeAccount?.email ?? activeAccount?.accountId ?? activeAccount?.projectId;
+		if (activeAccountLabel) {
+			lines.push(`  ${uiTheme.fg("accent", "in use by this session:")} ${activeAccountLabel}`);
+		}
 
 		const renderableGroups = Array.from(limitGroups.values()).map(group => {
 			const entries = group.limits.map((limit, index) => ({
@@ -1533,7 +1556,14 @@ function renderUsageReports(
 
 			const windowSuffix = formatWindowSuffix(group.label, group.windowLabel, uiTheme);
 			lines.push(`${statusIcon} ${uiTheme.bold(group.label)} ${windowSuffix}`.trim());
-			const accountLabels = formatAccountHeaderRow(sortedLimits, sortedReports, nowMs, sectionColumnWidth, uiTheme);
+			const accountLabels = formatAccountHeaderRow(
+				sortedLimits,
+				sortedReports,
+				nowMs,
+				sectionColumnWidth,
+				uiTheme,
+				activeAccount,
+			);
 			lines.push(`  ${accountLabels.join(" ")}`.trimEnd());
 			const bars = sortedLimits.map(limit =>
 				padColumn(renderUsageBar(limit, uiTheme, sectionColumnWidth), sectionColumnWidth),

packages/coding-agent/src/session/auth-storage.ts

@@ -12,6 +12,7 @@ export type {
 	AuthStorageOptions,
 	CredentialOrigin,
 	CredentialOriginKind,
+	OAuthAccountIdentity,
 	OAuthCredential,
 	SerializedAuthStorage,
 	SnapshotResponse,

packages/coding-agent/src/slash-commands/helpers/active-oauth-account.ts

@@ -0,0 +1,44 @@
+import type { UsageLimit, UsageReport } from "@oh-my-pi/pi-ai";
+import type { OAuthAccountIdentity } from "../../session/auth-storage";
+
+function normalizeIdentityValue(value: unknown): string | undefined {
+	return typeof value === "string" && value.trim() ? value.trim().toLowerCase() : undefined;
+}
+
+/**
+ * True when a single usage-limit column belongs to the given OAuth identity.
+ *
+ * Single definition of the matching rules for both `/usage` renderers:
+ * - `accountId` ↔ report metadata `accountId`/`account_id` or `limit.scope.accountId`
+ * - `email`     ↔ report metadata `email`
+ * - `projectId` ↔ report metadata `projectId` or `limit.scope.projectId`
+ *   (Google-style providers key usage on the GCP project, not an account id)
+ */
+export function limitMatchesActiveAccount(
+	report: UsageReport,
+	limit: UsageLimit,
+	identity: OAuthAccountIdentity | undefined,
+): boolean {
+	if (!identity) return false;
+	const metadata = report.metadata ?? {};
+	const activeAccountId = normalizeIdentityValue(identity.accountId);
+	if (activeAccountId) {
+		const reportAccountId = normalizeIdentityValue(metadata.accountId) ?? normalizeIdentityValue(metadata.account_id);
+		if (reportAccountId === activeAccountId) return true;
+		if (normalizeIdentityValue(limit.scope.accountId) === activeAccountId) return true;
+	}
+	const activeEmail = normalizeIdentityValue(identity.email);
+	if (activeEmail && normalizeIdentityValue(metadata.email) === activeEmail) return true;
+	const activeProjectId = normalizeIdentityValue(identity.projectId);
+	if (activeProjectId) {
+		if (normalizeIdentityValue(metadata.projectId) === activeProjectId) return true;
+		if (normalizeIdentityValue(limit.scope.projectId) === activeProjectId) return true;
+	}
+	return false;
+}
+
+/** True when any limit column in `report` belongs to the given OAuth identity. */
+export function reportMatchesActiveAccount(report: UsageReport, identity: OAuthAccountIdentity | undefined): boolean {
+	if (!identity) return false;
+	return report.limits.some(limit => limitMatchesActiveAccount(report, limit, identity));
+}

packages/coding-agent/src/slash-commands/helpers/usage-report.ts

@@ -1,5 +1,7 @@
 import type { UsageLimit, UsageReport } from "@oh-my-pi/pi-ai";
+import type { OAuthAccountIdentity } from "../../session/auth-storage";
 import type { SlashCommandRuntime } from "../types";
+import { reportMatchesActiveAccount } from "./active-oauth-account";
 import { formatDuration, renderAsciiBar } from "./format";
 
 function formatProviderName(provider: string): string {
@@ -31,7 +33,11 @@ function formatUsageReportAccount(report: UsageReport, limit: UsageLimit, index:
 	return `account ${index + 1}`;
 }
 
-function renderUsageReports(reports: UsageReport[], nowMs: number): string {
+function renderUsageReports(
+	reports: UsageReport[],
+	nowMs: number,
+	resolveActiveAccount?: (provider: string) => OAuthAccountIdentity | undefined,
+): string {
 	const latestFetchedAt = Math.max(...reports.map(report => report.fetchedAt ?? 0));
 	const lines = [`Usage${latestFetchedAt ? ` (${formatDuration(nowMs - latestFetchedAt)} ago)` : ""}`];
 	const grouped = new Map<string, UsageReport[]>();
@@ -45,7 +51,9 @@ function renderUsageReports(reports: UsageReport[], nowMs: number): string {
 		left.localeCompare(right),
 	)) {
 		lines.push("", formatProviderName(provider));
+		const activeAccount = resolveActiveAccount?.(provider);
 		for (const report of providerReports) {
+			const inUse = reportMatchesActiveAccount(report, activeAccount);
 			if (report.limits.length === 0) {
 				const email = typeof report.metadata?.email === "string" ? report.metadata.email : "account";
 				lines.push(`- ${email}: no limits reported`);
@@ -56,7 +64,9 @@ function renderUsageReports(reports: UsageReport[], nowMs: number): string {
 				const window = limit.window?.label ?? limit.scope.windowId;
 				const tier = limit.scope.tier ? ` (${limit.scope.tier})` : "";
 				lines.push(`- ${limit.label}${tier}${window ? ` — ${window}` : ""}`);
-				lines.push(`  ${formatUsageReportAccount(report, limit, index)}: ${formatUsageAmount(limit)}`);
+				lines.push(
+					`  ${formatUsageReportAccount(report, limit, index)}: ${formatUsageAmount(limit)}${inUse ? "  ← in use by this session" : ""}`,
+				);
 				lines.push(`  ${renderAsciiBar(limit.amount.usedFraction)}`);
 				if (limit.window?.resetsAt && limit.window.resetsAt > nowMs) {
 					lines.push(`  resets in ${formatDuration(limit.window.resetsAt - nowMs)}`);
@@ -79,7 +89,18 @@ export async function buildUsageReportText(runtime: SlashCommandRuntime): Promis
 	};
 	if (provider.fetchUsageReports) {
 		const reports = await provider.fetchUsageReports();
-		if (reports && reports.length > 0) return renderUsageReports(reports, Date.now());
+		if (reports && reports.length > 0) {
+			const currentProvider = runtime.session.model?.provider;
+			const activeAccount = currentProvider
+				? runtime.session.modelRegistry.authStorage.getOAuthAccountIdentity(
+						currentProvider,
+						runtime.session.sessionId,
+					)
+				: undefined;
+			return renderUsageReports(reports, Date.now(), providerId =>
+				providerId === currentProvider ? activeAccount : undefined,
+			);
+		}
 	}
 
 	const stats = runtime.session.sessionManager.getUsageStatistics();