From cd324f1bdcffa6d46bb894e84943b5956601be1a Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Sat, 26 Sep 2026 19:25:31 +0200 Subject: [PATCH] feat(usage): resolve reports from active session accounts Resolve source inputs at fetch time so model and credential changes select the account each session spends. Keep the legacy discovered-account RPC for existing clients. Sync OpenCode model events outside foreground turns so usage follows model switches. --- packages/client/src/daemon-client.ts | 17 ++++ packages/plugin/src/server/provider.ts | 23 +++++ packages/protocol/src/messages.ts | 15 ++++ packages/server/src/plugins/index.ts | 6 ++ .../plugins/opencode-go-usage-source/icon.svg | 1 + .../opencode-go-usage-source/index.server.ts | 14 +++ .../paseo-plugin.json | 1 + .../server/usage.test.ts | 34 +++++++ .../opencode-go-usage-source/server/usage.ts | 57 ++++++++++++ .../opencode-go-usage-source/shared/input.ts | 4 + .../src/server/agent/agent-sdk-types.ts | 2 + .../src/server/agent/plugin-provider.test.ts | 18 ++++ .../src/server/agent/plugin-provider.ts | 40 +++++++-- .../server/agent/providers/acp-agent.test.ts | 7 ++ .../src/server/agent/providers/acp-agent.ts | 4 + .../agent/providers/claude/agent.test.ts | 19 ++++ .../server/agent/providers/claude/agent.ts | 10 +++ .../providers/codex-app-server-agent.test.ts | 38 ++++++++ .../agent/providers/codex-app-server-agent.ts | 15 ++++ .../agent/providers/opencode-agent.test.ts | 59 +++++++++++++ .../server/agent/providers/opencode-agent.ts | 88 ++++++++++++++++++- .../server/agent/providers/pi/agent.test.ts | 53 ++++++++++- .../src/server/agent/providers/pi/agent.ts | 31 +++++++ .../authorization/operation-permissions.ts | 2 + packages/server/src/server/plugins/index.ts | 12 ++- .../session-usage-reference.e2e.test.ts | 33 +++++++ .../session-usage-reference/index.server.ts | 72 +++++++++++++++ .../session-usage-reference/paseo-plugin.json | 1 + .../plugins/usage-sources/index.test.ts | 18 ++++ .../src/server/plugins/usage-sources/index.ts | 27 +++++- packages/server/src/server/session.ts | 59 ++++++++++++- 31 files changed, 767 insertions(+), 13 deletions(-) create mode 100644 packages/server/src/plugins/opencode-go-usage-source/icon.svg create mode 100644 packages/server/src/plugins/opencode-go-usage-source/index.server.ts create mode 100644 packages/server/src/plugins/opencode-go-usage-source/paseo-plugin.json create mode 100644 packages/server/src/plugins/opencode-go-usage-source/server/usage.test.ts create mode 100644 packages/server/src/plugins/opencode-go-usage-source/server/usage.ts create mode 100644 packages/server/src/plugins/opencode-go-usage-source/shared/input.ts create mode 100644 packages/server/src/server/plugins/session-usage-reference.e2e.test.ts create mode 100644 packages/server/src/server/plugins/test-fixtures/session-usage-reference/index.server.ts create mode 100644 packages/server/src/server/plugins/test-fixtures/session-usage-reference/paseo-plugin.json diff --git a/packages/client/src/daemon-client.ts b/packages/client/src/daemon-client.ts index 6a909f5b6a..2a1a793b65 100644 --- a/packages/client/src/daemon-client.ts +++ b/packages/client/src/daemon-client.ts @@ -95,6 +95,7 @@ import type { ProviderDiagnosticResponseMessage, ProviderUsageListResponseMessage, UsageListReportsResponseMessage, + AgentGetUsageReportResponseMessage, DaemonGetStatusResponse, DaemonGetPairingOfferResponse, DaemonConfigReloadResponse, @@ -561,6 +562,7 @@ type RefreshProvidersSnapshotPayload = RefreshProvidersSnapshotResponseMessage[" type ProviderDiagnosticPayload = ProviderDiagnosticResponseMessage["payload"]; type ProviderUsageListPayload = ProviderUsageListResponseMessage["payload"]; type UsageListReportsPayload = UsageListReportsResponseMessage["payload"]; +type AgentGetUsageReportPayload = AgentGetUsageReportResponseMessage["payload"]; type DaemonStatusPayload = DaemonGetStatusResponse["payload"]; type DaemonPairingOfferPayload = DaemonGetPairingOfferResponse["payload"]; type DiagnosticsPayload = DiagnosticsResponse["payload"]; @@ -5250,6 +5252,21 @@ export class DaemonClient { }); } + async getAgentUsageReport(options: { + agentId: string; + requestId?: string; + forceRefresh?: boolean; + }): Promise { + return this.sendNamespacedCorrelatedSessionRequest({ + requestId: options.requestId, + message: { + type: "agent.get_usage_report.request", + agentId: options.agentId, + forceRefresh: options.forceRefresh, + }, + }); + } + async listCommands(options: ListCommandsOptions): Promise; async listCommands(agentId: string, requestId?: string): Promise; async listCommands( diff --git a/packages/plugin/src/server/provider.ts b/packages/plugin/src/server/provider.ts index c65a814753..a72477d19d 100644 --- a/packages/plugin/src/server/provider.ts +++ b/packages/plugin/src/server/provider.ts @@ -18,6 +18,7 @@ export const PROVIDER_CAPABILITIES = [ "session.revert.files", "session.subsession", "session.unarchive", + "session.usage_reference", "permission", "permission.tool_policy", "timeline.plugin", @@ -227,6 +228,7 @@ export type ProviderInput = } | { type: "session.prompt"; sessionId: string; prompt: ProviderPrompt } | { type: "session.interrupt"; requestId: string; sessionId: string } + | { type: "session.usage_reference"; requestId: string; sessionId: string } | { type: "session.permission"; sessionId: string; @@ -526,6 +528,11 @@ export type ProviderEvent = | { type: "catalog"; requestId: string; catalog: ProviderCatalog } | { type: "sessions"; requestId: string; sessions: ProviderSessionSummary[] } | { type: "request.completed"; requestId: string } + | { + type: "usage_reference"; + requestId: string; + reference: { source: string; input: JsonValue } | null; + } | { type: "request.failed"; requestId: string; error: ProviderError } | { type: "session.opened"; @@ -609,6 +616,8 @@ export function requiredProviderCapabilities(input: ProviderInput): readonly Pro case "session.interrupt": case "session.close": return []; + case "session.usage_reference": + return ["session.usage_reference"]; case "sessions": return ["session.list"]; case "session.open": { @@ -887,6 +896,13 @@ export const ProviderInputSchema: z.ZodType = z.discriminatedUnio z .object({ type: z.literal("session.interrupt"), requestId: idSchema, sessionId: idSchema }) .strict(), + z + .object({ + type: z.literal("session.usage_reference"), + requestId: idSchema, + sessionId: idSchema, + }) + .strict(), z .object({ type: z.literal("session.permission"), @@ -1272,6 +1288,13 @@ export const ProviderEventSchema: z.ZodType = z.discriminatedUnio }) .strip(), z.object({ type: z.literal("request.completed"), requestId: idSchema }).strip(), + z + .object({ + type: z.literal("usage_reference"), + requestId: idSchema, + reference: z.object({ source: idSchema, input: z.json() }).nullable(), + }) + .strip(), z .object({ type: z.literal("request.failed"), requestId: idSchema, error: providerErrorSchema }) .strip(), diff --git a/packages/protocol/src/messages.ts b/packages/protocol/src/messages.ts index 25e31983d2..84fe0bb2df 100644 --- a/packages/protocol/src/messages.ts +++ b/packages/protocol/src/messages.ts @@ -1779,6 +1779,12 @@ export const UsageListReportsRequestMessageSchema = z.object({ requestId: z.string(), forceRefresh: z.boolean().optional(), }); +export const AgentGetUsageReportRequestMessageSchema = z.object({ + type: z.literal("agent.get_usage_report.request"), + requestId: z.string(), + agentId: z.string(), + forceRefresh: z.boolean().optional(), +}); export const ResumeAgentRequestMessageSchema = z.object({ type: z.literal("resume_agent_request"), @@ -3239,6 +3245,7 @@ export const SessionInboundMessageSchema = z.discriminatedUnion("type", [ ProviderDiagnosticRequestMessageSchema, ProviderUsageListRequestMessageSchema, UsageListReportsRequestMessageSchema, + AgentGetUsageReportRequestMessageSchema, ResumeAgentRequestMessageSchema, ImportAgentRequestMessageSchema, RefreshAgentRequestMessageSchema, @@ -6255,6 +6262,10 @@ 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() }), +}); const AgentSlashCommandSchema = z.object({ name: z.string(), @@ -6942,6 +6953,7 @@ export const SessionOutboundMessageSchema = z.discriminatedUnion("type", [ ProviderDiagnosticResponseMessageSchema, ProviderUsageListResponseMessageSchema, UsageListReportsResponseMessageSchema, + AgentGetUsageReportResponseMessageSchema, ListCommandsResponseSchema, ListTerminalsResponseSchema, TerminalsChangedSchema, @@ -7123,6 +7135,9 @@ 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 ProviderUsageStatus = z.infer; export type ProviderUsage = z.infer; export type ProviderUsageWindow = z.infer; diff --git a/packages/server/src/plugins/index.ts b/packages/server/src/plugins/index.ts index 4bbf849975..fe9b8e824a 100644 --- a/packages/server/src/plugins/index.ts +++ b/packages/server/src/plugins/index.ts @@ -8,6 +8,7 @@ import grokUsageSource from "./grok-usage-source/index.server.js"; import kimiUsageSource from "./kimi-usage-source/index.server.js"; import minimaxUsageSource from "./minimax-usage-source/index.server.js"; import zaiUsageSource from "./zai-usage-source/index.server.js"; +import opencodeGoUsageSource from "./opencode-go-usage-source/index.server.js"; export interface InternalPlugin { id: string; @@ -56,4 +57,9 @@ export const INTERNAL_PLUGINS: InternalPlugin[] = [ directory: fileURLToPath(new URL("./zai-usage-source/", import.meta.url)), contribute: zaiUsageSource, }, + { + id: "opencode-go-usage-source", + directory: fileURLToPath(new URL("./opencode-go-usage-source/", import.meta.url)), + contribute: opencodeGoUsageSource, + }, ]; diff --git a/packages/server/src/plugins/opencode-go-usage-source/icon.svg b/packages/server/src/plugins/opencode-go-usage-source/icon.svg new file mode 100644 index 0000000000..2099d47ec4 --- /dev/null +++ b/packages/server/src/plugins/opencode-go-usage-source/icon.svg @@ -0,0 +1 @@ + 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 new file mode 100644 index 0000000000..c048c9dcb5 --- /dev/null +++ b/packages/server/src/plugins/opencode-go-usage-source/index.server.ts @@ -0,0 +1,14 @@ +import type { PluginServerContext } from "@getpaseo/plugin/server"; +import { inputSchema } from "./shared/input.js"; +import { fetchUsage } from "./server/usage.js"; + +export default function contribute(server: PluginServerContext) { + server.registerUsageSource({ + id: "opencode-go", + label: "OpenCode Go", + icon: "icon.svg", + input: inputSchema, + fetch: fetchUsage, + }); + return () => {}; +} diff --git a/packages/server/src/plugins/opencode-go-usage-source/paseo-plugin.json b/packages/server/src/plugins/opencode-go-usage-source/paseo-plugin.json new file mode 100644 index 0000000000..587e662d4f --- /dev/null +++ b/packages/server/src/plugins/opencode-go-usage-source/paseo-plugin.json @@ -0,0 +1 @@ +{ "id": "opencode-go-usage-source", "requirements": { "paseo": ">=0.9.2" } } 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 new file mode 100644 index 0000000000..57453396f2 --- /dev/null +++ b/packages/server/src/plugins/opencode-go-usage-source/server/usage.test.ts @@ -0,0 +1,34 @@ +import { expect, test } from "vitest"; +import { fetchUsage } from "./usage.js"; + +test("reads the OpenCode Go usage windows", async () => { + const requests: Array<{ url: string; authorization: string | null }> = []; + const fetchApi = async (url: string | URL | Request, init?: RequestInit) => { + const headers = new Headers(init?.headers); + requests.push({ url: String(url), authorization: headers.get("Authorization") }); + return Response.json({ + rollingUsage: { usagePercent: 21, resetInSec: 1800 }, + weeklyUsage: { usagePercent: 42 }, + monthlyUsage: { usagePercent: 63 }, + }); + }; + const report = await fetchUsage({ apiKey: "fixture-key" }, fetchApi as typeof fetch); + expect(requests).toEqual([ + { url: "https://opencode.ai/zen/go/v1/usage", authorization: "Bearer fixture-key" }, + ]); + expect(report.status).toBe("available"); + expect(report.windows.map((window) => [window.id, window.usedPct, window.headline])).toEqual([ + ["rolling", 21, true], + ["weekly", 42, undefined], + ["monthly", 63, undefined], + ]); +}); + +test("treats rejected credentials as unavailable", async () => { + const report = await fetchUsage( + { apiKey: "fixture-key" }, + async () => new Response(null, { status: 401 }), + ); + expect(report.status).toBe("unavailable"); + expect(report.windows).toEqual([]); +}); 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 new file mode 100644 index 0000000000..7dff014137 --- /dev/null +++ b/packages/server/src/plugins/opencode-go-usage-source/server/usage.ts @@ -0,0 +1,57 @@ +import { createHash } from "node:crypto"; +import { z } from "zod"; +import { + toneFromUsedPct, + windowFromUsedPct, + type UsageReport, +} from "@getpaseo/plugin/server/usage"; +import type { Input } from "../shared/input.js"; + +const windowSchema = z + .object({ usagePercent: z.number().finite(), resetInSec: z.number().finite().optional() }) + .passthrough(); +const responseSchema = z.object({ + rollingUsage: windowSchema.optional(), + weeklyUsage: windowSchema.optional(), + monthlyUsage: windowSchema.optional(), +}); + +export async function fetchUsage( + input: Input, + fetchApi: typeof fetch = fetch, +): Promise { + const key = createHash("sha256").update(input.apiKey).digest("hex"); + const response = await fetchApi("https://opencode.ai/zen/go/v1/usage", { + headers: { Authorization: `Bearer ${input.apiKey}`, Accept: "application/json" }, + signal: AbortSignal.timeout(15_000), + }); + if (response.status === 401 || response.status === 403) + return { account: { key }, status: "unavailable", windows: [] }; + if (!response.ok) throw new Error(`OpenCode Go usage API returned ${response.status}`); + const data = responseSchema.parse(await response.json()); + const now = Date.now(); + const windows = ( + [ + ["rolling", "Rolling", data.rollingUsage], + ["weekly", "Weekly", data.weeklyUsage], + ["monthly", "Monthly", data.monthlyUsage], + ] as const + ).flatMap(([id, label, value]) => + value + ? [ + windowFromUsedPct({ + id, + label, + utilizationPct: value.usagePercent, + resetsAt: + value.resetInSec === undefined + ? null + : new Date(now + value.resetInSec * 1000).toISOString(), + tone: toneFromUsedPct(value.usagePercent), + headline: id === "rolling", + }), + ] + : [], + ); + return { account: { key }, status: "available", planLabel: "Go", windows }; +} diff --git a/packages/server/src/plugins/opencode-go-usage-source/shared/input.ts b/packages/server/src/plugins/opencode-go-usage-source/shared/input.ts new file mode 100644 index 0000000000..3f8774be5b --- /dev/null +++ b/packages/server/src/plugins/opencode-go-usage-source/shared/input.ts @@ -0,0 +1,4 @@ +import { z } from "zod"; + +export const inputSchema = z.object({ apiKey: z.string().min(1) }).strict(); +export type Input = z.infer; diff --git a/packages/server/src/server/agent/agent-sdk-types.ts b/packages/server/src/server/agent/agent-sdk-types.ts index e2fb9779b9..64e5b11eca 100644 --- a/packages/server/src/server/agent/agent-sdk-types.ts +++ b/packages/server/src/server/agent/agent-sdk-types.ts @@ -668,6 +668,8 @@ export interface AgentSession { /** New provider-owned rows to commit on registration. streamHistory must also * replay them at their original timestamps; restored sessions omit old rows. */ readonly initialTimeline?: ImportedTimelineEntry[]; + /** Resolved at fetch time because model and credentials may change during a session. */ + getUsageReference?(): Promise<{ source: string; input: JsonValue } | null>; run(prompt: AgentPromptInput, options?: AgentRunOptions): Promise; startTurn(prompt: AgentPromptInput, options?: AgentRunOptions): Promise<{ turnId: string }>; steerActiveTurn?(prompt: AgentPromptInput, options: SteerActiveTurnOptions): Promise; diff --git a/packages/server/src/server/agent/plugin-provider.test.ts b/packages/server/src/server/agent/plugin-provider.test.ts index 47d3d67cf9..6aa32e2731 100644 --- a/packages/server/src/server/agent/plugin-provider.test.ts +++ b/packages/server/src/server/agent/plugin-provider.test.ts @@ -396,6 +396,24 @@ describe("PluginAgentClientRegistry", () => { await registry.shutdown(); }); + test("returns no usage reference when a plugin provider lacks the capability", async () => { + const harness = createProviderHarness(); + const registry = new PluginAgentClientRegistry(createTestLogger()); + registry.replace([harness.registration]); + const client = registry.clients()[harness.registration.id]!; + const session = await client.createSession({ + provider: harness.registration.id, + cwd: "/workspace", + }); + try { + expect(await session.getUsageReference?.()).toBeNull(); + expect(harness.inputs.some((input) => input.type === "session.usage_reference")).toBe(false); + } finally { + await session.close(); + registry.replace([]); + } + }); + test("terminalizes an active turn exactly once when its plugin provider is removed", async () => { const harness = createProviderHarness({ completeTurn: false }); const registry = new PluginAgentClientRegistry(createTestLogger()); diff --git a/packages/server/src/server/agent/plugin-provider.ts b/packages/server/src/server/agent/plugin-provider.ts index 1977c787e1..bd5a7d9986 100644 --- a/packages/server/src/server/agent/plugin-provider.ts +++ b/packages/server/src/server/agent/plugin-provider.ts @@ -105,6 +105,20 @@ function providerError(error: ProviderError): Error { }); } +function isProviderRequestReply( + event: ProviderEvent, +): event is Extract< + ProviderEvent, + { type: "request.completed" | "catalog" | "sessions" | "usage_reference" } +> { + return ( + event.type === "request.completed" || + event.type === "catalog" || + event.type === "sessions" || + event.type === "usage_reference" + ); +} + class ProviderRuntime { private connection: ProviderConnection | null = null; private connecting: Promise | null = null; @@ -356,11 +370,7 @@ class ProviderRuntime { this.failRequest(event); return; } - if ( - event.type === "request.completed" || - event.type === "catalog" || - event.type === "sessions" - ) { + if (isProviderRequestReply(event)) { this.finishRequest(event); return; } @@ -515,7 +525,10 @@ class ProviderRuntime { } private finishRequest( - event: Extract, + event: Extract< + ProviderEvent, + { type: "request.completed" | "catalog" | "sessions" | "usage_reference" } + >, ): void { const request = this.requests.get(event.requestId); if (!request) return; @@ -554,6 +567,17 @@ class ProviderRuntimeSession { return this.capabilities; } + async getUsageReference(): Promise<{ source: string; input: JsonValue } | null> { + if (!this.capabilities.includes("session.usage_reference")) return null; + const event = await this.runtime.complete({ + type: "session.usage_reference", + requestId: randomUUID(), + sessionId: this.providerSessionId, + }); + if (event.type !== "usage_reference") throw new Error("Invalid usage reference response"); + return event.reference; + } + onEvent(listener: (event: ProviderEvent) => void): () => void { this.listeners.add(listener); return () => this.listeners.delete(listener); @@ -1063,6 +1087,10 @@ class PluginAgentSession implements AgentSession { return agentCapabilities(this.bridge.negotiatedCapabilities); } + getUsageReference(): Promise<{ source: string; input: JsonValue } | null> { + return this.bridge.getUsageReference(); + } + get features(): AgentFeature[] { return this.bridge.config.settings.map((setting) => setting.type === "toggle" diff --git a/packages/server/src/server/agent/providers/acp-agent.test.ts b/packages/server/src/server/agent/providers/acp-agent.test.ts index 9d106a2c1a..0394900d3b 100644 --- a/packages/server/src/server/agent/providers/acp-agent.test.ts +++ b/packages/server/src/server/agent/providers/acp-agent.test.ts @@ -222,6 +222,13 @@ function createSessionWithConfig( ); } +test("ACP usage reference uses the provider ID", async () => { + for (const provider of ["copilot", "cursor", "kimi", "custom-source"]) { + const session = createSessionWithConfig({ provider }); + expect(await session.getUsageReference()).toEqual({ source: provider, input: {} }); + } +}); + function createKiroSession( options: { waitForInitialCommands?: boolean; initialCommandsWaitTimeoutMs?: number } = {}, logger: ReturnType = createTestLogger(), diff --git a/packages/server/src/server/agent/providers/acp-agent.ts b/packages/server/src/server/agent/providers/acp-agent.ts index d2e0a1beb6..10f384af84 100644 --- a/packages/server/src/server/agent/providers/acp-agent.ts +++ b/packages/server/src/server/agent/providers/acp-agent.ts @@ -1633,6 +1633,10 @@ export class ACPAgentSession implements AgentSession, ACPClient { readonly provider: string; readonly capabilities: AgentCapabilityFlags; + async getUsageReference() { + return { source: this.provider, input: {} }; + } + private readonly logger: Logger; private readonly runtimeSettings?: ProviderRuntimeSettings; private readonly defaultCommand: [string, ...string[]]; diff --git a/packages/server/src/server/agent/providers/claude/agent.test.ts b/packages/server/src/server/agent/providers/claude/agent.test.ts index c908e17cf5..c034fde9b9 100644 --- a/packages/server/src/server/agent/providers/claude/agent.test.ts +++ b/packages/server/src/server/agent/providers/claude/agent.test.ts @@ -527,6 +527,25 @@ describe("ClaudeAgentClient.fetchCatalog", () => { describe("ClaudeAgentClient binary resolution", () => { const logger = createTestLogger(); + test("Claude usage reference follows CLAUDE_CONFIG_DIR and excludes API overrides", async () => { + const client = new ClaudeAgentClient({ logger, resolveBinary: async () => "/test/claude/bin" }); + const session = await client.createSession( + { provider: "claude", cwd: process.cwd() }, + { env: { CLAUDE_CONFIG_DIR: "/accounts/second" } }, + ); + expect(await session.getUsageReference?.()).toEqual({ + source: "claude", + input: { configDir: "/accounts/second" }, + }); + await session.close(); + const custom = await client.createSession( + { provider: "claude", cwd: process.cwd() }, + { env: { ANTHROPIC_AUTH_TOKEN: "token" } }, + ); + expect(await custom.getUsageReference?.()).toBeNull(); + await custom.close(); + }); + test("resolves the installed Claude Code version", async () => { await expect(resolveClaudeCodeVersion()).resolves.toMatch(/^\d+\.\d+\.\d+$/); }); diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index 04fef2b1f6..38476490b3 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -3,6 +3,7 @@ import { randomUUID } from "node:crypto"; import fs from "node:fs"; import { promises } from "node:fs"; import path from "node:path"; +import { homedir } from "node:os"; import { type AgentDefinition, type CanUseTool, @@ -2041,6 +2042,15 @@ class ClaudeAgentSession implements AgentSession { readonly provider = "claude" as const; readonly capabilities = CLAUDE_CAPABILITIES; + async getUsageReference() { + const env = this.buildSdkEnv(); + if (env.ANTHROPIC_BASE_URL || env.ANTHROPIC_API_KEY || env.ANTHROPIC_AUTH_TOKEN) return null; + return { + source: "claude", + input: { configDir: env.CLAUDE_CONFIG_DIR || path.join(env.HOME || homedir(), ".claude") }, + }; + } + private readonly config: ClaudeAgentConfig; private readonly launchEnv?: Record; private readonly agentId?: string; diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts index 00dc26652c..78286c1988 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts +++ b/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts @@ -171,6 +171,44 @@ function createSession( return session; } +test("Codex usage reference follows CODEX_HOME and excludes custom base URLs", async () => { + const session = new CodexAppServerAgentSession( + createConfig(), + null, + createTestLogger(), + () => { + throw new Error("unused"); + }, + {}, + false, + false, + false, + undefined, + "interactive", + { CODEX_HOME: "/accounts/second" }, + ); + expect(await session.getUsageReference()).toEqual({ + source: "codex", + input: { codexHome: "/accounts/second" }, + }); + const custom = new CodexAppServerAgentSession( + createConfig(), + null, + createTestLogger(), + () => { + throw new Error("unused"); + }, + {}, + false, + false, + false, + undefined, + "interactive", + { CODEX_HOME: "/accounts/second", OPENAI_BASE_URL: "https://example.test" }, + ); + expect(await custom.getUsageReference()).toBeNull(); +}); + function createProviderWithFakeAppServer( appServer: FakeCodexAppServer, options?: { diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.ts index b7fc721f24..cd8dfc2d0e 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.ts +++ b/packages/server/src/server/agent/providers/codex-app-server-agent.ts @@ -3453,6 +3453,7 @@ export class CodexAppServerAgentSession implements AgentSession { private readonly autoReviewEnabled: boolean = false, private readonly agentId?: string, private readonly initialResumePurpose: "interactive" | "history" = "interactive", + private readonly usageEnv: NodeJS.ProcessEnv = process.env, ) { this.logger = logger.child({ module: "agent", @@ -3481,6 +3482,17 @@ export class CodexAppServerAgentSession implements AgentSession { } } + async getUsageReference() { + if (this.usageEnv.OPENAI_BASE_URL) return null; + return { + source: "codex", + input: { + codexHome: + this.usageEnv.CODEX_HOME || path.join(this.usageEnv.HOME || os.homedir(), ".codex"), + }, + }; + } + get id(): string | null { return this.currentThreadId; } @@ -7143,6 +7155,8 @@ export class CodexAppServerAgentClient implements AgentClient { goalsEnabled, autoReviewEnabled, launchContext?.agentId, + "interactive", + buildCodexAppServerEnv(this.runtimeSettings, launchContext?.env), ); await session.connect(); return session; @@ -7175,6 +7189,7 @@ export class CodexAppServerAgentClient implements AgentClient { autoReviewEnabled, launchContext?.agentId, options?.purpose ?? "interactive", + buildCodexAppServerEnv(this.runtimeSettings, launchContext?.env), ); await session.connect(); return session; 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 1ba8b282a9..058c56d8da 100644 --- a/packages/server/src/server/agent/providers/opencode-agent.test.ts +++ b/packages/server/src/server/agent/providers/opencode-agent.test.ts @@ -333,6 +333,65 @@ describe("OpenCodeAgentClient adapter smoke tests", () => { model: TEST_MODEL, }); + test("usage reference follows the active OpenCode model and OAuth account", async () => { + const cwd = tmpCwd(); + const runtime = new TestOpenCodeHarness(); + const openCode = new TestOpenCodeClient(); + runtime.enqueueClient(openCode); + const client = new OpenCodeAgentClient( + logger, + { + env: { + OPENCODE_AUTH_CONTENT: JSON.stringify({ + openai: { type: "oauth", access: "oauth-token", accountId: "acct-1" }, + "opencode-go": { type: "api", key: "go-key" }, + }), + }, + }, + { serverManager: runtime, createClient: runtime.createClient }, + ); + const session = await client.createSession(buildConfig(cwd)); + await session.setModel?.("openai/gpt-5"); + expect(await session.getUsageReference?.()).toEqual({ + source: "codex", + input: { accessToken: "oauth-token", accountId: "acct-1" }, + }); + await session.setModel?.("opencode-go/qwen"); + expect(await session.getUsageReference?.()).toEqual({ + source: "opencode-go", + input: { apiKey: "go-key" }, + }); + await session.setModel?.("other/model"); + expect(await session.getUsageReference?.()).toBeNull(); + const events: AgentStreamEvent[] = []; + session.subscribe((event) => events.push(event)); + const sessionId = (await session.getRuntimeInfo()).sessionId; + openCode.emitEvent({ + type: "message.updated", + properties: { + info: { + id: "msg_changed_model", + sessionID: sessionId, + role: "user", + model: { providerID: "openai", modelID: "gpt-5" }, + }, + }, + } as OpenCodeEvent); + await vi.waitFor(async () => + expect(await session.getUsageReference?.()).toEqual({ + source: "codex", + input: { accessToken: "oauth-token", accountId: "acct-1" }, + }), + ); + expect(events).toContainEqual({ + type: "model_changed", + provider: "opencode", + runtimeInfo: { provider: "opencode", sessionId, model: "openai/gpt-5", modeId: null }, + }); + 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 2276e85786..0de0d4cd33 100644 --- a/packages/server/src/server/agent/providers/opencode-agent.ts +++ b/packages/server/src/server/agent/providers/opencode-agent.ts @@ -12,6 +12,8 @@ import { type TextPartInput as OpenCodeTextPartInput, } from "@opencode-ai/sdk/v2/client"; import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; import { randomBytes } from "node:crypto"; import { isDeepStrictEqual } from "node:util"; import { createPathEquivalenceMatcher } from "../../../utils/path.js"; @@ -1500,6 +1502,7 @@ export class OpenCodeAgentClient implements AgentClient { false, unbindBridge, connectServer, + { ...process.env, ...this.runtimeSettings?.env, ...launchContext?.env }, ); } catch (error) { await connection.release(); @@ -1554,6 +1557,7 @@ export class OpenCodeAgentClient implements AgentClient { registeredAcquisition !== null, unbindBridge, connectServer, + { ...process.env, ...this.runtimeSettings?.env, ...launchContext?.env }, ); } catch (error) { await connection.release(); @@ -3279,7 +3283,7 @@ function isOpenCodeTerminalEvent(event: OpenCodeEvent, sessionId: string): boole } function isOpenCodeProviderInternalEvent(event: AgentStreamEvent): boolean { - return event.type === "provider_subagent"; + return event.type === "provider_subagent" || event.type === "model_changed"; } function readOpenCodeChildSessionInfo(value: unknown): OpenCodeChildSessionInfo | null { @@ -3360,6 +3364,56 @@ interface OpenCodeServerConnection { release: () => Promise; } +async function readOpenCodeUsageAuth( + env: NodeJS.ProcessEnv, +): Promise | null> { + try { + const content = + env.OPENCODE_AUTH_CONTENT ?? + (await fs.readFile( + path.join( + env.XDG_DATA_HOME || path.join(env.HOME || os.homedir(), ".local", "share"), + "opencode", + "auth.json", + ), + "utf8", + )); + const parsed: unknown = JSON.parse(content); + return parsed && typeof parsed === "object" && !Array.isArray(parsed) + ? (parsed as Record) + : null; + } catch { + return null; + } +} + +function resolveOpenCodeUsageReference( + model: string, + auth: Record | null, +): { source: string; input: import("@getpaseo/protocol/agent-types").JsonValue } | 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; + return { + source: "codex", + input: { + accessToken: credential.access, + ...(typeof credential.accountId === "string" ? { 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; + return { source: "opencode-go", input: { apiKey: credential.key } }; +} + class OpenCodeAgentSession implements AgentSession { readonly provider = "opencode" as const; readonly capabilities = OPENCODE_CAPABILITIES; @@ -3448,6 +3502,7 @@ class OpenCodeAgentSession implements AgentSession { private readonly externallyDriven = false, releaseBridge?: () => void, connectServer?: () => Promise, + private readonly usageEnv: NodeJS.ProcessEnv = process.env, ) { this.config = config; this.server = { client, events, url: serverUrl ?? "", release: releaseServer }; @@ -3552,6 +3607,16 @@ class OpenCodeAgentSession implements AgentSession { }; } + async getUsageReference(): Promise<{ + source: string; + input: import("@getpaseo/protocol/agent-types").JsonValue; + } | null> { + const model = this.config.model; + if (!model?.startsWith("openai/") && !model?.startsWith("opencode-go/")) return null; + const auth = await readOpenCodeUsageAuth(this.usageEnv); + return resolveOpenCodeUsageReference(model, auth); + } + async setModel(modelId: string | null): Promise { const normalizedModelId = typeof modelId === "string" && modelId.trim().length > 0 ? modelId : null; @@ -5346,6 +5411,7 @@ class OpenCodeAgentSession implements AgentSession { private async translateEvent(event: OpenCodeEvent): Promise { const eventSessionId = getOpenCodeEventSessionId(event); + const runtimeModelChanged = this.syncRuntimeModel(event, eventSessionId); if ( event.type !== "session.created" && eventSessionId && @@ -5362,6 +5428,12 @@ class OpenCodeAgentSession implements AgentSession { } } const translated = translateOpenCodeEvent(event, this.createTranslationState()); + if (runtimeModelChanged) + translated.push({ + type: "model_changed", + provider: this.provider, + runtimeInfo: await this.getRuntimeInfo(), + }); this.appendProviderSubagentEvents(event, translated); const events: AgentStreamEvent[] = []; @@ -5402,6 +5474,20 @@ class OpenCodeAgentSession implements AgentSession { return events; } + private syncRuntimeModel(event: OpenCodeEvent, eventSessionId: string | null): boolean { + if (event.type !== "message.updated" || eventSessionId !== this.sessionId) return false; + const info = event.properties.info; + let model: string | undefined; + if (info.role === "assistant") model = resolveOpenCodeModelLookupKeyFromAssistantMessage(info); + if (info.role === "user" && info.model) + model = buildOpenCodeModelLookupKey(info.model.providerID, info.model.modelID); + if (!model || model === this.config.model) return false; + this.config.model = model; + this.selectedModelContextWindowMaxTokens = + this.resolveConfiguredModelContextWindowMaxTokens(model); + return true; + } + private async tryAutoApproveToolPermission( request: AgentPermissionRequest, directory: string, 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 969b221f02..5163ef37f3 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -23,7 +23,7 @@ import { PiRpcAgentSession, transformPiModels, } from "./agent.js"; -import { FakePi } from "./test-utils/fake-pi.js"; +import { FakePi, type FakePiSession } from "./test-utils/fake-pi.js"; import { createPiExtensionHost } from "./extensions/index.js"; import { PiExtensionHost } from "./extensions/host.js"; import type { PiModel, PiThinkingLevel } from "./rpc-types.js"; @@ -202,6 +202,57 @@ async function createSession( return { pi, session, events }; } +test("Pi usage reference follows the current OAuth model and agent directory", async () => { + const agentDir = mkdtempSync(path.join(tmpdir(), "paseo-pi-usage-")); + try { + writeFileSync( + path.join(agentDir, "auth.json"), + JSON.stringify({ + "openai-codex": { + type: "oauth", + access: "codex-token", + refresh: "refresh", + expires: Date.now() + 60_000, + }, + anthropic: { + type: "oauth", + access: "claude-token", + refresh: "refresh", + expires: Date.now() + 60_000, + }, + }), + ); + const pi = new FakePi(); + let runtime!: FakePiSession; + pi.queueSessionSetup((sessionRuntime) => { + runtime = sessionRuntime; + runtime.state.model = { provider: "openai-codex", id: "gpt", name: "GPT" }; + }); + const session = await createClient(pi).createSession(createConfig(), { + env: { PI_CODING_AGENT_DIR: agentDir }, + }); + expect(await session.getUsageReference?.()).toEqual({ + source: "codex", + input: { accessToken: "codex-token" }, + }); + runtime.state.model = { provider: "anthropic", id: "claude", name: "Claude" }; + expect(await session.getUsageReference?.()).toEqual({ + source: "claude", + input: { accessToken: "claude-token" }, + }); + writeFileSync( + path.join(agentDir, "auth.json"), + JSON.stringify({ anthropic: { type: "api_key", key: "api-key" } }), + ); + expect(await session.getUsageReference?.()).toBeNull(); + runtime.state.model = { provider: "other", id: "other", name: "Other" }; + expect(await session.getUsageReference?.()).toBeNull(); + await session.close(); + } finally { + rmSync(agentDir, { recursive: true, force: true }); + } +}); + test("forwards launch-context env to the Pi process launch", async () => { const pi = new FakePi(); const client = createClient(pi); diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 5fd7b051ef..fdd1fc1480 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -242,6 +242,7 @@ interface PiRpcAgentSessionOptions { extensionTimeoutMs?: number; logger: Logger; usagePollScheduler?: PiUsagePollScheduler; + usageEnv?: Record; } interface PiResumeConfig { @@ -1122,6 +1123,32 @@ export class PiRpcAgentSession implements AgentSession { readonly provider: AgentProvider; readonly capabilities: AgentCapabilityFlags; + async getUsageReference() { + await this.refreshState(); + const provider = this.state.model?.provider; + let source: string | null = null; + if (provider === "openai-codex") source = "codex"; + 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 } }; + } catch { + return null; + } + } + private readonly subscribers = new Set<(event: AgentStreamEvent) => void>(); private readonly activeToolCalls = new Map(); private readonly pendingExtensionUiRequests = new Map(); @@ -1149,6 +1176,7 @@ export class PiRpcAgentSession implements AgentSession { private readonly currentModeId: string | null; private readonly logger: Logger; private readonly usagePoller: PiUsagePoller; + private readonly usageEnv?: Record; private closed = false; private readonly closeController = new AbortController(); private readonly pendingExtensionHydrations = new Set>(); @@ -1168,6 +1196,7 @@ export class PiRpcAgentSession implements AgentSession { this.extensionTimeoutMs = options.extensionTimeoutMs ?? DEFAULT_PI_EXTENSION_RESULT_TIMEOUT_MS; this.logger = options.logger; this.extensionHost = createPiExtensionHost(this.logger); + this.usageEnv = options.usageEnv; this.usagePoller = new PiUsagePoller({ scheduler: options.usagePollScheduler, readStats: () => this.runtimeSession.getSessionStats(), @@ -2461,6 +2490,7 @@ export class PiRpcAgentClient implements AgentClient { extensionTimeoutMs: this.providerParams.extensionTimeoutMs, logger: this.logger, usagePollScheduler: this.usagePollScheduler, + usageEnv: mcpEnv, }); } catch (error) { await runtimeSession.close().catch(() => undefined); @@ -2524,6 +2554,7 @@ export class PiRpcAgentClient implements AgentClient { extensionTimeoutMs: this.providerParams.extensionTimeoutMs, logger: this.logger, usagePollScheduler: this.usagePollScheduler, + usageEnv: mcpEnv, }); } catch (error) { await runtimeSession.close().catch(() => undefined); diff --git a/packages/server/src/server/authorization/operation-permissions.ts b/packages/server/src/server/authorization/operation-permissions.ts index 214544c39c..dffcc10d29 100644 --- a/packages/server/src/server/authorization/operation-permissions.ts +++ b/packages/server/src/server/authorization/operation-permissions.ts @@ -150,6 +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", provider_diagnostic_request: "daemon.read", pull_request_timeline_request: "workspace.read", "push.unregister.request": "workspace.read", @@ -375,6 +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", 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 86b3003164..f49bd95e48 100644 --- a/packages/server/src/server/plugins/index.ts +++ b/packages/server/src/server/plugins/index.ts @@ -142,10 +142,20 @@ export class PluginService { return [...this.providers.values()].sort((left, right) => left.id.localeCompare(right.id)); } - listUsageReports(options?: { forceRefresh?: boolean }) { + listUsageReports(options?: { + forceRefresh?: boolean; + references?: Array<{ source: string; input: unknown }>; + }) { return this.usageSources.listReports(options); } + fetchUsageReference( + reference: { source: string; input: unknown }, + options?: { forceRefresh?: boolean }, + ) { + return this.usageSources.fetchReference(reference, options); + } + listLegacyUsage() { return this.usageSources.listLegacyUsage(); } 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 new file mode 100644 index 0000000000..7eceecd767 --- /dev/null +++ b/packages/server/src/server/plugins/session-usage-reference.e2e.test.ts @@ -0,0 +1,33 @@ +import { fileURLToPath } from "node:url"; +import { expect, test } from "vitest"; +import { DaemonClient } from "../test-utils/daemon-client.js"; +import { createTestPaseoDaemon } from "../test-utils/paseo-daemon.js"; +import contribute from "./test-fixtures/session-usage-reference/index.server.js"; + +test("resolves a provider plugin session reference through its usage source", async () => { + const directory = fileURLToPath( + new URL("./test-fixtures/session-usage-reference/", import.meta.url), + ); + const daemon = await createTestPaseoDaemon({ + pluginsEnabled: false, + internalPlugins: [{ id: "fixture-session-usage-reference", directory, contribute }], + }); + const client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` }); + try { + await client.connect(); + const agent = await client.createAgent({ + provider: "fixture-session-provider", + cwd: directory, + }); + const result = await client.getAgentUsageReport({ agentId: agent.id }); + expect(result.entry).toMatchObject({ + sourceId: "fixture-session-usage", + report: { account: { key: "from-session" }, windows: [{ usedPct: 31 }] }, + }); + const reports = await client.listUsageReports(); + expect(reports.reports.map((entry) => entry.report.account.key)).toContain("from-session"); + } finally { + await client.close(); + await daemon.close(); + } +}, 60_000); 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 new file mode 100644 index 0000000000..2ba78a4785 --- /dev/null +++ b/packages/server/src/server/plugins/test-fixtures/session-usage-reference/index.server.ts @@ -0,0 +1,72 @@ +import type { PluginServerContext } from "@getpaseo/plugin/server"; +import { + negotiateProviderCapabilities, + type ProviderEvent, +} from "@getpaseo/plugin/server/provider"; +import { z } from "zod"; + +export default function contribute(server: PluginServerContext) { + server.registerUsageSource({ + 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 }, + status: "available", + windows: [{ id: "hour", label: "Hour", usedPct: 31, headline: true }], + }), + }); + server.registerProvider({ + id: "fixture-session-provider", + label: "Fixture session provider", + async connect(request) { + const capabilities = negotiateProviderCapabilities(request.capabilities, [ + "prompt.message", + "session.usage_reference", + ]); + const listeners = new Set<(event: ProviderEvent) => void>(); + const emit = (event: ProviderEvent) => { + for (const listener of listeners) listener(event); + }; + return { + version: 1, + capabilities, + async send(input) { + if (input.type === "catalog") + emit({ + type: "catalog", + requestId: input.requestId, + catalog: { models: [{ id: "fixture", label: "Fixture" }], modes: [] }, + }); + if (input.type === "session.open") { + emit({ + type: "session.opened", + requestId: input.requestId, + sessionId: input.sessionId, + capabilities, + restoration: "core", + cwd: input.config.cwd, + }); + emit({ type: "session.ready", requestId: input.requestId, sessionId: input.sessionId }); + } + if (input.type === "session.usage_reference") + emit({ + type: "usage_reference", + requestId: input.requestId, + reference: { source: "fixture-session-usage", input: { account: "from-session" } }, + }); + if (input.type === "session.close") + emit({ type: "request.completed", requestId: input.requestId }); + }, + onEvent(listener) { + listeners.add(listener); + return () => listeners.delete(listener); + }, + async close() { + listeners.clear(); + }, + }; + }, + }); + return () => {}; +} diff --git a/packages/server/src/server/plugins/test-fixtures/session-usage-reference/paseo-plugin.json b/packages/server/src/server/plugins/test-fixtures/session-usage-reference/paseo-plugin.json new file mode 100644 index 0000000000..3e4937e1e2 --- /dev/null +++ b/packages/server/src/server/plugins/test-fixtures/session-usage-reference/paseo-plugin.json @@ -0,0 +1 @@ +{ "id": "fixture-session-usage-reference", "requirements": { "paseo": ">=0.9.2" } } 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 274bfea02b..ce51810ce8 100644 --- a/packages/server/src/server/plugins/usage-sources/index.test.ts +++ b/packages/server/src/server/plugins/usage-sources/index.test.ts @@ -44,6 +44,24 @@ test("coalesces concurrent fetches and dedupes accounts", async () => { 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); diff --git a/packages/server/src/server/plugins/usage-sources/index.ts b/packages/server/src/server/plugins/usage-sources/index.ts index a3beeecd25..cb52f98328 100644 --- a/packages/server/src/server/plugins/usage-sources/index.ts +++ b/packages/server/src/server/plugins/usage-sources/index.ts @@ -34,7 +34,12 @@ export class UsageSourceRegistry { for (const key of this.cache.keys()) if (key.startsWith(`${id}:`)) this.cache.delete(key); } - async listReports(options: { forceRefresh?: boolean } = {}): Promise { + async listReports( + options: { + forceRefresh?: boolean; + references?: Array<{ source: string; input: unknown }>; + } = {}, + ): Promise { const discovered = await Promise.all( [...this.sources.values()].map(async (source) => { try { @@ -47,11 +52,29 @@ export class UsageSourceRegistry { }), ); const unique = new Map(); - for (const entry of discovered.flat()) + 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: { source: string; input: unknown }, + options: { forceRefresh?: boolean } = {}, + ): Promise { + return this.sources.has(reference.source) + ? this.fetch(reference.source, reference.input, options) + : Promise.resolve(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(); diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index 287990a497..9d6c1d79cf 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -509,7 +509,14 @@ export interface SessionOptions { subscribeSettings?(listener: (pluginId: string, settingsId: string) => void): () => void; catalog(): Array<{ id: string; clientBundle: string }>; invokePluginRpc(pluginId: string, method: string, input: unknown): Promise; - listUsageReports(options?: { forceRefresh?: boolean }): Promise; + listUsageReports(options?: { + forceRefresh?: boolean; + references?: Array<{ source: string; input: unknown }>; + }): Promise; + fetchUsageReference( + reference: { source: string; input: unknown }, + options?: { forceRefresh?: boolean }, + ): Promise; listLegacyUsage(): Promise<{ fetchedAt: string; providers: ProviderUsage[] }>; }; orchestrationSkills?: import("./orchestration-skills/index.js").OrchestrationSkills; @@ -3005,6 +3012,8 @@ export class Session { return this.providerCatalogSession.handleProviderUsageListRequest(msg); case "usage.list_reports.request": return this.handleUsageListReportsRequest(msg); + case "agent.get_usage_report.request": + return this.handleAgentGetUsageReportRequest(msg); default: return undefined; } @@ -3015,7 +3024,24 @@ export class Session { ): Promise { try { if (!this.pluginRuntime) throw new Error("Plugin runtime is unavailable"); - const reports = await this.pluginRuntime.listUsageReports({ forceRefresh: msg.forceRefresh }); + 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 }, @@ -3033,6 +3059,35 @@ export class Session { } } + 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":