@hieplp/pi-account-switcher 0.2.0

package/src/services/providers.ts

@@ -0,0 +1,123 @@
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import { type ProviderStore, useProviderStore } from "@/storage";
+import type { ProviderConfig, ProviderModelConfig } from "@/types";
+import { commonUtil } from "@/utils";
+
+export interface ProviderService {
+  load(): Promise<void>;
+  getProviders(): ProviderConfig[];
+  addProvider(provider: ProviderConfig): Promise<void>;
+  editProvider(original: ProviderConfig, updated: ProviderConfig): Promise<void>;
+  removeProvider(provider: ProviderConfig): Promise<void>;
+  registerProviders(providers: ProviderConfig[]): void;
+  registerProvider(provider: ProviderConfig): void;
+}
+
+export function useProviderService(pi: ExtensionAPI, path: string): ProviderService {
+  return new ProviderServiceImpl(pi, useProviderStore(path));
+}
+
+// ===============================================================================================
+// Provider Service
+// ===============================================================================================
+
+class ProviderServiceImpl implements ProviderService {
+  private providers: ProviderConfig[] = [];
+
+  constructor(
+    private readonly pi: ExtensionAPI,
+    private readonly store: ProviderStore,
+  ) {}
+
+  async load(): Promise<void> {
+    this.providers = await this.store.load();
+    this.registerProviders(this.providers);
+  }
+
+  getProviders(): ProviderConfig[] {
+    return this.providers;
+  }
+
+  async addProvider(provider: ProviderConfig): Promise<void> {
+    if (this.providers.some((p) => p.id === provider.id)) {
+      throw new Error(`Provider already exists: ${provider.id}`);
+    }
+    this.registerProvider(provider);
+    this.providers.push(provider);
+    await this.store.save(this.providers);
+  }
+
+  async editProvider(original: ProviderConfig, updated: ProviderConfig): Promise<void> {
+    const index = this.providers.findIndex((p) => p.id === original.id);
+    if (index === -1) throw new Error(`Provider not found: ${original.id}`);
+    if (updated.id !== original.id && this.providers.some((p) => p.id === updated.id)) {
+      throw new Error(`Provider already exists: ${updated.id}`);
+    }
+    this.providers[index] = updated;
+    await this.store.save(this.providers);
+    this.registerProvider(updated);
+  }
+
+  async removeProvider(provider: ProviderConfig): Promise<void> {
+    const index = this.providers.findIndex((p) => p.id === provider.id);
+    if (index === -1) throw new Error(`Provider not found: ${provider.id}`);
+    this.providers.splice(index, 1);
+    await this.store.save(this.providers);
+  }
+
+  registerProviders(providers: ProviderConfig[]): void {
+    providers.forEach((provider) => this.registerProvider(provider));
+  }
+
+  registerProvider(provider: ProviderConfig): void {
+    const config = this.toPiProvider(provider);
+    if (!config) return;
+    this.pi.registerProvider(
+      provider.id,
+      config as Parameters<ExtensionAPI["registerProvider"]>[1],
+    );
+  }
+
+  private toPiProvider(provider: ProviderConfig): Record<string, unknown> | undefined {
+    if (
+      !provider.baseUrl &&
+      !provider.api &&
+      !provider.apiKey &&
+      !provider.models &&
+      !provider.headers &&
+      !provider.authHeader &&
+      !provider.compat
+    ) {
+      return undefined;
+    }
+
+    return commonUtil.omitUndefined({
+      name: provider.name ?? provider.label,
+      baseUrl: provider.baseUrl,
+      apiKey: provider.apiKey,
+      api: provider.api,
+      headers: provider.headers,
+      authHeader: provider.authHeader,
+      models: provider.models?.map((model) => this.toPiModel(provider, model)),
+      modelOverrides: provider.modelOverrides,
+      compat: provider.compat,
+    });
+  }
+
+  private toPiModel(
+    provider: ProviderConfig,
+    model: ProviderModelConfig,
+  ): Record<string, unknown> | undefined {
+    return commonUtil.omitUndefined({
+      ...model,
+      api: model.api ?? provider.api,
+      name: model.name ?? model.id,
+      reasoning: model.reasoning ?? false,
+      input: model.input ?? ["text"],
+      contextWindow: model.contextWindow ?? 128000,
+      maxTokens: model.maxTokens ?? 16384,
+      cost: model.cost ?? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+      compat: model.compat ?? provider.compat,
+    });
+  }
+}

