From abf2646e188419f73ee4b15e2be24a912cca2a64 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Sat, 26 Sep 2026 18:28:50 +0200 Subject: [PATCH] feat(usage): add plugin usage sources and Codex discovery --- packages/client/src/daemon-client.ts | 12 ++ packages/plugin/package.json | 5 + packages/plugin/src/server/contracts.ts | 2 + packages/plugin/src/server/index.ts | 7 + packages/plugin/src/server/usage.ts | 51 +++++++ packages/protocol/src/messages.ts | 32 +++++ .../src/plugins/codex-usage-source/icon.svg | 1 + .../codex-usage-source/index.server.ts | 15 ++ .../codex-usage-source/paseo-plugin.json | 1 + .../codex-usage-source/server/usage.test.ts | 135 ++++++++++++++++++ .../codex-usage-source/server/usage.ts | 120 ++++++++++++++++ .../codex-usage-source/shared/input.ts | 8 ++ packages/server/src/plugins/index.ts | 10 +- .../authorization/operation-permissions.ts | 2 + packages/server/src/server/plugins/index.ts | 53 +++++++ .../server/plugins/plugin-process-protocol.ts | 32 +++++ .../src/server/plugins/plugin-process.ts | 103 ++++++++++--- .../server/plugins/plugin-sdk-specifiers.ts | 1 + packages/server/src/server/plugins/runtime.ts | 24 ++++ .../usage-source-directory/icon.svg | 1 + .../usage-source-directory/index.server.ts | 31 ++++ .../usage-source-directory/paseo-plugin.json | 1 + .../test-fixtures/usage-source/icon.svg | 1 + .../usage-source/index.server.ts | 31 ++++ .../usage-source/paseo-plugin.json | 1 + .../server/plugins/usage-sources.e2e.test.ts | 47 ++++++ .../plugins/usage-sources/index.test.ts | 69 +++++++++ .../src/server/plugins/usage-sources/index.ts | 96 +++++++++++++ packages/server/src/server/session.ts | 28 ++++ .../server/src/server/websocket-server.ts | 1 + 30 files changed, 900 insertions(+), 21 deletions(-) create mode 100644 packages/plugin/src/server/usage.ts create mode 100644 packages/server/src/plugins/codex-usage-source/icon.svg create mode 100644 packages/server/src/plugins/codex-usage-source/index.server.ts create mode 100644 packages/server/src/plugins/codex-usage-source/paseo-plugin.json create mode 100644 packages/server/src/plugins/codex-usage-source/server/usage.test.ts create mode 100644 packages/server/src/plugins/codex-usage-source/server/usage.ts create mode 100644 packages/server/src/plugins/codex-usage-source/shared/input.ts create mode 100644 packages/server/src/server/plugins/test-fixtures/usage-source-directory/icon.svg create mode 100644 packages/server/src/server/plugins/test-fixtures/usage-source-directory/index.server.ts create mode 100644 packages/server/src/server/plugins/test-fixtures/usage-source-directory/paseo-plugin.json create mode 100644 packages/server/src/server/plugins/test-fixtures/usage-source/icon.svg create mode 100644 packages/server/src/server/plugins/test-fixtures/usage-source/index.server.ts create mode 100644 packages/server/src/server/plugins/test-fixtures/usage-source/paseo-plugin.json create mode 100644 packages/server/src/server/plugins/usage-sources.e2e.test.ts create mode 100644 packages/server/src/server/plugins/usage-sources/index.test.ts create mode 100644 packages/server/src/server/plugins/usage-sources/index.ts diff --git a/packages/client/src/daemon-client.ts b/packages/client/src/daemon-client.ts index 46194ada89..6a909f5b6a 100644 --- a/packages/client/src/daemon-client.ts +++ b/packages/client/src/daemon-client.ts @@ -94,6 +94,7 @@ import type { RefreshProvidersSnapshotResponseMessage, ProviderDiagnosticResponseMessage, ProviderUsageListResponseMessage, + UsageListReportsResponseMessage, DaemonGetStatusResponse, DaemonGetPairingOfferResponse, DaemonConfigReloadResponse, @@ -559,6 +560,7 @@ type GetProvidersSnapshotPayload = GetProvidersSnapshotResponseMessage["payload" type RefreshProvidersSnapshotPayload = RefreshProvidersSnapshotResponseMessage["payload"]; type ProviderDiagnosticPayload = ProviderDiagnosticResponseMessage["payload"]; type ProviderUsageListPayload = ProviderUsageListResponseMessage["payload"]; +type UsageListReportsPayload = UsageListReportsResponseMessage["payload"]; type DaemonStatusPayload = DaemonGetStatusResponse["payload"]; type DaemonPairingOfferPayload = DaemonGetPairingOfferResponse["payload"]; type DiagnosticsPayload = DiagnosticsResponse["payload"]; @@ -5238,6 +5240,16 @@ export class DaemonClient { }); } + async listUsageReports(options?: { + requestId?: string; + forceRefresh?: boolean; + }): Promise { + return this.sendNamespacedCorrelatedSessionRequest({ + requestId: options?.requestId, + message: { type: "usage.list_reports.request", forceRefresh: options?.forceRefresh }, + }); + } + async listCommands(options: ListCommandsOptions): Promise; async listCommands(agentId: string, requestId?: string): Promise; async listCommands( diff --git a/packages/plugin/package.json b/packages/plugin/package.json index f6e397cc50..7782732e3e 100644 --- a/packages/plugin/package.json +++ b/packages/plugin/package.json @@ -23,6 +23,11 @@ "source": "./src/server/provider.ts", "default": "./dist/server/provider.js" }, + "./server/usage": { + "types": "./dist/server/usage.d.ts", + "source": "./src/server/usage.ts", + "default": "./dist/server/usage.js" + }, "./server/acp": { "types": "./dist/server/acp.d.ts", "source": "./src/server/acp.ts", diff --git a/packages/plugin/src/server/contracts.ts b/packages/plugin/src/server/contracts.ts index c1cf10fc15..984b9c117b 100644 --- a/packages/plugin/src/server/contracts.ts +++ b/packages/plugin/src/server/contracts.ts @@ -3,6 +3,7 @@ import type { ZodType, input as ZodInput, output as ZodOutput } from "zod"; import type { PluginRpcContract } from "../rpc.js"; import type { PluginCleanup } from "../contracts.js"; import type { ProviderRegistration } from "./provider.js"; +import type { UsageSourceRegistration } from "./usage.js"; import type { PluginLifecycleRegistration } from "./lifecycle.js"; export interface PluginHandlerContext { @@ -38,6 +39,7 @@ export interface PluginServerContext extends PluginLifecycleRegistration { ) => ZodInput | Promise>, ): void; registerProvider(provider: ProviderRegistration): void; + registerUsageSource(source: UsageSourceRegistration): void; } export type PluginServerContribution = (server: PluginServerContext) => PluginCleanup; diff --git a/packages/plugin/src/server/index.ts b/packages/plugin/src/server/index.ts index 9e1d78b8bb..dc211f61bf 100644 --- a/packages/plugin/src/server/index.ts +++ b/packages/plugin/src/server/index.ts @@ -1,3 +1,10 @@ +export type { + UsageSourceRegistration, + UsageReport, + UsageWindow, + UsageBalance, + UsageDetail, +} from "./usage.js"; export type { PluginHandlerContext, PluginServerContext, diff --git a/packages/plugin/src/server/usage.ts b/packages/plugin/src/server/usage.ts new file mode 100644 index 0000000000..af6d3bea4d --- /dev/null +++ b/packages/plugin/src/server/usage.ts @@ -0,0 +1,51 @@ +import type { ZodType } from "zod"; +import type { JsonValue } from "@getpaseo/protocol/agent-types"; + +export interface UsageWindow { + id: string; + label: string; + usedPct?: number | null; + remainingPct?: number | null; + resetsAt?: string | null; + runsOutAt?: string | null; + shortfallPct?: number | null; + tone?: "default" | "ok" | "warning" | "danger"; + headline?: boolean; +} + +export interface UsageBalance { + id: string; + label: string; + used?: number | null; + remaining?: number | null; + limit?: number | null; + unit: "usd" | "credits" | "requests" | "tokens"; + resetsAt?: string | null; + tone?: UsageWindow["tone"]; +} + +export interface UsageDetail { + id: string; + label: string; + value: string; + tone?: UsageWindow["tone"]; +} + +export interface UsageReport { + account: { key: string; label?: string }; + status: "available" | "unavailable" | "error"; + planLabel?: string; + windows: UsageWindow[]; + balances?: UsageBalance[]; + details?: UsageDetail[]; + error?: string; +} + +export interface UsageSourceRegistration { + id: string; + label: string; + icon?: string; + input: ZodType; + fetch(input: unknown): Promise; + discover?(): Promise; +} diff --git a/packages/protocol/src/messages.ts b/packages/protocol/src/messages.ts index db72cd3cb6..25e31983d2 100644 --- a/packages/protocol/src/messages.ts +++ b/packages/protocol/src/messages.ts @@ -1774,6 +1774,12 @@ export const ProviderUsageListRequestMessageSchema = z.object({ requestId: z.string(), }); +export const UsageListReportsRequestMessageSchema = z.object({ + type: z.literal("usage.list_reports.request"), + requestId: z.string(), + forceRefresh: z.boolean().optional(), +}); + export const ResumeAgentRequestMessageSchema = z.object({ type: z.literal("resume_agent_request"), handle: AgentPersistenceHandleSchema, @@ -3232,6 +3238,7 @@ export const SessionInboundMessageSchema = z.discriminatedUnion("type", [ RefreshProvidersSnapshotRequestMessageSchema, ProviderDiagnosticRequestMessageSchema, ProviderUsageListRequestMessageSchema, + UsageListReportsRequestMessageSchema, ResumeAgentRequestMessageSchema, ImportAgentRequestMessageSchema, RefreshAgentRequestMessageSchema, @@ -3534,6 +3541,7 @@ export const ServerInfoStatusPayloadSchema = z // COMPAT(hubAgentRpc): added in v0.8.0; remove gate after 2027-03-05. hubAgentRpc: z.boolean().optional(), providersSnapshot: z.boolean().optional(), + usageSources: z.boolean().optional(), // COMPAT(providersSnapshotCwd): added in v0.3.2, remove gate after 2027-02-10. providersSnapshotCwd: z.boolean().optional(), // COMPAT(directorySync): added in v0.3.x, remove gate after 2027-02-12. @@ -6228,6 +6236,26 @@ 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() })), + balances: z.array(ProviderUsageBalanceSchema).optional(), + details: z.array(ProviderUsageDetailSchema).optional(), + error: z.string().optional(), +}); +export const UsageReportEntrySchema = z.object({ + sourceId: z.string(), + sourceLabel: z.string(), + icon: z.string().optional(), + report: UsageReportSchema, +}); +export const UsageListReportsResponseMessageSchema = z.object({ + type: z.literal("usage.list_reports.response"), + payload: z.object({ requestId: z.string(), reports: z.array(UsageReportEntrySchema) }), +}); + const AgentSlashCommandSchema = z.object({ name: z.string(), description: z.string(), @@ -6913,6 +6941,7 @@ export const SessionOutboundMessageSchema = z.discriminatedUnion("type", [ RefreshProvidersSnapshotResponseMessageSchema, ProviderDiagnosticResponseMessageSchema, ProviderUsageListResponseMessageSchema, + UsageListReportsResponseMessageSchema, ListCommandsResponseSchema, ListTerminalsResponseSchema, TerminalsChangedSchema, @@ -7091,6 +7120,9 @@ export type ProviderDiagnosticResponseMessage = z.infer< typeof ProviderDiagnosticResponseMessageSchema >; export type ProviderUsageTone = z.infer; +export type UsageReport = z.infer; +export type UsageReportEntry = z.infer; +export type UsageListReportsResponseMessage = z.infer; export type ProviderUsageStatus = z.infer; export type ProviderUsage = z.infer; export type ProviderUsageWindow = z.infer; diff --git a/packages/server/src/plugins/codex-usage-source/icon.svg b/packages/server/src/plugins/codex-usage-source/icon.svg new file mode 100644 index 0000000000..d184309848 --- /dev/null +++ b/packages/server/src/plugins/codex-usage-source/icon.svg @@ -0,0 +1 @@ + diff --git a/packages/server/src/plugins/codex-usage-source/index.server.ts b/packages/server/src/plugins/codex-usage-source/index.server.ts new file mode 100644 index 0000000000..ddc416b0dc --- /dev/null +++ b/packages/server/src/plugins/codex-usage-source/index.server.ts @@ -0,0 +1,15 @@ +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: "codex", + label: "Codex", + icon: "icon.svg", + input: inputSchema, + discover: async () => [{}], + fetch: fetchUsage, + }); + return () => {}; +} diff --git a/packages/server/src/plugins/codex-usage-source/paseo-plugin.json b/packages/server/src/plugins/codex-usage-source/paseo-plugin.json new file mode 100644 index 0000000000..64470099fb --- /dev/null +++ b/packages/server/src/plugins/codex-usage-source/paseo-plugin.json @@ -0,0 +1 @@ +{ "id": "codex-usage-source", "requirements": { "paseo": ">=0.9.2" } } 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 new file mode 100644 index 0000000000..d0a415ba02 --- /dev/null +++ b/packages/server/src/plugins/codex-usage-source/server/usage.test.ts @@ -0,0 +1,135 @@ +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"; + +const originalHome = process.env["CODEX_HOME"]; +afterEach(() => { + if (originalHome === undefined) delete process.env["CODEX_HOME"]; + else process.env["CODEX_HOME"] = originalHome; +}); + +function response(headers: HeadersInit, accountId?: string): Promise { + const request = new Headers(headers); + expect(request.get("Authorization")).toMatch(/^Bearer fixture-/); + expect(request.get("ChatGPT-Account-Id")).toBe(accountId ?? null); + return Promise.resolve( + new Response( + JSON.stringify({ + plan_type: "plus", + rate_limit: { primary_window: { used_percent: 30, reset_at: 1700000000 } }, + }), + { status: 200 }, + ), + ); +} + +test("default input reads CODEX_HOME auth and preserves the usage request", async () => { + const home = await mkdtemp(join(tmpdir(), "usage-codex-")); + try { + process.env["CODEX_HOME"] = home; + await writeFile( + join(home, "auth.json"), + JSON.stringify({ + tokens: { access_token: "fixture-default", account_id: "account-default" }, + }), + ); + const report = await fetchUsage({}, (_url, init) => + response(init?.headers ?? {}, "account-default"), + ); + expect(report).toMatchObject({ + account: { key: "account-default" }, + status: "available", + planLabel: "plus", + windows: [{ id: "session", usedPct: 30, headline: true }], + }); + } finally { + await rm(home, { recursive: true, force: true }); + } +}); + +test("explicit codexHome reads only that auth file", async () => { + const home = await mkdtemp(join(tmpdir(), "usage-codex-")); + try { + await writeFile( + join(home, "auth.json"), + JSON.stringify({ tokens: { access_token: "fixture-home", account_id: "account-home" } }), + ); + const report = await fetchUsage({ codexHome: home }, (_url, init) => + response(init?.headers ?? {}, "account-home"), + ); + expect(report.account.key).toBe("account-home"); + } finally { + await rm(home, { recursive: true, force: true }); + } +}); + +test("explicit access token needs no auth file", async () => { + const report = await fetchUsage( + { accessToken: "fixture-supplied", accountId: "account-supplied" }, + (_url, init) => response(init?.headers ?? {}, "account-supplied"), + ); + expect(report.account.key).toBe("account-supplied"); +}); + +test("coerces credit balance and marks a 96 percent window dangerous", async () => { + const report = await fetchUsage( + { accessToken: "fixture-supplied" }, + async () => + new Response( + JSON.stringify({ + rate_limit: { + primary_window: { used_percent: 12 }, + secondary_window: { used_percent: 96 }, + }, + credits: { balance: "0" }, + }), + { status: 200 }, + ), + ); + expect(report.windows).toEqual( + expect.arrayContaining([ + expect.objectContaining({ id: "session", tone: "ok" }), + expect.objectContaining({ id: "weekly", tone: "danger" }), + ]), + ); + expect(report.balances).toEqual([expect.objectContaining({ remaining: 0, tone: "danger" })]); +}); + +test("HTML usage body is unavailable", async () => { + const report = await fetchUsage( + { accessToken: "fixture-supplied" }, + async () => new Response("Login", { status: 200 }), + ); + expect(report.status).toBe("unavailable"); +}); + +test("401 leaves auth.json byte for byte unchanged and makes no refresh request", async () => { + const home = await mkdtemp(join(tmpdir(), "usage-codex-")); + try { + const authPath = join(home, "auth.json"); + const before = JSON.stringify({ + OPENAI_API_KEY: null, + tokens: { + id_token: "fixture-id", + access_token: "fixture-stale", + refresh_token: "fixture-refresh", + account_id: "fixture-account", + }, + last_refresh: "2026-07-04T20:35:00Z", + }); + await writeFile(authPath, before); + let calls = 0; + const report = await fetchUsage({ codexHome: home }, async (url) => { + expect(String(url)).toBe("https://chatgpt.com/backend-api/wham/usage"); + calls++; + return new Response(null, { status: 401 }); + }); + expect(report.status).toBe("unavailable"); + expect(calls).toBe(1); + expect(await readFile(authPath, "utf8")).toBe(before); + } finally { + await rm(home, { recursive: true, force: true }); + } +}); diff --git a/packages/server/src/plugins/codex-usage-source/server/usage.ts b/packages/server/src/plugins/codex-usage-source/server/usage.ts new file mode 100644 index 0000000000..d444ed5570 --- /dev/null +++ b/packages/server/src/plugins/codex-usage-source/server/usage.ts @@ -0,0 +1,120 @@ +import { createHash } from "node:crypto"; +import { readFile } from "node:fs/promises"; +import { homedir } from "node:os"; +import { join } from "node:path"; +import type { UsageReport, UsageWindow } from "@getpaseo/plugin/server/usage"; +import { z } from "zod"; +import type { CodexUsageInput } from "../shared/input.js"; + +const authSchema = z.object({ + tokens: z + .object({ access_token: z.string().optional(), account_id: z.string().optional() }) + .optional(), +}); +const number = z.coerce.number().finite(); +const windowSchema = z.object({ used_percent: number.optional(), reset_at: number.optional() }); +const responseSchema = z.object({ + plan_type: z.string().optional(), + email: z.string().optional(), + rate_limit: z + .object({ primary_window: windowSchema.nullish(), secondary_window: windowSchema.nullish() }) + .nullish(), + code_review_rate_limit: z.object({ primary_window: windowSchema.nullish() }).nullish(), + credits: z.object({ balance: number.optional() }).nullish(), +}); + +async function readAuth( + input: CodexUsageInput, +): Promise<{ token: string; accountId?: string } | null> { + if ("accessToken" in input) return { token: input.accessToken, accountId: input.accountId }; + const candidates = + "codexHome" in input + ? [join(input.codexHome, "auth.json")] + : [ + ...(process.env["CODEX_HOME"] ? [join(process.env["CODEX_HOME"], "auth.json")] : []), + join(homedir(), ".config", "codex", "auth.json"), + join(homedir(), ".codex", "auth.json"), + ]; + for (const path of candidates) { + try { + const auth = authSchema.parse(JSON.parse(await readFile(path, "utf8"))); + if (auth.tokens?.access_token) + return { token: auth.tokens.access_token, accountId: auth.tokens.account_id }; + } catch { + continue; + } + } + return null; +} + +function usageWindow( + id: string, + label: string, + value: z.infer | null | undefined, + headline = false, +): UsageWindow | null { + if (!value) return null; + const usedPct = value.used_percent ?? 0; + let tone: UsageWindow["tone"] = "ok"; + if (usedPct > 90) tone = "danger"; + else if (usedPct >= 70) tone = "warning"; + return { + id, + label, + usedPct, + remainingPct: Math.max(0, 100 - usedPct), + resetsAt: value.reset_at != null ? new Date(value.reset_at * 1000).toISOString() : null, + tone, + ...(headline ? { headline: true } : {}), + }; +} + +export async function fetchUsage( + input: CodexUsageInput, + 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"); + const headers: Record = { + Authorization: `Bearer ${auth.token}`, + Accept: "application/json", + "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)", + }; + if (auth.accountId) headers["ChatGPT-Account-Id"] = auth.accountId; + const response = await fetchApi("https://chatgpt.com/backend-api/wham/usage", { + headers, + signal: AbortSignal.timeout(15_000), + }); + if (response.status === 401 || response.status === 403) + return { account: { key }, 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: [] }; + const usage = responseSchema.parse(JSON.parse(text)); + const windows = [ + usageWindow("session", "Session", usage.rate_limit?.primary_window, true), + usageWindow("weekly", "Weekly", usage.rate_limit?.secondary_window), + usageWindow("code_review", "Code review", usage.code_review_rate_limit?.primary_window), + ].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, + balances: + balance === undefined + ? [] + : [ + { + id: "credits", + label: "Credits", + remaining: balance, + unit: "usd", + tone: balance <= 0 ? "danger" : "ok", + }, + ], + details: [], + }; +} diff --git a/packages/server/src/plugins/codex-usage-source/shared/input.ts b/packages/server/src/plugins/codex-usage-source/shared/input.ts new file mode 100644 index 0000000000..0c54b51e69 --- /dev/null +++ b/packages/server/src/plugins/codex-usage-source/shared/input.ts @@ -0,0 +1,8 @@ +import { z } from "zod"; + +export const inputSchema = z.union([ + z.object({ codexHome: z.string().min(1) }).strict(), + z.object({ accessToken: z.string().min(1), accountId: z.string().optional() }).strict(), + z.object({}).strict(), +]); +export type CodexUsageInput = z.infer; diff --git a/packages/server/src/plugins/index.ts b/packages/server/src/plugins/index.ts index 6ccd5ee954..3dd7c1d778 100644 --- a/packages/server/src/plugins/index.ts +++ b/packages/server/src/plugins/index.ts @@ -1,4 +1,6 @@ import type { PluginServerContribution } from "@getpaseo/plugin/server"; +import { fileURLToPath } from "node:url"; +import codexUsageSource from "./codex-usage-source/index.server.js"; export interface InternalPlugin { id: string; @@ -6,4 +8,10 @@ export interface InternalPlugin { contribute: PluginServerContribution; } -export const INTERNAL_PLUGINS: InternalPlugin[] = []; +export const INTERNAL_PLUGINS: InternalPlugin[] = [ + { + id: "codex-usage-source", + directory: fileURLToPath(new URL("./codex-usage-source/", import.meta.url)), + contribute: codexUsageSource, + }, +]; diff --git a/packages/server/src/server/authorization/operation-permissions.ts b/packages/server/src/server/authorization/operation-permissions.ts index 7ce3947fa8..214544c39c 100644 --- a/packages/server/src/server/authorization/operation-permissions.ts +++ b/packages/server/src/server/authorization/operation-permissions.ts @@ -149,6 +149,7 @@ const INBOUND_PERMISSION = { "project.rename.request": "workspace.manage", project_icon_request: "workspace.read", "provider.usage.list.request": "daemon.read", + "usage.list_reports.request": "daemon.read", provider_diagnostic_request: "daemon.read", pull_request_timeline_request: "workspace.read", "push.unregister.request": "workspace.read", @@ -373,6 +374,7 @@ const OUTBOUND_PERMISSION = { "project.update": "workspace.read", project_icon_response: "workspace.read", "provider.usage.list.response": "daemon.read", + "usage.list_reports.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 287b33224c..2d49f45116 100644 --- a/packages/server/src/server/plugins/index.ts +++ b/packages/server/src/server/plugins/index.ts @@ -26,6 +26,8 @@ import { PluginRuntime } from "./runtime.js"; import type { InternalPlugin } from "../../plugins/index.js"; import type { PluginProviderMetadata } from "./plugin-process-protocol.js"; import { readPluginProviderIcon } from "./provider-icon.js"; +import { UsageSourceRegistry } from "./usage-sources/index.js"; +import type { PluginUsageSourceMetadata } from "./plugin-process-protocol.js"; const BUILTIN_PROVIDER_ID_SET: ReadonlySet = new Set(BUILTIN_PROVIDER_IDS); @@ -37,6 +39,9 @@ interface PluginRuntimePort { getLogs(pluginId: string): PluginLogEntry[]; clearLogs(pluginId: string): void; getProviderRegistrations?(pluginId: string): readonly PluginProviderMetadata[]; + getUsageSourceRegistrations?(pluginId: string): readonly PluginUsageSourceMetadata[]; + fetchUsage?: PluginRuntime["fetchUsage"]; + discoverUsage?: PluginRuntime["discoverUsage"]; connectProvider: PluginRuntime["connectProvider"]; getProviderCatalogCacheKey?: PluginRuntime["getProviderCatalogCacheKey"]; validatePlugin?(path: string): Promise; @@ -72,6 +77,8 @@ export class PluginService { private readonly errors = new Map(); private readonly listeners = new Set<(pluginId: string) => void>(); private readonly providers = new Map(); + private readonly usageSources = new UsageSourceRegistry(); + private readonly usageSourceIdsByPlugin = new Map(); private readonly providerIdsByPlugin = new Map(); private readonly providerListeners = new Set<() => void>(); private lifecycle = Promise.resolve(); @@ -100,6 +107,7 @@ export class PluginService { ); this.runtime.subscribe((pluginId, error) => { this.removeProviderRegistrations(pluginId); + this.removeUsageSources(pluginId); if (error) this.errors.set(pluginId, error); this.notify(pluginId); }); @@ -134,6 +142,10 @@ export class PluginService { return [...this.providers.values()].sort((left, right) => left.id.localeCompare(right.id)); } + listUsageReports(options?: { forceRefresh?: boolean }) { + return this.usageSources.listReports(options); + } + subscribeProviderRegistrations(listener: () => void): () => void { this.providerListeners.add(listener); return () => this.providerListeners.delete(listener); @@ -148,6 +160,7 @@ export class PluginService { try { await this.runtime.startInternalPlugin?.(plugin); await this.publishProviderRegistrations(plugin.id, plugin.directory); + this.publishUsageSources(plugin.id); } catch (error) { this.logger.error({ err: error, pluginId: plugin.id }, "Failed to start internal plugin"); } @@ -498,9 +511,11 @@ export class PluginService { await this.runtime.startPlugin(pluginId, sourcePath, () => this.canPublish(pluginId)); try { await this.publishProviderRegistrations(pluginId, sourcePath); + this.publishUsageSources(pluginId); } catch (error) { try { this.removeProviderRegistrations(pluginId); + this.removeUsageSources(pluginId); } finally { await this.runtime.stopPluginById(pluginId); } @@ -510,6 +525,7 @@ export class PluginService { private stopPlugin(pluginId: string): Promise { this.removeProviderRegistrations(pluginId); + this.removeUsageSources(pluginId); return this.runtime.stopPluginById(pluginId); } @@ -524,9 +540,46 @@ export class PluginService { for (const pluginId of this.providerIdsByPlugin.keys()) { this.removeProviderRegistrations(pluginId); } + for (const pluginId of this.usageSourceIdsByPlugin.keys()) this.removeUsageSources(pluginId); await this.runtime.stopAll(); } + private publishUsageSources(pluginId: string): void { + const metadata = this.runtime.getUsageSourceRegistrations?.(pluginId) ?? []; + const registered: string[] = []; + try { + for (const source of metadata) { + this.usageSources.register({ + id: source.id, + label: source.label, + icon: source.icon, + discover: async () => { + if (!source.discover) return []; + const result = await this.runtime.discoverUsage?.(pluginId, source.id); + if (!Array.isArray(result)) + throw new Error(`Invalid usage discovery from ${source.id}`); + return result; + }, + fetch: (input) => { + if (!this.runtime.fetchUsage) throw new Error("Usage source runtime unavailable"); + return this.runtime.fetchUsage(pluginId, source.id, input); + }, + }); + registered.push(source.id); + } + } catch (error) { + for (const id of registered) this.usageSources.unregister(id); + throw error; + } + this.usageSourceIdsByPlugin.set(pluginId, registered); + } + + private removeUsageSources(pluginId: string): void { + for (const id of this.usageSourceIdsByPlugin.get(pluginId) ?? []) + this.usageSources.unregister(id); + this.usageSourceIdsByPlugin.delete(pluginId); + } + private async publishProviderRegistrations( pluginId: string, pluginDirectory: string, diff --git a/packages/server/src/server/plugins/plugin-process-protocol.ts b/packages/server/src/server/plugins/plugin-process-protocol.ts index 0066fad500..d1e3bd9aa0 100644 --- a/packages/server/src/server/plugins/plugin-process-protocol.ts +++ b/packages/server/src/server/plugins/plugin-process-protocol.ts @@ -15,12 +15,20 @@ export interface PluginProviderMetadata { iconPath?: string; } +export interface PluginUsageSourceMetadata { + id: string; + label: string; + icon?: string; + discover: boolean; +} + export type PluginProcessRequest = | { type: "initialize"; pluginId: string; bundle: string; appVersion: string; + pluginDirectory: string; settingsDirectory?: string; } | { @@ -31,6 +39,8 @@ export type PluginProcessRequest = } | { type: "hook"; requestId: string; kind: "event" | "before"; name: string; input: unknown } | { type: "hook.cancel"; requestId: string } + | { type: "usage.fetch"; requestId: string; sourceId: string; input: unknown } + | { type: "usage.discover"; requestId: string; sourceId: string } | { type: "invoke"; requestId: string; method: string; input: unknown } | { type: "provider.connect"; @@ -56,6 +66,7 @@ export type PluginProcessMessage = type: "ready"; methods: string[]; providers: PluginProviderMetadata[]; + usageSources?: PluginUsageSourceMetadata[]; hooks?: { events: string[]; before: string[] }; } | { type: "result"; requestId: string; output: unknown } @@ -91,6 +102,14 @@ const providerMetadataSchema = z hasCatalogCacheKey: z.boolean().optional(), }) .strict(); +const usageSourceMetadataSchema = z + .object({ + id: z.string().min(1), + label: z.string().min(1), + icon: z.string().optional(), + discover: z.boolean(), + }) + .strict(); const providerConnectRequestSchema = z .object({ versions: z.array(z.number().int().positive()), @@ -111,6 +130,7 @@ export const PluginProcessRequestSchema: z.ZodType = z.dis pluginId: z.string().min(1), bundle: z.string(), appVersion: z.string(), + pluginDirectory: z.string(), settingsDirectory: z.string().optional(), }) .strict(), @@ -141,6 +161,17 @@ 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.fetch"), + requestId: z.string(), + sourceId: z.string(), + input: z.unknown(), + }) + .strict(), + z + .object({ type: z.literal("usage.discover"), requestId: z.string(), sourceId: z.string() }) + .strict(), z .object({ type: z.literal("invoke"), @@ -182,6 +213,7 @@ export const PluginProcessMessageSchema: z.ZodType = z.dis type: z.literal("ready"), methods: z.array(z.string()), providers: z.array(providerMetadataSchema), + usageSources: z.array(usageSourceMetadataSchema).optional(), hooks: hooksSchema.optional(), }) .strict(), diff --git a/packages/server/src/server/plugins/plugin-process.ts b/packages/server/src/server/plugins/plugin-process.ts index 6839b60665..c3d84daa55 100644 --- a/packages/server/src/server/plugins/plugin-process.ts +++ b/packages/server/src/server/plugins/plugin-process.ts @@ -8,6 +8,8 @@ import { createRequire } from "node:module"; import * as pluginSharedRuntime from "@getpaseo/plugin"; import * as pluginProviderRuntime from "@getpaseo/plugin/server/provider"; import * as pluginAcpRuntime from "@getpaseo/plugin/server/acp"; +import * as pluginUsageRuntime from "@getpaseo/plugin/server/usage"; +import type { UsageSourceRegistration } from "@getpaseo/plugin/server/usage"; import type { SettingsDefinition, PluginRpcContract } from "@getpaseo/plugin"; import type { PluginHandlerContext, PluginServerContribution } from "@getpaseo/plugin/server"; import { fileURLToPath } from "node:url"; @@ -24,6 +26,7 @@ import { isPluginClientOnlySdkSpecifier } from "./plugin-sdk-specifiers.js"; import { createPluginClientId } from "./plugin-session-identity.js"; import { PluginSettingsStore } from "./settings/index.js"; +import { readPluginProviderIcon } from "./provider-icon.js"; const nodeRequire = createRequire(import.meta.url); function runtimeRequire(name: string): unknown { @@ -34,6 +37,7 @@ function runtimeRequire(name: string): unknown { if (name === "@getpaseo/plugin/server") return {}; if (name === "@getpaseo/plugin/server/provider") return pluginProviderRuntime; if (name === "@getpaseo/plugin/server/acp") return pluginAcpRuntime; + if (name === "@getpaseo/plugin/server/usage") return pluginUsageRuntime; if (name === "@getpaseo/plugin/client/host") throw new Error(`${name} is private to the app host`); return nodeRequire(name); @@ -93,6 +97,7 @@ export function createPluginWorker(options: { }); const handlers = new Map(); const providers = new Map(); + const usageSources = new Map(); const providerConnections = new Map< string, { connection: ProviderConnection; unsubscribe: () => void; closing?: Promise } @@ -156,6 +161,21 @@ export function createPluginWorker(options: { providers.set(id, { ...provider, id }); } + function registerUsageSource(source: UsageSourceRegistration): void { + const id = source.id.trim(); + if ( + !/^[a-z][a-z0-9._-]*$/.test(id) || + !source.label.trim() || + typeof source.fetch !== "function" || + !source.input || + typeof source.input.parseAsync !== "function" + ) { + throw new Error(`Invalid usage source: ${source.id}`); + } + if (usageSources.has(id)) throw new Error(`Duplicate usage source: ${id}`); + usageSources.set(id, { ...source, id }); + } + function providerMetadata(provider: ProviderRegistration) { return { id: provider.id, @@ -285,6 +305,7 @@ export function createPluginWorker(options: { handle: (contract, handler) => register(contract, (input, context) => handler(contract.input.parse(input), context)), registerProvider, + registerUsageSource, registerSettings, on: hooks.on, before: hooks.before, @@ -293,6 +314,18 @@ export function createPluginWorker(options: { throw new Error("Plugin contribution must return a cleanup function"); } cleanup = contributedCleanup; + const usageSourceMetadata = await Promise.all( + [...usageSources.values()] + .sort((a, b) => a.id.localeCompare(b.id)) + .map(async (source) => ({ + id: source.id, + label: source.label, + icon: source.icon + ? await readPluginProviderIcon(message.pluginDirectory, source.icon) + : undefined, + discover: !!source.discover, + })), + ); send({ type: "ready", methods: [...handlers.keys()].sort(), @@ -300,6 +333,7 @@ export function createPluginWorker(options: { providers: [...providers.values()] .sort((left, right) => left.id.localeCompare(right.id)) .map(providerMetadata), + usageSources: usageSourceMetadata, }); } @@ -327,7 +361,48 @@ export function createPluginWorker(options: { channel.disconnect(); } - channel.onMessage((rawMessage: unknown) => { + function handleUsageRequest( + message: Extract, + ): void { + void (async () => { + const source = usageSources.get(message.sourceId); + if (!source) throw new Error(`Unknown usage source: ${message.sourceId}`); + 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)); + })().then( + (output) => send({ type: "result", requestId: message.requestId, output }), + (error) => send({ type: "error", requestId: message.requestId, error: describeError(error) }), + ); + } + + function rejectWhileStopping(message: PluginProcessRequest): void { + if ( + message.type === "provider.catalog_key" || + message.type === "usage.fetch" || + message.type === "usage.discover" + ) { + send({ type: "error", requestId: message.requestId, error: "Plugin is stopping" }); + } else if (message.type === "provider.connect") { + send({ + type: "provider.connect_failed", + connectionId: message.connectionId, + error: "Plugin is stopping", + }); + } else if (message.type === "provider.send") { + send({ + type: "provider.rejected", + connectionId: message.connectionId, + acceptanceId: message.acceptanceId, + error: "Plugin is stopping", + }); + } else if (message.type === "provider.close") { + send({ type: "provider.closed", connectionId: message.connectionId }); + } + } + + function handleMessage(rawMessage: unknown): void { const parsed = PluginProcessRequestSchema.safeParse(rawMessage); if (!parsed.success) { const value = rawMessage as { connectionId?: unknown; acceptanceId?: unknown } | null; @@ -365,24 +440,7 @@ export function createPluginWorker(options: { return; } if (stopping) { - if (message.type === "provider.catalog_key") { - send({ type: "error", requestId: message.requestId, error: "Plugin is stopping" }); - } else if (message.type === "provider.connect") { - send({ - type: "provider.connect_failed", - connectionId: message.connectionId, - error: "Plugin is stopping", - }); - } else if (message.type === "provider.send") { - send({ - type: "provider.rejected", - connectionId: message.connectionId, - acceptanceId: message.acceptanceId, - error: "Plugin is stopping", - }); - } else if (message.type === "provider.close") { - send({ type: "provider.closed", connectionId: message.connectionId }); - } + rejectWhileStopping(message); return; } if (message.type === "provider.catalog_key") { @@ -398,6 +456,10 @@ export function createPluginWorker(options: { ); return; } + if (message.type === "usage.fetch" || message.type === "usage.discover") { + handleUsageRequest(message); + return; + } if (message.type === "provider.connect") { void connectProvider(message).catch((error) => { if (stopping) return; @@ -451,7 +513,8 @@ export function createPluginWorker(options: { (error) => send({ type: "error", requestId: message.requestId, error: describeError(error) }), ); - }); + } + channel.onMessage(handleMessage); function handleHookMessage( message: Extract, diff --git a/packages/server/src/server/plugins/plugin-sdk-specifiers.ts b/packages/server/src/server/plugins/plugin-sdk-specifiers.ts index 2c17db0483..079f1c501d 100644 --- a/packages/server/src/server/plugins/plugin-sdk-specifiers.ts +++ b/packages/server/src/server/plugins/plugin-sdk-specifiers.ts @@ -8,6 +8,7 @@ export const PLUGIN_CLIENT_ONLY_SDK_SPECIFIERS = [ const PLUGIN_SERVER_ONLY_SDK_SPECIFIERS = [ "@getpaseo/plugin/server", "@getpaseo/plugin/server/provider", + "@getpaseo/plugin/server/usage", "@getpaseo/plugin/server/acp", ] as const; diff --git a/packages/server/src/server/plugins/runtime.ts b/packages/server/src/server/plugins/runtime.ts index af04af06b6..5a645c94b7 100644 --- a/packages/server/src/server/plugins/runtime.ts +++ b/packages/server/src/server/plugins/runtime.ts @@ -24,6 +24,7 @@ import type { PluginProcessMessage, PluginProcessRequest, PluginProviderMetadata, + PluginUsageSourceMetadata, } from "./plugin-process-protocol.js"; import { PluginProcessMessageSchema } from "./plugin-process-protocol.js"; import { PluginSessionSocket } from "./session-socket.js"; @@ -67,6 +68,7 @@ interface LoadedPlugin { methods: ReadonlySet; hooks: { events: string[]; before: string[] }; providers: readonly PluginProviderMetadata[]; + usageSources: readonly PluginUsageSourceMetadata[]; child: PluginChild | null; outputCapture: PluginOutputCapture | null; pending: Map; @@ -346,6 +348,7 @@ export class PluginRuntime { assertPluginCompatibility({ ...manifest, version: this.daemonVersion, runtime: "daemon" }); const loaded = await this.launchPlugin({ pluginId: input.id, + pluginDirectory: directory, requirements: manifest.requirements, child: new InternalPluginChild(input.contribute), bundle: "", @@ -382,6 +385,22 @@ export class PluginRuntime { return this.plugins.get(pluginId)?.providers ?? []; } + getUsageSourceRegistrations(pluginId: string): readonly PluginUsageSourceMetadata[] { + return this.plugins.get(pluginId)?.usageSources ?? []; + } + + fetchUsage(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.fetch", requestId: randomUUID(), sourceId, input }); + } + + discoverUsage(pluginId: string, sourceId: string): Promise { + const loaded = this.plugins.get(pluginId); + if (!loaded) throw new Error(`Plugin is not available: ${pluginId}`); + return this.request(loaded, { type: "usage.discover", requestId: randomUUID(), sourceId }); + } + async connectProvider( pluginId: string, providerId: string, @@ -589,6 +608,7 @@ export class PluginRuntime { methods: new Set(), hooks: { events: [], before: [] }, providers: [], + usageSources: [], child: null, outputCapture: null, pending: new Map(), @@ -600,6 +620,7 @@ export class PluginRuntime { } return this.launchPlugin({ pluginId, + pluginDirectory: directory, requirements: manifest.requirements, child: this.spawnChild(), bundle: serverBundle, @@ -609,6 +630,7 @@ export class PluginRuntime { private async launchPlugin(input: { pluginId: string; + pluginDirectory: string; requirements: PluginRequirements | undefined; child: PluginChild; bundle: string; @@ -692,6 +714,7 @@ export class PluginRuntime { void send(child, { type: "initialize", pluginId, + pluginDirectory: input.pluginDirectory, appVersion: this.daemonVersion, bundle, settingsDirectory: this.dependencies.settingsDirectory @@ -713,6 +736,7 @@ export class PluginRuntime { methods: new Set(ready.methods), hooks: ready.hooks ?? { events: [], before: [] }, providers: ready.providers ?? [], + usageSources: ready.usageSources ?? [], child, outputCapture, pending, diff --git a/packages/server/src/server/plugins/test-fixtures/usage-source-directory/icon.svg b/packages/server/src/server/plugins/test-fixtures/usage-source-directory/icon.svg new file mode 100644 index 0000000000..b7d6b25aef --- /dev/null +++ b/packages/server/src/server/plugins/test-fixtures/usage-source-directory/icon.svg @@ -0,0 +1 @@ + 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 new file mode 100644 index 0000000000..62102d07c1 --- /dev/null +++ b/packages/server/src/server/plugins/test-fixtures/usage-source-directory/index.server.ts @@ -0,0 +1,31 @@ +import type { PluginServerContext } from "@getpaseo/plugin/server"; +import { z } from "zod"; + +let fetches = 0; +export default function contribute(server: PluginServerContext) { + server.registerUsageSource({ + id: "fixture-directory", + label: "Fixture", + icon: "icon.svg", + input: z.object({ account: z.string() }).strict(), + discover: async () => { + const inputs: Array> = [ + { account: "one" }, + { account: "bad", extra: true }, + { account: "throws" }, + ]; + return inputs; + }, + 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 }], + }; + }, + }); + return () => {}; +} diff --git a/packages/server/src/server/plugins/test-fixtures/usage-source-directory/paseo-plugin.json b/packages/server/src/server/plugins/test-fixtures/usage-source-directory/paseo-plugin.json new file mode 100644 index 0000000000..273f46f948 --- /dev/null +++ b/packages/server/src/server/plugins/test-fixtures/usage-source-directory/paseo-plugin.json @@ -0,0 +1 @@ +{ "id": "usage-source-fixture", "requirements": { "paseo": ">=0.8.0" } } diff --git a/packages/server/src/server/plugins/test-fixtures/usage-source/icon.svg b/packages/server/src/server/plugins/test-fixtures/usage-source/icon.svg new file mode 100644 index 0000000000..b7d6b25aef --- /dev/null +++ b/packages/server/src/server/plugins/test-fixtures/usage-source/icon.svg @@ -0,0 +1 @@ + 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 new file mode 100644 index 0000000000..b76cd67b42 --- /dev/null +++ b/packages/server/src/server/plugins/test-fixtures/usage-source/index.server.ts @@ -0,0 +1,31 @@ +import type { PluginServerContext } from "@getpaseo/plugin/server"; +import { z } from "zod"; + +let fetches = 0; +export default function contribute(server: PluginServerContext) { + server.registerUsageSource({ + id: "fixture", + label: "Fixture", + icon: "icon.svg", + input: z.object({ account: z.string() }).strict(), + discover: async () => { + const inputs: Array> = [ + { account: "one" }, + { account: "bad", extra: true }, + { account: "throws" }, + ]; + return inputs; + }, + 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 }], + }; + }, + }); + return () => {}; +} diff --git a/packages/server/src/server/plugins/test-fixtures/usage-source/paseo-plugin.json b/packages/server/src/server/plugins/test-fixtures/usage-source/paseo-plugin.json new file mode 100644 index 0000000000..273f46f948 --- /dev/null +++ b/packages/server/src/server/plugins/test-fixtures/usage-source/paseo-plugin.json @@ -0,0 +1 @@ +{ "id": "usage-source-fixture", "requirements": { "paseo": ">=0.8.0" } } diff --git a/packages/server/src/server/plugins/usage-sources.e2e.test.ts b/packages/server/src/server/plugins/usage-sources.e2e.test.ts new file mode 100644 index 0000000000..d5e63fc795 --- /dev/null +++ b/packages/server/src/server/plugins/usage-sources.e2e.test.ts @@ -0,0 +1,47 @@ +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/usage-source/index.server.js"; + +const directory = fileURLToPath(new URL("./test-fixtures/usage-source/", import.meta.url)); +const subprocessDirectory = fileURLToPath( + new URL("./test-fixtures/usage-source-directory/", import.meta.url), +); + +test("lists internal and subprocess usage; validates input and isolates fetch errors", async () => { + const daemon = await createTestPaseoDaemon({ + daemonVersion: "0.9.2", + pluginsEnabled: false, + internalPlugins: [{ id: "fixture-internal", directory, contribute }], + }); + const client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws`, appVersion: "0.9.2" }); + try { + await client.connect(); + 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, + ).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, + ).toBe(1); + expect( + (await client.listUsageReports({ forceRefresh: true })).reports.find( + (entry) => entry.report.account.key === "one", + )?.report.windows[0]?.usedPct, + ).toBe(2); + await client.patchDaemonConfig({ pluginsEnabled: true }); + await client.installDirectoryPlugin(subprocessDirectory, "fixture-directory"); + const both = await client.listUsageReports({ forceRefresh: true }); + expect(both.reports).toHaveLength(6); + await client.patchDaemonConfig({ pluginsEnabled: false }); + await expect.poll(async () => (await client.listUsageReports()).reports.length).toBe(3); + } finally { + await client.close(); + await daemon.close(); + } +}, 60_000); diff --git a/packages/server/src/server/plugins/usage-sources/index.test.ts b/packages/server/src/server/plugins/usage-sources/index.test.ts new file mode 100644 index 0000000000..2aa0e1f511 --- /dev/null +++ b/packages/server/src/server/plugins/usage-sources/index.test.ts @@ -0,0 +1,69 @@ +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 }], +}); + +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; + 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("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", + ]); +}); diff --git a/packages/server/src/server/plugins/usage-sources/index.ts b/packages/server/src/server/plugins/usage-sources/index.ts new file mode 100644 index 0000000000..7571b11da5 --- /dev/null +++ b/packages/server/src/server/plugins/usage-sources/index.ts @@ -0,0 +1,96 @@ +import { createHash } from "node:crypto"; +import { UsageReportSchema, type UsageReportEntry } from "@getpaseo/protocol/messages"; + +export interface UsageSource { + id: string; + label: string; + icon?: string; + discover(): Promise; + fetch(input: unknown): Promise; +} + +/** Owns running sources, account deduplication, and the five-minute fetch cache. */ +export class UsageSourceRegistry { + private readonly sources = new Map(); + private readonly cache = new Map(); + private readonly pending = new Map>(); + + constructor( + private readonly now: () => number = Date.now, + private readonly ttlMs = 300_000, + ) {} + + register(source: UsageSource): void { + if (this.sources.has(source.id)) throw new Error(`Duplicate usage source: ${source.id}`); + this.sources.set(source.id, source); + } + + unregister(id: string): void { + this.sources.delete(id); + for (const key of this.cache.keys()) if (key.startsWith(`${id}:`)) this.cache.delete(key); + } + + async listReports(options: { forceRefresh?: boolean } = {}): 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))); + } catch (error) { + return [this.errorEntry(source, error)]; + } + }), + ); + const unique = new Map(); + for (const entry of discovered.flat()) + unique.set(`${entry.sourceId}:${entry.report.account.key}`, entry); + return [...unique.values()]; + } + + 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) + return Promise.resolve(cached.entry); + const pending = this.pending.get(key); + 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 }; + } catch (error) { + entry = this.errorEntry(source, error, hash); + } + this.cache.set(key, { at: this.now(), entry }); + return entry; + })(); + this.pending.set(key, request); + void request.finally(() => { + if (this.pending.get(key) === request) this.pending.delete(key); + }); + return request; + } + + private errorEntry(source: UsageSource, error: unknown, key = source.id): UsageReportEntry { + return { + sourceId: source.id, + sourceLabel: source.label, + icon: source.icon, + 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 c69033ab3f..9347827231 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -506,6 +506,9 @@ 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; }; orchestrationSkills?: import("./orchestration-skills/index.js").OrchestrationSkills; mcpBaseUrl?: string | null; @@ -2997,11 +3000,36 @@ export class Session { return this.providerCatalogSession.handleProviderDiagnosticRequest(msg); case "provider.usage.list.request": return this.providerCatalogSession.handleProviderUsageListRequest(msg); + case "usage.list_reports.request": + return this.handleUsageListReportsRequest(msg); default: return undefined; } } + private async handleUsageListReportsRequest( + msg: Extract, + ): Promise { + try { + const reports = + (await this.pluginRuntime?.listUsageReports?.({ forceRefresh: msg.forceRefresh })) ?? []; + 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 dispatchTerminalMessage(msg: SessionInboundMessage): Promise | undefined { switch (msg.type) { case "start_workspace_script_request": diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index 7c533ce23d..e2acb7038e 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -1786,6 +1786,7 @@ export class VoiceAssistantWebSocketServer { desktopManaged: this.daemonRuntimeConfig?.desktopManaged === true, ...(this.serverCapabilities ? { capabilities: this.serverCapabilities } : {}), features: { + usageSources: true, ownedSubscriptions: true, agentRequestReceipts: true, workspaceRequestReceipts: true,