From b0d2597715e355c737ca4603ba34f55bc0e466f7 Mon Sep 17 00:00:00 2001 From: can1357 Date: Fri, 12 Jun 2026 04:23:54 +0200 Subject: [PATCH] refactor(coding-agent): share active-account matching across /usage renderers - Extracted `limitMatchesActiveAccount`/`reportMatchesActiveAccount` into `slash-commands/helpers/active-oauth-account.ts` as the single definition of the report-to-account matching rules, including projectId matching against `limit.scope.projectId`/metadata. - Dropped the duplicated `ActiveAccountIdentity`/`OAuthAccessResolver` shims, `as unknown` session casts, the dead `getOAuthAccountId` fallback, and the email-vs-scope-accountId comparison from `command-controller.ts` and `usage-report.ts`. - Replaced the async per-provider `resolveActiveAccountsForReports` map with one synchronous typed `authStorage.getOAuthAccountIdentity()` call per render, gated to the session's current provider. - Re-exported `OAuthAccountIdentity` from `session/auth-storage.ts` and added `active-oauth-account.test.ts` covering the matching rules. --- .../modes/controllers/command-controller.ts | 82 +++---------------- .../coding-agent/src/session/auth-storage.ts | 1 + .../helpers/active-oauth-account.ts | 44 ++++++++++ .../slash-commands/helpers/usage-report.ts | 71 +++------------- .../test/active-oauth-account.test.ts | 82 +++++++++++++++++++ 5 files changed, 151 insertions(+), 129 deletions(-) create mode 100644 packages/coding-agent/src/slash-commands/helpers/active-oauth-account.ts create mode 100644 packages/coding-agent/test/active-oauth-account.test.ts diff --git a/packages/coding-agent/src/modes/controllers/command-controller.ts b/packages/coding-agent/src/modes/controllers/command-controller.ts index e24110d8f..94545e416 100644 --- a/packages/coding-agent/src/modes/controllers/command-controller.ts +++ b/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,10 +405,15 @@ export class CommandController { } const availableWidth = Math.max(40, (this.ctx.ui.terminal.columns ?? 100) - 2); - const activeAccounts = await resolveActiveAccountsForReports(this.ctx.session, usageReports); 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 ? activeAccounts.get(provider) : undefined, + provider === currentProvider ? activeAccount : undefined, ); this.ctx.present([new Spacer(1), new Text(output, 1, 0)]); } @@ -1309,22 +1315,18 @@ function formatResetShort(limit: UsageLimit, nowMs: number): string | undefined return formatDuration(resetsAt - nowMs); } -function formatActiveAccountLabel(activeAccount: ActiveAccountIdentity | undefined): string | undefined { - return activeAccount?.email ?? activeAccount?.accountId; -} - function formatAccountHeaderRow( limits: UsageLimit[], reports: UsageReport[], nowMs: number, columnWidth: number, uiTheme: typeof theme, - activeAccount?: ActiveAccountIdentity, + activeAccount?: OAuthAccountIdentity, ): string[] { const parts = limits.map((limit, index) => { const reset = formatResetShort(limit, nowMs); const report = reports[index]; - const active = reportMatchesActiveAccount(report, limit, activeAccount); + const active = report !== undefined && limitMatchesActiveAccount(report, limit, activeAccount); const label = formatAccountLabel(limit, report, index); return { label: active ? `● ${label}` : label, @@ -1471,7 +1473,7 @@ function renderUsageReports( uiTheme: typeof theme, nowMs: number, availableWidth: number, - resolveActiveAccount?: (provider: string) => ActiveAccountIdentity | undefined, + resolveActiveAccount?: (provider: string) => OAuthAccountIdentity | undefined, ): string { const lines: string[] = []; const latestFetchedAt = Math.max(...reports.map(report => report.fetchedAt ?? 0)); @@ -1521,7 +1523,7 @@ function renderUsageReports( } lines.push(uiTheme.bold(uiTheme.fg("accent", providerName))); - const activeAccountLabel = formatActiveAccountLabel(activeAccount); + const activeAccountLabel = activeAccount?.email ?? activeAccount?.accountId ?? activeAccount?.projectId; if (activeAccountLabel) { lines.push(` ${uiTheme.fg("accent", "in use by this session:")} ${activeAccountLabel}`); } @@ -1592,61 +1594,3 @@ function renderUsageReports( return lines.join("\n"); } - -type ActiveAccountIdentity = { - accountId?: string; - email?: string; -}; - -type OAuthAccessResolver = { - getOAuthAccountId?: (provider: string, sessionId?: string) => string | undefined; - getOAuthAccountIdentity?: (provider: string, sessionId?: string) => ActiveAccountIdentity | undefined; -}; - -function normalizeIdentityValue(value: unknown): string | undefined { - return typeof value === "string" && value.trim() ? value.trim().toLowerCase() : undefined; -} - -function reportMatchesActiveAccount( - report: UsageReport | undefined, - limit: UsageLimit, - activeAccount: ActiveAccountIdentity | undefined, -): boolean { - if (!report || !activeAccount) return false; - const activeAccountId = normalizeIdentityValue(activeAccount.accountId); - const activeEmail = normalizeIdentityValue(activeAccount.email); - const metadata = report.metadata ?? {}; - const reportAccountId = - normalizeIdentityValue(metadata.accountId) ?? normalizeIdentityValue(metadata.account_id) ?? undefined; - const reportEmail = normalizeIdentityValue(metadata.email); - const scopeAccountId = normalizeIdentityValue(limit.scope.accountId); - return Boolean( - (activeAccountId && (reportAccountId === activeAccountId || scopeAccountId === activeAccountId)) || - (activeEmail && (reportEmail === activeEmail || scopeAccountId === activeEmail)), - ); -} - -async function resolveActiveAccountsForReports( - sessionValue: unknown, - reports: UsageReport[], -): Promise> { - const session = sessionValue as { - sessionId?: string; - modelRegistry?: { authStorage?: OAuthAccessResolver }; - }; - const authStorage = session.modelRegistry?.authStorage; - if (!authStorage) return new Map(); - const providers = [...new Set(reports.map(report => report.provider))]; - const entries = await Promise.all( - providers.map(provider => { - const identity = authStorage.getOAuthAccountIdentity?.(provider, session.sessionId); - const accountId = identity?.accountId ?? authStorage.getOAuthAccountId?.(provider, session.sessionId); - const activeIdentity: ActiveAccountIdentity = { - ...(accountId ? { accountId } : {}), - ...(identity?.email ? { email: identity.email } : {}), - }; - return [provider, activeIdentity] as const; - }), - ); - return new Map(entries.filter(([, identity]) => identity.accountId || identity.email)); -} diff --git a/packages/coding-agent/src/session/auth-storage.ts b/packages/coding-agent/src/session/auth-storage.ts index a0a8c0133..0b8725443 100644 --- a/packages/coding-agent/src/session/auth-storage.ts +++ b/packages/coding-agent/src/session/auth-storage.ts @@ -12,6 +12,7 @@ export type { AuthStorageOptions, CredentialOrigin, CredentialOriginKind, + OAuthAccountIdentity, OAuthCredential, SerializedAuthStorage, SnapshotResponse, diff --git a/packages/coding-agent/src/slash-commands/helpers/active-oauth-account.ts b/packages/coding-agent/src/slash-commands/helpers/active-oauth-account.ts new file mode 100644 index 000000000..b32350c28 --- /dev/null +++ b/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)); +} diff --git a/packages/coding-agent/src/slash-commands/helpers/usage-report.ts b/packages/coding-agent/src/slash-commands/helpers/usage-report.ts index 702fa5d6f..1e12c978b 100644 --- a/packages/coding-agent/src/slash-commands/helpers/usage-report.ts +++ b/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,66 +33,10 @@ function formatUsageReportAccount(report: UsageReport, limit: UsageLimit, index: return `account ${index + 1}`; } -type ActiveAccountIdentity = { - accountId?: string; - email?: string; -}; - -type OAuthAccessResolver = { - getOAuthAccountId?: (provider: string, sessionId?: string) => string | undefined; - getOAuthAccountIdentity?: (provider: string, sessionId?: string) => ActiveAccountIdentity | undefined; -}; - -function normalizeIdentityValue(value: unknown): string | undefined { - return typeof value === "string" && value.trim() ? value.trim().toLowerCase() : undefined; -} - -/** True when the report belongs to the given OAuth account identity. */ -function isReportForAccount(report: UsageReport, activeAccount: ActiveAccountIdentity): boolean { - const activeAccountId = normalizeIdentityValue(activeAccount.accountId); - const activeEmail = normalizeIdentityValue(activeAccount.email); - const metadata = report.metadata ?? {}; - const reportAccountId = - normalizeIdentityValue(metadata.accountId) ?? normalizeIdentityValue(metadata.account_id) ?? undefined; - const reportEmail = normalizeIdentityValue(metadata.email); - return report.limits.some(limit => { - const scopeAccountId = normalizeIdentityValue(limit.scope.accountId); - return Boolean( - (activeAccountId && (reportAccountId === activeAccountId || scopeAccountId === activeAccountId)) || - (activeEmail && (reportEmail === activeEmail || scopeAccountId === activeEmail)), - ); - }); -} - -async function resolveActiveAccountsForReports( - sessionValue: unknown, - reports: UsageReport[], -): Promise> { - const session = sessionValue as { - sessionId?: string; - modelRegistry?: { authStorage?: OAuthAccessResolver }; - }; - const authStorage = session.modelRegistry?.authStorage; - if (!authStorage) return new Map(); - const providers = [...new Set(reports.map(report => report.provider))]; - const entries = await Promise.all( - providers.map(provider => { - const identity = authStorage.getOAuthAccountIdentity?.(provider, session.sessionId); - const accountId = identity?.accountId ?? authStorage.getOAuthAccountId?.(provider, session.sessionId); - const activeIdentity: ActiveAccountIdentity = { - ...(accountId ? { accountId } : {}), - ...(identity?.email ? { email: identity.email } : {}), - }; - return [provider, activeIdentity] as const; - }), - ); - return new Map(entries.filter(([, identity]) => identity.accountId || identity.email)); -} - function renderUsageReports( reports: UsageReport[], nowMs: number, - resolveActiveAccount?: (provider: string) => ActiveAccountIdentity | undefined, + resolveActiveAccount?: (provider: string) => OAuthAccountIdentity | undefined, ): string { const latestFetchedAt = Math.max(...reports.map(report => report.fetchedAt ?? 0)); const lines = [`Usage${latestFetchedAt ? ` (${formatDuration(nowMs - latestFetchedAt)} ago)` : ""}`]; @@ -107,7 +53,7 @@ function renderUsageReports( lines.push("", formatProviderName(provider)); const activeAccount = resolveActiveAccount?.(provider); for (const report of providerReports) { - const inUse = activeAccount !== undefined && isReportForAccount(report, activeAccount); + 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`); @@ -144,10 +90,15 @@ export async function buildUsageReportText(runtime: SlashCommandRuntime): Promis if (provider.fetchUsageReports) { const reports = await provider.fetchUsageReports(); if (reports && reports.length > 0) { - const activeAccounts = await resolveActiveAccountsForReports(runtime.session, reports); 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 ? activeAccounts.get(providerId) : undefined, + providerId === currentProvider ? activeAccount : undefined, ); } } diff --git a/packages/coding-agent/test/active-oauth-account.test.ts b/packages/coding-agent/test/active-oauth-account.test.ts new file mode 100644 index 000000000..05275a450 --- /dev/null +++ b/packages/coding-agent/test/active-oauth-account.test.ts @@ -0,0 +1,82 @@ +import { describe, expect, test } from "bun:test"; +import type { UsageLimit, UsageReport } from "@oh-my-pi/pi-ai"; +import { + limitMatchesActiveAccount, + reportMatchesActiveAccount, +} from "../src/slash-commands/helpers/active-oauth-account"; + +function makeLimit(scope: Partial = {}): UsageLimit { + return { + id: "limit-1", + label: "Requests", + scope: { provider: "anthropic", ...scope }, + amount: { usedFraction: 0.5, unit: "percent" }, + }; +} + +function makeReport(overrides: Partial = {}): UsageReport { + return { + provider: "anthropic", + fetchedAt: Date.now(), + limits: [makeLimit()], + ...overrides, + }; +} + +describe("limitMatchesActiveAccount", () => { + test("matches accountId against report metadata (camel and snake case) and limit scope", () => { + const identity = { accountId: "ACC-1" }; + expect(limitMatchesActiveAccount(makeReport({ metadata: { accountId: "acc-1" } }), makeLimit(), identity)).toBe( + true, + ); + expect(limitMatchesActiveAccount(makeReport({ metadata: { account_id: "acc-1" } }), makeLimit(), identity)).toBe( + true, + ); + expect(limitMatchesActiveAccount(makeReport(), makeLimit({ accountId: "acc-1" }), identity)).toBe(true); + expect(limitMatchesActiveAccount(makeReport({ metadata: { accountId: "acc-2" } }), makeLimit(), identity)).toBe( + false, + ); + }); + + test("matches email against report metadata only — never against scope accountId", () => { + const identity = { email: "user@example.com" }; + expect( + limitMatchesActiveAccount(makeReport({ metadata: { email: "User@Example.com" } }), makeLimit(), identity), + ).toBe(true); + // An email must not match an opaque account-id slot that happens to hold the same string. + expect(limitMatchesActiveAccount(makeReport(), makeLimit({ accountId: "user@example.com" }), identity)).toBe( + false, + ); + }); + + test("matches projectId for Google-style providers via scope or metadata", () => { + const identity = { projectId: "gcp-proj-1" }; + expect(limitMatchesActiveAccount(makeReport(), makeLimit({ projectId: "gcp-proj-1" }), identity)).toBe(true); + expect( + limitMatchesActiveAccount(makeReport({ metadata: { projectId: "gcp-proj-1" } }), makeLimit(), identity), + ).toBe(true); + expect(limitMatchesActiveAccount(makeReport(), makeLimit({ projectId: "gcp-proj-2" }), identity)).toBe(false); + }); + + test("returns false without an identity or with an empty identity", () => { + expect(limitMatchesActiveAccount(makeReport({ metadata: { email: "a@b.c" } }), makeLimit(), undefined)).toBe( + false, + ); + expect(limitMatchesActiveAccount(makeReport({ metadata: { email: "a@b.c" } }), makeLimit(), {})).toBe(false); + }); +}); + +describe("reportMatchesActiveAccount", () => { + test("matches when any limit column belongs to the identity", () => { + const report = makeReport({ + limits: [makeLimit({ accountId: "other" }), makeLimit({ accountId: "acc-1" })], + }); + expect(reportMatchesActiveAccount(report, { accountId: "acc-1" })).toBe(true); + expect(reportMatchesActiveAccount(report, { accountId: "acc-3" })).toBe(false); + }); + + test("does not match a report with no limits", () => { + const report = makeReport({ limits: [], metadata: { email: "user@example.com" } }); + expect(reportMatchesActiveAccount(report, { email: "user@example.com" })).toBe(false); + }); +});