package/src/storage/accounts.ts

@@ -0,0 +1,109 @@
+import { readFile } from "node:fs/promises";
+import { configSchema } from "@/schemas";
+import type { AccountConfig, AccountSwitcherConfig } from "@/types";
+import { errorUtil, fileUtil } from "@/utils";
+
+export interface AccountStore {
+  load(): Promise<AccountConfig[]>;
+  addAccount(account: AccountConfig): Promise<AccountConfig[]>;
+  replaceAccount(originalId: string, account: AccountConfig): Promise<AccountConfig[]>;
+  removeAccount(id: string): Promise<AccountConfig[]>;
+}
+
+export function useAccountStore(path: string) {
+  return new AccountStoreImpl(path);
+}
+
+// ===============================================================================================
+// Account Store
+// ===============================================================================================
+
+class AccountStoreImpl implements AccountStore {
+  constructor(private readonly path: string) {}
+
+  async load(): Promise<AccountConfig[]> {
+    try {
+      const raw = await readFile(this.path, "utf8");
+      const parsed = configSchema.parse(JSON.parse(raw));
+
+      const { accounts } = parsed;
+      accountValidator.assertNoDuplicateAccounts(accounts);
+      accountValidator.assertAccountsHaveCredentials(accounts);
+
+      return accounts;
+    } catch (error) {
+      if (fileUtil.isMissingFileError(error)) {
+        return [];
+      }
+      throw new Error(`Failed to load accounts at ${this.path}: ${errorUtil.format(error)}`);
+    }
+  }
+
+  async addAccount(account: AccountConfig): Promise<AccountConfig[]> {
+    const accounts = await this.load();
+    accountValidator.assertAccountIdAvailable(accounts, account.id);
+    const next = [...accounts, account];
+    await this.save(next);
+    return next;
+  }
+
+  async replaceAccount(originalId: string, account: AccountConfig): Promise<AccountConfig[]> {
+    const accounts = await this.load();
+    const index = accounts.findIndex((a) => a.id === originalId);
+    if (index === -1) throw new Error(`Account not found: ${originalId}`);
+    accountValidator.assertAccountIdAvailable(accounts, account.id, originalId);
+    const next = accounts.filter((a) => a.id !== originalId);
+    next.splice(index, 0, account);
+    await this.save(next);
+    return next;
+  }
+
+  async removeAccount(id: string): Promise<AccountConfig[]> {
+    const accounts = await this.load();
+    const next = accounts.filter((a) => a.id !== id);
+    if (next.length === accounts.length) throw new Error(`Account not found: ${id}`);
+    await this.save(next);
+    return next;
+  }
+
+  private async save(accounts: AccountConfig[]): Promise<void> {
+    const config: AccountSwitcherConfig = { accounts };
+    await fileUtil.writePrivateJson(this.path, config);
+  }
+}
+
+// ===============================================================================================
+// Account Validator
+// ===============================================================================================
+
+const accountValidator = {
+  assertNoDuplicateAccounts: (accounts: AccountConfig[]): void => {
+    const seen = new Set<string>();
+    const duplicates = new Set<string>();
+    for (const { id } of accounts) {
+      if (seen.has(id)) duplicates.add(id);
+      else seen.add(id);
+    }
+    if (duplicates.size > 0) {
+      throw new Error(`Duplicate account ids: ${Array.from(duplicates).sort().join(", ")}`);
+    }
+  },
+
+  assertAccountIdAvailable: (accounts: AccountConfig[], id: string, originalId?: string): void => {
+    if (id === originalId) return;
+    if (accounts.some((account) => account.id === id)) {
+      throw new Error(`Account already exists: ${id}`);
+    }
+  },
+
+  assertAccountsHaveCredentials: (accounts: AccountConfig[]): void => {
+    const missing = accounts
+      .filter(
+        (a) => (!a.env || Object.keys(a.env).length === 0) && !a.providerApiKey && !a.usesProviderApiKey && !a.piAuth,
+      )
+      .map((a) => a.id);
+    if (missing.length > 0) {
+      throw new Error(`Accounts missing credentials: ${missing.sort().join(", ")}`);
+    }
+  },
+};

package/src/storage/index.ts

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

package/src/storage/paths.ts

@@ -0,0 +1,8 @@
+import { homedir } from "node:os";
+import { join } from "node:path";
+
+export const APP_DIR = join(homedir(), ".pi", "account-switcher");
+export const CONFIG_PATH = join(APP_DIR, "accounts.json");
+export const PROVIDERS_PATH = join(APP_DIR, "providers.json");
+export const STATE_PATH = join(APP_DIR, "state.json");
+export const PI_AUTH_PATH = join(homedir(), ".pi", "agent", "auth.json");

