From 88b7e096a6808b73bafa3759da3ece5514966caf Mon Sep 17 00:00:00 2001 From: Qichang Zheng Date: Sun, 18 Oct 2026 08:33:56 +0800 Subject: [PATCH] refactor(usage): name the provider usage types and pass options objects Follow the server conventions: the provider environments list returns a named ProviderUsageEnvironment, and the registry's discover() and mergeReports() and the e2e injectedUsageSources() helper take a single options object. No behavior change. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/server/agent/agent-sdk-types.ts | 10 ++++ .../server/agent/provider-snapshot-manager.ts | 12 +--- .../server/plugins/usage-sources.e2e.test.ts | 52 +++++++++++------- .../src/server/plugins/usage-sources/index.ts | 55 +++++++++++-------- 4 files changed, 76 insertions(+), 53 deletions(-) diff --git a/packages/server/src/server/agent/agent-sdk-types.ts b/packages/server/src/server/agent/agent-sdk-types.ts index 34d36c9e16..81bc6d9003 100644 --- a/packages/server/src/server/agent/agent-sdk-types.ts +++ b/packages/server/src/server/agent/agent-sdk-types.ts @@ -661,6 +661,16 @@ export interface AgentUsageSession { sessionKey: string; } +/** An enabled provider that sets its own env, as its sessions report it for usage. */ +export interface ProviderUsageEnvironment { + /** The provider a custom provider extends, or its own id. */ + provider: string; + /** The provider's label, used to name the logins found through it. */ + label: string; + /** The environment the provider launches with. */ + env: Record; +} + export interface AgentSession { usageSession?(): AgentUsageSession | null; readonly provider: AgentProvider; diff --git a/packages/server/src/server/agent/provider-snapshot-manager.ts b/packages/server/src/server/agent/provider-snapshot-manager.ts index db1d748172..b3e130cdcb 100644 --- a/packages/server/src/server/agent/provider-snapshot-manager.ts +++ b/packages/server/src/server/agent/provider-snapshot-manager.ts @@ -19,6 +19,7 @@ import { type AgentProvider, type FetchCatalogOptions, type ProviderSnapshotEntry, + type ProviderUsageEnvironment, } from "./agent-sdk-types.js"; import { raceProviderRefreshAbort, @@ -361,15 +362,8 @@ export class ProviderSnapshotManager { return this.createAgentManagerState(this.generation.definitions, this.providerClients); } - /** - * Enabled providers that set their own env, as their sessions report them for usage: the - * provider a custom provider extends, its own label, and the environment it launches with. - */ - listProviderEnvironments(): Array<{ - provider: AgentProvider; - label: string; - env: Record; - }> { + /** Enabled providers that set their own env, as their sessions report them for usage. */ + listProviderEnvironments(): ProviderUsageEnvironment[] { return Object.entries(this.generation.definitions).flatMap(([provider, definition]) => { const runtimeSettings = definition.configuration?.runtimeSettings; if (!definition.enabled || !runtimeSettings?.env) return []; 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 ea2c015ad8..d21e36ea3f 100644 --- a/packages/server/src/server/plugins/usage-sources.e2e.test.ts +++ b/packages/server/src/server/plugins/usage-sources.e2e.test.ts @@ -18,11 +18,23 @@ async function json(file: string, value: unknown) { } /** The real Claude and Codex readers, compiled as built-ins, with home, env and HTTP injected. */ -async function injectedUsageSources( - root: string, - lookup: { home: string; env: Record }, - fetchApi: string, -): Promise { +interface UsageLookup { + home: string; + env: Record; +} + +interface InjectedUsageSourcesOptions { + root: string; + lookup: UsageLookup; + /** Source of an `async function fetchApi(url, init)` the readers call instead of the network. */ + fetchApi: string; +} + +async function injectedUsageSources({ + root, + lookup, + fetchApi, +}: InjectedUsageSourcesOptions): Promise { const plugins = path.join(root, "plugins"); for (const source of ["codex", "claude"]) { const directory = path.join(plugins, `${source}-usage-source`); @@ -233,10 +245,10 @@ test("built-in sources deduplicate cross-harness accounts and fall back without JSON.stringify({ access: "fixture-omp-claude", expires: Date.now() + 60_000 }), ); db.close(); - const builtinPlugins = await injectedUsageSources( - root, - { home, env }, - `async function fetchApi(url, init) { + const builtinPlugins = await injectedUsageSources({ + root: root, + lookup: { home, env }, + fetchApi: `async function fetchApi(url, init) { const token = new Headers(init.headers).get('Authorization'); if (String(url).endsWith('/profile')) { const second = token === 'Bearer fixture-pi-claude'; @@ -246,7 +258,7 @@ test("built-in sources deduplicate cross-harness accounts and fall back without const used = { 'Bearer fixture-codex': 11, 'Bearer fixture-opencode': 22, 'Bearer fixture-pi-codex': 33, 'Bearer fixture-omp-codex': 44, 'Bearer fixture-omp-claude': 66, 'Bearer fixture-pi-claude': 77 }[token] || 55; return new Response(JSON.stringify({ rate_limit: { primary_window: { used_percent: used } }, five_hour: { utilization: used } })); }`, - ); + }); const daemon = await createTestPaseoDaemon({ pluginsEnabled: false, builtinPlugins }); const client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` }); try { @@ -336,10 +348,10 @@ test("#3589: host usage lists each enabled provider's own login without an agent await json(path.join(dir("codex-work"), "auth.json"), { tokens: { access_token: "fixture-codex-work", account_id: "chatgpt-work" }, }); - const builtinPlugins = await injectedUsageSources( - root, - { home, env: {} }, - `async function fetchApi(url, init) { + const builtinPlugins = await injectedUsageSources({ + root: root, + lookup: { home, env: {} }, + fetchApi: `async function fetchApi(url, init) { const token = new Headers(init.headers).get('Authorization'); if (String(url).endsWith('/profile')) { const uuid = { 'Bearer fixture-work': 'work', 'Bearer fixture-off': 'off' }[token] || 'default'; @@ -347,7 +359,7 @@ test("#3589: host usage lists each enabled provider's own login without an agent } return new Response(JSON.stringify({ rate_limit: { primary_window: { used_percent: 1 } }, five_hour: { utilization: 1 } })); }`, - ); + }); const daemon = await createTestPaseoDaemon({ pluginsEnabled: false, builtinPlugins, @@ -427,16 +439,16 @@ test("#3589: a provider whose agents talk to another API host has no login in th await json(path.join(root, "claude-work", ".credentials.json"), { claudeAiOauth: { accessToken: "fixture-work" }, }); - const builtinPlugins = await injectedUsageSources( - root, - { home: path.join(root, "home"), env: {} }, - `async function fetchApi(url) { + const builtinPlugins = await injectedUsageSources({ + root: root, + lookup: { home: path.join(root, "home"), env: {} }, + fetchApi: `async function fetchApi(url) { if (String(url).endsWith('/profile')) { return new Response(JSON.stringify({ account: { uuid: 'work' }, organization: { uuid: 'org' } })); } return new Response(JSON.stringify({ rate_limit: { primary_window: { used_percent: 1 } }, five_hour: { utilization: 1 } })); }`, - ); + }); const daemon = await createTestPaseoDaemon({ pluginsEnabled: false, builtinPlugins, diff --git a/packages/server/src/server/plugins/usage-sources/index.ts b/packages/server/src/server/plugins/usage-sources/index.ts index f04092c722..8f88fd0caf 100644 --- a/packages/server/src/server/plugins/usage-sources/index.ts +++ b/packages/server/src/server/plugins/usage-sources/index.ts @@ -1,5 +1,5 @@ import { createHash } from "node:crypto"; -import type { AgentUsageSession } from "../../agent/agent-sdk-types.js"; +import type { AgentUsageSession, ProviderUsageEnvironment } from "../../agent/agent-sdk-types.js"; import type { UsageScope } from "@getpaseo/plugin/server/usage"; import { z } from "zod"; import type { Logger } from "pino"; @@ -33,11 +33,19 @@ interface KnownReport { export interface AgentUsageLookup { hasAgent(id: string): boolean; usageSession(id: string): AgentUsageSession | null; - /** - * Enabled providers that set their own env, with the label to show for their logins. Host - * listings discover these like a session. - */ - providerEnvironments?(): Array & { label: string }>; + /** Enabled providers that set their own env. Host listings discover these like a session. */ + providerEnvironments?(): ProviderUsageEnvironment[]; +} + +interface DiscoverOptions { + scope: UsageScope; + /** Names the logins found in this scope; defaults to the account's harness or the source label. */ + harness?: string; +} + +interface MergeReportsOptions { + into: Map; + reports: Iterable<[string, KnownReport]>; } export interface ListUsageReportsOptions { @@ -105,7 +113,7 @@ export class UsageSourceRegistry { // its agents runs. The host's own logins come first in each account's fallback order. // A login is named after the provider it was found through: a login that can't be read // (an expired token has no account) is otherwise one more card that looks like the rest. - const scopes: Array<{ scope: UsageScope; harness?: string }> = [ + const scopes: DiscoverOptions[] = [ { scope: { kind: "global" } }, ...(this.agents.providerEnvironments?.() ?? []).map(({ provider, label, env }) => ({ scope: { kind: "session" as const, provider, env }, @@ -113,10 +121,8 @@ export class UsageSourceRegistry { })), ]; const defaults = new Map(); - for (const discovered of await Promise.all( - scopes.map(({ scope, harness }) => this.discover(scope, harness)), - )) - mergeReports(defaults, discovered); + for (const discovered of await Promise.all(scopes.map((entry) => this.discover(entry)))) + mergeReports({ into: defaults, reports: discovered }); this.defaults = defaults; this.mergeKnown(); ids = [...this.known.keys()]; @@ -153,10 +159,12 @@ export class UsageSourceRegistry { previous.model === session.model; if (sameScope) return [...previous.reports.keys()]; const reports = await this.discover({ - kind: "session", - provider: session.provider, - model: session.model, - env: session.env, + scope: { + kind: "session", + provider: session.provider, + model: session.model, + env: session.env, + }, }); // A query that finishes after a resume must not publish the previous launch's mapping. if (this.agents.usageSession(agentId)?.sessionKey !== session.sessionKey) return []; @@ -176,14 +184,16 @@ export class UsageSourceRegistry { this.defaults, ...[...this.byAgent.values()].map((agent) => agent.reports), ]) { - mergeReports( - this.known, - [...reports].filter(([, report]) => this.sources.get(report.source.id) === report.source), - ); + mergeReports({ + into: this.known, + reports: [...reports].filter( + ([, report]) => this.sources.get(report.source.id) === report.source, + ), + }); } } - private async discover(scope: UsageScope, harness?: string): Promise> { + private async discover({ scope, harness }: DiscoverOptions): Promise> { const discovered = await Promise.all( [...this.sources.values()].map(async (source) => { const reports = new Map(); @@ -355,10 +365,7 @@ function legacyError(report: UsageReport, now: number): string | null { } // An account found in several scopes is one report; its logins keep discovery order, once each. -function mergeReports( - into: Map, - reports: Iterable<[string, KnownReport]>, -): void { +function mergeReports({ into, reports }: MergeReportsOptions): void { for (const [id, report] of reports) { const known = into.get(id); if (!known) {