@hieplp/pi-account-switcher 0.2.0

package/src/constants/providers.ts

@@ -0,0 +1,36 @@
+export const PROVIDER_API_TYPES = [
+  "openai-completions",
+  "openai-responses",
+  "anthropic-messages",
+  "azure-openai-responses",
+  "openai-codex-responses",
+  "mistral-conversations",
+  "google-generative-ai",
+  "google-vertex",
+  "bedrock-converse-stream",
+] as const;
+
+export const OAUTH_PROVIDER_IDS = [
+  "anthropic",
+  "openai-codex",
+  "github-copilot",
+  "google-antigravity",
+  "custom",
+] as const;
+
+export const BUILT_IN_PROVIDER_IDS = ["anthropic", "openai", "openai-codex", "google", "xai", "openrouter"] as const;
+
+export const PROVIDER_ALIASES: Record<string, string> = {
+  claude: "anthropic",
+  codex: "openai-codex",
+  gemini: "google",
+};
+
+export const PROVIDER_ENV_KEYS: Record<string, string[]> = {
+  anthropic: ["ANTHROPIC_API_KEY"],
+  openai: ["OPENAI_API_KEY"],
+  "openai-codex": ["OPENAI_API_KEY"],
+  google: ["GOOGLE_API_KEY", "GEMINI_API_KEY"],
+  xai: ["XAI_API_KEY"],
+  openrouter: ["OPENROUTER_API_KEY"],
+};

package/src/extension.ts

@@ -0,0 +1,21 @@
+import { createJiti } from "@mariozechner/jiti";
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import { dirname } from "node:path";
+import { fileURLToPath } from "node:url";
+
+export default async function accountSwitcherBootstrap(pi: ExtensionAPI) {
+  const srcDir = dirname(fileURLToPath(import.meta.url));
+  const jiti = createJiti(import.meta.url, {
+    alias: {
+      "@": srcDir,
+    },
+  });
+
+  const extension = await jiti.import<
+    (pi: ExtensionAPI) => void | Promise<void>
+  >("./index", {
+    default: true,
+  });
+
+  await extension(pi);
+}

package/src/index.ts

@@ -0,0 +1,20 @@
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type { AccountSwitcherContext } from "@/types";
+import { useAccountSwitcher, type AccountSwitcher } from "./runtime";
+import { registerAllCommands } from "./commands";
+
+async function accountSwitcher(pi: ExtensionAPI) {
+  const runtime: AccountSwitcher = useAccountSwitcher(pi);
+
+  pi.on("session_start", async (_, ctx) => {
+    await runtime.init(ctx as AccountSwitcherContext);
+  });
+
+  pi.on("model_select", async (event, ctx) => {
+    await runtime.onModelSelect(event.model.provider, ctx as AccountSwitcherContext);
+  });
+
+  registerAllCommands(pi, runtime);
+}
+
+export default accountSwitcher;

package/src/runtime/account-switcher-runtime.ts

