refactor(usage): own usage RPCs and validate session auth

2026-11-06 08:24:46 +00:00 · 2026-09-26 20:52:57 +02:00
parent 0612bf9348
commit 42be71e5a7
9 changed files with 329 additions and 173 deletions
@@ -392,6 +392,35 @@ describe("OpenCodeAgentClient adapter smoke tests", () => {
    rmSync(cwd, { recursive: true, force: true });
  });

+  test("usage reference ignores another auth entry with a different type", async () => {
+    const cwd = tmpCwd();
+    const runtime = new TestOpenCodeHarness();
+    runtime.enqueueClient(new TestOpenCodeClient());
+    const client = new OpenCodeAgentClient(
+      logger,
+      {
+        env: {
+          OPENCODE_AUTH_CONTENT: JSON.stringify({
+            openai: { type: "api", key: "other-key" },
+            "opencode-go": { type: "api", key: "go-key" },
+          }),
+        },
+      },
+      { serverManager: runtime, createClient: runtime.createClient },
+    );
+    const session = await client.createSession(buildConfig(cwd));
+    try {
+      await session.setModel?.("opencode-go/qwen");
+      expect(await session.getUsageReference?.()).toEqual({
+        source: "opencode-go",
+        input: { apiKey: "go-key" },
+      });
+    } finally {
+      await session.close();
+      rmSync(cwd, { recursive: true, force: true });
+    }
+  });
+
  test("creates a session with valid id and provider", async () => {
    const cwd = tmpCwd();
    const runtime = new TestOpenCodeHarness();
@@ -3365,9 +3365,28 @@ interface OpenCodeServerConnection {
  release: () => Promise<void>;
 }

-async function readOpenCodeUsageAuth(
-  env: NodeJS.ProcessEnv,
-): Promise<Record<string, unknown> | null> {
+const openCodeOAuthAuthSchema = z
+  .object({
+    openai: z
+      .object({
+        type: z.literal("oauth"),
+        access: z.string().min(1),
+        accountId: z.string().optional(),
+      })
+      .passthrough()
+      .optional(),
+  })
+  .passthrough();
+const openCodeGoAuthSchema = z
+  .object({
+    "opencode-go": z
+      .object({ type: z.literal("api"), key: z.string().min(1) })
+      .passthrough()
+      .optional(),
+  })
+  .passthrough();
+
+async function readOpenCodeUsageAuth(env: NodeJS.ProcessEnv): Promise<unknown> {
  try {
    const content =
      env.OPENCODE_AUTH_CONTENT ??
@@ -3379,39 +3398,31 @@ async function readOpenCodeUsageAuth(
        ),
        "utf8",
      ));
-    const parsed: unknown = JSON.parse(content);
-    return parsed && typeof parsed === "object" && !Array.isArray(parsed)
-      ? (parsed as Record<string, unknown>)
-      : null;
+    return JSON.parse(content);
  } catch {
    return null;
  }
 }

-function resolveOpenCodeUsageReference(
-  model: string,
-  auth: Record<string, unknown> | null,
-): UsageReference | null {
+function resolveOpenCodeUsageReference(model: string, auth: unknown): UsageReference | null {
  if (!auth) return null;
  if (model.startsWith("openai/")) {
-    const entry = auth.openai;
-    if (!entry || typeof entry !== "object") return null;
-    const credential = entry as Record<string, unknown>;
-    if (credential.type !== "oauth" || typeof credential.access !== "string" || !credential.access)
-      return null;
+    const parsed = openCodeOAuthAuthSchema.safeParse(auth);
+    if (!parsed.success) return null;
+    const credential = parsed.data.openai;
+    if (!credential) return null;
    return {
      source: "codex",
      input: {
        accessToken: credential.access,
-        ...(typeof credential.accountId === "string" ? { accountId: credential.accountId } : {}),
+        ...(credential.accountId ? { accountId: credential.accountId } : {}),
      },
    };
  }
-  const entry = auth["opencode-go"];
-  if (!entry || typeof entry !== "object") return null;
-  const credential = entry as Record<string, unknown>;
-  if (credential.type !== "api" || typeof credential.key !== "string" || !credential.key)
-    return null;
+  const parsed = openCodeGoAuthSchema.safeParse(auth);
+  if (!parsed.success) return null;
+  const credential = parsed.data["opencode-go"];
+  if (!credential) return null;
  return { source: "opencode-go", input: { apiKey: credential.key } };
 }

@@ -241,6 +241,19 @@ test("Pi usage reference follows the current OAuth model and agent directory", a
      source: "claude",
      input: { accessToken: "claude-token" },
    });
+    writeFileSync(
+      path.join(agentDir, "auth.json"),
+      JSON.stringify({
+        "openai-codex": { type: "oauth", access: "codex-token" },
+        anthropic: { type: "api_key", key: "api-key" },
+      }),
+    );
+    runtime.state.model = { provider: "openai-codex", id: "gpt", name: "GPT" };
+    expect(await session.getUsageReference?.()).toEqual({
+      source: "codex",
+      input: { accessToken: "codex-token" },
+    });
+    runtime.state.model = { provider: "anthropic", id: "claude", name: "Claude" };
    writeFileSync(
      path.join(agentDir, "auth.json"),
      JSON.stringify({ anthropic: { type: "api_key", key: "api-key" } }),
@@ -1,5 +1,6 @@
 import { randomUUID } from "node:crypto";
 import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
+import { readFile } from "node:fs/promises";
 import { homedir, tmpdir } from "node:os";
 import { join, resolve as resolvePath } from "node:path";
 import type { Logger } from "pino";
@@ -529,6 +530,27 @@ function resolvePiAgentDir(env: Record<string, string> | undefined): string {
  return resolvePath(configured);
 }

+const piCodexUsageAuthSchema = z
+  .object({
+    "openai-codex": z
+      .object({
+        type: z.literal("oauth"),
+        access: z.string().min(1),
+        accountId: z.string().optional(),
+      })
+      .passthrough()
+      .optional(),
+  })
+  .passthrough();
+const piClaudeUsageAuthSchema = z
+  .object({
+    anthropic: z
+      .object({ type: z.literal("oauth"), access: z.string().min(1) })
+      .passthrough()
+      .optional(),
+  })
+  .passthrough();
+
 function readPiGlobalMcpConfig(env: Record<string, string> | undefined): Record<string, unknown> {
  const globalConfigPath = join(resolvePiAgentDir(env), "mcp.json");
  if (!existsSync(globalConfigPath)) {
@@ -1131,27 +1153,22 @@ export class PiRpcAgentSession implements AgentSession {
    if (provider === "anthropic") source = "claude";
    if (!source || !provider) return null;
    try {
-      const credentials = JSON.parse(
-        readFileSync(join(resolvePiAgentDir(this.usageEnv), "auth.json"), "utf8"),
-      ) as Record<string, unknown>;
-      const entry = credentials[provider];
-      if (!entry || typeof entry !== "object") return null;
-      const credential = entry as Record<string, unknown>;
-      if (
-        credential.type !== "oauth" ||
-        typeof credential.access !== "string" ||
-        !credential.access
-      )
-        return null;
-      return {
-        source,
-        input: {
-          accessToken: credential.access,
-          ...(source === "codex" && typeof credential.accountId === "string" && credential.accountId
-            ? { accountId: credential.accountId }
-            : {}),
-        },
-      };
+      const auth: unknown = JSON.parse(
+        await readFile(join(resolvePiAgentDir(this.usageEnv), "auth.json"), "utf8"),
+      );
+      if (provider === "openai-codex") {
+        const credential = piCodexUsageAuthSchema.parse(auth)["openai-codex"];
+        if (!credential) return null;
+        return {
+          source,
+          input: {
+            accessToken: credential.access,
+            ...(credential.accountId ? { accountId: credential.accountId } : {}),
+          },
+        };
+      }
+      const credential = piClaudeUsageAuthSchema.parse(auth).anthropic;
+      return credential ? { source, input: { accessToken: credential.access } } : null;
    } catch {
      return null;
    }
@@ -180,6 +180,7 @@ import {
 } from "./session/checkout/git-metadata-generator.js";
 import { ScheduleSession } from "./session/schedule/schedule-session.js";
 import { ProviderCatalogSession } from "./session/provider/provider-catalog-session.js";
+import { UsageSession } from "./session/usage/usage-session.js";
 import { WorkspaceFilesSession } from "./session/files/workspace-files-session.js";
 import { AgentConfigSession } from "./session/agent-config/agent-config-session.js";
 import { ProjectConfigSession } from "./session/project-config/project-config-session.js";
@@ -790,6 +791,7 @@ export class Session {
  private readonly checkoutSession: CheckoutSession;
  private readonly scheduleSession: ScheduleSession;
  private readonly providerCatalogSession: ProviderCatalogSession;
+  private readonly usageSession: UsageSession;
  private readonly workspaceFilesSession: WorkspaceFilesSession;
  private readonly agentConfigSession: AgentConfigSession;
  private readonly projectConfigSession: ProjectConfigSession;
@@ -995,10 +997,13 @@ export class Session {
        listDraftFeatures: (config) => this.agentManager.listDraftFeatures(config),
      },
      providerSnapshotManager,
-      listLegacyUsage: async () => {
-        if (!pluginRuntime) throw new Error("Plugin runtime is unavailable");
-        return pluginRuntime.listLegacyUsage();
-      },
+      logger: this.sessionLogger,
+    });
+    this.usageSession = new UsageSession({
+      emit: (msg) => this.emit(msg),
+      listAgents: () => this.agentManager.listAgents(),
+      getAgent: (agentId) => this.agentManager.getAgent(agentId),
+      runtime: pluginRuntime,
      logger: this.sessionLogger,
    });
    this.agentConfigSession = new AgentConfigSession({
@@ -3010,85 +3015,16 @@ export class Session {
      case "provider_diagnostic_request":
        return this.providerCatalogSession.handleProviderDiagnosticRequest(msg);
      case "provider.usage.list.request":
-        return this.providerCatalogSession.handleProviderUsageListRequest(msg);
+        return this.usageSession.handleLegacyList(msg);
      case "usage.list_reports.request":
-        return this.handleUsageListReportsRequest(msg);
+        return this.usageSession.handleListReports(msg);
      case "agent.get_usage_report.request":
-        return this.handleAgentGetUsageReportRequest(msg);
+        return this.usageSession.handleGetAgentReport(msg);
      default:
        return undefined;
    }
  }

-  private async handleUsageListReportsRequest(
-    msg: Extract<SessionInboundMessage, { type: "usage.list_reports.request" }>,
-  ): Promise<void> {
-    try {
-      if (!this.pluginRuntime) throw new Error("Plugin runtime is unavailable");
-      const references = (
-        await Promise.all(
-          this.agentManager
-            .listAgents()
-            .map(async (agent) => agent.session?.getUsageReference?.().catch(() => null) ?? null),
-        )
-      ).filter(
-        (
-          reference,
-        ): reference is {
-          source: string;
-          input: import("@getpaseo/protocol/agent-types").JsonValue;
-        } => reference !== null,
-      );
-      const reports = await this.pluginRuntime.listUsageReports({
-        forceRefresh: msg.forceRefresh,
-        references,
-      });
-      this.emit({
-        type: "usage.list_reports.response",
-        payload: { requestId: msg.requestId, reports },
-      });
-    } catch (error) {
-      this.emit({
-        type: "rpc_error",
-        payload: {
-          requestId: msg.requestId,
-          requestType: msg.type,
-          error: error instanceof Error ? error.message : String(error),
-          code: "usage_list_reports_failed",
-        },
-      });
-    }
-  }
-
-  private async handleAgentGetUsageReportRequest(
-    msg: Extract<SessionInboundMessage, { type: "agent.get_usage_report.request" }>,
-  ): Promise<void> {
-    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<void> | undefined {
    switch (msg.type) {
      case "start_workspace_script_request":
@@ -23,7 +23,6 @@ interface MakeOptions {
  supportsCustomModeIcons?: boolean;
  supportsCompactProviderSnapshots?: boolean;
  snapshot?: Partial<ProviderSnapshotManager>;
-  usage?: () => Promise<{ fetchedAt: string; providers: [] }>;
  host?: Partial<ProviderCatalogSessionHost>;
 }

@@ -72,8 +71,6 @@ function makeSubsystem(options: MakeOptions = {}) {
  const subsystem = new ProviderCatalogSession({
    host,
    providerSnapshotManager,
-    listLegacyUsage:
-      options.usage ?? (async () => ({ fetchedAt: "2026-01-01T00:00:00.000Z", providers: [] })),
    logger: pino({ level: "silent" }),
  });
  function pushSnapshotChange(
@@ -299,23 +296,6 @@ describe("ProviderCatalogSession", () => {
    });
  });

-  it("surfaces a usage-list failure as an rpc_error envelope", async () => {
-    const { subsystem, emitted } = makeSubsystem({
-      usage: async () => {
-        throw new Error("quota service down");
-      },
-    });
-
-    await subsystem.handleProviderUsageListRequest({
-      type: "provider.usage.list.request",
-      requestId: "u1",
-    });
-
-    const err = findByType(emitted, "rpc_error");
-    expect(err?.payload.code).toBe("provider_usage_list_failed");
-    expect(err?.payload.requestId).toBe("u1");
-  });
-
  it("surfaces a feature-list failure inline, not as an rpc_error", async () => {
    const { subsystem, emitted } = makeSubsystem({
      host: {
@@ -386,7 +366,6 @@ it("announces shared content without retransmitting models or hashing discovery
    new ProviderCatalogSession({
      providerSnapshotManager: manager,
      logger: pino({ level: "silent" }),
-      listLegacyUsage: async () => ({ fetchedAt: "2026-01-01T00:00:00.000Z", providers: [] }),
      host: {
        emit(message) {
          emitted.push(message);
@@ -19,7 +19,6 @@ import {
  type ProviderSnapshotEntry,
 } from "../../agent/agent-sdk-types.js";
 import type { ProviderAvailability } from "../../agent/agent-manager.js";
-import type { ProviderUsage } from "@getpaseo/protocol/messages";
 import { expandTilde } from "../../../utils/path.js";

 // COMPAT(customModeIcons): the only mode icons known to clients before v0.1.84. Any
@@ -56,13 +55,12 @@ export interface ProviderCatalogSessionHost {
 export interface ProviderCatalogSessionOptions {
  host: ProviderCatalogSessionHost;
  providerSnapshotManager: ProviderSnapshotManager;
-  listLegacyUsage(): Promise<{ fetchedAt: string; providers: ProviderUsage[] }>;
  logger: pino.Logger;
 }

 /**
 * A client's provider catalog surface: model / mode / feature listing, the providers
- * snapshot push + pull, provider diagnostics, and usage. The snapshot PUSH (start) and
+ * snapshot push + pull, and provider diagnostics. The snapshot PUSH (start) and
 * every PULL handler gate visibility and downgrade mode icons through the SAME predicates,
 * so an older client sees one consistent provider set across both paths — the COMPAT
 * invariant the shell could only enforce by code proximity before this carve.
@@ -70,14 +68,12 @@ export interface ProviderCatalogSessionOptions {
 export class ProviderCatalogSession {
  private readonly host: ProviderCatalogSessionHost;
  private readonly providerSnapshotManager: ProviderSnapshotManager;
-  private readonly listLegacyUsage: ProviderCatalogSessionOptions["listLegacyUsage"];
  private readonly logger: pino.Logger;
  private unsubscribeSnapshotEvents: (() => void) | null = null;

  constructor(options: ProviderCatalogSessionOptions) {
    this.host = options.host;
    this.providerSnapshotManager = options.providerSnapshotManager;
-    this.listLegacyUsage = options.listLegacyUsage;
    this.logger = options.logger;
  }

@@ -486,34 +482,6 @@ export class ProviderCatalogSession {
      });
    }
  }
-
-  async handleProviderUsageListRequest(
-    msg: Extract<SessionInboundMessage, { type: "provider.usage.list.request" }>,
-  ): Promise<void> {
-    try {
-      const usage = await this.listLegacyUsage();
-      this.host.emit({
-        type: "provider.usage.list.response",
-        payload: {
-          requestId: msg.requestId,
-          fetchedAt: usage.fetchedAt,
-          providers: usage.providers,
-        },
-      });
-    } catch (error) {
-      const err = error instanceof Error ? error : new Error(String(error));
-      this.logger.error({ err }, "Failed to list provider usage");
-      this.host.emit({
-        type: "rpc_error",
-        payload: {
-          requestId: msg.requestId,
-          requestType: msg.type,
-          error: `Failed to list provider usage: ${err.message}`,
-          code: "provider_usage_list_failed",
-        },
-      });
-    }
-  }
 }

 function resolveCatalogRequestCwd(cwd?: string | null): string | undefined {
@@ -0,0 +1,78 @@
+import pino from "pino";
+import { expect, test } from "vitest";
+import type { AgentSession } from "../../agent/agent-sdk-types.js";
+import type { SessionOutboundMessage } from "../../messages.js";
+import { UsageSession } from "./usage-session.js";
+
+test("collects live references for usage reports and resolves an agent report", async () => {
+  const emitted: SessionOutboundMessage[] = [];
+  const references: unknown[] = [];
+  const reference = { source: "fixture", input: { account: "one" } };
+  const agent = {
+    session: {
+      getUsageReference: async () => reference,
+    } as AgentSession,
+  };
+  const entry = {
+    sourceId: "fixture",
+    sourceLabel: "Fixture",
+    report: { account: { key: "one" }, status: "available" as const, windows: [] },
+  };
+  const usage = new UsageSession({
+    emit: (message) => emitted.push(message),
+    listAgents: () => [agent, { session: null }],
+    getAgent: (id) => (id === "one" ? agent : null),
+    runtime: {
+      async listUsageReports(options) {
+        references.push(options.references);
+        return [entry];
+      },
+      async fetchUsageReference(value) {
+        references.push(value);
+        return entry;
+      },
+      async listLegacyUsage() {
+        return { fetchedAt: "2026-01-01T00:00:00.000Z", providers: [] };
+      },
+    },
+    logger: pino({ level: "silent" }),
+  });
+
+  await usage.handleListReports({ type: "usage.list_reports.request", requestId: "list" });
+  await usage.handleGetAgentReport({
+    type: "agent.get_usage_report.request",
+    requestId: "one",
+    agentId: "one",
+  });
+  expect(references).toEqual([[reference], reference]);
+  expect(emitted.map((message) => message.type)).toEqual([
+    "usage.list_reports.response",
+    "agent.get_usage_report.response",
+  ]);
+});
+
+test("surfaces a legacy usage-list failure as an rpc_error envelope", async () => {
+  const emitted: SessionOutboundMessage[] = [];
+  const usage = new UsageSession({
+    emit: (message) => emitted.push(message),
+    listAgents: () => [],
+    getAgent: () => null,
+    runtime: {
+      async listUsageReports() {
+        return [];
+      },
+      async fetchUsageReference() {
+        return null;
+      },
+      async listLegacyUsage(): Promise<never> {
+        throw new Error("quota service down");
+      },
+    },
+    logger: pino({ level: "silent" }),
+  });
+  await usage.handleLegacyList({ type: "provider.usage.list.request", requestId: "u1" });
+  expect(emitted[0]).toMatchObject({
+    type: "rpc_error",
+    payload: { requestId: "u1", code: "provider_usage_list_failed" },
+  });
+});
@@ -0,0 +1,125 @@
+import type pino from "pino";
+import type { ProviderUsage, UsageReportEntry } from "@getpaseo/protocol/messages";
+import type { AgentSession, UsageReference } from "../../agent/agent-sdk-types.js";
+import type { SessionInboundMessage, SessionOutboundMessage } from "../../messages.js";
+
+interface UsageAgent {
+  session: AgentSession | null;
+}
+
+export interface UsageSessionOptions {
+  emit(message: SessionOutboundMessage): void;
+  listAgents(): UsageAgent[];
+  getAgent(agentId: string): UsageAgent | null;
+  runtime?: {
+    listUsageReports(options: {
+      forceRefresh?: boolean;
+      references: UsageReference[];
+    }): Promise<UsageReportEntry[]>;
+    fetchUsageReference(
+      reference: UsageReference,
+      options: { forceRefresh?: boolean },
+    ): Promise<UsageReportEntry | null>;
+    listLegacyUsage(): Promise<{ fetchedAt: string; providers: ProviderUsage[] }>;
+  };
+  logger: pino.Logger;
+}
+
+export class UsageSession {
+  constructor(private readonly options: UsageSessionOptions) {}
+
+  async handleListReports(
+    msg: Extract<SessionInboundMessage, { type: "usage.list_reports.request" }>,
+  ): Promise<void> {
+    try {
+      if (!this.options.runtime) throw new Error("Plugin runtime is unavailable");
+      const references = (
+        await Promise.all(
+          this.options
+            .listAgents()
+            .map(async (agent) => agent.session?.getUsageReference?.().catch(() => null) ?? null),
+        )
+      ).filter((reference): reference is UsageReference => reference !== null);
+      const reports = await this.options.runtime.listUsageReports({
+        forceRefresh: msg.forceRefresh,
+        references,
+      });
+      this.options.emit({
+        type: "usage.list_reports.response",
+        payload: { requestId: msg.requestId, reports },
+      });
+    } catch (error) {
+      this.emitError(msg, error, "usage_list_reports_failed");
+    }
+  }
+
+  async handleGetAgentReport(
+    msg: Extract<SessionInboundMessage, { type: "agent.get_usage_report.request" }>,
+  ): Promise<void> {
+    try {
+      if (!this.options.runtime) throw new Error("Plugin runtime is unavailable");
+      const reference =
+        (await this.options.getAgent(msg.agentId)?.session?.getUsageReference?.()) ?? null;
+      const entry = reference
+        ? await this.options.runtime.fetchUsageReference(reference, {
+            forceRefresh: msg.forceRefresh,
+          })
+        : null;
+      this.options.emit({
+        type: "agent.get_usage_report.response",
+        payload: { requestId: msg.requestId, entry },
+      });
+    } catch (error) {
+      this.emitError(msg, error, "agent_get_usage_report_failed");
+    }
+  }
+
+  // COMPAT(providerUsageList): added in v0.9.3, remove after 2027-03-26.
+  async handleLegacyList(
+    msg: Extract<SessionInboundMessage, { type: "provider.usage.list.request" }>,
+  ): Promise<void> {
+    try {
+      if (!this.options.runtime) throw new Error("Plugin runtime is unavailable");
+      const usage = await this.options.runtime.listLegacyUsage();
+      this.options.emit({
+        type: "provider.usage.list.response",
+        payload: {
+          requestId: msg.requestId,
+          fetchedAt: usage.fetchedAt,
+          providers: usage.providers,
+        },
+      });
+    } catch (error) {
+      const err = error instanceof Error ? error : new Error(String(error));
+      this.options.logger.error({ err }, "Failed to list provider usage");
+      this.options.emit({
+        type: "rpc_error",
+        payload: {
+          requestId: msg.requestId,
+          requestType: msg.type,
+          error: `Failed to list provider usage: ${err.message}`,
+          code: "provider_usage_list_failed",
+        },
+      });
+    }
+  }
+
+  private emitError(
+    msg: Extract<
+      SessionInboundMessage,
+      { type: "usage.list_reports.request" | "agent.get_usage_report.request" }
+    >,
+    error: unknown,
+    code: string,
+  ): void {
+    this.options.emit({
+      type: "rpc_error",
+      payload: {
+        requestId: msg.requestId,
+        requestType: msg.type,
+        error: error instanceof Error ? error.message : String(error),
+        code,
+      },
+    });
+  }
+}