takomi 2.0.6 → 2.1.0

package/.pi/extensions/oauth-router/oauth-flow.ts

@@ -0,0 +1,154 @@
+import { spawn } from "node:child_process";
+import { randomUUID } from "node:crypto";
+import type { OAuthCredentials } from "@mariozechner/pi-ai";
+import { getOAuthProvider } from "@mariozechner/pi-ai/oauth";
+import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
+import type { RouterUpstreamConfig, StoredRouterAccount } from "./types.ts";
+
+function now() {
+  return Date.now();
+}
+
+function normalizeCredentials(credentials: OAuthCredentials) {
+  const { access, refresh, expires, ...meta } = credentials;
+  return {
+    access,
+    refresh,
+    expires,
+    meta,
+  };
+}
+
+export function openUrlInBrowser(url: string) {
+  const platform = process.platform;
+
+  try {
+    if (platform === "win32") {
+      const child = spawn("rundll32.exe", ["url.dll,FileProtocolHandler", url], {
+        detached: true,
+        stdio: "ignore",
+      });
+      child.unref();
+      return;
+    }
+
+    if (platform === "darwin") {
+      const child = spawn("open", [url], { detached: true, stdio: "ignore" });
+      child.unref();
+      return;
+    }
+
+    const child = spawn("xdg-open", [url], { detached: true, stdio: "ignore" });
+    child.unref();
+  } catch {
+    // Best effort only.
+  }
+}
+
+async function promptRequired(ctx: ExtensionContext, message: string, placeholder?: string): Promise<string> {
+  const response = await ctx.ui.input(message, placeholder);
+  if (response === undefined) throw new Error("Cancelled by user");
+  return response;
+}
+
+export async function createAccountFromUpstream(
+  upstream: RouterUpstreamConfig,
+  label: string,
+  ctx: ExtensionContext,
+): Promise<StoredRouterAccount> {
+  const createdAt = now();
+
+  if (upstream.authMode === "api-key") {
+    const token = await promptRequired(ctx, `Enter API key or bearer token for ${upstream.label}:`);
+    return {
+      id: `acct_${randomUUID().slice(0, 8)}`,
+      label,
+      provider: "api-key",
+      upstreamId: upstream.id,
+      access: token.trim(),
+      refresh: "",
+      expires: Number.MAX_SAFE_INTEGER,
+      enabled: true,
+      weight: 1,
+      createdAt,
+      updatedAt: createdAt,
+      meta: {},
+    };
+  }
+
+  if (!upstream.oauthProviderId) {
+    throw new Error(`Upstream ${upstream.id} is missing oauthProviderId`);
+  }
+
+  const provider = getOAuthProvider(upstream.oauthProviderId);
+  if (!provider) {
+    throw new Error(`OAuth provider not available: ${upstream.oauthProviderId}`);
+  }
+
+  const credentials = await provider.login({
+    onAuth(info) {
+      openUrlInBrowser(info.url);
+      ctx.ui.notify(`${provider.name}: ${info.instructions ?? "Finish login in your browser."}`, "info");
+      ctx.ui.notify(info.url, "info");
+    },
+    onPrompt(prompt) {
+      return promptRequired(ctx, prompt.message, prompt.placeholder);
+    },
+    onProgress(message) {
+      ctx.ui.notify(message, "info");
+    },
+  });
+
+  const normalized = normalizeCredentials(credentials);
+
+  return {
+    id: `acct_${randomUUID().slice(0, 8)}`,
+    label,
+    provider: upstream.oauthProviderId,
+    upstreamId: upstream.id,
+    access: normalized.access,
+    refresh: normalized.refresh,
+    expires: normalized.expires,
+    enabled: true,
+    weight: 1,
+    createdAt,
+    updatedAt: createdAt,
+    meta: normalized.meta,
+  };
+}
+
+function toCredentials(account: StoredRouterAccount): OAuthCredentials {
+  return {
+    access: account.access,
+    refresh: account.refresh,
+    expires: account.expires,
+    ...(account.meta ?? {}),
+  };
+}
+
+export async function refreshAccountCredentials(account: StoredRouterAccount): Promise<StoredRouterAccount> {
+  if (account.provider === "api-key") return account;
+
+  const provider = getOAuthProvider(account.provider);
+  if (!provider) throw new Error(`OAuth provider not available: ${account.provider}`);
+
+  const refreshed = await provider.refreshToken(toCredentials(account));
+  const normalized = normalizeCredentials(refreshed);
+
+  return {
+    ...account,
+    access: normalized.access,
+    refresh: normalized.refresh,
+    expires: normalized.expires,
+    meta: normalized.meta,
+    updatedAt: now(),
+  };
+}
+
+export async function getApiKeyForAccount(account: StoredRouterAccount): Promise<string> {
+  if (account.provider === "api-key") return account.access;
+
+  const provider = getOAuthProvider(account.provider);
+  if (!provider) throw new Error(`OAuth provider not available: ${account.provider}`);
+  return provider.getApiKey(toCredentials(account));
+}

package/.pi/extensions/oauth-router/oauth-store.ts

