pi-opencode-go-provider 1.0.25 → 1.1.0

package/README.md

@@ -21,7 +21,7 @@ _Go-optimized endpoints for lower latency — 14+ models for [pi](https://github
 - **Cost Tracking** with per-model pricing for budget management
 - **Reasoning Models** with thinking level maps for proper effort control
 - **Prompt-cache session affinity** — sends `x-opencode-session` and `x-opencode-client` so OpenCode Go can pin a session to the same cache node
-- **Usage widget** — shows how much of the OpenCode Go 5h / 7d / 30d budgets you have spent, below the editor (footer status line outside the TUI)
+- **Usage widget** — shows how much of the OpenCode Go 5h / 7d / 30d budgets you have spent, below the editor (footer status line outside the TUI). Each account carries its own budgets, so when [pi-multiprovider](https://github.com/monotykamary/pi-multiprovider) 0.8.0+ pools several `opencode-go` accounts the widget bills the session's active one and repaints on every switch or resume.
 
 ## Installation
 
@@ -159,6 +159,17 @@ model changes. Settings are read from `~/.pi/agent/opencode-go-provider.json`:
 `placement` accepts `belowEditor` (default) or `aboveEditor`. Every key is
 optional; `/opencode-go-usage on|off` writes only `enabled`.
 
+### Pooled accounts
+
+When [pi-multiprovider](https://github.com/monotykamary/pi-multiprovider)
+0.8.0+ pools several `opencode-go` accounts, usage reads bill the session's
+active account instead of Pi's default credential: every account has its own
+5h / 7d / 30d budgets. The widget also repaints when the account changes,
+including when a resumed session restores the account last chosen with
+`/switch-account`, rather than showing the previous account until the next
+poll. Without pi-multiprovider nothing changes: the widget bills the key from
+the resolution order above.
+
 ## Authentication
 
 The opencode-go API key can be configured in multiple ways. Credentials are resolved in this order:

package/index.ts

@@ -33,6 +33,12 @@ import { USAGE_WIDGET_KEY, readUsageConfig, writeUsageConfig } from "./config.ts
 import { sanitizeStatusText, truncateToWidth } from "./format.ts";
 import { usageSegments, type UsageSegment, type UsageSeverity } from "./usage.ts";
 import { UsageController } from "./usage-controller.ts";
+import {
+  isMultiproviderService,
+  MULTIPROVIDER_SERVICE_EVENT,
+  setActiveMultiproviderService,
+  type MultiproviderService,
+} from "./multiprovider.ts";
 import modelsData from "./models.json" with { type: "json" };
 import customModelsData from "./custom-models.json" with { type: "json" };
 import patchData from "./patch.json" with { type: "json" };
@@ -421,6 +427,34 @@ export default function (pi: ExtensionAPI) {
   let usageConfig = readUsageConfig();
   let usageWidgetInstalled = false;
   const usageController = new UsageController(() => usageConfig, updateUsageWidget);
+  let usageContext: ExtensionContext | undefined;
+
+  // Follow pi-multiprovider's active pooled account. Usage is per-account, so a
+  // switch — and a resume, which replays the account the session last switched
+  // to — must repaint the widget instead of waiting for the next poll. Without
+  // pi-multiprovider nothing here activates and resolution stays unchanged.
+  let multiproviderService: MultiproviderService | undefined;
+  let unsubscribeMultiprovider: (() => void) | undefined;
+  const refreshUsageForActiveAccount = (ctx: ExtensionContext | undefined): void => {
+    if (ctx === undefined) return;
+    void usageController.refresh(ctx, { force: true });
+  };
+  if (typeof pi.events?.on === "function") {
+    pi.events.on(MULTIPROVIDER_SERVICE_EVENT, (value: unknown) => {
+      if (!isMultiproviderService(value)) return;
+      if (value !== multiproviderService) {
+        unsubscribeMultiprovider?.();
+        multiproviderService = value;
+        setActiveMultiproviderService(value);
+        unsubscribeMultiprovider = value.onActiveAccountChanged(PROVIDER_ID, (event) => {
+          refreshUsageForActiveAccount(usageContext ?? event.ctx);
+        });
+      }
+      // The event re-fires at every session start with the same stable object,
+      // so this also catches a service that appeared mid-session.
+      refreshUsageForActiveAccount(usageContext);
+    });
+  }
 
   const USAGE_SEVERITY_COLORS: Record<UsageSeverity, ThemeColor> = {
     ok: "success",
@@ -510,6 +544,7 @@ export default function (pi: ExtensionAPI) {
   });
 
   pi.on("session_start", async (_event, ctx) => {
+    usageContext = ctx;
     usageController.start(ctx);
     revalidateAbort?.abort();
     revalidateAbort = new AbortController();
@@ -551,6 +586,11 @@ export default function (pi: ExtensionAPI) {
 
   pi.on("session_shutdown", () => {
     revalidateAbort?.abort();
+    usageContext = undefined;
+    unsubscribeMultiprovider?.();
+    unsubscribeMultiprovider = undefined;
+    multiproviderService = undefined;
+    setActiveMultiproviderService(undefined);
     usageController.shutdown();
   });
 }

package/multiprovider.ts

@@ -0,0 +1,74 @@
+/**
+ * Soft bridge to pi-multiprovider.
+ *
+ * When pi-multiprovider pools several opencode-go accounts it announces a
+ * service on Pi's event bus and notifies followers whenever the session's
+ * active account changes — including a resume, which replays the account the
+ * session last switched to. Usage is account-scoped: each account has its own
+ * 5h / 7d / 30d budget, so the widget has to bill the active pooled account
+ * instead of whichever credential Pi resolves on its own. Without that
+ * extension everything here stays inert and resolution is unchanged.
+ */
+
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+
+export const MULTIPROVIDER_SERVICE_EVENT = "pi-multiprovider:service";
+
+export type MultiproviderActiveAccount = {
+  id: string;
+  label: string;
+  authKind: string;
+};
+
+export type MultiproviderAccountAuth = {
+  accessToken: string;
+  label: string;
+  source?: string;
+};
+
+export type MultiproviderAccountChangedEvent = {
+  providerId: string;
+  account: MultiproviderActiveAccount | undefined;
+  ctx: ExtensionContext;
+};
+
+export type MultiproviderServiceContext = Pick<
+  ExtensionContext,
+  "modelRegistry" | "model" | "sessionManager"
+>;
+
+export type MultiproviderService = {
+  getActiveAccount(
+    providerId: string,
+    ctx: MultiproviderServiceContext,
+  ): Promise<MultiproviderActiveAccount | undefined>;
+  resolveActiveAccountAuth(
+    providerId: string,
+    ctx: MultiproviderServiceContext,
+    signal?: AbortSignal,
+  ): Promise<MultiproviderAccountAuth | undefined>;
+  onActiveAccountChanged(
+    providerId: string,
+    callback: (event: MultiproviderAccountChangedEvent) => void,
+  ): () => void;
+};
+
+export function isMultiproviderService(value: unknown): value is MultiproviderService {
+  if (typeof value !== "object" || value === null) return false;
+  const candidate = value as Partial<MultiproviderService>;
+  return (
+    typeof candidate.getActiveAccount === "function" &&
+    typeof candidate.resolveActiveAccountAuth === "function" &&
+    typeof candidate.onActiveAccountChanged === "function"
+  );
+}
+
+let activeService: MultiproviderService | undefined;
+
+export function setActiveMultiproviderService(service: MultiproviderService | undefined): void {
+  activeService = service;
+}
+
+export function getActiveMultiproviderService(): MultiproviderService | undefined {
+  return activeService;
+}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-opencode-go-provider",
-  "version": "1.0.25",
+  "version": "1.1.0",
   "description": "Opencode Go provider extension for pi - Fast, efficient GLM, Kimi, and MiniMax models through the opencode.ai API",
   "type": "module",
   "main": "index.ts",

package/tests/multiprovider.test.ts

@@ -0,0 +1,132 @@
+import assert from "node:assert/strict";
+import { test } from "node:test";
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { DEFAULT_USAGE_CONFIG, PROVIDER_ID } from "../config.ts";
+import {
+  isMultiproviderService,
+  setActiveMultiproviderService,
+  type MultiproviderAccountAuth,
+  type MultiproviderService,
+} from "../multiprovider.ts";
+import { UsageController } from "../usage-controller.ts";
+
+function fakeService(
+  resolve: () => Promise<MultiproviderAccountAuth | undefined>,
+): { value: MultiproviderService; resolveActiveAccountAuth: () => Promise<MultiproviderAccountAuth | undefined> } {
+  const resolveActiveAccountAuth = async () => resolve();
+  const value: MultiproviderService = {
+    getActiveAccount: async () => undefined,
+    resolveActiveAccountAuth,
+    onActiveAccountChanged: () => () => {},
+  };
+  return { value, resolveActiveAccountAuth };
+}
+
+function sessionContext(registryKey: string | undefined): ExtensionContext {
+  return {
+    model: { provider: PROVIDER_ID, id: "kimi-k3" },
+    hasUI: true,
+    signal: undefined,
+    ui: { notify: () => {} },
+    modelRegistry: { getApiKeyForProvider: async () => registryKey },
+    sessionManager: { getSessionId: () => "session-1" },
+  } as unknown as ExtensionContext;
+}
+
+function makeController(): UsageController {
+  return new UsageController(() => DEFAULT_USAGE_CONFIG, () => {});
+}
+
+/** Records the bearer token each usage request billed. */
+function stubUsageFetch(): { headers: string[]; restore: () => void } {
+  const original = globalThis.fetch;
+  const headers: string[] = [];
+  globalThis.fetch = (async (_input: unknown, init?: RequestInit) => {
+    const requestHeaders = (init?.headers ?? {}) as Record<string, string>;
+    headers.push(String(requestHeaders.authorization ?? ""));
+    return new Response(
+      JSON.stringify({
+        usage: {
+          rolling: { status: "ok", percent: 10, resetsAt: null },
+          weekly: { status: "ok", percent: 20, resetsAt: null },
+          monthly: { status: "ok", percent: 30, resetsAt: null },
+        },
+      }),
+    );
+  }) as typeof fetch;
+  return {
+    headers,
+    restore() {
+      globalThis.fetch = original;
+    },
+  };
+}
+
+test("detects the pi-multiprovider service payload", () => {
+  assert.equal(isMultiproviderService(fakeService(async () => undefined).value), true);
+  assert.equal(isMultiproviderService(undefined), false);
+  assert.equal(isMultiproviderService({ getActiveAccount: () => {} }), false);
+  assert.equal(
+    isMultiproviderService({
+      getActiveAccount: () => {},
+      resolveActiveAccountAuth: () => {},
+      onActiveAccountChanged: null,
+    }),
+    false,
+  );
+});
+
+test("bills the session's pooled account when pi-multiprovider pins one", async () => {
+  const fetchStub = stubUsageFetch();
+  const service = fakeService(async () => ({ accessToken: "pooled-key", label: "Work" }));
+  setActiveMultiproviderService(service.value);
+  try {
+    const controller = makeController();
+    await controller.refresh(sessionContext("registry-key"), { force: true });
+    assert.deepEqual(fetchStub.headers, ["Bearer pooled-key"]);
+    assert.ok(controller.snapshot);
+  } finally {
+    setActiveMultiproviderService(undefined);
+    fetchStub.restore();
+  }
+});
+
+test("falls back to the registry key when the session has no pooled account", async () => {
+  const fetchStub = stubUsageFetch();
+  const service = fakeService(async () => undefined);
+  setActiveMultiproviderService(service.value);
+  try {
+    const controller = makeController();
+    await controller.refresh(sessionContext("registry-key"), { force: true });
+    assert.deepEqual(fetchStub.headers, ["Bearer registry-key"]);
+  } finally {
+    setActiveMultiproviderService(undefined);
+    fetchStub.restore();
+  }
+});
+
+test("falls back to the registry key when the bridge rejects", async () => {
+  const fetchStub = stubUsageFetch();
+  const service = fakeService(() => Promise.reject(new Error("store locked")));
+  setActiveMultiproviderService(service.value);
+  try {
+    const controller = makeController();
+    await controller.refresh(sessionContext("registry-key"), { force: true });
+    assert.deepEqual(fetchStub.headers, ["Bearer registry-key"]);
+  } finally {
+    setActiveMultiproviderService(undefined);
+    fetchStub.restore();
+  }
+});
+
+test("keeps its own resolution when pi-multiprovider is absent", async () => {
+  const fetchStub = stubUsageFetch();
+  setActiveMultiproviderService(undefined);
+  try {
+    const controller = makeController();
+    await controller.refresh(sessionContext("registry-key"), { force: true });
+    assert.deepEqual(fetchStub.headers, ["Bearer registry-key"]);
+  } finally {
+    fetchStub.restore();
+  }
+});

package/usage-controller.ts

@@ -8,6 +8,7 @@
 
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 import { PROVIDER_ID, configPath, type UsageConfig } from "./config.ts";
+import { getActiveMultiproviderService } from "./multiprovider.ts";
 import {
   USAGE_LIMITS_NOTE,
   USAGE_URL,
@@ -156,7 +157,27 @@ export class UsageController {
     }
   }
 
+  /**
+   * The session's active pooled account, when pi-multiprovider pools
+   * opencode-go. Every account carries its own budget, so a switched or
+   * restored account must bill itself instead of Pi's default credential.
+   */
+  private async resolvePooledApiKey(ctx: ExtensionContext): Promise<string | undefined> {
+    const service = getActiveMultiproviderService();
+    if (service === undefined) return undefined;
+    try {
+      const resolved = await service.resolveActiveAccountAuth(PROVIDER_ID, ctx);
+      const token = resolved?.accessToken.trim();
+      return token ? token : undefined;
+    } catch {
+      // A failing bridge must never block the default resolution below.
+      return undefined;
+    }
+  }
+
   private async resolveApiKey(ctx: ExtensionContext): Promise<string | undefined> {
+    const pooled = await this.resolvePooledApiKey(ctx);
+    if (pooled !== undefined) return pooled;
     let key: string | undefined;
     try {
       key = await ctx.modelRegistry.getApiKeyForProvider(PROVIDER_ID);