diff --git a/packages/client/src/daemon-client.ts b/packages/client/src/daemon-client.ts index 2a1a793b65..c691898267 100644 --- a/packages/client/src/daemon-client.ts +++ b/packages/client/src/daemon-client.ts @@ -95,7 +95,7 @@ import type { ProviderDiagnosticResponseMessage, ProviderUsageListResponseMessage, UsageListReportsResponseMessage, - AgentGetUsageReportResponseMessage, + AgentResolveUsageReportResponseMessage, DaemonGetStatusResponse, DaemonGetPairingOfferResponse, DaemonConfigReloadResponse, @@ -562,7 +562,7 @@ type RefreshProvidersSnapshotPayload = RefreshProvidersSnapshotResponseMessage[" type ProviderDiagnosticPayload = ProviderDiagnosticResponseMessage["payload"]; type ProviderUsageListPayload = ProviderUsageListResponseMessage["payload"]; type UsageListReportsPayload = UsageListReportsResponseMessage["payload"]; -type AgentGetUsageReportPayload = AgentGetUsageReportResponseMessage["payload"]; +type AgentResolveUsageReportPayload = AgentResolveUsageReportResponseMessage["payload"]; type DaemonStatusPayload = DaemonGetStatusResponse["payload"]; type DaemonPairingOfferPayload = DaemonGetPairingOfferResponse["payload"]; type DiagnosticsPayload = DiagnosticsResponse["payload"]; @@ -5245,24 +5245,27 @@ export class DaemonClient { async listUsageReports(options?: { requestId?: string; forceRefresh?: boolean; + reportIds?: string[]; }): Promise { return this.sendNamespacedCorrelatedSessionRequest({ requestId: options?.requestId, - message: { type: "usage.list_reports.request", forceRefresh: options?.forceRefresh }, + message: { + type: "usage.list_reports.request", + forceRefresh: options?.forceRefresh, + reportIds: options?.reportIds, + }, }); } - async getAgentUsageReport(options: { + async resolveAgentUsageReport(options: { agentId: string; requestId?: string; - forceRefresh?: boolean; - }): Promise { + }): Promise { return this.sendNamespacedCorrelatedSessionRequest({ requestId: options.requestId, message: { - type: "agent.get_usage_report.request", + type: "agent.resolve_usage_report.request", agentId: options.agentId, - forceRefresh: options.forceRefresh, }, }); } diff --git a/packages/plugin/src/server/usage.ts b/packages/plugin/src/server/usage.ts index 26dd22763a..3e34e7acf1 100644 --- a/packages/plugin/src/server/usage.ts +++ b/packages/plugin/src/server/usage.ts @@ -1,3 +1,4 @@ +import { createHash } from "node:crypto"; import type { ZodType } from "zod"; import type { JsonValue } from "@getpaseo/protocol/agent-types"; @@ -32,7 +33,6 @@ export interface UsageDetail { } export interface UsageReport { - account: { key: string; label?: string }; status: "available" | "unavailable" | "error"; planLabel?: string; windows: UsageWindow[]; @@ -46,6 +46,8 @@ export interface UsageSourceRegistration { label: string; icon?: string; input: ZodType; + /** Stable account identity, resolved without fetching usage. */ + identify(input: unknown): Promise<{ key: string; label?: string } | null>; fetch(input: unknown): Promise; discover?(): Promise; } @@ -108,9 +110,12 @@ export function usedPctOf( return (used / limit) * 100; } -export function unavailableUsage(accountKey = "default"): UsageReport { +export function hashAccountKey(value: string): string { + return createHash("sha256").update(value).digest("hex"); +} + +export function unavailableUsage(): UsageReport { return { - account: { key: accountKey }, status: "unavailable", windows: [], balances: [], diff --git a/packages/protocol/src/messages.ts b/packages/protocol/src/messages.ts index 84fe0bb2df..a716d45893 100644 --- a/packages/protocol/src/messages.ts +++ b/packages/protocol/src/messages.ts @@ -1777,13 +1777,13 @@ export const ProviderUsageListRequestMessageSchema = z.object({ export const UsageListReportsRequestMessageSchema = z.object({ type: z.literal("usage.list_reports.request"), requestId: z.string(), + reportIds: z.array(z.string()).optional(), forceRefresh: z.boolean().optional(), }); -export const AgentGetUsageReportRequestMessageSchema = z.object({ - type: z.literal("agent.get_usage_report.request"), +export const AgentResolveUsageReportRequestMessageSchema = z.object({ + type: z.literal("agent.resolve_usage_report.request"), requestId: z.string(), agentId: z.string(), - forceRefresh: z.boolean().optional(), }); export const ResumeAgentRequestMessageSchema = z.object({ @@ -3245,7 +3245,7 @@ export const SessionInboundMessageSchema = z.discriminatedUnion("type", [ ProviderDiagnosticRequestMessageSchema, ProviderUsageListRequestMessageSchema, UsageListReportsRequestMessageSchema, - AgentGetUsageReportRequestMessageSchema, + AgentResolveUsageReportRequestMessageSchema, ResumeAgentRequestMessageSchema, ImportAgentRequestMessageSchema, RefreshAgentRequestMessageSchema, @@ -6244,7 +6244,6 @@ export const ProviderUsageListResponseMessageSchema = z.object({ }); export const UsageReportSchema = z.object({ - account: z.object({ key: z.string(), label: z.string().optional() }), status: ProviderUsageStatusSchema, planLabel: z.string().optional(), windows: z.array(ProviderUsageWindowSchema.extend({ headline: z.boolean().optional() })), @@ -6253,6 +6252,9 @@ export const UsageReportSchema = z.object({ error: z.string().optional(), }); export const UsageReportEntrySchema = z.object({ + id: z.string(), + account: z.object({ label: z.string().optional() }), + fetchedAt: z.string(), sourceId: z.string(), sourceLabel: z.string(), icon: z.string().optional(), @@ -6262,9 +6264,9 @@ export const UsageListReportsResponseMessageSchema = z.object({ type: z.literal("usage.list_reports.response"), payload: z.object({ requestId: z.string(), reports: z.array(UsageReportEntrySchema) }), }); -export const AgentGetUsageReportResponseMessageSchema = z.object({ - type: z.literal("agent.get_usage_report.response"), - payload: z.object({ requestId: z.string(), entry: UsageReportEntrySchema.nullable() }), +export const AgentResolveUsageReportResponseMessageSchema = z.object({ + type: z.literal("agent.resolve_usage_report.response"), + payload: z.object({ requestId: z.string(), reportId: z.string().nullable() }), }); const AgentSlashCommandSchema = z.object({ @@ -6953,7 +6955,7 @@ export const SessionOutboundMessageSchema = z.discriminatedUnion("type", [ ProviderDiagnosticResponseMessageSchema, ProviderUsageListResponseMessageSchema, UsageListReportsResponseMessageSchema, - AgentGetUsageReportResponseMessageSchema, + AgentResolveUsageReportResponseMessageSchema, ListCommandsResponseSchema, ListTerminalsResponseSchema, TerminalsChangedSchema, @@ -7135,8 +7137,8 @@ export type ProviderUsageTone = z.infer; export type UsageReport = z.infer; export type UsageReportEntry = z.infer; export type UsageListReportsResponseMessage = z.infer; -export type AgentGetUsageReportResponseMessage = z.infer< - typeof AgentGetUsageReportResponseMessageSchema +export type AgentResolveUsageReportResponseMessage = z.infer< + typeof AgentResolveUsageReportResponseMessageSchema >; export type ProviderUsageStatus = z.infer; export type ProviderUsage = z.infer; diff --git a/packages/server/src/plugins/claude-usage-source/index.server.ts b/packages/server/src/plugins/claude-usage-source/index.server.ts index 636f9d286b..ce5b16c4f7 100644 --- a/packages/server/src/plugins/claude-usage-source/index.server.ts +++ b/packages/server/src/plugins/claude-usage-source/index.server.ts @@ -1,6 +1,6 @@ import type { PluginServerContext } from "@getpaseo/plugin/server"; import { inputSchema } from "./shared/input.js"; -import { fetchUsage } from "./server/usage.js"; +import { fetchUsage, identify } from "./server/usage.js"; export default function contribute(server: PluginServerContext) { server.registerUsageSource({ @@ -9,6 +9,7 @@ export default function contribute(server: PluginServerContext) { icon: "icon.svg", input: inputSchema, discover: async () => [{}], + identify, fetch: fetchUsage, }); return () => {}; diff --git a/packages/server/src/plugins/claude-usage-source/server/usage.test.ts b/packages/server/src/plugins/claude-usage-source/server/usage.test.ts index 34b561ff92..b2c84af7d7 100644 --- a/packages/server/src/plugins/claude-usage-source/server/usage.test.ts +++ b/packages/server/src/plugins/claude-usage-source/server/usage.test.ts @@ -599,3 +599,51 @@ describe("Claude usage source scoped limit reconciliation", () => { ]); }); }); + +it("identify reads account and organization from the selected Claude config", async () => { + const { mkdtemp, writeFile, rm } = await import("node:fs/promises"); + const { identify } = await import("./usage.js"); + const directory = await mkdtemp(join(tmpdir(), "claude-identity-")); + try { + await writeFile( + join(directory, ".claude.json"), + JSON.stringify({ + oauthAccount: { + accountUuid: "account-uuid", + organizationUuid: "org-uuid", + emailAddress: "test@example.com", + }, + }), + ); + expect(await identify({ configDir: directory })).toEqual({ + key: "account-uuid.org-uuid", + label: "test@example.com", + }); + } finally { + await rm(directory, { recursive: true, force: true }); + } +}); + +it("token-only Claude inputs use a cached OAuth profile identity", async () => { + const { identify } = await import("./usage.js"); + let calls = 0; + const fetchProfile: typeof fetch = async () => { + calls++; + return new Response( + JSON.stringify({ + account: { uuid: "account-uuid", email: "test@example.com" }, + organization: { uuid: "org-uuid" }, + }), + { status: 200 }, + ); + }; + expect(await identify({ accessToken: "fixture-claude-token" }, fetchProfile)).toEqual({ + key: "account-uuid.org-uuid", + label: "test@example.com", + }); + expect(await identify({ accessToken: "fixture-claude-token" }, fetchProfile)).toEqual({ + key: "account-uuid.org-uuid", + label: "test@example.com", + }); + expect(calls).toBe(1); +}); diff --git a/packages/server/src/plugins/claude-usage-source/server/usage.ts b/packages/server/src/plugins/claude-usage-source/server/usage.ts index 6a7e8d25b3..9be43dd6de 100644 --- a/packages/server/src/plugins/claude-usage-source/server/usage.ts +++ b/packages/server/src/plugins/claude-usage-source/server/usage.ts @@ -1,12 +1,12 @@ import type { UsageInput } from "../shared/input.js"; import { execFile } from "node:child_process"; -import { createHash } from "node:crypto"; import { existsSync, promises as fs } from "node:fs"; import { homedir, userInfo } from "node:os"; import { join } from "node:path"; import { promisify } from "node:util"; import { z } from "zod"; import { + hashAccountKey, toneFromUsedPct, unavailableUsage, windowFromUsedPct, @@ -416,13 +416,12 @@ export async function fetchUsage( } const { oauth } = credentials; - const accountKey = createHash("sha256").update(oauth.accessToken).digest("hex"); const plan = buildClaudePlan(oauth.subscriptionType, oauth.rateLimitTier); const resp = await callClaudeApi(oauth.accessToken); if (resp === "NEEDS_AUTH") { // Read-only on credentials; the Claude CLI owns refresh. See docs/providers.md. - return unavailableUsage(accountKey); + return unavailableUsage(); } const scoped = reconcileScopedLimits( @@ -450,7 +449,6 @@ export async function fetchUsage( } return { - account: { key: accountKey }, status: "available", planLabel: plan ?? undefined, windows, @@ -458,3 +456,63 @@ export async function fetchUsage( details, }; } + +// The OAuth usage endpoint meters the active organization selected by the token. +const ClaudeAccountSchema = z.object({ + accountUuid: z.string().min(1), + organizationUuid: z.string().min(1), + emailAddress: z.string().optional(), +}); +const profileCache = new Map>(); + +export async function identify(input: UsageInput, fetchApi: typeof fetch = fetch) { + const directory = "configDir" in input ? input.configDir : homedir(); + if (!("accessToken" in input)) { + try { + const config = z + .object({ oauthAccount: ClaudeAccountSchema }) + .parse(JSON.parse(await fs.readFile(join(directory, ".claude.json"), "utf8"))); + const account = config.oauthAccount; + return { + key: `${account.accountUuid}.${account.organizationUuid}`, + ...(account.emailAddress ? { label: account.emailAddress } : {}), + }; + } catch { + return null; + } + } + const token = input.accessToken; + const tokenHash = hashAccountKey(token); + const pending = profileCache.get(tokenHash); + if (pending) return pending; + const request = (async () => { + const response = await fetchApi("https://api.anthropic.com/api/oauth/profile", { + signal: AbortSignal.timeout(15_000), + headers: { + Authorization: `Bearer ${token}`, + Accept: "application/json", + "anthropic-beta": CLAUDE_OAUTH_BETA, + }, + }); + if (response.status === 401 || response.status === 403) return null; + if (!response.ok) throw new Error(`Claude profile API returned ${response.status}`); + const profile = z + .object({ + account: z.object({ uuid: z.string().min(1), email: z.string().optional() }), + organization: z.object({ uuid: z.string().min(1) }), + }) + .parse(await response.json()); + const account = profile.account; + return { + key: `${account.uuid}.${profile.organization.uuid}`, + ...(account.email ? { label: account.email } : {}), + }; + })(); + profileCache.set(tokenHash, request); + try { + return await request; + } catch (error) { + profileCache.delete(tokenHash); + throw error; + } +} diff --git a/packages/server/src/plugins/codex-usage-source/index.server.ts b/packages/server/src/plugins/codex-usage-source/index.server.ts index ddc416b0dc..3a7c472784 100644 --- a/packages/server/src/plugins/codex-usage-source/index.server.ts +++ b/packages/server/src/plugins/codex-usage-source/index.server.ts @@ -1,6 +1,6 @@ import type { PluginServerContext } from "@getpaseo/plugin/server"; import { inputSchema } from "./shared/input.js"; -import { fetchUsage } from "./server/usage.js"; +import { fetchUsage, identify } from "./server/usage.js"; export default function contribute(server: PluginServerContext) { server.registerUsageSource({ @@ -9,6 +9,7 @@ export default function contribute(server: PluginServerContext) { icon: "icon.svg", input: inputSchema, discover: async () => [{}], + identify, fetch: fetchUsage, }); return () => {}; diff --git a/packages/server/src/plugins/codex-usage-source/server/usage.test.ts b/packages/server/src/plugins/codex-usage-source/server/usage.test.ts index d0a415ba02..a898e2770f 100644 --- a/packages/server/src/plugins/codex-usage-source/server/usage.test.ts +++ b/packages/server/src/plugins/codex-usage-source/server/usage.test.ts @@ -2,7 +2,7 @@ import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { afterEach, expect, test } from "vitest"; -import { fetchUsage } from "./usage.js"; +import { fetchUsage, identify } from "./usage.js"; const originalHome = process.env["CODEX_HOME"]; afterEach(() => { @@ -39,7 +39,6 @@ test("default input reads CODEX_HOME auth and preserves the usage request", asyn response(init?.headers ?? {}, "account-default"), ); expect(report).toMatchObject({ - account: { key: "account-default" }, status: "available", planLabel: "plus", windows: [{ id: "session", usedPct: 30, headline: true }], @@ -56,21 +55,23 @@ test("explicit codexHome reads only that auth file", async () => { join(home, "auth.json"), JSON.stringify({ tokens: { access_token: "fixture-home", account_id: "account-home" } }), ); - const report = await fetchUsage({ codexHome: home }, (_url, init) => + await fetchUsage({ codexHome: home }, (_url, init) => response(init?.headers ?? {}, "account-home"), ); - expect(report.account.key).toBe("account-home"); + expect(await identify({ codexHome: home })).toEqual({ key: "account-home" }); } finally { await rm(home, { recursive: true, force: true }); } }); test("explicit access token needs no auth file", async () => { - const report = await fetchUsage( + await fetchUsage( { accessToken: "fixture-supplied", accountId: "account-supplied" }, (_url, init) => response(init?.headers ?? {}, "account-supplied"), ); - expect(report.account.key).toBe("account-supplied"); + expect( + await identify({ accessToken: "fixture-supplied", accountId: "account-supplied" }), + ).toEqual({ key: "account-supplied" }); }); test("coerces credit balance and marks a 96 percent window dangerous", async () => { @@ -133,3 +134,11 @@ test("401 leaves auth.json byte for byte unchanged and makes no refresh request" await rm(home, { recursive: true, force: true }); } }); + +test("token-only account claim survives token rotation", async () => { + const token = (suffix: string) => + `header.${Buffer.from(JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: "work-id" }, suffix })).toString("base64url")}.signature`; + expect(await identify({ accessToken: token("first") })).toEqual({ key: "work-id" }); + expect(await identify({ accessToken: token("second") })).toEqual({ key: "work-id" }); + expect(await identify({ accessToken: "opaque-token" })).toBeNull(); +}); diff --git a/packages/server/src/plugins/codex-usage-source/server/usage.ts b/packages/server/src/plugins/codex-usage-source/server/usage.ts index c04ab28fef..16aae8345e 100644 --- a/packages/server/src/plugins/codex-usage-source/server/usage.ts +++ b/packages/server/src/plugins/codex-usage-source/server/usage.ts @@ -1,4 +1,3 @@ -import { createHash } from "node:crypto"; import { readFile } from "node:fs/promises"; import { homedir } from "node:os"; import { join } from "node:path"; @@ -29,7 +28,7 @@ const responseSchema = z.object({ credits: z.object({ balance: number.optional() }).nullish(), }); -async function readAuth( +export async function readAuth( input: CodexUsageInput, ): Promise<{ token: string; accountId?: string } | null> { if ("accessToken" in input) return { token: input.accessToken, accountId: input.accountId }; @@ -76,8 +75,7 @@ export async function fetchUsage( fetchApi: typeof fetch = fetch, ): Promise { const auth = await readAuth(input); - if (!auth) return { account: { key: "default" }, status: "unavailable", windows: [] }; - const key = auth.accountId ?? createHash("sha256").update(auth.token).digest("hex"); + if (!auth) return { status: "unavailable", windows: [] }; const headers: Record = { Authorization: `Bearer ${auth.token}`, Accept: "application/json", @@ -89,10 +87,10 @@ export async function fetchUsage( signal: AbortSignal.timeout(15_000), }); if (response.status === 401 || response.status === 403) - return { account: { key }, status: "unavailable", windows: [] }; + return { status: "unavailable", windows: [] }; if (!response.ok) throw new Error(`Codex usage API returned ${response.status}`); const text = await response.text(); - if (text.trim().startsWith("<")) return { account: { key }, status: "unavailable", windows: [] }; + if (text.trim().startsWith("<")) return { status: "unavailable", windows: [] }; const usage = responseSchema.parse(JSON.parse(text)); const windows = [ usageWindow("session", "Session", usage.rate_limit?.primary_window, true), @@ -101,7 +99,6 @@ export async function fetchUsage( ].filter((window): window is UsageWindow => window !== null); const balance = usage.credits?.balance; return { - account: { key, ...(usage.email ? { label: usage.email } : {}) }, status: "available", planLabel: usage.plan_type, windows, @@ -120,3 +117,34 @@ export async function fetchUsage( details: [], }; } + +/** The account claim is decoded locally; the token itself never becomes identity. */ +export async function identify(input: CodexUsageInput) { + const auth = await readAuth(input); + if (!auth) return null; + const claims = (() => { + try { + const payload = auth.token.split(".")[1]; + return payload + ? (JSON.parse(Buffer.from(payload, "base64url").toString("utf8")) as Record< + string, + unknown + >) + : null; + } catch { + return null; + } + })(); + const nested = claims?.["https://api.openai.com/auth"]; + const claim = + nested && typeof nested === "object" + ? (nested as Record)["chatgpt_account_id"] + : null; + const key = + auth.accountId ?? + (typeof claim === "string" ? claim : null) ?? + (typeof claims?.["chatgpt_account_id"] === "string" + ? (claims["chatgpt_account_id"] as string) + : null); + return key ? { key } : null; +} diff --git a/packages/server/src/plugins/copilot-usage-source/index.server.ts b/packages/server/src/plugins/copilot-usage-source/index.server.ts index 2227f972de..79da4a578e 100644 --- a/packages/server/src/plugins/copilot-usage-source/index.server.ts +++ b/packages/server/src/plugins/copilot-usage-source/index.server.ts @@ -1,6 +1,6 @@ import type { PluginServerContext } from "@getpaseo/plugin/server"; import { inputSchema } from "./shared/input.js"; -import { fetchUsage } from "./server/usage.js"; +import { fetchUsage, identify } from "./server/usage.js"; export default function contribute(server: PluginServerContext) { server.registerUsageSource({ @@ -9,6 +9,7 @@ export default function contribute(server: PluginServerContext) { icon: "icon.svg", input: inputSchema, discover: async () => [{}], + identify, fetch: fetchUsage, }); return () => {}; diff --git a/packages/server/src/plugins/copilot-usage-source/server/usage.test.ts b/packages/server/src/plugins/copilot-usage-source/server/usage.test.ts index ed3eb6b0ac..13c8cf700d 100644 --- a/packages/server/src/plugins/copilot-usage-source/server/usage.test.ts +++ b/packages/server/src/plugins/copilot-usage-source/server/usage.test.ts @@ -2,7 +2,7 @@ import { mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { fetchUsage } from "./usage.js"; +import { identify, fetchUsage } from "./usage.js"; import type { UsageReport } from "@getpaseo/plugin/server/usage"; function mockFetch(handlers: Map Response>): typeof fetch { @@ -111,3 +111,14 @@ describe("copilot usage source", () => { }); }); }); + +it("identify uses the configured copilot account key", async () => { + const previous = process.env["COPILOT_TOKEN"]; + try { + process.env["COPILOT_TOKEN"] = "fixture-token"; + expect(await identify()).toEqual({ key: "default" }); + } finally { + if (previous === undefined) delete process.env["COPILOT_TOKEN"]; + else process.env["COPILOT_TOKEN"] = previous; + } +}); diff --git a/packages/server/src/plugins/copilot-usage-source/server/usage.ts b/packages/server/src/plugins/copilot-usage-source/server/usage.ts index 06eaa35719..27a5c84e35 100644 --- a/packages/server/src/plugins/copilot-usage-source/server/usage.ts +++ b/packages/server/src/plugins/copilot-usage-source/server/usage.ts @@ -75,7 +75,6 @@ export async function fetchUsage( : []; return { - account: { key: "default" }, status: "available", planLabel: resp.copilot_plan || undefined, windows: [], @@ -83,3 +82,12 @@ export async function fetchUsage( details, }; } + +export async function identify() { + const token = + process.env["COPILOT_TOKEN"] || + process.env["GITHUB_TOKEN"] || + process.env["GITHUB_PAT"] || + (await readGithubCliToken()); + return token ? { key: "default" } : null; +} diff --git a/packages/server/src/plugins/cursor-usage-source/index.server.ts b/packages/server/src/plugins/cursor-usage-source/index.server.ts index 9397dd07c4..a4c1947cb4 100644 --- a/packages/server/src/plugins/cursor-usage-source/index.server.ts +++ b/packages/server/src/plugins/cursor-usage-source/index.server.ts @@ -1,6 +1,6 @@ import type { PluginServerContext } from "@getpaseo/plugin/server"; import { inputSchema } from "./shared/input.js"; -import { fetchUsage } from "./server/usage.js"; +import { fetchUsage, identify } from "./server/usage.js"; export default function contribute(server: PluginServerContext) { server.registerUsageSource({ @@ -9,6 +9,7 @@ export default function contribute(server: PluginServerContext) { icon: "icon.svg", input: inputSchema, discover: async () => [{}], + identify, fetch: fetchUsage, }); return () => {}; diff --git a/packages/server/src/plugins/cursor-usage-source/server/usage.test.ts b/packages/server/src/plugins/cursor-usage-source/server/usage.test.ts index df9cf4bac5..c13cbc5a6f 100644 --- a/packages/server/src/plugins/cursor-usage-source/server/usage.test.ts +++ b/packages/server/src/plugins/cursor-usage-source/server/usage.test.ts @@ -3,7 +3,7 @@ import { createRequire } from "node:module"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { fetchUsage } from "./usage.js"; +import { identify, fetchUsage } from "./usage.js"; import type { UsageReport } from "@getpaseo/plugin/server/usage"; // node:sqlite has no @types/node@20 typings; require it with a narrow local type. @@ -258,3 +258,14 @@ describe("cursor usage source", () => { expect(cursor.status).toBe("unavailable"); }); }); + +it("identify uses the configured cursor account key", async () => { + const previous = process.env["CURSOR_ACCESS_TOKEN"]; + try { + process.env["CURSOR_ACCESS_TOKEN"] = "fixture-token"; + expect(await identify()).toEqual({ key: "default" }); + } finally { + if (previous === undefined) delete process.env["CURSOR_ACCESS_TOKEN"]; + else process.env["CURSOR_ACCESS_TOKEN"] = previous; + } +}); diff --git a/packages/server/src/plugins/cursor-usage-source/server/usage.ts b/packages/server/src/plugins/cursor-usage-source/server/usage.ts index b1fbf9830e..15a9f6e7b3 100644 --- a/packages/server/src/plugins/cursor-usage-source/server/usage.ts +++ b/packages/server/src/plugins/cursor-usage-source/server/usage.ts @@ -222,7 +222,6 @@ export async function fetchUsage( } return { - account: { key: "default" }, status: "available", planLabel: undefined, windows: [], @@ -230,3 +229,13 @@ export async function fetchUsage( details: [], }; } + +export async function identify() { + const home = homedir(); + const token = + process.env["CURSOR_ACCESS_TOKEN"] || + process.env["CURSOR_TOKEN"] || + (await readCursorTokenFromSqlite(home)) || + (await readCursorTokenFromAuthJson(home)); + return token ? { key: "default" } : null; +} diff --git a/packages/server/src/plugins/grok-usage-source/index.server.ts b/packages/server/src/plugins/grok-usage-source/index.server.ts index c7f148fdb5..3c87562300 100644 --- a/packages/server/src/plugins/grok-usage-source/index.server.ts +++ b/packages/server/src/plugins/grok-usage-source/index.server.ts @@ -1,6 +1,6 @@ import type { PluginServerContext } from "@getpaseo/plugin/server"; import { inputSchema } from "./shared/input.js"; -import { fetchUsage } from "./server/usage.js"; +import { fetchUsage, identify } from "./server/usage.js"; export default function contribute(server: PluginServerContext) { server.registerUsageSource({ @@ -9,6 +9,7 @@ export default function contribute(server: PluginServerContext) { icon: "icon.svg", input: inputSchema, discover: async () => [{}], + identify, fetch: fetchUsage, }); return () => {}; diff --git a/packages/server/src/plugins/grok-usage-source/server/usage.test.ts b/packages/server/src/plugins/grok-usage-source/server/usage.test.ts index eb03b8f7e6..9b4d30668f 100644 --- a/packages/server/src/plugins/grok-usage-source/server/usage.test.ts +++ b/packages/server/src/plugins/grok-usage-source/server/usage.test.ts @@ -2,7 +2,7 @@ import { mkdirSync, writeFileSync, mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { fetchUsage } from "./usage.js"; +import { identify, fetchUsage } from "./usage.js"; import type { UsageReport } from "@getpaseo/plugin/server/usage"; function writeGrokAuth(home: string, auth: Record): void { @@ -266,3 +266,14 @@ describe("grok usage source", () => { }); }); }); + +it("identify uses the configured grok account key", async () => { + const previous = process.env["GROK_TOKEN"]; + try { + process.env["GROK_TOKEN"] = "fixture-token"; + expect(await identify()).toEqual({ key: "default" }); + } finally { + if (previous === undefined) delete process.env["GROK_TOKEN"]; + else process.env["GROK_TOKEN"] = previous; + } +}); diff --git a/packages/server/src/plugins/grok-usage-source/server/usage.ts b/packages/server/src/plugins/grok-usage-source/server/usage.ts index 9cd86cefd2..251674f597 100644 --- a/packages/server/src/plugins/grok-usage-source/server/usage.ts +++ b/packages/server/src/plugins/grok-usage-source/server/usage.ts @@ -147,7 +147,6 @@ export async function fetchUsage( if (window) window.headline = true; return { - account: { key: "default" }, status: "available", planLabel: undefined, windows: window ? [window] : [], @@ -155,3 +154,13 @@ export async function fetchUsage( details: [], }; } + +export async function identify() { + if (process.env["GROK_API_KEY"] || process.env["GROK_TOKEN"]) return { key: "default" }; + try { + const auth = JSON.parse(await fs.readFile(join(homedir(), ".grok", "auth.json"), "utf8")); + return extractGrokTokenFromAuth(auth) ? { key: "default" } : null; + } catch { + return null; + } +} diff --git a/packages/server/src/plugins/kimi-usage-source/index.server.ts b/packages/server/src/plugins/kimi-usage-source/index.server.ts index 45e7486c7b..a1df077baa 100644 --- a/packages/server/src/plugins/kimi-usage-source/index.server.ts +++ b/packages/server/src/plugins/kimi-usage-source/index.server.ts @@ -1,6 +1,6 @@ import type { PluginServerContext } from "@getpaseo/plugin/server"; import { inputSchema } from "./shared/input.js"; -import { fetchUsage } from "./server/usage.js"; +import { fetchUsage, identify } from "./server/usage.js"; export default function contribute(server: PluginServerContext) { server.registerUsageSource({ @@ -9,6 +9,7 @@ export default function contribute(server: PluginServerContext) { icon: "icon.svg", input: inputSchema, discover: async () => [{}], + identify, fetch: fetchUsage, }); return () => {}; diff --git a/packages/server/src/plugins/kimi-usage-source/server/usage.test.ts b/packages/server/src/plugins/kimi-usage-source/server/usage.test.ts index d772517034..3b3039a7c2 100644 --- a/packages/server/src/plugins/kimi-usage-source/server/usage.test.ts +++ b/packages/server/src/plugins/kimi-usage-source/server/usage.test.ts @@ -2,7 +2,7 @@ import { mkdirSync, readFileSync, writeFileSync, mkdtempSync, rmSync } from "nod import { tmpdir } from "node:os"; import { join } from "node:path"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { fetchUsage } from "./usage.js"; +import { identify, fetchUsage } from "./usage.js"; import type { UsageReport } from "@getpaseo/plugin/server/usage"; function kimiCredentialPath(dir: string): string { @@ -368,3 +368,14 @@ describe("Kimi usage source usage windows", () => { ]); }); }); + +it("identify uses the configured kimi account key", async () => { + const previous = process.env["KIMI_TOKEN"]; + try { + process.env["KIMI_TOKEN"] = "fixture-token"; + expect(await identify()).toEqual({ key: "default" }); + } finally { + if (previous === undefined) delete process.env["KIMI_TOKEN"]; + else process.env["KIMI_TOKEN"] = previous; + } +}); diff --git a/packages/server/src/plugins/kimi-usage-source/server/usage.ts b/packages/server/src/plugins/kimi-usage-source/server/usage.ts index 1cf4e32653..bd20560ae7 100644 --- a/packages/server/src/plugins/kimi-usage-source/server/usage.ts +++ b/packages/server/src/plugins/kimi-usage-source/server/usage.ts @@ -315,7 +315,6 @@ export async function fetchUsage( if (windows[0]) windows[0].headline = true; return { - account: { key: "default" }, status: "available", planLabel: undefined, windows, @@ -323,3 +322,25 @@ export async function fetchUsage( details: [], }; } + +export async function identify() { + if (process.env["KIMI_TOKEN"] || process.env["KIMI_API_KEY"]) return { key: "default" }; + const home = homedir(); + const paths = [ + join( + process.env["KIMI_CODE_HOME"] || join(home, ".kimi-code"), + "credentials", + "kimi-code.json", + ), + join(home, ".kimi", "credentials", "kimi-code.json"), + ]; + for (const path of paths) { + try { + const auth = KimiAuthSchema.parse(JSON.parse(await fs.readFile(path, "utf8"))); + if (auth.access_token) return { key: "default" }; + } catch { + continue; + } + } + return null; +} diff --git a/packages/server/src/plugins/minimax-usage-source/index.server.ts b/packages/server/src/plugins/minimax-usage-source/index.server.ts index f4b11cbc47..95d470867b 100644 --- a/packages/server/src/plugins/minimax-usage-source/index.server.ts +++ b/packages/server/src/plugins/minimax-usage-source/index.server.ts @@ -1,6 +1,6 @@ import type { PluginServerContext } from "@getpaseo/plugin/server"; import { inputSchema } from "./shared/input.js"; -import { fetchUsage } from "./server/usage.js"; +import { fetchUsage, identify } from "./server/usage.js"; export default function contribute(server: PluginServerContext) { server.registerUsageSource({ @@ -9,6 +9,7 @@ export default function contribute(server: PluginServerContext) { icon: "icon.svg", input: inputSchema, discover: async () => [{}], + identify, fetch: fetchUsage, }); return () => {}; diff --git a/packages/server/src/plugins/minimax-usage-source/server/usage.test.ts b/packages/server/src/plugins/minimax-usage-source/server/usage.test.ts index 620326f0be..b7660d7d68 100644 --- a/packages/server/src/plugins/minimax-usage-source/server/usage.test.ts +++ b/packages/server/src/plugins/minimax-usage-source/server/usage.test.ts @@ -2,7 +2,7 @@ import { mkdirSync, writeFileSync, mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { fetchUsage } from "./usage.js"; +import { identify, fetchUsage } from "./usage.js"; import type { UsageReport } from "@getpaseo/plugin/server/usage"; function writeMiniMaxConfig(dir: string, payload: Record): void { @@ -313,3 +313,14 @@ describe("minimax usage source", () => { }); }); }); + +it("identify uses the configured minimax account key", async () => { + const previous = process.env["MINIMAX_API_KEY"]; + try { + process.env["MINIMAX_API_KEY"] = "fixture-token"; + expect(await identify()).toEqual({ key: "default" }); + } finally { + if (previous === undefined) delete process.env["MINIMAX_API_KEY"]; + else process.env["MINIMAX_API_KEY"] = previous; + } +}); diff --git a/packages/server/src/plugins/minimax-usage-source/server/usage.ts b/packages/server/src/plugins/minimax-usage-source/server/usage.ts index f87110e79e..9c0b019e0f 100644 --- a/packages/server/src/plugins/minimax-usage-source/server/usage.ts +++ b/packages/server/src/plugins/minimax-usage-source/server/usage.ts @@ -258,7 +258,6 @@ export async function fetchUsage( if (windows[0]) windows[0].headline = true; return { - account: { key: "default" }, status: windows.length > 0 ? "available" : "unavailable", planLabel: undefined, windows, @@ -266,3 +265,25 @@ export async function fetchUsage( details: [], }; } + +export async function identify() { + if (process.env["MINIMAX_API_KEY"]) return { key: "default" }; + const home = join(homedir(), ".mmx"); + try { + const credentials = MiniMaxCredentialsSchema.parse( + JSON.parse(await fs.readFile(join(home, "credentials.json"), "utf8")), + ); + if (credentials.access_token) return { key: "default" }; + } catch { + /* Check config next. */ + } + try { + const config = MiniMaxConfigSchema.parse( + JSON.parse(await fs.readFile(join(home, "config.json"), "utf8")), + ); + if (config.api_key || config.oauth?.access_token) return { key: "default" }; + } catch { + /* No configured account. */ + } + return null; +} diff --git a/packages/server/src/plugins/opencode-go-usage-source/index.server.ts b/packages/server/src/plugins/opencode-go-usage-source/index.server.ts index b011eaf2f4..cb47cc1ac6 100644 --- a/packages/server/src/plugins/opencode-go-usage-source/index.server.ts +++ b/packages/server/src/plugins/opencode-go-usage-source/index.server.ts @@ -1,6 +1,6 @@ import type { PluginServerContext } from "@getpaseo/plugin/server"; import { inputSchema } from "./shared/input.js"; -import { discover, fetchUsage } from "./server/usage.js"; +import { discover, fetchUsage, identify } from "./server/usage.js"; export default function contribute(server: PluginServerContext) { server.registerUsageSource({ @@ -9,6 +9,7 @@ export default function contribute(server: PluginServerContext) { icon: "icon.svg", input: inputSchema, discover, + identify, fetch: fetchUsage, }); return () => {}; diff --git a/packages/server/src/plugins/opencode-go-usage-source/server/usage.test.ts b/packages/server/src/plugins/opencode-go-usage-source/server/usage.test.ts index bb33576b49..8072f68a00 100644 --- a/packages/server/src/plugins/opencode-go-usage-source/server/usage.test.ts +++ b/packages/server/src/plugins/opencode-go-usage-source/server/usage.test.ts @@ -2,7 +2,7 @@ import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { expect, test } from "vitest"; -import { discover, fetchUsage } from "./usage.js"; +import { discover, fetchUsage, identify } from "./usage.js"; const upstreamResponse = { usage: { @@ -60,7 +60,7 @@ test("omits default discovery when auth.json lacks an OpenCode Go API key", asyn }, path, ); - expect(report).toEqual({ account: { key: "default" }, status: "unavailable", windows: [] }); + expect(report).toEqual({ status: "unavailable", windows: [] }); } finally { await rm(directory, { recursive: true, force: true }); } @@ -74,3 +74,10 @@ test.each([401, 403])("maps HTTP %i to unavailable", async (status) => { expect(report.status).toBe("unavailable"); expect(report.windows).toEqual([]); }); + +test("identifies OpenCode Go key without a usage request", async () => { + const identity = await identify({ apiKey: "fixture-secret" }); + expect(identity?.key).toMatch(/^[a-f0-9]{64}$/); + expect(identity?.key).not.toBe("fixture-secret"); + expect(await identify({ apiKey: "fixture-secret" })).toEqual(identity); +}); diff --git a/packages/server/src/plugins/opencode-go-usage-source/server/usage.ts b/packages/server/src/plugins/opencode-go-usage-source/server/usage.ts index 348f9c2d9e..0a761fdb99 100644 --- a/packages/server/src/plugins/opencode-go-usage-source/server/usage.ts +++ b/packages/server/src/plugins/opencode-go-usage-source/server/usage.ts @@ -1,9 +1,9 @@ -import { createHash } from "node:crypto"; import { readFile } from "node:fs/promises"; import { homedir } from "node:os"; import { join } from "node:path"; import { z } from "zod"; import { + hashAccountKey, toneFromUsedPct, windowFromUsedPct, type UsageReport, @@ -35,7 +35,7 @@ function authPath(env: NodeJS.ProcessEnv = process.env): string { ); } -async function readDefaultKey(path = authPath()): Promise { +export async function readDefaultKey(path = authPath()): Promise { try { const auth = authSchema.parse(JSON.parse(await readFile(path, "utf8"))); return auth["opencode-go"]?.key ?? null; @@ -54,14 +54,13 @@ export async function fetchUsage( path = authPath(), ): Promise { const apiKey = "apiKey" in input ? input.apiKey : await readDefaultKey(path); - if (!apiKey) return { account: { key: "default" }, status: "unavailable", windows: [] }; - const key = createHash("sha256").update(apiKey).digest("hex"); + if (!apiKey) return { status: "unavailable", windows: [] }; const response = await fetchApi("https://opencode.ai/zen/go/v1/usage", { headers: { Authorization: `Bearer ${apiKey}`, Accept: "application/json" }, signal: AbortSignal.timeout(15_000), }); if (response.status === 401 || response.status === 403) - return { account: { key }, status: "unavailable", windows: [] }; + return { status: "unavailable", windows: [] }; if (!response.ok) throw new Error(`OpenCode Go usage API returned ${response.status}`); const data = responseSchema.parse(await response.json()); const windows = ( @@ -80,5 +79,10 @@ export async function fetchUsage( headline: id === "rolling", }), ); - return { account: { key }, status: "available", planLabel: "Go", windows }; + return { status: "available", planLabel: "Go", windows }; +} + +export async function identify(input: Input, path = authPath()) { + const key = "apiKey" in input ? input.apiKey : await readDefaultKey(path); + return key ? { key: hashAccountKey(key) } : null; } diff --git a/packages/server/src/plugins/zai-usage-source/index.server.ts b/packages/server/src/plugins/zai-usage-source/index.server.ts index 95b86aefd7..b4fd86943a 100644 --- a/packages/server/src/plugins/zai-usage-source/index.server.ts +++ b/packages/server/src/plugins/zai-usage-source/index.server.ts @@ -1,6 +1,6 @@ import type { PluginServerContext } from "@getpaseo/plugin/server"; import { inputSchema } from "./shared/input.js"; -import { fetchUsage } from "./server/usage.js"; +import { fetchUsage, identify } from "./server/usage.js"; export default function contribute(server: PluginServerContext) { server.registerUsageSource({ @@ -9,6 +9,7 @@ export default function contribute(server: PluginServerContext) { icon: "icon.svg", input: inputSchema, discover: async () => [{}], + identify, fetch: fetchUsage, }); return () => {}; diff --git a/packages/server/src/plugins/zai-usage-source/server/usage.test.ts b/packages/server/src/plugins/zai-usage-source/server/usage.test.ts index 3f8fd0082d..9988640e32 100644 --- a/packages/server/src/plugins/zai-usage-source/server/usage.test.ts +++ b/packages/server/src/plugins/zai-usage-source/server/usage.test.ts @@ -2,7 +2,7 @@ import { mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { fetchUsage } from "./usage.js"; +import { identify, fetchUsage } from "./usage.js"; import type { UsageReport } from "@getpaseo/plugin/server/usage"; function mockFetch(handlers: Map Response>): typeof fetch { @@ -117,3 +117,14 @@ describe("zai usage source", () => { }); }); }); + +it("identify uses the configured zai account key", async () => { + const previous = process.env["ZAI_API_KEY"]; + try { + process.env["ZAI_API_KEY"] = "fixture-token"; + expect(await identify()).toEqual({ key: "default" }); + } finally { + if (previous === undefined) delete process.env["ZAI_API_KEY"]; + else process.env["ZAI_API_KEY"] = previous; + } +}); diff --git a/packages/server/src/plugins/zai-usage-source/server/usage.ts b/packages/server/src/plugins/zai-usage-source/server/usage.ts index 8d23f88f1a..432990f374 100644 --- a/packages/server/src/plugins/zai-usage-source/server/usage.ts +++ b/packages/server/src/plugins/zai-usage-source/server/usage.ts @@ -57,7 +57,6 @@ export async function fetchUsage( } return { - account: { key: "default" }, status: "available", planLabel: sub.productName || undefined, windows: [], @@ -65,3 +64,7 @@ export async function fetchUsage( details, }; } + +export async function identify() { + return process.env["ZAI_API_KEY"] || process.env["GLM_API_KEY"] ? { key: "default" } : null; +} diff --git a/packages/server/src/server/authorization/operation-permissions.ts b/packages/server/src/server/authorization/operation-permissions.ts index dffcc10d29..fa260011c3 100644 --- a/packages/server/src/server/authorization/operation-permissions.ts +++ b/packages/server/src/server/authorization/operation-permissions.ts @@ -150,7 +150,7 @@ const INBOUND_PERMISSION = { project_icon_request: "workspace.read", "provider.usage.list.request": "daemon.read", "usage.list_reports.request": "daemon.read", - "agent.get_usage_report.request": "daemon.read", + "agent.resolve_usage_report.request": "daemon.read", provider_diagnostic_request: "daemon.read", pull_request_timeline_request: "workspace.read", "push.unregister.request": "workspace.read", @@ -376,7 +376,7 @@ const OUTBOUND_PERMISSION = { project_icon_response: "workspace.read", "provider.usage.list.response": "daemon.read", "usage.list_reports.response": "daemon.read", - "agent.get_usage_report.response": "daemon.read", + "agent.resolve_usage_report.response": "daemon.read", provider_diagnostic_response: "daemon.read", providers_snapshot_update: ["daemon.read", "hub.execute"], pull_request_timeline_response: "workspace.read", diff --git a/packages/server/src/server/plugins/index.ts b/packages/server/src/server/plugins/index.ts index 2b0c943cf5..09247488dc 100644 --- a/packages/server/src/server/plugins/index.ts +++ b/packages/server/src/server/plugins/index.ts @@ -41,6 +41,7 @@ interface PluginRuntimePort { clearLogs(pluginId: string): void; getProviderRegistrations?(pluginId: string): readonly PluginProviderMetadata[]; getUsageSourceRegistrations(pluginId: string): readonly PluginUsageSourceMetadata[]; + identifyUsage: PluginRuntime["identifyUsage"]; fetchUsage: PluginRuntime["fetchUsage"]; discoverUsage: PluginRuntime["discoverUsage"]; connectProvider: PluginRuntime["connectProvider"]; @@ -143,12 +144,16 @@ export class PluginService { return [...this.providers.values()].sort((left, right) => left.id.localeCompare(right.id)); } - listUsageReports(options?: { forceRefresh?: boolean; references?: UsageReference[] }) { + listUsageReports(options?: { + forceRefresh?: boolean; + reportIds?: string[]; + references?: UsageReference[]; + }) { return this.usageSources.listReports(options); } - fetchUsageReference(reference: UsageReference, options?: { forceRefresh?: boolean }) { - return this.usageSources.fetchReference(reference, options); + resolveUsageReference(reference: UsageReference) { + return this.usageSources.resolveReference(reference); } listLegacyUsage() { @@ -569,6 +574,11 @@ export class PluginService { throw new Error(`Invalid usage discovery from ${source.id}`); return result; }, + identify: (input) => + this.runtime.identifyUsage(pluginId, source.id, input) as Promise<{ + key: string; + label?: string; + } | null>, fetch: (input) => { return this.runtime.fetchUsage(pluginId, source.id, input); }, diff --git a/packages/server/src/server/plugins/plugin-process-protocol.ts b/packages/server/src/server/plugins/plugin-process-protocol.ts index d1e3bd9aa0..16ee512c37 100644 --- a/packages/server/src/server/plugins/plugin-process-protocol.ts +++ b/packages/server/src/server/plugins/plugin-process-protocol.ts @@ -39,6 +39,7 @@ export type PluginProcessRequest = } | { type: "hook"; requestId: string; kind: "event" | "before"; name: string; input: unknown } | { type: "hook.cancel"; requestId: string } + | { type: "usage.identify"; requestId: string; sourceId: string; input: unknown } | { type: "usage.fetch"; requestId: string; sourceId: string; input: unknown } | { type: "usage.discover"; requestId: string; sourceId: string } | { type: "invoke"; requestId: string; method: string; input: unknown } @@ -161,6 +162,14 @@ export const PluginProcessRequestSchema: z.ZodType = z.dis }) .strict(), z.object({ type: z.literal("hook.cancel"), requestId: z.string() }).strict(), + z + .object({ + type: z.literal("usage.identify"), + requestId: z.string(), + sourceId: z.string(), + input: z.unknown(), + }) + .strict(), z .object({ type: z.literal("usage.fetch"), diff --git a/packages/server/src/server/plugins/plugin-process.ts b/packages/server/src/server/plugins/plugin-process.ts index c3d84daa55..b848c89789 100644 --- a/packages/server/src/server/plugins/plugin-process.ts +++ b/packages/server/src/server/plugins/plugin-process.ts @@ -362,7 +362,10 @@ export function createPluginWorker(options: { } function handleUsageRequest( - message: Extract, + message: Extract< + PluginProcessRequest, + { type: "usage.identify" | "usage.fetch" | "usage.discover" } + >, ): void { void (async () => { const source = usageSources.get(message.sourceId); @@ -370,7 +373,11 @@ export function createPluginWorker(options: { if (message.type === "usage.discover") return jsonTransportValue(source.discover ? await source.discover() : []); const input = await source.input.parseAsync(message.input); - return jsonTransportValue(await source.fetch(input)); + return jsonTransportValue( + message.type === "usage.identify" + ? await source.identify(input) + : await source.fetch(input), + ); })().then( (output) => send({ type: "result", requestId: message.requestId, output }), (error) => send({ type: "error", requestId: message.requestId, error: describeError(error) }), @@ -380,6 +387,7 @@ export function createPluginWorker(options: { function rejectWhileStopping(message: PluginProcessRequest): void { if ( message.type === "provider.catalog_key" || + message.type === "usage.identify" || message.type === "usage.fetch" || message.type === "usage.discover" ) { @@ -456,7 +464,11 @@ export function createPluginWorker(options: { ); return; } - if (message.type === "usage.fetch" || message.type === "usage.discover") { + if ( + message.type === "usage.identify" || + message.type === "usage.fetch" || + message.type === "usage.discover" + ) { handleUsageRequest(message); return; } diff --git a/packages/server/src/server/plugins/runtime.ts b/packages/server/src/server/plugins/runtime.ts index 5a645c94b7..ea6b1ae554 100644 --- a/packages/server/src/server/plugins/runtime.ts +++ b/packages/server/src/server/plugins/runtime.ts @@ -389,6 +389,17 @@ export class PluginRuntime { return this.plugins.get(pluginId)?.usageSources ?? []; } + identifyUsage(pluginId: string, sourceId: string, input: unknown): Promise { + const loaded = this.plugins.get(pluginId); + if (!loaded) throw new Error(`Plugin is not available: ${pluginId}`); + return this.request(loaded, { + type: "usage.identify", + requestId: randomUUID(), + sourceId, + input, + }); + } + fetchUsage(pluginId: string, sourceId: string, input: unknown): Promise { const loaded = this.plugins.get(pluginId); if (!loaded) throw new Error(`Plugin is not available: ${pluginId}`); diff --git a/packages/server/src/server/plugins/session-usage-reference.e2e.test.ts b/packages/server/src/server/plugins/session-usage-reference.e2e.test.ts index f35cc1e021..ae69a0a60d 100644 --- a/packages/server/src/server/plugins/session-usage-reference.e2e.test.ts +++ b/packages/server/src/server/plugins/session-usage-reference.e2e.test.ts @@ -25,26 +25,32 @@ test("resolves a provider plugin session reference through its usage source", as provider: "fixture-session-provider", cwd: directory, }); - const result = await client.getAgentUsageReport({ agentId: agent.id }); - expect(result.entry).toMatchObject({ + const result = await client.resolveAgentUsageReport({ agentId: agent.id }); + expect(result.reportId).toBe("fixture-session-usage:from-session"); + expect( + (await client.listUsageReports({ reportIds: [result.reportId!] })).reports[0], + ).toMatchObject({ + id: result.reportId, sourceId: "fixture-session-usage", - report: { account: { key: "from-session" }, windows: [{ usedPct: 31 }] }, + report: { windows: [{ usedPct: 31 }] }, }); const reports = await client.listUsageReports(); - expect(reports.reports.map((entry) => entry.report.account.key)).toContain("from-session"); + expect(reports.reports.map((entry) => entry.id)).toContain( + "fixture-session-usage:from-session", + ); const missing = await client.createAgent({ provider: "fixture-session-provider", cwd: directory, model: "missing", }); - expect((await client.getAgentUsageReport({ agentId: missing.id })).entry).toBeNull(); + expect((await client.resolveAgentUsageReport({ agentId: missing.id })).reportId).toBeNull(); } finally { await client.close(); await daemon.close(); } }, 60_000); -test("agent.get_usage_report resolves source IDs from default built-in and ACP sessions", async () => { +test("agent.resolve_usage_report resolves source IDs from default built-in and ACP sessions", async () => { const directory = fileURLToPath( new URL("./test-fixtures/session-usage-reference/", import.meta.url), ); @@ -131,8 +137,8 @@ test("agent.get_usage_report resolves source IDs from default built-in and ACP s await client.connect(); for (const provider of providers) { const agent = await client.createAgent({ provider, cwd: directory }); - const result = await client.getAgentUsageReport({ agentId: agent.id }); - expect(result.entry?.sourceId, provider).toBe(provider); + const result = await client.resolveAgentUsageReport({ agentId: agent.id }); + expect(result.reportId, provider).toBe(`${provider}:default`); } } finally { await client.close(); diff --git a/packages/server/src/server/plugins/test-fixtures/session-usage-reference/index.server.ts b/packages/server/src/server/plugins/test-fixtures/session-usage-reference/index.server.ts index c74504a958..c5aeaede16 100644 --- a/packages/server/src/server/plugins/test-fixtures/session-usage-reference/index.server.ts +++ b/packages/server/src/server/plugins/test-fixtures/session-usage-reference/index.server.ts @@ -14,8 +14,8 @@ export default function contribute(server: PluginServerContext) { id, label: id, input, + identify: async () => ({ key: "default" }), fetch: async () => ({ - account: { key: "default" }, status: "available", windows: [{ id: "hour", label: "Hour", usedPct: 31, headline: true }], }), @@ -25,8 +25,8 @@ export default function contribute(server: PluginServerContext) { id: "fixture-session-usage", label: "Fixture session usage", input: z.object({ account: z.string() }).strict(), - fetch: async (input) => ({ - account: { key: (input as { account: string }).account }, + identify: async (input) => ({ key: (input as { account: string }).account }), + fetch: async () => ({ status: "available", windows: [{ id: "hour", label: "Hour", usedPct: 31, headline: true }], }), diff --git a/packages/server/src/server/plugins/test-fixtures/usage-source-directory/index.server.ts b/packages/server/src/server/plugins/test-fixtures/usage-source-directory/index.server.ts index 62102d07c1..48a4ee5eec 100644 --- a/packages/server/src/server/plugins/test-fixtures/usage-source-directory/index.server.ts +++ b/packages/server/src/server/plugins/test-fixtures/usage-source-directory/index.server.ts @@ -16,12 +16,12 @@ export default function contribute(server: PluginServerContext) { ]; return inputs; }, + identify: async (input) => ({ key: (input as { account: string }).account }), fetch: async (input) => { const account = (input as { account: string }).account; if (account === "throws") throw new Error("fixture failure"); fetches++; return { - account: { key: account }, status: "available", windows: [{ id: "count", label: "Count", usedPct: fetches, headline: true }], }; diff --git a/packages/server/src/server/plugins/test-fixtures/usage-source/index.server.ts b/packages/server/src/server/plugins/test-fixtures/usage-source/index.server.ts index b76cd67b42..a8f1f28211 100644 --- a/packages/server/src/server/plugins/test-fixtures/usage-source/index.server.ts +++ b/packages/server/src/server/plugins/test-fixtures/usage-source/index.server.ts @@ -16,12 +16,12 @@ export default function contribute(server: PluginServerContext) { ]; return inputs; }, + identify: async (input) => ({ key: (input as { account: string }).account }), fetch: async (input) => { const account = (input as { account: string }).account; if (account === "throws") throw new Error("fixture failure"); fetches++; return { - account: { key: account }, status: "available", windows: [{ id: "count", label: "Count", usedPct: fetches, headline: true }], }; diff --git a/packages/server/src/server/plugins/usage-sources.e2e.test.ts b/packages/server/src/server/plugins/usage-sources.e2e.test.ts index d5e63fc795..588cbc7e78 100644 --- a/packages/server/src/server/plugins/usage-sources.e2e.test.ts +++ b/packages/server/src/server/plugins/usage-sources.e2e.test.ts @@ -21,19 +21,23 @@ test("lists internal and subprocess usage; validates input and isolates fetch er const first = await client.listUsageReports(); expect(first.reports).toHaveLength(3); expect(first.reports[0]?.icon).toContain(" entry.report.account.key === "one")?.report.windows[0]?.usedPct, + first.reports.find((entry) => entry.id === "fixture:one")?.report.windows[0]?.usedPct, ).toBe(1); expect(first.reports.filter((entry) => entry.report.status === "error")).toHaveLength(2); expect( - (await client.listUsageReports()).reports.find((entry) => entry.report.account.key === "one") - ?.report.windows[0]?.usedPct, + (await client.listUsageReports()).reports.find((entry) => entry.id === "fixture:one")?.report + .windows[0]?.usedPct, ).toBe(1); expect( - (await client.listUsageReports({ forceRefresh: true })).reports.find( - (entry) => entry.report.account.key === "one", - )?.report.windows[0]?.usedPct, + ( + await client.listUsageReports({ reportIds: ["fixture:one"], forceRefresh: true }) + ).reports.find((entry) => entry.id === "fixture:one")?.report.windows[0]?.usedPct, ).toBe(2); + expect( + (await client.listUsageReports({ reportIds: ["fixture:throws"] })).reports[0]?.report.status, + ).toBe("error"); await client.patchDaemonConfig({ pluginsEnabled: true }); await client.installDirectoryPlugin(subprocessDirectory, "fixture-directory"); const both = await client.listUsageReports({ forceRefresh: true }); diff --git a/packages/server/src/server/plugins/usage-sources/index.test.ts b/packages/server/src/server/plugins/usage-sources/index.test.ts index ce51810ce8..cbfde562b5 100644 --- a/packages/server/src/server/plugins/usage-sources/index.test.ts +++ b/packages/server/src/server/plugins/usage-sources/index.test.ts @@ -1,148 +1,181 @@ import { expect, test } from "vitest"; import { UsageSourceRegistry } from "./index.js"; -const report = (key: string, count: number) => ({ - account: { key }, - status: "available" as const, - windows: [{ id: "count", label: "Count", usedPct: count }], -}); +function source(input: { + id: string; + discover?: () => Promise; + identify?: (value: unknown) => Promise<{ key: string } | null>; + fetch?: (value: unknown) => Promise; +}) { + return { + id: input.id, + label: input.id, + discover: input.discover ?? (async () => []), + identify: + input.identify ?? + (async (value: unknown) => ({ key: (value as { account: string }).account })), + fetch: input.fetch ?? (async () => ({ status: "available", windows: [] })), + }; +} -test("caches by source and input for five minutes and forceRefresh bypasses it", async () => { - let now = 0; - let count = 0; - const registry = new UsageSourceRegistry(() => now); - registry.register({ - id: "fixture", - label: "Fixture", - discover: async () => [{ account: "a" }], - fetch: async () => report("a", ++count), - }); - expect((await registry.listReports())[0]?.report.windows[0]?.usedPct).toBe(1); - expect((await registry.listReports())[0]?.report.windows[0]?.usedPct).toBe(1); - expect((await registry.listReports({ forceRefresh: true }))[0]?.report.windows[0]?.usedPct).toBe( - 2, - ); - now = 300_001; - expect((await registry.listReports())[0]?.report.windows[0]?.usedPct).toBe(3); -}); - -test("coalesces concurrent fetches and dedupes accounts", async () => { - let count = 0; +test("identifies accounts without fetching and preserves IDs across token rotation", async () => { const registry = new UsageSourceRegistry(); - registry.register({ - id: "fixture", - label: "Fixture", - discover: async () => [{ a: 1 }, { a: 2 }], - fetch: async () => report("same", ++count), - }); - const [first, second] = await Promise.all([ - registry.fetch("fixture", { a: 1 }), - registry.fetch("fixture", { a: 1 }), - ]); - expect(first).toBe(second); - expect(count).toBe(1); - expect(await registry.listReports()).toHaveLength(1); -}); - -test("includes live session references and ignores unregistered sources", async () => { - const registry = new UsageSourceRegistry(); - registry.register({ - id: "fixture", - label: "Fixture", - discover: async () => [{ account: "default" }], - fetch: async (input) => report((input as { account: string }).account, 25), - }); - const reports = await registry.listReports({ - references: [ - { source: "fixture", input: { account: "second" } }, - { source: "missing", input: {} }, - ], - }); - expect(reports.map((entry) => entry.report.account.key)).toEqual(["default", "second"]); - expect(await registry.fetchReference({ source: "missing", input: {} })).toBeNull(); -}); - -test("drops expired cached inputs on a later write", async () => { - let now = 0; - const registry = new UsageSourceRegistry(() => now, 100); - registry.register({ - id: "fixture", - label: "Fixture", - discover: async () => [], - fetch: async (input) => report(String(input), 1), - }); - await registry.fetch("fixture", "old"); - now = 101; - await registry.fetch("fixture", "new"); - const cache = Reflect.get(registry, "cache") as Map; - expect(cache.size).toBe(1); -}); - -test("turns fetch and discovery failures into error reports", async () => { - const registry = new UsageSourceRegistry(); - registry.register({ - id: "failure", - label: "Failure", - discover: async () => [{}], - fetch: async () => { - throw new Error("failed"); - }, - }); - registry.register({ - id: "discovery", - label: "Discovery", - discover: async () => { - throw new Error("discovery failed"); - }, - fetch: async () => report("never", 0), - }); - expect((await registry.listReports()).map((entry) => entry.report.status)).toEqual([ - "error", - "error", - ]); -}); - -test("maps discovered reports to the legacy provider usage response", async () => { - const registry = new UsageSourceRegistry(() => Date.parse("2026-06-19T00:00:00.000Z")); - registry.register({ - id: "fixture", - label: "Fixture Source", - discover: async () => [{}], - fetch: async () => ({ - account: { key: "one" }, - status: "available", - windows: [{ id: "weekly", label: "Weekly", usedPct: 42, headline: true }], + let fetches = 0; + let token = "old"; + registry.register( + source({ + id: "codex", + discover: async () => [{ account: "work", token }], + fetch: async (input) => { + fetches++; + return { + status: "available", + windows: [{ id: "token", label: (input as { token: string }).token }], + }; + }, }), + ); + expect( + await registry.resolveReference({ source: "codex", input: { account: "work", token } }), + ).toBe("codex:work"); + expect(fetches).toBe(0); + token = "new"; + expect( + await registry.resolveReference({ source: "codex", input: { account: "work", token } }), + ).toBe("codex:work"); + expect( + (await registry.listReports({ reportIds: ["codex:work"] }))[0]?.report.windows[0]?.label, + ).toBe("new"); + expect(fetches).toBe(1); +}); + +test("coalesces per account, caches errors, and refreshes only requested IDs", async () => { + let now = 0; + const counts = new Map(); + const registry = new UsageSourceRegistry(() => now); + registry.register( + source({ + id: "source", + discover: async () => [{ account: "a" }, { account: "b" }], + fetch: async (input) => { + const account = (input as { account: string }).account; + counts.set(account, (counts.get(account) ?? 0) + 1); + if (account === "b") throw new Error("failed"); + return { status: "available", windows: [] }; + }, + }), + ); + const first = await registry.listReports(); + expect(first.map((entry) => entry.id)).toEqual(["source:a", "source:b"]); + expect(first[1]?.report.status).toBe("error"); + expect(await registry.listReports()).toEqual(first); + expect(counts.get("a")).toBe(1); + expect(counts.get("b")).toBe(1); + now = 1000; + const refreshed = await registry.listReports({ + reportIds: ["source:a", "missing:id"], + forceRefresh: true, }); - registry.register({ - id: "missing", - label: "Missing Source", - discover: async () => [{}], - fetch: async () => ({ account: { key: "default" }, status: "unavailable", windows: [] }), + expect(refreshed.map((entry) => entry.id)).toEqual(["source:a"]); + expect(refreshed[0]?.fetchedAt).not.toBe(first[0]?.fetchedAt); + expect(counts.get("a")).toBe(2); + expect(counts.get("b")).toBe(1); + now = 301_001; + await registry.listReports({ reportIds: ["source:a"] }); + expect(counts.get("a")).toBe(3); +}); + +test("invalid keys become source errors and missing identities produce no report", async () => { + const registry = new UsageSourceRegistry(); + registry.register( + source({ + id: "source", + discover: async () => [{ account: "bad:key" }, { account: "none" }], + identify: async (input) => { + const account = (input as { account: string }).account; + return account === "none" ? null : { key: account }; + }, + }), + ); + const reports = await registry.listReports(); + expect(reports.map((entry) => entry.id)).toEqual(["source:error"]); + expect(reports[0]?.report.status).toBe("error"); + expect( + await registry.resolveReference({ source: "source", input: { account: "none" } }), + ).toBeNull(); +}); + +test("legacy listing uses oldest fetchedAt", async () => { + let now = 1000; + const registry = new UsageSourceRegistry(() => now); + registry.register( + source({ id: "source", discover: async () => [{ account: "a" }, { account: "b" }] }), + ); + await registry.listReports({ references: [{ source: "source", input: { account: "a" } }] }); + now = 2000; + await registry.listReports({ reportIds: ["source:b"], forceRefresh: true }); + expect((await registry.listLegacyUsage()).fetchedAt).toBe(new Date(1000).toISOString()); +}); + +test("two Codex homes and a token route resolve to their vendor account IDs", async () => { + const { mkdtemp, writeFile, rm } = await import("node:fs/promises"); + const { tmpdir } = await import("node:os"); + const { join } = await import("node:path"); + const { identify } = await import("../../../plugins/codex-usage-source/server/usage.js"); + const personal = await mkdtemp(join(tmpdir(), "usage-personal-")); + const work = await mkdtemp(join(tmpdir(), "usage-work-")); + try { + await writeFile( + join(personal, "auth.json"), + JSON.stringify({ tokens: { account_id: "personal-id", access_token: "personal-token" } }), + ); + await writeFile( + join(work, "auth.json"), + JSON.stringify({ tokens: { account_id: "work-id", access_token: "work-token" } }), + ); + const registry = new UsageSourceRegistry(); + registry.register( + source({ + id: "codex", + discover: async () => [{ codexHome: personal }, { codexHome: work }], + identify: async (input) => identify(input as Parameters[0]), + }), + ); + const token = `header.${Buffer.from(JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: "work-id" } })).toString("base64url")}.signature`; + expect( + await registry.resolveReference({ source: "codex", input: { accessToken: token } }), + ).toBe("codex:work-id"); + expect((await registry.listReports()).map((entry) => entry.id)).toEqual([ + "codex:personal-id", + "codex:work-id", + ]); + } finally { + await rm(personal, { recursive: true, force: true }); + await rm(work, { recursive: true, force: true }); + } +}); + +test("concurrent requests for the same ID share one vendor fetch", async () => { + let finish!: (report: unknown) => void; + const response = new Promise((resolve) => { + finish = resolve; }); - expect(await registry.listLegacyUsage()).toEqual({ - fetchedAt: "2026-06-19T00:00:00.000Z", - providers: [ - { - providerId: "fixture", - displayName: "Fixture Source", - status: "available", - planLabel: null, - windows: [{ id: "weekly", label: "Weekly", usedPct: 42, headline: true }], - balances: [], - details: [], - error: null, + let fetches = 0; + const registry = new UsageSourceRegistry(); + registry.register( + source({ + id: "coalesced", + fetch: async () => { + fetches++; + return response; }, - { - providerId: "missing", - displayName: "Missing Source", - status: "unavailable", - planLabel: null, - windows: [], - balances: [], - details: [], - error: null, - }, - ], - }); + }), + ); + await registry.resolveReference({ source: "coalesced", input: { account: "one" } }); + const first = registry.listReports({ reportIds: ["coalesced:one"] }); + const second = registry.listReports({ reportIds: ["coalesced:one"], forceRefresh: true }); + finish({ status: "available", windows: [] }); + const [one, two] = await Promise.all([first, second]); + expect(one[0]).toBe(two[0]); + expect(fetches).toBe(1); }); diff --git a/packages/server/src/server/plugins/usage-sources/index.ts b/packages/server/src/server/plugins/usage-sources/index.ts index 686a3e509a..1734817104 100644 --- a/packages/server/src/server/plugins/usage-sources/index.ts +++ b/packages/server/src/server/plugins/usage-sources/index.ts @@ -1,8 +1,7 @@ -import { createHash } from "node:crypto"; import { UsageReportSchema, - type UsageReportEntry, type ProviderUsage, + type UsageReportEntry, } from "@getpaseo/protocol/messages"; import type { UsageReference } from "../../agent/agent-sdk-types.js"; @@ -11,12 +10,20 @@ export interface UsageSource { label: string; icon?: string; discover(): Promise; + identify(input: unknown): Promise<{ key: string; label?: string } | null>; fetch(input: unknown): Promise; } -/** Owns running sources, account deduplication, and the five-minute fetch cache. */ +interface KnownReport { + source: UsageSource; + input: unknown; + label?: string; +} + +/** Owns account identity, the latest input for each report, and the five-minute fetch cache. */ export class UsageSourceRegistry { private readonly sources = new Map(); + private readonly known = new Map(); private readonly cache = new Map(); private readonly pending = new Map>(); @@ -32,55 +39,78 @@ export class UsageSourceRegistry { unregister(id: string): void { this.sources.delete(id); + for (const key of this.known.keys()) if (key.startsWith(`${id}:`)) this.known.delete(key); for (const key of this.cache.keys()) if (key.startsWith(`${id}:`)) this.cache.delete(key); } + async resolveReference(reference: UsageReference): Promise { + const source = this.sources.get(reference.source); + if (!source) return null; + return this.identify(source, reference.input); + } + + private async identify(source: UsageSource, input: unknown): Promise { + try { + const account = await source.identify(input); + if (account === null) return null; + if ( + typeof account !== "object" || + (account.label !== undefined && typeof account.label !== "string") + ) + throw new Error(`Invalid account identity from ${source.id}`); + if (typeof account.key !== "string" || !/^[A-Za-z0-9._-]{1,128}$/.test(account.key)) + throw new Error(`Invalid account key from ${source.id}`); + const id = `${source.id}:${account.key}`; + this.known.set(id, { source, input, label: account.label }); + return id; + } catch (error) { + const id = `${source.id}:error`; + this.cache.set(id, { at: this.now(), entry: this.errorEntry(source, id, error) }); + return id; + } + } + async listReports( - options: { - forceRefresh?: boolean; - references?: UsageReference[]; - } = {}, + options: { forceRefresh?: boolean; reportIds?: string[]; references?: UsageReference[] } = {}, ): Promise { + const ids = options.reportIds ?? (await this.discoverReportIds(options.references ?? [])); + return Promise.all( + [...new Set(ids)] + .filter((id) => this.known.has(id) || this.cache.has(id)) + .map((id) => this.fetchId(id, options.forceRefresh)), + ); + } + + private async discoverReportIds(references: UsageReference[]): Promise { const discovered = await Promise.all( [...this.sources.values()].map(async (source) => { try { const inputs = await source.discover(); if (!Array.isArray(inputs)) throw new Error("Usage discovery must return an array"); - return await Promise.all(inputs.map((input) => this.fetch(source.id, input, options))); + return (await Promise.all(inputs.map((input) => this.identify(source, input)))).filter( + (id): id is string => id !== null, + ); } catch (error) { - return [this.errorEntry(source, error)]; + const id = `${source.id}:error`; + this.cache.set(id, { at: this.now(), entry: this.errorEntry(source, id, error) }); + return [id]; } }), ); - const unique = new Map(); - const referenced = await Promise.all( - (options.references ?? []).map(async (reference) => { - if (!this.sources.has(reference.source)) return null; - return this.fetch(reference.source, reference.input, options); - }), - ); - for (const entry of [ - ...discovered.flat(), - ...referenced.filter((result): result is UsageReportEntry => result !== null), - ]) - unique.set(`${entry.sourceId}:${entry.report.account.key}`, entry); - return [...unique.values()]; - } - - fetchReference( - reference: UsageReference, - options: { forceRefresh?: boolean } = {}, - ): Promise { - return this.sources.has(reference.source) - ? this.fetch(reference.source, reference.input, options) - : Promise.resolve(null); + const live = await Promise.all(references.map((reference) => this.resolveReference(reference))); + return [...discovered.flat(), ...live.filter((id): id is string => id !== null)]; } // COMPAT(providerUsageList): added in v0.9.3, remove after 2027-03-26. async listLegacyUsage(): Promise<{ fetchedAt: string; providers: ProviderUsage[] }> { const reports = await this.listReports(); return { - fetchedAt: new Date(this.now()).toISOString(), + fetchedAt: reports.length + ? reports.reduce( + (oldest, entry) => (entry.fetchedAt < oldest ? entry.fetchedAt : oldest), + reports[0]!.fetchedAt, + ) + : new Date(this.now()).toISOString(), providers: reports.map((entry) => ({ providerId: entry.sourceId, displayName: entry.sourceLabel, @@ -94,49 +124,55 @@ export class UsageSourceRegistry { }; } - fetch( - sourceId: string, - input: unknown, - options: { forceRefresh?: boolean } = {}, - ): Promise { - const source = this.sources.get(sourceId); - if (!source) throw new Error(`Unknown usage source: ${sourceId}`); - const hash = createHash("sha256").update(JSON.stringify(input)).digest("hex"); - const key = `${sourceId}:${hash}`; - const cached = this.cache.get(key); - if (!options.forceRefresh && cached && this.now() - cached.at < this.ttlMs) + private fetchId(id: string, forceRefresh = false): Promise { + const known = this.known.get(id); + const cached = this.cache.get(id); + if (!known) return Promise.resolve(cached!.entry); + if (!forceRefresh && cached && this.now() - cached.at < this.ttlMs) return Promise.resolve(cached.entry); - const pending = this.pending.get(key); + const pending = this.pending.get(id); if (pending) return pending; const request = (async () => { let entry: UsageReportEntry; try { - const report = UsageReportSchema.parse(await source.fetch(input)); - entry = { sourceId, sourceLabel: source.label, icon: source.icon, report }; + const report = UsageReportSchema.parse(await known.source.fetch(known.input)); + entry = { + id, + sourceId: known.source.id, + sourceLabel: known.source.label, + icon: known.source.icon, + account: { label: known.label }, + fetchedAt: new Date(this.now()).toISOString(), + report, + }; } catch (error) { - entry = this.errorEntry(source, error, hash); + entry = this.errorEntry(known.source, id, error, known.label); } - const now = this.now(); - for (const [cachedKey, stored] of this.cache) { - if (now - stored.at >= this.ttlMs) this.cache.delete(cachedKey); - } - this.cache.set(key, { at: now, entry }); + this.cache.set(id, { at: this.now(), entry }); return entry; })(); - this.pending.set(key, request); + this.pending.set(id, request); void request.finally(() => { - if (this.pending.get(key) === request) this.pending.delete(key); + if (this.pending.get(id) === request) this.pending.delete(id); }); return request; } - private errorEntry(source: UsageSource, error: unknown, key = source.id): UsageReportEntry { + private errorEntry( + source: UsageSource, + id: string, + error: unknown, + label?: string, + fetchedAt = new Date(this.now()).toISOString(), + ): UsageReportEntry { return { + id, sourceId: source.id, sourceLabel: source.label, icon: source.icon, + account: { label }, + fetchedAt, report: { - account: { key }, status: "error", windows: [], error: error instanceof Error ? error.message : String(error), diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index 5549ecde69..f83040c8c9 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -513,12 +513,10 @@ export interface SessionOptions { invokePluginRpc(pluginId: string, method: string, input: unknown): Promise; listUsageReports(options?: { forceRefresh?: boolean; + reportIds?: string[]; references?: UsageReference[]; }): Promise; - fetchUsageReference( - reference: UsageReference, - options?: { forceRefresh?: boolean }, - ): Promise; + resolveUsageReference(reference: UsageReference): Promise; listLegacyUsage(): Promise<{ fetchedAt: string; providers: ProviderUsage[] }>; }; orchestrationSkills?: import("./orchestration-skills/index.js").OrchestrationSkills; @@ -3018,8 +3016,8 @@ export class Session { return this.usageSession.handleLegacyList(msg); case "usage.list_reports.request": return this.usageSession.handleListReports(msg); - case "agent.get_usage_report.request": - return this.usageSession.handleGetAgentReport(msg); + case "agent.resolve_usage_report.request": + return this.usageSession.handleResolveAgentReport(msg); default: return undefined; } diff --git a/packages/server/src/server/session/usage/usage-session.test.ts b/packages/server/src/server/session/usage/usage-session.test.ts index a7583cea90..7914af3872 100644 --- a/packages/server/src/server/session/usage/usage-session.test.ts +++ b/packages/server/src/server/session/usage/usage-session.test.ts @@ -14,9 +14,12 @@ test("collects live references for usage reports and resolves an agent report", } as AgentSession, }; const entry = { + id: "fixture:one", + account: {}, + fetchedAt: "2026-01-01T00:00:00.000Z", sourceId: "fixture", sourceLabel: "Fixture", - report: { account: { key: "one" }, status: "available" as const, windows: [] }, + report: { status: "available" as const, windows: [] }, }; const usage = new UsageSession({ emit: (message) => emitted.push(message), @@ -27,9 +30,9 @@ test("collects live references for usage reports and resolves an agent report", references.push(options.references); return [entry]; }, - async fetchUsageReference(value) { + async resolveUsageReference(value) { references.push(value); - return entry; + return entry.id; }, async listLegacyUsage() { return { fetchedAt: "2026-01-01T00:00:00.000Z", providers: [] }; @@ -39,15 +42,15 @@ test("collects live references for usage reports and resolves an agent report", }); await usage.handleListReports({ type: "usage.list_reports.request", requestId: "list" }); - await usage.handleGetAgentReport({ - type: "agent.get_usage_report.request", + await usage.handleResolveAgentReport({ + type: "agent.resolve_usage_report.request", requestId: "one", agentId: "one", }); expect(references).toEqual([[reference], reference]); expect(emitted.map((message) => message.type)).toEqual([ "usage.list_reports.response", - "agent.get_usage_report.response", + "agent.resolve_usage_report.response", ]); }); @@ -61,7 +64,7 @@ test("surfaces a legacy usage-list failure as an rpc_error envelope", async () = async listUsageReports() { return []; }, - async fetchUsageReference() { + async resolveUsageReference() { return null; }, async listLegacyUsage(): Promise { @@ -76,3 +79,24 @@ test("surfaces a legacy usage-list failure as an rpc_error envelope", async () = payload: { requestId: "u1", code: "provider_usage_list_failed" }, }); }); + +test("unknown agent returns agent_not_found", async () => { + const emitted: SessionOutboundMessage[] = []; + const usage = new UsageSession({ + emit: (message) => emitted.push(message), + listAgents: () => [], + getAgent: () => null, + runtime: { + listUsageReports: async () => [], + resolveUsageReference: async () => null, + listLegacyUsage: async () => ({ fetchedAt: "", providers: [] }), + }, + logger: pino({ level: "silent" }), + }); + await usage.handleResolveAgentReport({ + type: "agent.resolve_usage_report.request", + requestId: "missing", + agentId: "missing", + }); + expect(emitted[0]).toMatchObject({ type: "rpc_error", payload: { code: "agent_not_found" } }); +}); diff --git a/packages/server/src/server/session/usage/usage-session.ts b/packages/server/src/server/session/usage/usage-session.ts index 49ba32455d..6665a7454f 100644 --- a/packages/server/src/server/session/usage/usage-session.ts +++ b/packages/server/src/server/session/usage/usage-session.ts @@ -14,12 +14,10 @@ export interface UsageSessionOptions { runtime?: { listUsageReports(options: { forceRefresh?: boolean; + reportIds?: string[]; references: UsageReference[]; }): Promise; - fetchUsageReference( - reference: UsageReference, - options: { forceRefresh?: boolean }, - ): Promise; + resolveUsageReference(reference: UsageReference): Promise; listLegacyUsage(): Promise<{ fetchedAt: string; providers: ProviderUsage[] }>; }; logger: pino.Logger; @@ -33,15 +31,21 @@ export class UsageSession { ): Promise { try { if (!this.options.runtime) throw new Error("Plugin runtime is unavailable"); - const references = ( - await Promise.all( - this.options - .listAgents() - .map(async (agent) => agent.session?.getUsageReference?.().catch(() => null) ?? null), - ) - ).filter((reference): reference is UsageReference => reference !== null); + const references = + msg.reportIds === undefined + ? ( + await Promise.all( + this.options + .listAgents() + .map( + async (agent) => agent.session?.getUsageReference?.().catch(() => null) ?? null, + ), + ) + ).filter((reference): reference is UsageReference => reference !== null) + : []; const reports = await this.options.runtime.listUsageReports({ forceRefresh: msg.forceRefresh, + reportIds: msg.reportIds, references, }); this.options.emit({ @@ -53,24 +57,26 @@ export class UsageSession { } } - async handleGetAgentReport( - msg: Extract, + async handleResolveAgentReport( + msg: Extract, ): Promise { try { if (!this.options.runtime) throw new Error("Plugin runtime is unavailable"); - const reference = - (await this.options.getAgent(msg.agentId)?.session?.getUsageReference?.()) ?? null; - const entry = reference - ? await this.options.runtime.fetchUsageReference(reference, { - forceRefresh: msg.forceRefresh, - }) + const agent = this.options.getAgent(msg.agentId); + if (!agent) { + this.emitError(msg, new Error(`Agent not found: ${msg.agentId}`), "agent_not_found"); + return; + } + const reference = (await agent.session?.getUsageReference?.()) ?? null; + const reportId = reference + ? await this.options.runtime.resolveUsageReference(reference) : null; this.options.emit({ - type: "agent.get_usage_report.response", - payload: { requestId: msg.requestId, entry }, + type: "agent.resolve_usage_report.response", + payload: { requestId: msg.requestId, reportId }, }); } catch (error) { - this.emitError(msg, error, "agent_get_usage_report_failed"); + this.emitError(msg, error, "agent_resolve_usage_report_failed"); } } @@ -107,7 +113,7 @@ export class UsageSession { private emitError( msg: Extract< SessionInboundMessage, - { type: "usage.list_reports.request" | "agent.get_usage_report.request" } + { type: "usage.list_reports.request" | "agent.resolve_usage_report.request" } >, error: unknown, code: string,