@@ -0,0 +1,194 @@
+import AccountSwitcher from "./account-switcher";
+import { ACCOUNTS_PATH, PROVIDERS_PATH, STATE_PATH } from "@/constants";
+import type { Api, Model } from "@mariozechner/pi-ai";
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type { AccountConfig, AccountSwitcherContext, PiAuthEntry, ProviderConfig } from "@/types";
+import type { AccountService, ModelService, PiAuthService, ProviderService } from "@/services";
+import { useAccountService, useModelService, usePiAuthService, useProviderService } from "@/services";
+import { accountUtil, modelUtil, providerUtil, uiUtil } from "@/utils";
+
+function resolveAuthProvider(account: AccountConfig, providers: ProviderConfig[]): string {
+  if (account.piAuth?.provider) return account.piAuth.provider;
+  const provider = providerUtil.findProvider(account.provider, providers);
+  return provider?.piAuthProvider ?? providerUtil.normalizeProvider(account.provider);
+}
+
+export default class AccountSwitcherRuntime implements AccountSwitcher {
+  private accountService: AccountService;
+  private modelService: ModelService;
+  private piAuthService: PiAuthService;
+  private providerService: ProviderService;
+
+  constructor(private readonly pi: Pick<ExtensionAPI, "registerProvider" | "setModel">) {
+    this.providerService = useProviderService(this.pi as ExtensionAPI, PROVIDERS_PATH);
+    this.accountService = useAccountService(ACCOUNTS_PATH, STATE_PATH);
+    this.modelService = useModelService(this.pi);
+    this.piAuthService = usePiAuthService();
+  }
+
+  // ===============================================================================================
+  // Core
+  // ===============================================================================================
+
+  async init(ctx: AccountSwitcherContext): Promise<void> {
+    await this.load();
+
+    const active = this.accountService.getActiveAccount();
+    uiUtil.setAccountStatus(ctx.ui, active?.label);
+
+    // Re-apply saved account credentials so env vars and OAuth auth storage are
+    // populated on session start, not only after the first explicit switch.
+    if (active) {
+      const providers = this.providerService.getProviders();
+      await this.applyProviderApiKey(active, providers);
+      await accountUtil.applyAccountEnv(active, ctx.modelRegistry, resolveAuthProvider(active, providers));
+    }
+
+    // Restore the last active model. modelRegistry.find returns undefined if the
+    // model is no longer available (e.g. provider was removed), in which case we
+    // leave Pi's default model selection untouched.
+    const modelState = this.accountService.getActiveModelState();
+    if (modelState) {
+      const model = ctx.modelRegistry.find(modelState.provider, modelState.id);
+      if (model) await this.modelService.applyModel(model, ctx);
+    }
+  }
+
+  async load(): Promise<void> {
+    await this.accountService.load();
+    await this.providerService.load();
+  }
+
+  async onModelSelect(provider: string, ctx: AccountSwitcherContext): Promise<void> {
+    const matchingAccount = this.findAccountsByProvider(provider)[0];
+    const activeAccount = this.accountService.getActiveAccount();
+    if (matchingAccount && matchingAccount.id !== activeAccount?.id) {
+      await this.activateAccount(matchingAccount, ctx);
+    }
+  }
+
+  // ===============================================================================================
+  // Pi Auth
+  // ===============================================================================================
+
+  async getPiAuthEntry(provider: string): Promise<PiAuthEntry | undefined> {
+    return this.piAuthService.getEntry(provider);
+  }
+
+  isOAuthEntry(entry: PiAuthEntry | undefined): boolean {
+    return this.piAuthService.isOAuthEntry(entry);
+  }
+
+  // ===============================================================================================
+  // Account
+  // ===============================================================================================
+
+  getAccounts(): AccountConfig[] {
+    return this.accountService.getAccounts();
+  }
+
+  findAccountById(id: string): AccountConfig | undefined {
+    return this.accountService.getAccounts().find((a) => a.id === id);
+  }
+
+  findAccountsByProvider(provider: string): AccountConfig[] {
+    const providers = this.providerService.getProviders();
+    const normalized = providerUtil.normalizeProviderWithCustom(provider, providers);
+    return this.accountService.getAccounts().filter((a) => {
+      const accountProvider = providerUtil.normalizeProviderWithCustom(resolveAuthProvider(a, providers), providers);
+      return accountProvider === normalized;
+    });
+  }
+
+  getActiveAccount(): AccountConfig | undefined {
+    return this.accountService.getActiveAccount();
+  }
+
+  async addAccount(account: AccountConfig): Promise<void> {
+    return this.accountService.addAccount(account);
+  }
+
+  async editAccount(original: AccountConfig, updated: AccountConfig): Promise<void> {
+    return this.accountService.editAccount(original, updated);
+  }
+
+  async removeAccount(account: AccountConfig): Promise<void> {
+    return this.accountService.removeAccount(account);
+  }
+
+  async activateAccount(account: AccountConfig, ctx: AccountSwitcherContext): Promise<string> {
+    const providers = this.providerService.getProviders();
+    const providerApiKey = await this.applyProviderApiKey(account, providers);
+    const result = await this.accountService.activateAccount(account, ctx, resolveAuthProvider(account, providers));
+
+    // piAuth accounts authenticate via a separate provider (e.g. github-copilot),
+    // so use that for model lookup rather than the account's own provider field.
+    const accountProvider = providerUtil.normalizeProviderWithCustom(
+      resolveAuthProvider(account, providers),
+      providers,
+    );
+    const currentProvider = ctx.model
+      ? providerUtil.normalizeProviderWithCustom(ctx.model.provider, providers)
+      : undefined;
+
+    // Skip model selection if the active model already belongs to the same provider.
+    if (accountProvider !== currentProvider) {
+      const model = await modelUtil.pickModel(ctx, account, providers);
+      if (model) await this.applyModel(model, ctx);
+    }
+
+    return providerApiKey ? `provider apiKey (${providerApiKey})` : result;
+  }
+
+  private async applyProviderApiKey(account: AccountConfig, providers: ProviderConfig[]): Promise<string | undefined> {
+    if (!account.providerApiKey && !account.usesProviderApiKey) return undefined;
+
+    const provider = providerUtil.findProvider(account.provider, providers);
+    if (!provider) throw new Error(`Custom provider not found for account ${account.id}: ${account.provider}`);
+
+    if (account.providerApiKey) {
+      const apiKey = await accountUtil.resolveSecret(account.providerApiKey);
+      if (!apiKey) throw new Error(`Resolved empty providerApiKey for account ${account.id}`);
+      this.providerService.registerProvider({ ...provider, apiKey });
+      return provider.id;
+    }
+
+    this.providerService.registerProvider(provider);
+    return provider.id;
+  }
+
+  // ===============================================================================================
+  // Model
+  // ===============================================================================================
+
+  async applyModel(model: Model<Api>, ctx: AccountSwitcherContext): Promise<void> {
+    await this.modelService.applyModel(model, ctx);
+    await this.accountService.saveActiveModel(model.id, model.provider);
+  }
+
+  // ===============================================================================================
+  // Provider
+  // ===============================================================================================
+
+  getProviders(): ProviderConfig[] {
+    return this.providerService.getProviders();
+  }
+
+  async addProvider(provider: ProviderConfig): Promise<void> {
+    return this.providerService.addProvider(provider);
+  }
+
+  async editProvider(original: ProviderConfig, updated: ProviderConfig): Promise<void> {
+    return this.providerService.editProvider(original, updated);
+  }
+
+  async removeProvider(provider: ProviderConfig): Promise<void> {
+    const providers = this.providerService.getProviders();
+    const dependents = this.accountService.findAccountsByProvider(provider.id, providers);
+    if (dependents.length > 0) {
+      const names = dependents.map((a) => `"${a.label ?? a.id}"`).join(", ");
+      throw new Error(`Cannot remove: ${names} ${dependents.length === 1 ? "uses" : "use"} this provider`);
+    }
+    return this.providerService.removeProvider(provider);
+  }
+}

