@hieplp/pi-account-switcher 0.2.3 → 0.3.0

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

@@ -1,11 +1,12 @@
-import AccountSwitcher from "./account-switcher";
+import type AccountSwitcher from "./account-switcher";
+import { createHash } from "node:crypto";
 import { ACCOUNTS_PATH, PROVIDERS_PATH, STATE_PATH } from "@/constants";
 import type { Api, Model } from "@earendil-works/pi-ai";
 import type { ExtensionAPI } from "@earendil-works/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";
+import { accountUtil, findLongestMatchingDir, modelUtil, providerUtil, uiUtil } from "@/utils";
 
 function resolveAuthProvider(account: AccountConfig, providers: ProviderConfig[]): string {
   if (account.piAuth?.provider) return account.piAuth.provider;
@@ -22,58 +23,117 @@ export default class AccountSwitcherRuntime implements AccountSwitcher {
   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);
+  private lastStatusLabel: string | undefined;
+  private sessionKey: string | undefined;
+
+  constructor(
+    private readonly pi: Pick<ExtensionAPI, "registerProvider" | "setModel">,
+    private readonly paths?: { accounts: string; providers: string; state: string },
+  ) {
+    this.providerService = useProviderService(this.pi as ExtensionAPI, paths?.providers ?? PROVIDERS_PATH);
+    this.accountService = useAccountService(paths?.accounts ?? ACCOUNTS_PATH, paths?.state ?? STATE_PATH);
     this.modelService = useModelService(this.pi);
     this.piAuthService = usePiAuthService();
   }
 
+  /** Derive a stable session key from the session manager. */
+  private getSessionKey(ctx: AccountSwitcherContext): string {
+    try {
+      const sm = (ctx as unknown as Record<string, unknown>).sessionManager as Record<string, unknown> | undefined;
+      const sessionFile = typeof sm?.getSessionFile === "function" ? (sm.getSessionFile as () => string)() : undefined;
+      if (sessionFile) return createHash("sha256").update(sessionFile).digest("hex").slice(0, 12);
+    } catch {
+      /* fall through */
+    }
+    return "default";
+  }
+
+  /** Find the account whose dirs contain the longest prefix of cwd. */
+  private findAccountForCwd(cwd: string | undefined): AccountConfig | undefined {
+    if (!cwd) return undefined;
+    const id = findLongestMatchingDir(this.accountService.getAccounts(), cwd);
+    return id ? this.accountService.getAccounts().find((a) => a.id === id) : undefined;
+  }
+
   // ===============================================================================================
   // Core
   // ===============================================================================================
 
   async init(ctx: AccountSwitcherContext): Promise<void> {
+    this.sessionKey = this.getSessionKey(ctx);
+    this.accountService.setSessionKey(this.sessionKey);
     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));
+    // Cascade:
+    // 0. PI_ACCOUNT_SWITCHER_ACTIVE_ID env var (from parent process)
+    // 1. Session key state (handled inside accountService.load())
+    // 2. CWD-based auto-select via dirs
+    // 3. defaultAccountId from config
+    let selected: AccountConfig | undefined;
+
+    // Step 0: env var from parent process (for subagent inheritance)
+    // PI_ACCOUNT_SWITCHER_NEXT_ID is a one-shot override (consumed after first read)
+    // PI_ACCOUNT_SWITCHER_ACTIVE_ID is the persistent inheritance from the parent
+    const nextId = process.env.PI_ACCOUNT_SWITCHER_NEXT_ID;
+    if (nextId) {
+      delete process.env.PI_ACCOUNT_SWITCHER_NEXT_ID;
+      selected = this.accountService.getAccounts().find((a) => a.id === nextId);
+    }
+    if (!selected) {
+      const envId = process.env.PI_ACCOUNT_SWITCHER_ACTIVE_ID;
+      if (envId) {
+        selected = this.accountService.getAccounts().find((a) => a.id === envId);
+      }
     }
 