package/src/storage/pi-auth.ts

@@ -0,0 +1,37 @@
+import { readFile } from "node:fs/promises";
+import { PI_AUTH_PATH } from "@/constants";
+import type { PiAuthEntry } from "@/types";
+import { fileUtil } from "@/utils";
+
+type PiAuthFile = Record<string, PiAuthEntry>;
+
+export interface PiAuthStore {
+  getEntry(provider: string): Promise<PiAuthEntry | undefined>;
+}
+
+export function usePiAuthStore(path = PI_AUTH_PATH): PiAuthStore {
+  return new PiAuthStoreImpl(path);
+}
+
+export function isOAuthEntry(entry: PiAuthEntry | undefined): boolean {
+  return entry?.type === "oauth";
+}
+
+class PiAuthStoreImpl implements PiAuthStore {
+  constructor(private readonly path: string) {}
+
+  async getEntry(provider: string): Promise<PiAuthEntry | undefined> {
+    const auth = await this.load();
+    return auth[provider];
+  }
+
+  private async load(): Promise<PiAuthFile> {
+    try {
+      const raw = await readFile(this.path, "utf8");
+      return JSON.parse(raw) as PiAuthFile;
+    } catch (error) {
+      if (fileUtil.isMissingFileError(error)) return {};
+      throw error;
+    }
+  }
+}

package/src/storage/providers.ts

@@ -0,0 +1,85 @@
+import type { ProviderConfig } from "@/types";
+import { commonUtil, errorUtil, fileUtil, providerUtil } from "@/utils";
+import { readFile } from "node:fs/promises";
+import { providerCatalogArraySchema, providerCatalogRecordSchema } from "@/schemas";
+
+export interface ProviderStore {
+  load(): Promise<ProviderConfig[]>;
+  save(providers: ProviderConfig[]): Promise<void>;
+}
+
+export function useProviderStore(path: string) {
+  return new ProviderStoreImpl(path);
+}
+
+// ===============================================================================================
+// Provider Store
+// ===============================================================================================
+
+class ProviderStoreImpl implements ProviderStore {
+  constructor(private readonly path: string) {}
+
+  async save(providers: ProviderConfig[]): Promise<void> {
+    await fileUtil.writePrivateJson(this.path, { providers });
+  }
+
+  async load(): Promise<ProviderConfig[]> {
+    try {
+      const raw = await readFile(this.path, "utf8");
+      const json = JSON.parse(raw) as unknown;
+      const result = providerCatalogRecordSchema.safeParse(json);
+      if (result.success) return JsonParser.parseRecord(result.data.providers);
+
+      const arrayResult = providerCatalogArraySchema.safeParse(json);
+      if (arrayResult.success) return JsonParser.parseArray(arrayResult.data.providers);
+
+      throw new Error("providers.json must contain either a providers object or providers array");
+    } catch (error) {
+      if (fileUtil.isMissingFileError(error)) {
+        return [];
+      }
+      throw new Error(`Failed to load account switcher providers at ${this.path}: ${errorUtil.format(error)}`);
+    }
+  }
+}
+
+// ===============================================================================================
+// Json Parser
+// ===============================================================================================
+
+type ProviderRecord = Record<string, Omit<ProviderConfig, "id"> & { id?: string }>;
+
+class JsonParser {
+  static parseRecord(providers: ProviderRecord): ProviderConfig[] {
+    return Object.entries(providers).map(([id, provider]) => this.normalize({ ...provider, id } as ProviderConfig));
+  }
+
+  static parseArray(providers: ProviderConfig[]): ProviderConfig[] {
+    return providers.map((provider) => this.normalize(provider));
+  }
+
+  private static normalize(provider: ProviderConfig): ProviderConfig {
+    const id = providerUtil.normalizeProvider(provider.id);
+
+    const envKeys = commonUtil.unique([
+      ...(provider.envKeys ?? []),
+      ...(provider.apiKey && commonUtil.isLikelyEnvKey(provider.apiKey) ? [provider.apiKey] : []),
+    ]);
+
+    const aliases = commonUtil
+      .unique((provider.aliases ?? []).map(providerUtil.normalizeProvider))
+      .filter((alias) => alias !== id);
+
+    const api =
+      provider.api ?? (provider.baseUrl || provider.apiKey || provider.models ? "openai-completions" : undefined);
+
+    return {
+      ...provider,
+      id,
+      label: provider.label ?? provider.name,
+      api,
+      aliases,
+      envKeys,
+    };
+  }
+}