package/src/runtime/account-switcher.ts

@@ -0,0 +1,32 @@
+import type { Api, Model } from "@mariozechner/pi-ai";
+import type { AccountConfig, AccountSwitcherContext, PiAuthEntry, ProviderConfig } from "@/types";
+
+export default interface AccountSwitcher {
+  // Core
+  init(ctx: AccountSwitcherContext): Promise<void>;
+  load(): Promise<void>;
+  onModelSelect(provider: string, ctx: AccountSwitcherContext): Promise<void>;
+
+  // Pi Auth
+  getPiAuthEntry(provider: string): Promise<PiAuthEntry | undefined>;
+  isOAuthEntry(entry: PiAuthEntry | undefined): boolean;
+
+  // Account
+  getAccounts(): AccountConfig[];
+  findAccountById(id: string): AccountConfig | undefined;
+  findAccountsByProvider(provider: string): AccountConfig[];
+  getActiveAccount(): AccountConfig | undefined;
+  addAccount(account: AccountConfig): Promise<void>;
+  editAccount(original: AccountConfig, updated: AccountConfig): Promise<void>;
+  removeAccount(account: AccountConfig): Promise<void>;
+  activateAccount(account: AccountConfig, ctx: AccountSwitcherContext): Promise<string>;
+
+  // Model
+  applyModel(model: Model<Api>, ctx: AccountSwitcherContext): Promise<void>;
+
+  // Provider
+  getProviders(): ProviderConfig[];
+  addProvider(config: ProviderConfig): Promise<void>;
+  editProvider(original: ProviderConfig, updated: ProviderConfig): Promise<void>;
+  removeProvider(provider: ProviderConfig): Promise<void>;
+}