@@ -0,0 +1,121 @@
+import { existsSync, readFileSync } from "node:fs";
+import { CREDENTIALS_PATH, writeJsonFile } from "./config.ts";
+import type { RouterCredentialStore, StoredRouterAccount } from "./types.ts";
+
+const EMPTY_STORE: RouterCredentialStore = {
+  version: 1,
+  accounts: [],
+};
+
+function normalizeAccount(account: StoredRouterAccount): StoredRouterAccount {
+  return {
+    ...account,
+    enabled: account.enabled !== false,
+    weight: Number.isFinite(account.weight) && account.weight > 0 ? Math.floor(account.weight) : 1,
+    refresh: account.refresh ?? "",
+    access: account.access ?? "",
+    expires: Number.isFinite(account.expires) ? account.expires : Number.MAX_SAFE_INTEGER,
+    createdAt: Number.isFinite(account.createdAt) ? account.createdAt : Date.now(),
+    updatedAt: Number.isFinite(account.updatedAt) ? account.updatedAt : Date.now(),
+    meta: account.meta ?? {},
+  };
+}
+
+export function redactToken(token: string): string {
+  if (!token) return "<empty>";
+  if (token.length <= 8) return "********";
+  return `${token.slice(0, 4)}…${token.slice(-4)}`;
+}
+
+export function summarizeAccount(account: StoredRouterAccount) {
+  return {
+    id: account.id,
+    label: account.label,
+    provider: account.provider,
+    upstreamId: account.upstreamId,
+    enabled: account.enabled,
+    weight: account.weight,
+    expires: account.expires,
+    createdAt: account.createdAt,
+    updatedAt: account.updatedAt,
+    access: redactToken(account.access),
+    refresh: redactToken(account.refresh),
+    meta: account.meta ?? {},
+  };
+}
+
+export class RouterAccountStore {
+  private data: RouterCredentialStore;
+
+  constructor() {
+    this.data = this.load();
+  }
+
+  private load(): RouterCredentialStore {
+    if (!existsSync(CREDENTIALS_PATH)) {
+      writeJsonFile(CREDENTIALS_PATH, EMPTY_STORE, true);
+      return { ...EMPTY_STORE, accounts: [] };
+    }
+
+    try {
+      const parsed = JSON.parse(readFileSync(CREDENTIALS_PATH, "utf8")) as Partial<RouterCredentialStore>;
+      return {
+        version: 1,
+        accounts: Array.isArray(parsed.accounts) ? parsed.accounts.map((account) => normalizeAccount(account)) : [],
+      };
+    } catch {
+      writeJsonFile(CREDENTIALS_PATH, EMPTY_STORE, true);
+      return { ...EMPTY_STORE, accounts: [] };
+    }
+  }
+
+  private save() {
+    writeJsonFile(CREDENTIALS_PATH, this.data, true);
+  }
+
+  reload() {
+    this.data = this.load();
+  }
+
+  list(): StoredRouterAccount[] {
+    return [...this.data.accounts].sort((a, b) => a.createdAt - b.createdAt);
+  }
+
+  get(id: string): StoredRouterAccount | undefined {
+    return this.data.accounts.find((account) => account.id === id);
+  }
+
+  add(account: StoredRouterAccount) {
+    this.data.accounts.push(normalizeAccount(account));
+    this.save();
+  }
+
+  update(account: StoredRouterAccount) {
+    const index = this.data.accounts.findIndex((item) => item.id === account.id);
+    if (index === -1) throw new Error(`Unknown account: ${account.id}`);
+    this.data.accounts[index] = normalizeAccount(account);
+    this.save();
+  }
+
+  remove(id: string) {
+    const next = this.data.accounts.filter((account) => account.id !== id);
+    this.data.accounts = next;
+    this.save();
+  }
+
+  setEnabled(id: string, enabled: boolean) {
+    const account = this.get(id);
+    if (!account) throw new Error(`Unknown account: ${id}`);
+    account.enabled = enabled;
+    account.updatedAt = Date.now();
+    this.update(account);
+  }
+
+  setWeight(id: string, weight: number) {
+    const account = this.get(id);
+    if (!account) throw new Error(`Unknown account: ${id}`);
+    account.weight = Math.max(1, Math.floor(weight));
+    account.updatedAt = Date.now();
+    this.update(account);
+  }
+}

package/.pi/extensions/oauth-router/package.json

@@ -0,0 +1,14 @@
+{
+  "name": "oauth-router",
+  "private": true,
+  "type": "module",
+  "description": "Pi extension for multi-account OAuth-aware routing across OpenAI-compatible upstreams",
+  "pi": {
+    "extensions": [
+      "./index.ts"
+    ]
+  },
+  "scripts": {
+    "verify": "python scripts/vibe-verify.py"
+  }
+}

package/.pi/extensions/oauth-router/policies.ts

@@ -0,0 +1,27 @@
+import type { EligibleAccount, RoutingPolicyName } from "./types.ts";
+
+function sortStable(accounts: EligibleAccount[]): EligibleAccount[] {
+  return [...accounts].sort((a, b) => {
+    if (a.account.createdAt !== b.account.createdAt) return a.account.createdAt - b.account.createdAt;
+    return a.account.id.localeCompare(b.account.id);
+  });
+}
+
+function expandByWeight(accounts: EligibleAccount[]): EligibleAccount[] {
+  return sortStable(accounts).flatMap((entry) => Array.from({ length: Math.max(1, entry.account.weight || 1) }, () => entry));
+}
+
+export function chooseEligibleAccount(
+  policy: RoutingPolicyName,
+  eligible: EligibleAccount[],
+  cursor: number,
+): { selected?: EligibleAccount; nextCursor: number } {
+  if (eligible.length === 0) return { selected: undefined, nextCursor: 0 };
+
+  const ordered = policy === "weighted-round-robin" ? expandByWeight(eligible) : sortStable(eligible);
+  const index = ((cursor % ordered.length) + ordered.length) % ordered.length;
+  return {
+    selected: ordered[index],
+    nextCursor: (index + 1) % ordered.length,
+  };
+}