package/src/storage/state.ts

@@ -0,0 +1,43 @@
+import { readFile } from "node:fs/promises";
+import z from "zod";
+import { STATE_PATH } from "@/constants";
+import { fileUtil } from "@/utils";
+
+const appStateSchema = z.object({
+  activeAccountId: z.string().optional(),
+  activeModelId: z.string().optional(),
+  activeModelProvider: z.string().optional(),
+});
+
+export interface AppState {
+  activeAccountId?: string;
+  activeModelId?: string;
+  activeModelProvider?: string;
+}
+
+export interface StateStore {
+  load(): Promise<AppState>;
+  save(state: AppState): Promise<void>;
+}
+
+export function useStateStore(path = STATE_PATH): StateStore {
+  return new StateStoreImpl(path);
+}
+
+class StateStoreImpl implements StateStore {
+  constructor(private readonly path: string) {}
+
+  async load(): Promise<AppState> {
+    try {
+      const raw = await readFile(this.path, "utf8");
+      return appStateSchema.parse(JSON.parse(raw));
+    } catch (error) {
+      if (fileUtil.isMissingFileError(error)) return {};
+      throw error;
+    }
+  }
+
+  async save(state: AppState): Promise<void> {
+    await fileUtil.writePrivateJson(this.path, state);
+  }
+}

package/src/types/accounts.ts

@@ -0,0 +1,37 @@
+import { ProviderId } from "./providers";
+
+export type SecretSource =
+  | string
+  | { type: "literal"; value: string }
+  | { type: "env"; name: string }
+  | { type: "file"; path: string }
+  | { type: "command"; command: string }
+  | { type: "op"; reference: string };
+
+export type PiAuthEntry =
+  | { type: "api_key"; key: string }
+  | ({
+      type: "oauth";
+      refresh: string;
+      access: string;
+      expires: number;
+    } & Record<string, unknown>);
+
+export interface AccountConfig {
+  id: string;
+  label: string;
+  provider: ProviderId;
+  /** API-key/env based account. */
+  env?: Record<string, SecretSource>;
+  /** Per-account override for a custom Pi model provider apiKey. */
+  providerApiKey?: SecretSource;
+  /** Uses apiKey from custom provider metadata. */
+  usesProviderApiKey?: boolean;
+  /** Optional model to switch to when this account is activated. */
+  model?: string;
+  /** Captured Pi /login credentials for built-in OAuth/subscription providers. */
+  piAuth?: {
+    provider: ProviderId;
+    entry: PiAuthEntry;
+  };
+}

package/src/types/config.ts

@@ -0,0 +1,6 @@
+import { AccountConfig } from "./accounts";
+
+export interface AccountSwitcherConfig {
+  accounts: AccountConfig[];
+  switchMode?: "env";
+}

package/src/types/context.ts

@@ -0,0 +1,3 @@
+import type { ExtensionCommandContext } from "@mariozechner/pi-coding-agent";
+
+export type AccountSwitcherContext = ExtensionCommandContext;

package/src/types/index.ts

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

package/src/types/providers.ts

@@ -0,0 +1,53 @@
+export type ProviderId = string;
+
+export type ProviderApi =
+  | "anthropic-messages"
+  | "openai-completions"
+  | "openai-responses"
+  | "azure-openai-responses"
+  | "openai-codex-responses"
+  | "mistral-conversations"
+  | "google-generative-ai"
+  | "google-vertex"
+  | "bedrock-converse-stream"
+  | string;
+
+export interface ProviderModelConfig {
+  id: string;
+  name?: string;
+  api?: ProviderApi;
+  baseUrl?: string;
+  reasoning?: boolean;
+  input?: ("text" | "image")[];
+  contextWindow?: number;
+  maxTokens?: number;
+  cost?: {
+    input: number;
+    output: number;
+    cacheRead: number;
+    cacheWrite: number;
+  };
+  compat?: Record<string, unknown>;
+  thinkingLevelMap?: Record<string, string | null>;
+  headers?: Record<string, string>;
+}
+
+export interface ProviderConfig {
+  id: ProviderId;
+  label?: string;
+  /** Alias for Pi provider display name when exported/registered. */
+  name?: string;
+  envKeys?: string[];
+  aliases?: string[];
+  /** Raw Pi auth provider id, when different from this provider's account id. */
+  piAuthProvider?: ProviderId;
+  /** Pi custom model provider config fields. */
+  baseUrl?: string;
+  api?: ProviderApi;
+  apiKey?: string;
+  headers?: Record<string, string>;
+  authHeader?: boolean;
+  compat?: Record<string, unknown>;
+  models?: ProviderModelConfig[];
+  modelOverrides?: Record<string, Partial<ProviderModelConfig>>;
+}