package/src/runtime/index.ts

@@ -0,0 +1,11 @@
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type AccountSwitcher from "./account-switcher";
+import AccountSwitcherRuntime from "./account-switcher-runtime";
+
+function useAccountSwitcher(
+  pi: Pick<ExtensionAPI, "registerProvider" | "setModel">,
+): AccountSwitcher {
+  return new AccountSwitcherRuntime(pi);
+}
+
+export { useAccountSwitcher, type AccountSwitcher };

package/src/schemas/accounts.ts

@@ -0,0 +1,50 @@
+import { z } from "zod";
+
+export const secretSourceSchema = z.union([
+  z.string().min(1),
+  z.object({ type: z.literal("literal"), value: z.string().min(1) }),
+  z.object({ type: z.literal("env"), name: z.string().min(1) }),
+  z.object({ type: z.literal("file"), path: z.string().min(1) }),
+  z.object({ type: z.literal("command"), command: z.string().min(1) }),
+  z.object({ type: z.literal("op"), reference: z.string().min(1) }),
+]);
+
+export const piAuthEntrySchema = z.union([
+  z.object({ type: z.literal("api_key"), key: z.string().min(1) }),
+  z
+    .object({
+      type: z.literal("oauth"),
+      refresh: z.string().min(1),
+      access: z.string().min(1),
+      expires: z.number(),
+    })
+    .passthrough(),
+]);
+
+export const accountSchema = z
+  .object({
+    id: z.string().min(1),
+    label: z.string().min(1),
+    provider: z.string().min(1),
+    model: z.string().min(1).optional(),
+    env: z.record(z.string().min(1), secretSourceSchema).optional(),
+    providerApiKey: secretSourceSchema.optional(),
+    usesProviderApiKey: z.boolean().optional(),
+    piAuth: z
+      .object({
+        provider: z.string().min(1),
+        entry: piAuthEntrySchema,
+      })
+      .optional(),
+  })
+  .refine(
+    (account) =>
+      (account.env && Object.keys(account.env).length > 0) ||
+      account.providerApiKey ||
+      account.usesProviderApiKey ||
+      account.piAuth,
+    {
+      message:
+        "Account must define env credentials, providerApiKey, provider apiKey, or piAuth credentials",
+    },
+  );

package/src/schemas/common.ts

@@ -0,0 +1,3 @@
+import z from "zod";
+
+export const jsonRecordSchema = z.record(z.string(), z.unknown());

package/src/schemas/config.ts

@@ -0,0 +1,7 @@
+import z from "zod";
+import { accountSchema } from "./accounts";
+
+export const configSchema = z.object({
+  accounts: z.array(accountSchema).default([]),
+  switchMode: z.literal("env").optional().default("env"),
+});

package/src/schemas/index.ts

@@ -0,0 +1,4 @@
+export * from "./common";
+export * from "./accounts";
+export * from "./config";
+export * from "./providers";

package/src/schemas/providers.ts