-    // 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 providers = this.providerService.getProviders();
-      const activeProvider = active ? resolveAccountProvider(active, providers) : undefined;
-      const savedProvider = providerUtil.normalizeProviderWithCustom(modelState.provider, providers);
-
-      // Only restore a saved model when it belongs to the active account's provider.
-      // Otherwise Pi can start with credentials for one provider and a model from another.
-      if (!activeProvider || savedProvider === activeProvider) {
-        const model = ctx.modelRegistry.find(modelState.provider, modelState.id);
-        if (model) await this.modelService.applyModel(model, ctx);
+    if (!selected) {
+      // Step 1: session key state (restored by accountService.load())
+      selected = this.accountService.getActiveAccount();
+    }
+    if (!selected) {
+      // Step 2: CWD-based auto-select via dirs
+      selected = this.findAccountForCwd(ctx.cwd);
+    }
+    if (!selected) {
+      // Step 3: defaultAccountId from config
+      const defaultId = await this.accountService.getDefaultAccountId();
+      if (defaultId) {
+        selected = this.accountService.getAccounts().find((a) => a.id === defaultId);
       }
     }
+    if (selected) {
+      const providers = this.providerService.getProviders();
+      await this.applyProviderApiKey(selected, providers);
+      await this.accountService.activateAccount(selected, ctx, resolveAuthProvider(selected, providers));
+    }
+
+    uiUtil.setAccountStatus(ctx.ui, selected?.label);
+  }
+
+  refreshStatus(ctx: AccountSwitcherContext): void {
+    const active = this.accountService.getActiveAccount();
+    const label = active?.label;
+    if (label === this.lastStatusLabel) return;
+    this.lastStatusLabel = label;
+    uiUtil.setAccountStatus(ctx.ui, label);
   }
 
   async load(): Promise<void> {
-    await this.accountService.load();
     await this.providerService.load();
+    await this.accountService.load();
   }
 
   async onModelSelect(provider: string, ctx: AccountSwitcherContext): Promise<void> {
-    const matchingAccount = this.findAccountsByProvider(provider)[0];
+    const providers = this.providerService.getProviders();
+    const normalizedProvider = providerUtil.normalizeProviderWithCustom(provider, providers);
     const activeAccount = this.accountService.getActiveAccount();
+
+    // If the active account already belongs to this provider, keep it.
+    if (activeAccount && resolveAccountProvider(activeAccount, providers) === normalizedProvider) {
+      return;
+    }
+
+    const matchingAccount = this.findAccountsByProvider(provider)[0];
     if (matchingAccount && matchingAccount.id !== activeAccount?.id) {
       await this.activateAccount(matchingAccount, ctx);
     }
@@ -130,6 +190,9 @@ export default class AccountSwitcherRuntime implements AccountSwitcher {
     const providerApiKey = await this.applyProviderApiKey(account, providers);
     const result = await this.accountService.activateAccount(account, ctx, resolveAuthProvider(account, providers));
 
+    // Persist the active account ID for subagent (cross-process) inheritance
+    process.env.PI_ACCOUNT_SWITCHER_ACTIVE_ID = account.id;
+
     // 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 = resolveAccountProvider(account, providers);
@@ -141,6 +204,11 @@ export default class AccountSwitcherRuntime implements AccountSwitcher {
     if (accountProvider !== currentProvider) {
       const model = await modelUtil.pickModel(ctx, account, providers, accountProvider);
       if (model) await this.applyModel(model, ctx);
+    } else {
+      // Same provider — persist current model for full session tracking
+      if (ctx.model) {
+        await this.accountService.saveActiveModel(ctx.model.id, ctx.model.provider);
+      }
     }
 
     return providerApiKey ? `provider apiKey (${providerApiKey})` : result;

package/src/runtime/account-switcher.ts

@@ -4,6 +4,7 @@ import type { AccountConfig, AccountSwitcherContext, PiAuthEntry, ProviderConfig
 export default interface AccountSwitcher {
   // Core
   init(ctx: AccountSwitcherContext): Promise<void>;
+  refreshStatus(ctx: AccountSwitcherContext): void;
   load(): Promise<void>;
   onModelSelect(provider: string, ctx: AccountSwitcherContext): Promise<void>;
 
@@ -27,7 +28,7 @@ export default interface AccountSwitcher {
   // Provider
   getProviders(): ProviderConfig[];
   registerProvider(provider: ProviderConfig): void;
-  addProvider(config: ProviderConfig): Promise<void>;
+  addProvider(provider: ProviderConfig): Promise<void>;
   editProvider(original: ProviderConfig, updated: ProviderConfig): Promise<void>;
   removeProvider(provider: ProviderConfig): Promise<void>;
 }

package/src/schemas/accounts.ts

@@ -30,6 +30,7 @@ export const accountSchema = z
     env: z.record(z.string().min(1), secretSourceSchema).optional(),
     providerApiKey: secretSourceSchema.optional(),
     usesProviderApiKey: z.boolean().optional(),
+    dirs: z.array(z.string()).optional(),
     piAuth: z
       .object({
         provider: z.string().min(1),

package/src/schemas/config.ts

@@ -4,4 +4,6 @@ import { accountSchema } from "./accounts";
 export const configSchema = z.object({
   accounts: z.array(accountSchema).default([]),
   switchMode: z.literal("env").optional().default("env"),
-});
+  defaultAccountId: z.string().optional(),
+  stateCleanupDays: z.number().int().min(1).optional().default(30),
+});

package/src/services/accounts.ts

@@ -4,6 +4,7 @@ import { accountUtil, providerUtil, uiUtil } from "@/utils";
 
 export interface AccountService {
   load(): Promise<void>;
+  setSessionKey(sessionKey: string): void;
   getAccounts(): AccountConfig[];
   findAccountsByProvider(provider: string, providers: ProviderConfig[]): AccountConfig[];
   getActiveAccount(): AccountConfig | undefined;
@@ -13,6 +14,8 @@ export interface AccountService {
   activateAccount(account: AccountConfig, ctx: AccountSwitcherContext, authProvider?: string): Promise<string>;
   getActiveModelState(): { id: string; provider: string } | undefined;
   saveActiveModel(id: string, provider: string): Promise<void>;
+  setDefaultAccountId(id: string): Promise<void>;
+  getDefaultAccountId(): Promise<string | undefined>;
 }
 
 export function useAccountService(accountsPath: string, statePath?: string): AccountService {
@@ -28,18 +31,51 @@ class AccountServiceImpl implements AccountService {
   private activeAccountId: string | undefined;
   private activeModelId: string | undefined;
   private activeModelProvider: string | undefined;
+  private sessionKey: string | undefined;
 
   constructor(
     private readonly store: AccountStore,
     private readonly stateStore: StateStore,
   ) {}
 
+  setSessionKey(sessionKey: string): void {
+    this.sessionKey = sessionKey;
+  }
+
   async load(): Promise<void> {
+    // Apply config TTL to state store before loading
+    const config = await this.store.loadConfig();
+    this.stateStore.setCleanupDays(config.stateCleanupDays ?? 30);
+
     this.accounts = await this.store.load();
-    const state = await this.stateStore.load();
-    this.activeAccountId = state.activeAccountId;
-    this.activeModelId = state.activeModelId;
-    this.activeModelProvider = state.activeModelProvider;
+    const key = this.sessionKey ?? "default";
+    const state = await this.stateStore.loadSession(key);
+
+    // Session key state only — no fallback to defaultAccountId here.
+    // The full cascade (session → dirs → defaultAccountId) is in runtime.init()
+    // so dir-based matching runs before we fall back to the default.
+    if (state.activeAccountId) {
+      this.activeAccountId = state.activeAccountId;
+      this.activeModelId = state.activeModelId;
+      this.activeModelProvider = state.activeModelProvider;
+    }
+
+    // Legacy cleanup (always runs): if sessions.default still exists from
+    // old format migration, promote its account to config-level default
+    // (if not set yet) and clear the key. This is dead data — sessions
+    // should only have per-session keys once migration completes.
+    if (await this.stateStore.sessionExists("default")) {
+      const legacyDefault = await this.stateStore.loadSession("default");
+      if (legacyDefault.activeAccountId && !(await this.getDefaultAccountId())) {
+        await this.setDefaultAccountId(legacyDefault.activeAccountId);
+      }
+      await this.stateStore.deleteSession("default");
+      if (!this.activeAccountId && legacyDefault.activeAccountId) {
+        this.activeAccountId = legacyDefault.activeAccountId;
+        this.activeModelId = legacyDefault.activeModelId;
+        this.activeModelProvider = legacyDefault.activeModelProvider;
+      }
+    }
   }
 
   getAccounts(): AccountConfig[] {
@@ -100,17 +136,28 @@ class AccountServiceImpl implements AccountService {
       applied = accountUtil.applyResolvedAccountEnv(account, resolved, ctx.modelRegistry, authProvider);
     }
     this.activeAccountId = account.id;
+    // Persist active account ID for subagent (cross-process) inheritance
+    process.env.PI_ACCOUNT_SWITCHER_ACTIVE_ID = account.id;
     await this.flushState();
     uiUtil.setAccountStatus(ctx.ui, account.label);
     if (account.piAuth) return "via OAuth";
     return applied.length > 0 ? applied.join(", ") : "";
   }
 
+  async setDefaultAccountId(id: string): Promise<void> {
+    await this.store.setDefaultAccountId(id);
+  }
+
+  async getDefaultAccountId(): Promise<string | undefined> {
+    const config = await this.store.loadConfig();
+    return config.defaultAccountId;
+  }
+
   private async flushState(): Promise<void> {
-    await this.stateStore.save({
+    await this.stateStore.saveSession(this.sessionKey ?? "default", {
       activeAccountId: this.activeAccountId,
       activeModelId: this.activeModelId,
       activeModelProvider: this.activeModelProvider,
     });
   }
-}
+}

package/src/storage/accounts.ts

@@ -5,9 +5,11 @@ import { errorUtil, fileUtil } from "@/utils";
 
 export interface AccountStore {
   load(): Promise<AccountConfig[]>;
+  loadConfig(): Promise<AccountSwitcherConfig>;
   addAccount(account: AccountConfig): Promise<AccountConfig[]>;
   replaceAccount(originalId: string, account: AccountConfig): Promise<AccountConfig[]>;
   removeAccount(id: string): Promise<AccountConfig[]>;
+  setDefaultAccountId(id: string): Promise<void>;
 }
 
 export function useAccountStore(path: string) {
@@ -22,18 +24,20 @@ class AccountStoreImpl implements AccountStore {
   constructor(private readonly path: string) {}
 
   async load(): Promise<AccountConfig[]> {
+    const config = await this.loadConfig();
+    return config.accounts;
+  }
+
+  async loadConfig(): Promise<AccountSwitcherConfig> {
     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;
+      accountValidator.assertNoDuplicateAccounts(parsed.accounts);
+      accountValidator.assertAccountsHaveCredentials(parsed.accounts);
+      return parsed;
     } catch (error) {
       if (fileUtil.isMissingFileError(error)) {
-        return [];
+        return { accounts: [] };
       }
       throw new Error(`Failed to load accounts at ${this.path}: ${errorUtil.format(error)}`);
     }
@@ -67,9 +71,23 @@ class AccountStoreImpl implements AccountStore {
   }
 
   private async save(accounts: AccountConfig[]): Promise<void> {
-    const config: AccountSwitcherConfig = { accounts };
+    const { defaultAccountId, stateCleanupDays } = await this.loadConfig();
+    const config: AccountSwitcherConfig = {
+      accounts,
+      ...(defaultAccountId ? { defaultAccountId } : {}),
+      ...(stateCleanupDays !== undefined ? { stateCleanupDays } : {}),
+    };
     await fileUtil.writePrivateJson(this.path, config);
   }
+
+  async setDefaultAccountId(id: string): Promise<void> {
+    const config = await this.loadConfig();
+    if (!config.accounts.some((a) => a.id === id)) {
+      throw new Error(`Account not found: ${id}`);
+    }
+    const updated: AccountSwitcherConfig = { ...config, defaultAccountId: id };
+    await fileUtil.writePrivateJson(this.path, updated);
+  }
 }
 
 // ===============================================================================================
@@ -106,4 +124,4 @@ const accountValidator = {
       throw new Error(`Accounts missing credentials: ${missing.sort().join(", ")}`);
     }
   },
-};
+};

package/src/storage/state.ts

@@ -3,36 +3,77 @@ import z from "zod";
 import { STATE_PATH } from "@/constants";
 import { fileUtil } from "@/utils";
 
-const appStateSchema = z.object({
+const sessionStateSchema = z.object({
   activeAccountId: z.string().optional(),
   activeModelId: z.string().optional(),
   activeModelProvider: z.string().optional(),
+  lastActive: z.string().optional(),
 });
 
-export interface AppState {
+const appStateSchema = z.preprocess(
+  (raw) => {
+    // Migrate legacy flat format: { activeAccountId, ... } → { sessions: { default: { ... } } }
+    if (
+      raw &&
+      typeof raw === "object" &&
+      !("sessions" in raw) &&
+      ("activeAccountId" in raw || "activeModelId" in raw)
+    ) {
+      return { sessions: { default: raw } };
+    }
+    return raw;
+  },
+  z.object({
+    sessions: z.record(z.string(), sessionStateSchema).default({}),
+  }),
+);
+
+export interface SessionState {
   activeAccountId?: string;
   activeModelId?: string;
   activeModelProvider?: string;
+  lastActive?: string;
+}
+
+export interface AppState {
+  sessions: Record<string, SessionState>;
 }
 
 export interface StateStore {
   load(): Promise<AppState>;
   save(state: AppState): Promise<void>;
+  loadSession(sessionKey: string): Promise<SessionState>;
+  saveSession(sessionKey: string, state: SessionState): Promise<void>;
+  deleteSession(sessionKey: string): Promise<void>;
+  sessionExists(sessionKey: string): Promise<boolean>;
+  setCleanupDays(days: number): void;
 }
 
 export function useStateStore(path = STATE_PATH): StateStore {
   return new StateStoreImpl(path);
 }
 
+// ===============================================================================================
+// State Store Implementation
+// ===============================================================================================
+
+const HARD_CAP = 500;
+
 class StateStoreImpl implements StateStore {
+  private cleanupDays = 30;
+
   constructor(private readonly path: string) {}
 
+  setCleanupDays(days: number): void {
+    this.cleanupDays = days;
+  }
+
   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 {};
+      if (fileUtil.isMissingFileError(error)) return { sessions: {} };
       throw error;
     }
   }
@@ -40,4 +81,81 @@ class StateStoreImpl implements StateStore {
   async save(state: AppState): Promise<void> {
     await fileUtil.writePrivateJson(this.path, state);
   }
-}
+
+  async loadSession(sessionKey: string): Promise<SessionState> {
+    const appState = await this.load();
+    return appState.sessions[sessionKey] ?? {};
+  }
+
+  async saveSession(sessionKey: string, state: SessionState): Promise<void> {
+    const appState = await this.load();
+
+    // Migration: timestamp any pre-upgrade entries so they become TTL-eligible
+    for (const [key, entry] of Object.entries(appState.sessions)) {
+      if (!entry.lastActive) {
+        appState.sessions[key] = { ...entry, lastActive: new Date().toISOString() };
+      }
+    }
+
+    // Stamp lastActive on the session being written
+    const stampedState: SessionState = { ...state, lastActive: new Date().toISOString() };
+    appState.sessions[sessionKey] = stampedState;
+
+    // Prune before saving
+    const pruned = this.pruneSessions(appState.sessions);
+    appState.sessions = pruned;
+
+    await this.save(appState);
+  }
+
+  async deleteSession(sessionKey: string): Promise<void> {
+    const appState = await this.load();
+    delete appState.sessions[sessionKey];
+    await this.save(appState);
+  }
+
+  async sessionExists(sessionKey: string): Promise<boolean> {
+    const appState = await this.load();
+    return sessionKey in appState.sessions;
+  }
+
+  /**
+   * Prune sessions in two passes:
+   * 1. TTL eviction: remove entries where lastActive exists AND entry is older than cleanupDays
+   * 2. Hard cap: if entries > 500, keep only the 500 most recent
+   *
+   * Entries without lastActive are preserved (pre-upgrade safety).
+   * Entries with invalid lastActive are treated as "recent enough" (not evicted).
+   */
+  private pruneSessions(sessions: Record<string, SessionState>): Record<string, SessionState> {
+    const entries = Object.entries(sessions);
+    const now = Date.now();
+    const ttlMs = this.cleanupDays * 86400000;
+
+    // TTL pass — keep entries that are either:
+    // - without lastActive (pre-upgrade safety)
+    // - with valid lastActive that is within TTL
+    const ttlFiltered = entries.filter(([, state]) => {
+      if (!state.lastActive) return true; // Preserve pre-upgrade entries
+      const ts = new Date(state.lastActive).getTime();
+      if (isNaN(ts)) return true; // Preserve invalid timestamps
+      return now - ts <= ttlMs;
+    });
+
+    // Hard cap pass — keep only the 500 most recent
+    if (ttlFiltered.length <= HARD_CAP) {
+      return Object.fromEntries(ttlFiltered);
+    }
+
+    // Sort by lastActive ascending (oldest first), treating undefined as 0
+    const sorted = [...ttlFiltered].sort(([, a], [, b]) => {
+      const aTs = a.lastActive ? new Date(a.lastActive).getTime() : 0;
+      const bTs = b.lastActive ? new Date(b.lastActive).getTime() : 0;
+      return aTs - bTs;
+    });
+
+    // Keep the last 500 (most recent)
+    const trimmed = sorted.slice(-HARD_CAP);
+    return Object.fromEntries(trimmed);
+  }
+}

package/src/types/accounts.ts

@@ -29,6 +29,8 @@ export interface AccountConfig {
   usesProviderApiKey?: boolean;
   /** Optional model to switch to when this account is activated. */
   model?: string;
+  /** Directories where this account should be the default (CWD-based auto-select). */
+  dirs?: string[];
   /** Captured Pi /login credentials for built-in OAuth/subscription providers. */
   piAuth?: {
     provider: ProviderId;

package/src/types/config.ts

@@ -3,4 +3,8 @@ import { AccountConfig } from "./accounts";
 export interface AccountSwitcherConfig {
   accounts: AccountConfig[];
   switchMode?: "env";
-}
+  /** Config-level fallback account when no session state or dirs match. */
+  defaultAccountId?: string;
+  /** Number of days after which inactive sessions are eligible for GC. Default: 30. */
+  stateCleanupDays?: number;
+}

package/src/utils/accounts.ts

@@ -1,4 +1,3 @@
-import * as piAi from "@earendil-works/pi-ai";
 import { readFile } from "node:fs/promises";
 import type { ModelRegistry } from "@earendil-works/pi-coding-agent";
 import type { AccountConfig, SecretSource } from "@/types";
@@ -8,6 +7,8 @@ import { providerUtil } from "./providers";
 
 export const accountUtil = {
   clearAccountEnv: async (account: AccountConfig, modelRegistry?: ModelRegistry): Promise<void> => {
+    // Clear the cross-process inheritance env var
+    delete process.env.PI_ACCOUNT_SWITCHER_ACTIVE_ID;
     const authProvider = account.piAuth?.provider ?? providerUtil.normalizeProvider(account.provider);
     if (!account.piAuth && account.env) {
       for (const envName of Object.keys(account.env)) {
@@ -89,11 +90,69 @@ export const accountUtil = {
   },
 };
 
+function normalizeDir(dir: string): string {
+  return dir.replace(/\/$/, "");
+}
+
+/**
+ * Check if an account has a specific directory.
+ * Normalizes trailing slashes before comparison.
+ */
+export function hasDir<T extends { dirs?: string[] }>(account: T, dir: string): boolean {
+  const dirs = account.dirs;
+  if (!dirs || dirs.length === 0) return false;
+  const normalized = normalizeDir(dir);
+  return dirs.some((d) => normalizeDir(d) === normalized);
+}
+
+/**
+ * Add a directory to an account.
+ * Returns a new AccountConfig with the dir added, or null if the dir already exists.
+ * Dirs are kept sorted for consistent display.
+ */
+export function addDirToAccount<T extends { id: string; label: string; provider: string; dirs?: string[] }>(
+  account: T,
+  dir: string,
+): T | null {
+  if (hasDir(account, dir)) return null;
+
+  const existing = account.dirs ?? [];
+  const newDirs = [...existing, dir].sort();
+  return { ...account, dirs: newDirs };
+}
+
+/**
+ * Remove a directory from an account.
+ * Returns a new AccountConfig with the dir removed, or null if the dir does not exist.
+ */
+export function removeDirFromAccount<T extends { id: string; label: string; provider: string; dirs?: string[] }>(
+  account: T,
+  dir: string,
+): T | null {
+  const dirs = account.dirs;
+  if (!dirs || dirs.length === 0) return null;
+
+  const normalized = normalizeDir(dir);
+  const filtered = dirs.filter((d) => normalizeDir(d) !== normalized);
+
+  if (filtered.length === dirs.length) return null;
+  if (filtered.length === 0) return { ...account, dirs: undefined };
+  return { ...account, dirs: filtered };
+}
+
 function closeCachedSessions(): void {
-  const helpers = piAi as {
-    cleanupSessionResources?: () => void;
-    closeOpenAICodexWebSocketSessions?: () => void;
-  };
-  helpers.cleanupSessionResources?.();
-  helpers.closeOpenAICodexWebSocketSessions?.();
+  // Dynamic import so the module is not required at load time — @earendil-works/pi-ai
+  // is a peerDependency provided by the pi agent host, not bundled with this package.
+  import("@earendil-works/pi-ai")
+    .then((piAi) => {
+      const helpers = piAi as {
+        cleanupSessionResources?: () => void;
+        closeOpenAICodexWebSocketSessions?: () => void;
+      };
+      helpers.cleanupSessionResources?.();
+      helpers.closeOpenAICodexWebSocketSessions?.();
+    })
+    .catch(() => {
+      // pi-ai not available in this environment — skip session cleanup
+    });
 }