package/src/utils/accounts.ts

@@ -0,0 +1,99 @@
+import * as piAi from "@mariozechner/pi-ai";
+import { readFile } from "node:fs/promises";
+import type { ModelRegistry } from "@mariozechner/pi-coding-agent";
+import type { AccountConfig, SecretSource } from "@/types";
+import { commonUtil } from "./common";
+import { fileUtil } from "./files";
+import { providerUtil } from "./providers";
+
+export const accountUtil = {
+  clearAccountEnv: async (account: AccountConfig, modelRegistry?: ModelRegistry): Promise<void> => {
+    const authProvider = account.piAuth?.provider ?? providerUtil.normalizeProvider(account.provider);
+    if (!account.piAuth && account.env) {
+      for (const envName of Object.keys(account.env)) {
+        delete process.env[envName];
+      }
+    }
+    modelRegistry?.authStorage.removeRuntimeApiKey(authProvider);
+  },
+
+  applyAccountEnv: async (
+    account: AccountConfig,
+    modelRegistry?: ModelRegistry,
+    authProviderOverride?: string,
+  ): Promise<string[]> => {
+    if (account.piAuth) {
+      const authProvider = authProviderOverride ?? account.piAuth.provider;
+      modelRegistry?.authStorage.set(authProvider, account.piAuth.entry);
+      modelRegistry?.authStorage.reload();
+      closeCachedSessions();
+      return [];
+    }
+
+    const resolved = await accountUtil.resolveAccountEnv(account);
+    return accountUtil.applyResolvedAccountEnv(account, resolved, modelRegistry, authProviderOverride);
+  },
+
+  resolveAccountEnv: async (account: AccountConfig): Promise<Array<[string, string]>> => {
+    if (!account.env) return [];
+
+    const resolvedEntries: Array<[string, string]> = [];
+    for (const [envName, source] of Object.entries(account.env)) {
+      const value = await accountUtil.resolveSecret(source);
+      if (!value) throw new Error(`Resolved empty value for ${envName} in account ${account.id}`);
+      resolvedEntries.push([envName, value]);
+    }
+    return resolvedEntries;
+  },
+
+  applyResolvedAccountEnv: (
+    account: AccountConfig,
+    resolvedEntries: Array<[string, string]>,
+    modelRegistry?: ModelRegistry,
+    authProviderOverride?: string,
+  ): string[] => {
+    const authProvider = authProviderOverride ?? providerUtil.normalizeProvider(account.provider);
+    const applied: string[] = [];
+    for (const [envName, value] of resolvedEntries) {
+      process.env[envName] = value;
+      applied.push(envName);
+    }
+
+    const firstValue = resolvedEntries[0]?.[1];
+    if (firstValue) modelRegistry?.authStorage.setRuntimeApiKey(authProvider, firstValue);
+    else modelRegistry?.authStorage.removeRuntimeApiKey(authProvider);
+
+    return applied;
+  },
+
+  resolveSecret: async (source: SecretSource): Promise<string> => {
+    if (typeof source === "string") {
+      if (source.startsWith("op://")) return commonUtil.runOpRead(source);
+      return source;
+    }
+    switch (source.type) {
+      case "literal":
+        return source.value;
+      case "env": {
+        const value = process.env[source.name];
+        if (!value) throw new Error(`Environment variable ${source.name} is not set`);
+        return value;
+      }
+      case "file":
+        return (await readFile(fileUtil.expandHome(source.path), "utf8")).trim();
+      case "command":
+        return commonUtil.runCommand(source.command);
+      case "op":
+        return commonUtil.runOpRead(source.reference);
+    }
+  },
+};
+
+function closeCachedSessions(): void {
+  const helpers = piAi as {
+    cleanupSessionResources?: () => void;
+    closeOpenAICodexWebSocketSessions?: () => void;
+  };
+  helpers.cleanupSessionResources?.();
+  helpers.closeOpenAICodexWebSocketSessions?.();
+}