@@ -0,0 +1,57 @@
+import z from "zod";
+import { jsonRecordSchema } from "@/schemas/common";
+
+export const providerModelSchema = z
+  .object({
+    id: z.string().min(1),
+    name: z.string().min(1).optional(),
+    api: z.string().min(1).optional(),
+    baseUrl: z.string().min(1).optional(),
+    reasoning: z.boolean().optional(),
+    input: z.array(z.enum(["text", "image"])).optional(),
+    contextWindow: z.number().int().positive().optional(),
+    maxTokens: z.number().int().positive().optional(),
+    cost: z
+      .object({
+        input: z.number(),
+        output: z.number(),
+        cacheRead: z.number(),
+        cacheWrite: z.number(),
+      })
+      .optional(),
+    compat: jsonRecordSchema.optional(),
+    thinkingLevelMap: z
+      .record(z.string(), z.union([z.string(), z.null()]))
+      .optional(),
+    headers: z.record(z.string(), z.string()).optional(),
+  })
+  .passthrough();
+
+export const providerSchema = z
+  .object({
+    id: z.string().min(1).optional(),
+    label: z.string().min(1).optional(),
+    name: z.string().min(1).optional(),
+    envKeys: z.array(z.string().min(1)).optional(),
+    aliases: z.array(z.string().min(1)).optional(),
+    piAuthProvider: z.string().min(1).optional(),
+    baseUrl: z.string().min(1).optional(),
+    api: z.string().min(1).optional(),
+    apiKey: z.string().min(1).optional(),
+    headers: z.record(z.string(), z.string()).optional(),
+    authHeader: z.boolean().optional(),
+    compat: jsonRecordSchema.optional(),
+    models: z.array(providerModelSchema).optional(),
+    modelOverrides: z
+      .record(z.string(), providerModelSchema.partial())
+      .optional(),
+  })
+  .passthrough();
+
+export const providerCatalogRecordSchema = z.object({
+  providers: z.record(z.string().min(1), providerSchema).default({}),
+});
+
+export const providerCatalogArraySchema = z.object({
+  providers: z.array(providerSchema.extend({ id: z.string().min(1) })).default([]),
+});

package/src/services/accounts.ts

@@ -0,0 +1,116 @@
+import { type AccountStore, useAccountStore, type StateStore, useStateStore } from "@/storage";
+import type { AccountConfig, AccountSwitcherContext, ProviderConfig } from "@/types";
+import { accountUtil, providerUtil, uiUtil } from "@/utils";
+
+export interface AccountService {
+  load(): Promise<void>;
+  getAccounts(): AccountConfig[];
+  findAccountsByProvider(provider: string, providers: ProviderConfig[]): AccountConfig[];
+  getActiveAccount(): AccountConfig | undefined;
+  addAccount(account: AccountConfig): Promise<void>;
+  editAccount(original: AccountConfig, updated: AccountConfig): Promise<void>;
+  removeAccount(account: AccountConfig): Promise<void>;
+  activateAccount(account: AccountConfig, ctx: AccountSwitcherContext, authProvider?: string): Promise<string>;
+  getActiveModelState(): { id: string; provider: string } | undefined;
+  saveActiveModel(id: string, provider: string): Promise<void>;
+}
+
+export function useAccountService(accountsPath: string, statePath?: string): AccountService {
+  return new AccountServiceImpl(useAccountStore(accountsPath), useStateStore(statePath));
+}
+
+// ===============================================================================================
+// Account Service
+// ===============================================================================================
+
+class AccountServiceImpl implements AccountService {
+  private accounts: AccountConfig[] = [];
+  private activeAccountId: string | undefined;
+  private activeModelId: string | undefined;
+  private activeModelProvider: string | undefined;
+
+  constructor(
+    private readonly store: AccountStore,
+    private readonly stateStore: StateStore,
+  ) {}
+
+  async load(): Promise<void> {
+    this.accounts = await this.store.load();
+    const state = await this.stateStore.load();
+    this.activeAccountId = state.activeAccountId;
+    this.activeModelId = state.activeModelId;
+    this.activeModelProvider = state.activeModelProvider;
+  }
+
+  getAccounts(): AccountConfig[] {
+    return this.accounts;
+  }
+
+  findAccountsByProvider(provider: string, providers: ProviderConfig[]): AccountConfig[] {
+    const normalized = providerUtil.normalizeProviderWithCustom(provider, providers);
+    return this.accounts.filter(
+      (account) => providerUtil.normalizeProviderWithCustom(account.provider, providers) === normalized,
+    );
+  }
+
+  getActiveAccount(): AccountConfig | undefined {
+    return this.accounts.find((a) => a.id === this.activeAccountId);
+  }
+
+  getActiveModelState(): { id: string; provider: string } | undefined {
+    if (!this.activeModelId || !this.activeModelProvider) return undefined;
+    return { id: this.activeModelId, provider: this.activeModelProvider };
+  }
+
+  async saveActiveModel(id: string, provider: string): Promise<void> {
+    this.activeModelId = id;
+    this.activeModelProvider = provider;
+    await this.flushState();
+  }
+
+  async addAccount(account: AccountConfig): Promise<void> {
+    this.accounts = await this.store.addAccount(account);
+  }
+
+  async editAccount(original: AccountConfig, updated: AccountConfig): Promise<void> {
+    this.accounts = await this.store.replaceAccount(original.id, updated);
+    if (this.activeAccountId === original.id) {
+      this.activeAccountId = updated.id;
+      await this.flushState();
+    }
+  }
+
+  async removeAccount(account: AccountConfig): Promise<void> {
+    this.accounts = await this.store.removeAccount(account.id);
+    if (this.activeAccountId === account.id) {
+      this.activeAccountId = undefined;
+      await this.flushState();
+    }
+  }
+
+  async activateAccount(account: AccountConfig, ctx: AccountSwitcherContext, authProvider?: string): Promise<string> {
+    const previous = this.getActiveAccount();
+    let applied: string[] = [];
+    if (account.piAuth) {
+      if (previous) await accountUtil.clearAccountEnv(previous, ctx.modelRegistry);
+      applied = await accountUtil.applyAccountEnv(account, ctx.modelRegistry, authProvider);
+    } else {
+      const resolved = await accountUtil.resolveAccountEnv(account);
+      if (previous) await accountUtil.clearAccountEnv(previous, ctx.modelRegistry);
+      applied = accountUtil.applyResolvedAccountEnv(account, resolved, ctx.modelRegistry, authProvider);
+    }
+    this.activeAccountId = account.id;
+    await this.flushState();
+    uiUtil.setAccountStatus(ctx.ui, account.label);
+    if (account.piAuth) return "via OAuth";
+    return applied.length > 0 ? applied.join(", ") : "";
+  }
+
+  private async flushState(): Promise<void> {
+    await this.stateStore.save({
+      activeAccountId: this.activeAccountId,
+      activeModelId: this.activeModelId,
+      activeModelProvider: this.activeModelProvider,
+    });
+  }
+}

