From 42be71e5a7d9976057695dcdb54524ad7c52c885 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Sat, 26 Sep 2026 19:50:19 +0200 Subject: [PATCH] refactor(usage): own usage RPCs and validate session auth --- .../agent/providers/opencode-agent.test.ts | 29 ++++ .../server/agent/providers/opencode-agent.ts | 55 +++++--- .../server/agent/providers/pi/agent.test.ts | 13 ++ .../src/server/agent/providers/pi/agent.ts | 59 ++++++--- packages/server/src/server/session.ts | 88 ++---------- .../provider/provider-catalog-session.test.ts | 21 --- .../provider/provider-catalog-session.ts | 34 +---- .../session/usage/usage-session.test.ts | 78 +++++++++++ .../src/server/session/usage/usage-session.ts | 125 ++++++++++++++++++ 9 files changed, 329 insertions(+), 173 deletions(-) create mode 100644 packages/server/src/server/session/usage/usage-session.test.ts create mode 100644 packages/server/src/server/session/usage/usage-session.ts diff --git a/packages/server/src/server/agent/providers/opencode-agent.test.ts b/packages/server/src/server/agent/providers/opencode-agent.test.ts index 058c56d8da..9406ddbce4 100644 --- a/packages/server/src/server/agent/providers/opencode-agent.test.ts +++ b/packages/server/src/server/agent/providers/opencode-agent.test.ts @@ -392,6 +392,35 @@ describe("OpenCodeAgentClient adapter smoke tests", () => { rmSync(cwd, { recursive: true, force: true }); }); + test("usage reference ignores another auth entry with a different type", async () => { + const cwd = tmpCwd(); + const runtime = new TestOpenCodeHarness(); + runtime.enqueueClient(new TestOpenCodeClient()); + const client = new OpenCodeAgentClient( + logger, + { + env: { + OPENCODE_AUTH_CONTENT: JSON.stringify({ + openai: { type: "api", key: "other-key" }, + "opencode-go": { type: "api", key: "go-key" }, + }), + }, + }, + { serverManager: runtime, createClient: runtime.createClient }, + ); + const session = await client.createSession(buildConfig(cwd)); + try { + await session.setModel?.("opencode-go/qwen"); + expect(await session.getUsageReference?.()).toEqual({ + source: "opencode-go", + input: { apiKey: "go-key" }, + }); + } finally { + await session.close(); + rmSync(cwd, { recursive: true, force: true }); + } + }); + test("creates a session with valid id and provider", async () => { const cwd = tmpCwd(); const runtime = new TestOpenCodeHarness(); diff --git a/packages/server/src/server/agent/providers/opencode-agent.ts b/packages/server/src/server/agent/providers/opencode-agent.ts index e3100025a5..e92a424ee6 100644 --- a/packages/server/src/server/agent/providers/opencode-agent.ts +++ b/packages/server/src/server/agent/providers/opencode-agent.ts @@ -3365,9 +3365,28 @@ interface OpenCodeServerConnection { release: () => Promise; } -async function readOpenCodeUsageAuth( - env: NodeJS.ProcessEnv, -): Promise | null> { +const openCodeOAuthAuthSchema = z + .object({ + openai: z + .object({ + type: z.literal("oauth"), + access: z.string().min(1), + accountId: z.string().optional(), + }) + .passthrough() + .optional(), + }) + .passthrough(); +const openCodeGoAuthSchema = z + .object({ + "opencode-go": z + .object({ type: z.literal("api"), key: z.string().min(1) }) + .passthrough() + .optional(), + }) + .passthrough(); + +async function readOpenCodeUsageAuth(env: NodeJS.ProcessEnv): Promise { try { const content = env.OPENCODE_AUTH_CONTENT ?? @@ -3379,39 +3398,31 @@ async function readOpenCodeUsageAuth( ), "utf8", )); - const parsed: unknown = JSON.parse(content); - return parsed && typeof parsed === "object" && !Array.isArray(parsed) - ? (parsed as Record) - : null; + return JSON.parse(content); } catch { return null; } } -function resolveOpenCodeUsageReference( - model: string, - auth: Record | null, -): UsageReference | null { +function resolveOpenCodeUsageReference(model: string, auth: unknown): UsageReference | null { if (!auth) return null; if (model.startsWith("openai/")) { - const entry = auth.openai; - if (!entry || typeof entry !== "object") return null; - const credential = entry as Record; - if (credential.type !== "oauth" || typeof credential.access !== "string" || !credential.access) - return null; + const parsed = openCodeOAuthAuthSchema.safeParse(auth); + if (!parsed.success) return null; + const credential = parsed.data.openai; + if (!credential) return null; return { source: "codex", input: { accessToken: credential.access, - ...(typeof credential.accountId === "string" ? { accountId: credential.accountId } : {}), + ...(credential.accountId ? { accountId: credential.accountId } : {}), }, }; } - const entry = auth["opencode-go"]; - if (!entry || typeof entry !== "object") return null; - const credential = entry as Record; - if (credential.type !== "api" || typeof credential.key !== "string" || !credential.key) - return null; + const parsed = openCodeGoAuthSchema.safeParse(auth); + if (!parsed.success) return null; + const credential = parsed.data["opencode-go"]; + if (!credential) return null; return { source: "opencode-go", input: { apiKey: credential.key } }; } diff --git a/packages/server/src/server/agent/providers/pi/agent.test.ts b/packages/server/src/server/agent/providers/pi/agent.test.ts index a2fecf7092..7308185db0 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -241,6 +241,19 @@ test("Pi usage reference follows the current OAuth model and agent directory", a source: "claude", input: { accessToken: "claude-token" }, }); + writeFileSync( + path.join(agentDir, "auth.json"), + JSON.stringify({ + "openai-codex": { type: "oauth", access: "codex-token" }, + anthropic: { type: "api_key", key: "api-key" }, + }), + ); + runtime.state.model = { provider: "openai-codex", id: "gpt", name: "GPT" }; + expect(await session.getUsageReference?.()).toEqual({ + source: "codex", + input: { accessToken: "codex-token" }, + }); + runtime.state.model = { provider: "anthropic", id: "claude", name: "Claude" }; writeFileSync( path.join(agentDir, "auth.json"), JSON.stringify({ anthropic: { type: "api_key", key: "api-key" } }), diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 5ee0e6a997..5d7966367e 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -1,5 +1,6 @@ import { randomUUID } from "node:crypto"; import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { readFile } from "node:fs/promises"; import { homedir, tmpdir } from "node:os"; import { join, resolve as resolvePath } from "node:path"; import type { Logger } from "pino"; @@ -529,6 +530,27 @@ function resolvePiAgentDir(env: Record | undefined): string { return resolvePath(configured); } +const piCodexUsageAuthSchema = z + .object({ + "openai-codex": z + .object({ + type: z.literal("oauth"), + access: z.string().min(1), + accountId: z.string().optional(), + }) + .passthrough() + .optional(), + }) + .passthrough(); +const piClaudeUsageAuthSchema = z + .object({ + anthropic: z + .object({ type: z.literal("oauth"), access: z.string().min(1) }) + .passthrough() + .optional(), + }) + .passthrough(); + function readPiGlobalMcpConfig(env: Record | undefined): Record { const globalConfigPath = join(resolvePiAgentDir(env), "mcp.json"); if (!existsSync(globalConfigPath)) { @@ -1131,27 +1153,22 @@ export class PiRpcAgentSession implements AgentSession { if (provider === "anthropic") source = "claude"; if (!source || !provider) return null; try { - const credentials = JSON.parse( - readFileSync(join(resolvePiAgentDir(this.usageEnv), "auth.json"), "utf8"), - ) as Record; - const entry = credentials[provider]; - if (!entry || typeof entry !== "object") return null; - const credential = entry as Record; - if ( - credential.type !== "oauth" || - typeof credential.access !== "string" || - !credential.access - ) - return null; - return { - source, - input: { - accessToken: credential.access, - ...(source === "codex" && typeof credential.accountId === "string" && credential.accountId - ? { accountId: credential.accountId } - : {}), - }, - }; + const auth: unknown = JSON.parse( + await readFile(join(resolvePiAgentDir(this.usageEnv), "auth.json"), "utf8"), + ); + if (provider === "openai-codex") { + const credential = piCodexUsageAuthSchema.parse(auth)["openai-codex"]; + if (!credential) return null; + return { + source, + input: { + accessToken: credential.access, + ...(credential.accountId ? { accountId: credential.accountId } : {}), + }, + }; + } + const credential = piClaudeUsageAuthSchema.parse(auth).anthropic; + return credential ? { source, input: { accessToken: credential.access } } : null; } catch { return null; } diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index 893021ec67..5549ecde69 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -180,6 +180,7 @@ import { } from "./session/checkout/git-metadata-generator.js"; import { ScheduleSession } from "./session/schedule/schedule-session.js"; import { ProviderCatalogSession } from "./session/provider/provider-catalog-session.js"; +import { UsageSession } from "./session/usage/usage-session.js"; import { WorkspaceFilesSession } from "./session/files/workspace-files-session.js"; import { AgentConfigSession } from "./session/agent-config/agent-config-session.js"; import { ProjectConfigSession } from "./session/project-config/project-config-session.js"; @@ -790,6 +791,7 @@ export class Session { private readonly checkoutSession: CheckoutSession; private readonly scheduleSession: ScheduleSession; private readonly providerCatalogSession: ProviderCatalogSession; + private readonly usageSession: UsageSession; private readonly workspaceFilesSession: WorkspaceFilesSession; private readonly agentConfigSession: AgentConfigSession; private readonly projectConfigSession: ProjectConfigSession; @@ -995,10 +997,13 @@ export class Session { listDraftFeatures: (config) => this.agentManager.listDraftFeatures(config), }, providerSnapshotManager, - listLegacyUsage: async () => { - if (!pluginRuntime) throw new Error("Plugin runtime is unavailable"); - return pluginRuntime.listLegacyUsage(); - }, + logger: this.sessionLogger, + }); + this.usageSession = new UsageSession({ + emit: (msg) => this.emit(msg), + listAgents: () => this.agentManager.listAgents(), + getAgent: (agentId) => this.agentManager.getAgent(agentId), + runtime: pluginRuntime, logger: this.sessionLogger, }); this.agentConfigSession = new AgentConfigSession({ @@ -3010,85 +3015,16 @@ export class Session { case "provider_diagnostic_request": return this.providerCatalogSession.handleProviderDiagnosticRequest(msg); case "provider.usage.list.request": - return this.providerCatalogSession.handleProviderUsageListRequest(msg); + return this.usageSession.handleLegacyList(msg); case "usage.list_reports.request": - return this.handleUsageListReportsRequest(msg); + return this.usageSession.handleListReports(msg); case "agent.get_usage_report.request": - return this.handleAgentGetUsageReportRequest(msg); + return this.usageSession.handleGetAgentReport(msg); default: return undefined; } } - private async handleUsageListReportsRequest( - msg: Extract, - ): Promise { - try { - if (!this.pluginRuntime) throw new Error("Plugin runtime is unavailable"); - const references = ( - await Promise.all( - this.agentManager - .listAgents() - .map(async (agent) => agent.session?.getUsageReference?.().catch(() => null) ?? null), - ) - ).filter( - ( - reference, - ): reference is { - source: string; - input: import("@getpaseo/protocol/agent-types").JsonValue; - } => reference !== null, - ); - const reports = await this.pluginRuntime.listUsageReports({ - forceRefresh: msg.forceRefresh, - references, - }); - this.emit({ - type: "usage.list_reports.response", - payload: { requestId: msg.requestId, reports }, - }); - } catch (error) { - this.emit({ - type: "rpc_error", - payload: { - requestId: msg.requestId, - requestType: msg.type, - error: error instanceof Error ? error.message : String(error), - code: "usage_list_reports_failed", - }, - }); - } - } - - private async handleAgentGetUsageReportRequest( - msg: Extract, - ): Promise { - try { - if (!this.pluginRuntime) throw new Error("Plugin runtime is unavailable"); - const agent = this.agentManager.getAgent(msg.agentId); - const reference = (await agent?.session?.getUsageReference?.()) ?? null; - const entry = reference - ? await this.pluginRuntime.fetchUsageReference(reference, { - forceRefresh: msg.forceRefresh, - }) - : null; - this.emit({ - type: "agent.get_usage_report.response", - payload: { requestId: msg.requestId, entry }, - }); - } catch (error) { - this.emit({ - type: "rpc_error", - payload: { - requestId: msg.requestId, - requestType: msg.type, - error: error instanceof Error ? error.message : String(error), - code: "agent_get_usage_report_failed", - }, - }); - } - } - private dispatchTerminalMessage(msg: SessionInboundMessage): Promise | undefined { switch (msg.type) { case "start_workspace_script_request": diff --git a/packages/server/src/server/session/provider/provider-catalog-session.test.ts b/packages/server/src/server/session/provider/provider-catalog-session.test.ts index eff3af67b8..c1cec0b42a 100644 --- a/packages/server/src/server/session/provider/provider-catalog-session.test.ts +++ b/packages/server/src/server/session/provider/provider-catalog-session.test.ts @@ -23,7 +23,6 @@ interface MakeOptions { supportsCustomModeIcons?: boolean; supportsCompactProviderSnapshots?: boolean; snapshot?: Partial; - usage?: () => Promise<{ fetchedAt: string; providers: [] }>; host?: Partial; } @@ -72,8 +71,6 @@ function makeSubsystem(options: MakeOptions = {}) { const subsystem = new ProviderCatalogSession({ host, providerSnapshotManager, - listLegacyUsage: - options.usage ?? (async () => ({ fetchedAt: "2026-01-01T00:00:00.000Z", providers: [] })), logger: pino({ level: "silent" }), }); function pushSnapshotChange( @@ -299,23 +296,6 @@ describe("ProviderCatalogSession", () => { }); }); - it("surfaces a usage-list failure as an rpc_error envelope", async () => { - const { subsystem, emitted } = makeSubsystem({ - usage: async () => { - throw new Error("quota service down"); - }, - }); - - await subsystem.handleProviderUsageListRequest({ - type: "provider.usage.list.request", - requestId: "u1", - }); - - const err = findByType(emitted, "rpc_error"); - expect(err?.payload.code).toBe("provider_usage_list_failed"); - expect(err?.payload.requestId).toBe("u1"); - }); - it("surfaces a feature-list failure inline, not as an rpc_error", async () => { const { subsystem, emitted } = makeSubsystem({ host: { @@ -386,7 +366,6 @@ it("announces shared content without retransmitting models or hashing discovery new ProviderCatalogSession({ providerSnapshotManager: manager, logger: pino({ level: "silent" }), - listLegacyUsage: async () => ({ fetchedAt: "2026-01-01T00:00:00.000Z", providers: [] }), host: { emit(message) { emitted.push(message); diff --git a/packages/server/src/server/session/provider/provider-catalog-session.ts b/packages/server/src/server/session/provider/provider-catalog-session.ts index a57cce5bbe..af04199bca 100644 --- a/packages/server/src/server/session/provider/provider-catalog-session.ts +++ b/packages/server/src/server/session/provider/provider-catalog-session.ts @@ -19,7 +19,6 @@ import { type ProviderSnapshotEntry, } from "../../agent/agent-sdk-types.js"; import type { ProviderAvailability } from "../../agent/agent-manager.js"; -import type { ProviderUsage } from "@getpaseo/protocol/messages"; import { expandTilde } from "../../../utils/path.js"; // COMPAT(customModeIcons): the only mode icons known to clients before v0.1.84. Any @@ -56,13 +55,12 @@ export interface ProviderCatalogSessionHost { export interface ProviderCatalogSessionOptions { host: ProviderCatalogSessionHost; providerSnapshotManager: ProviderSnapshotManager; - listLegacyUsage(): Promise<{ fetchedAt: string; providers: ProviderUsage[] }>; logger: pino.Logger; } /** * A client's provider catalog surface: model / mode / feature listing, the providers - * snapshot push + pull, provider diagnostics, and usage. The snapshot PUSH (start) and + * snapshot push + pull, and provider diagnostics. The snapshot PUSH (start) and * every PULL handler gate visibility and downgrade mode icons through the SAME predicates, * so an older client sees one consistent provider set across both paths — the COMPAT * invariant the shell could only enforce by code proximity before this carve. @@ -70,14 +68,12 @@ export interface ProviderCatalogSessionOptions { export class ProviderCatalogSession { private readonly host: ProviderCatalogSessionHost; private readonly providerSnapshotManager: ProviderSnapshotManager; - private readonly listLegacyUsage: ProviderCatalogSessionOptions["listLegacyUsage"]; private readonly logger: pino.Logger; private unsubscribeSnapshotEvents: (() => void) | null = null; constructor(options: ProviderCatalogSessionOptions) { this.host = options.host; this.providerSnapshotManager = options.providerSnapshotManager; - this.listLegacyUsage = options.listLegacyUsage; this.logger = options.logger; } @@ -486,34 +482,6 @@ export class ProviderCatalogSession { }); } } - - async handleProviderUsageListRequest( - msg: Extract, - ): Promise { - try { - const usage = await this.listLegacyUsage(); - this.host.emit({ - type: "provider.usage.list.response", - payload: { - requestId: msg.requestId, - fetchedAt: usage.fetchedAt, - providers: usage.providers, - }, - }); - } catch (error) { - const err = error instanceof Error ? error : new Error(String(error)); - this.logger.error({ err }, "Failed to list provider usage"); - this.host.emit({ - type: "rpc_error", - payload: { - requestId: msg.requestId, - requestType: msg.type, - error: `Failed to list provider usage: ${err.message}`, - code: "provider_usage_list_failed", - }, - }); - } - } } function resolveCatalogRequestCwd(cwd?: string | null): string | 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 new file mode 100644 index 0000000000..a7583cea90 --- /dev/null +++ b/packages/server/src/server/session/usage/usage-session.test.ts @@ -0,0 +1,78 @@ +import pino from "pino"; +import { expect, test } from "vitest"; +import type { AgentSession } from "../../agent/agent-sdk-types.js"; +import type { SessionOutboundMessage } from "../../messages.js"; +import { UsageSession } from "./usage-session.js"; + +test("collects live references for usage reports and resolves an agent report", async () => { + const emitted: SessionOutboundMessage[] = []; + const references: unknown[] = []; + const reference = { source: "fixture", input: { account: "one" } }; + const agent = { + session: { + getUsageReference: async () => reference, + } as AgentSession, + }; + const entry = { + sourceId: "fixture", + sourceLabel: "Fixture", + report: { account: { key: "one" }, status: "available" as const, windows: [] }, + }; + const usage = new UsageSession({ + emit: (message) => emitted.push(message), + listAgents: () => [agent, { session: null }], + getAgent: (id) => (id === "one" ? agent : null), + runtime: { + async listUsageReports(options) { + references.push(options.references); + return [entry]; + }, + async fetchUsageReference(value) { + references.push(value); + return entry; + }, + async listLegacyUsage() { + return { fetchedAt: "2026-01-01T00:00:00.000Z", providers: [] }; + }, + }, + logger: pino({ level: "silent" }), + }); + + await usage.handleListReports({ type: "usage.list_reports.request", requestId: "list" }); + await usage.handleGetAgentReport({ + type: "agent.get_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", + ]); +}); + +test("surfaces a legacy usage-list failure as an rpc_error envelope", async () => { + const emitted: SessionOutboundMessage[] = []; + const usage = new UsageSession({ + emit: (message) => emitted.push(message), + listAgents: () => [], + getAgent: () => null, + runtime: { + async listUsageReports() { + return []; + }, + async fetchUsageReference() { + return null; + }, + async listLegacyUsage(): Promise { + throw new Error("quota service down"); + }, + }, + logger: pino({ level: "silent" }), + }); + await usage.handleLegacyList({ type: "provider.usage.list.request", requestId: "u1" }); + expect(emitted[0]).toMatchObject({ + type: "rpc_error", + payload: { requestId: "u1", code: "provider_usage_list_failed" }, + }); +}); diff --git a/packages/server/src/server/session/usage/usage-session.ts b/packages/server/src/server/session/usage/usage-session.ts new file mode 100644 index 0000000000..49ba32455d --- /dev/null +++ b/packages/server/src/server/session/usage/usage-session.ts @@ -0,0 +1,125 @@ +import type pino from "pino"; +import type { ProviderUsage, UsageReportEntry } from "@getpaseo/protocol/messages"; +import type { AgentSession, UsageReference } from "../../agent/agent-sdk-types.js"; +import type { SessionInboundMessage, SessionOutboundMessage } from "../../messages.js"; + +interface UsageAgent { + session: AgentSession | null; +} + +export interface UsageSessionOptions { + emit(message: SessionOutboundMessage): void; + listAgents(): UsageAgent[]; + getAgent(agentId: string): UsageAgent | null; + runtime?: { + listUsageReports(options: { + forceRefresh?: boolean; + references: UsageReference[]; + }): Promise; + fetchUsageReference( + reference: UsageReference, + options: { forceRefresh?: boolean }, + ): Promise; + listLegacyUsage(): Promise<{ fetchedAt: string; providers: ProviderUsage[] }>; + }; + logger: pino.Logger; +} + +export class UsageSession { + constructor(private readonly options: UsageSessionOptions) {} + + async handleListReports( + msg: Extract, + ): 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 reports = await this.options.runtime.listUsageReports({ + forceRefresh: msg.forceRefresh, + references, + }); + this.options.emit({ + type: "usage.list_reports.response", + payload: { requestId: msg.requestId, reports }, + }); + } catch (error) { + this.emitError(msg, error, "usage_list_reports_failed"); + } + } + + async handleGetAgentReport( + 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, + }) + : null; + this.options.emit({ + type: "agent.get_usage_report.response", + payload: { requestId: msg.requestId, entry }, + }); + } catch (error) { + this.emitError(msg, error, "agent_get_usage_report_failed"); + } + } + + // COMPAT(providerUsageList): added in v0.9.3, remove after 2027-03-26. + async handleLegacyList( + msg: Extract, + ): Promise { + try { + if (!this.options.runtime) throw new Error("Plugin runtime is unavailable"); + const usage = await this.options.runtime.listLegacyUsage(); + this.options.emit({ + type: "provider.usage.list.response", + payload: { + requestId: msg.requestId, + fetchedAt: usage.fetchedAt, + providers: usage.providers, + }, + }); + } catch (error) { + const err = error instanceof Error ? error : new Error(String(error)); + this.options.logger.error({ err }, "Failed to list provider usage"); + this.options.emit({ + type: "rpc_error", + payload: { + requestId: msg.requestId, + requestType: msg.type, + error: `Failed to list provider usage: ${err.message}`, + code: "provider_usage_list_failed", + }, + }); + } + } + + private emitError( + msg: Extract< + SessionInboundMessage, + { type: "usage.list_reports.request" | "agent.get_usage_report.request" } + >, + error: unknown, + code: string, + ): void { + this.options.emit({ + type: "rpc_error", + payload: { + requestId: msg.requestId, + requestType: msg.type, + error: error instanceof Error ? error.message : String(error), + code, + }, + }); + } +}