package/src/services/index.ts

@@ -0,0 +1,4 @@
+export * from "./accounts";
+export * from "./models";
+export * from "./pi-auth";
+export * from "./providers";

package/src/services/models.ts

@@ -0,0 +1,27 @@
+import type { Api, Model } from "@mariozechner/pi-ai";
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type { AccountSwitcherContext } from "@/types";
+
+type ProviderModel = Model<Api>;
+
+export interface ModelService {
+  applyModel(model: ProviderModel, ctx: AccountSwitcherContext): Promise<void>;
+}
+
+export function useModelService(pi: Pick<ExtensionAPI, "setModel">): ModelService {
+  return new ModelServiceImpl(pi);
+}
+
+class ModelServiceImpl implements ModelService {
+  constructor(private readonly pi: Pick<ExtensionAPI, "setModel">) {}
+
+  async applyModel(model: ProviderModel, ctx: AccountSwitcherContext): Promise<void> {
+    const ok = await this.pi.setModel(model);
+    if (!ok) {
+      ctx.ui.notify(
+        `Account switched, but Pi refused model ${model.provider}/${model.id}. Check credentials.`,
+        "warning",
+      );
+    }
+  }
+}

package/src/services/pi-auth.ts

@@ -0,0 +1,23 @@
+import { type PiAuthStore, usePiAuthStore, isOAuthEntry } from "@/storage";
+import type { PiAuthEntry } from "@/types";
+
+export interface PiAuthService {
+  getEntry(provider: string): Promise<PiAuthEntry | undefined>;
+  isOAuthEntry(entry: PiAuthEntry | undefined): boolean;
+}
+
+export function usePiAuthService(path?: string): PiAuthService {
+  return new PiAuthServiceImpl(usePiAuthStore(path));
+}
+
+class PiAuthServiceImpl implements PiAuthService {
+  constructor(private readonly store: PiAuthStore) {}
+
+  async getEntry(provider: string): Promise<PiAuthEntry | undefined> {
+    return this.store.getEntry(provider);
+  }
+
+  isOAuthEntry(entry: PiAuthEntry | undefined): boolean {
+    return isOAuthEntry(entry);
+  }
+}